mtmd: use ggml_rope_set_offset (#27521 )

* mtmd: use ggml_rope_set_offset * add comment
common: add json.h abstraction (#27511 )
2026-08-23 13:12:34 +02:00 · 2026-08-22 16:33:47 +02:00 · 2026-08-22 16:28:28 +02:00 · 2026-08-22 16:16:06 +02:00 · 2026-08-22 15:53:56 +02:00 · 2026-08-22 13:28:30 +03:00
113 changed files with 14368 additions and 13023 deletions
@@ -1,23 +1,88 @@
 # note: place this as the last step of the job, so the new cache is saved by "Post ccache" right after the old one is cleared
 name: "ccache-clear"
-description: "Delete all GitHub Actions caches matching a key prefix"
+description: "Delete GitHub Actions caches matching a key prefix, oldest first"
 inputs:
  key:
    description: "Cache key prefix to match and delete"
    required: true
+  older:
+    description: "Only delete caches created more than this long ago (e.g. 90m, 1h, 1d). By default all matching caches are deleted"
+    required: false
+    default: ""
+  min:
+    description: "Stop deleting if fewer than this many caches would remain (e.g. 1). By default there is no minimum"
+    required: false
+    default: "0"
+  dry-run:
+    description: "Only print the caches that would be deleted, without deleting them"
+    required: false
+    default: "false"

 runs:
  using: "composite"
  steps:
    - name: Clear caches
      shell: bash
+      env:
+        CLEAR_KEY: ${{ inputs.key }}
+        CLEAR_OLDER: ${{ inputs.older }}
+        CLEAR_MIN: ${{ inputs.min }}
+        CLEAR_DRY_RUN: ${{ inputs.dry-run }}
      run: |
-        CACHES=$(gh cache list --key "ccache-${{ inputs.key }}" --json id,key --jq '.[] | "\(.id) \(.key)"' 2>/dev/null)
+        # Convert a duration (e.g. 90m, 1h, 1d, plain seconds) to seconds
+        to_seconds() {
+          local val="$1"
+          [[ "$val" =~ ^[0-9]+$ ]] && { echo "$val"; return 0; }
+          local num="${val%?}" unit="${val: -1}" mult
+          [[ "$num" =~ ^[0-9]+$ ]] || return 1
+          case "$unit" in
+            s) mult=1 ;;
+            m) mult=60 ;;
+            h) mult=3600 ;;
+            d) mult=86400 ;;
+            *) return 1 ;;
+          esac
+          echo $((num * mult))
+        }
+
+        [[ "$CLEAR_MIN" =~ ^[0-9]+$ ]] || { echo "Invalid min value: $CLEAR_MIN" >&2; exit 1; }
+        [[ "$CLEAR_DRY_RUN" =~ ^(true|false)$ ]] || { echo "Invalid dry-run value: $CLEAR_DRY_RUN" >&2; exit 1; }
+
+        CACHES=$(gh cache list --key "ccache-$CLEAR_KEY" --json id,key,createdAt --jq '.[] | [.createdAt, .id, .key] | @tsv' 2>/dev/null | LC_ALL=C sort)
        if [ -z "$CACHES" ]; then
-          echo "No caches found with key prefix: ${{ inputs.key }}"
+          echo "No caches found with key prefix: $CLEAR_KEY"
          exit 0
        fi
-        while read -r id key; do
-          echo "Deleting cache: $id ($key)"
-          gh cache delete "$id"
+
+        TOTAL=$(( $(wc -l <<< "$CACHES") ))
+
+        echo "Found $TOTAL cache(s) with key prefix: $CLEAR_KEY (oldest first):"
+        while IFS=$'\t' read -r CREATED ID KEY; do
+          printf '  %s  %s  %s\n' "$CREATED" "$ID" "$KEY"
+        done <<< "$CACHES"
+
+        CUTOFF=""
+        if [ -n "$CLEAR_OLDER" ]; then
+          OLDER_SECONDS=$(to_seconds "$CLEAR_OLDER") || { echo "Invalid older value: $CLEAR_OLDER (expected e.g. 90m, 1h, 1d)" >&2; exit 1; }
+          CUTOFF=$(( $(date +%s) - OLDER_SECONDS ))
+        fi
+
+        # Caches are sorted oldest first
+        DELETED=0
+        while IFS=$'\t' read -r CREATED ID KEY; do
+          if [ -n "$CUTOFF" ] && [ "$(date -d "$CREATED" +%s)" -ge "$CUTOFF" ]; then
+            echo "Rest are not older than $CLEAR_OLDER, stopping"
+            break
+          fi
+          if [ $((TOTAL - DELETED - 1)) -lt "$CLEAR_MIN" ]; then
+            echo "Keeping at least $CLEAR_MIN cache(s), stopping"
+            break
+          fi
+          if [ "$CLEAR_DRY_RUN" = "true" ]; then
+            echo "Would delete cache: $ID ($KEY)"
+          else
+            echo "Deleting cache: $ID ($KEY)"
+            gh cache delete "$ID"
+          fi
+          DELETED=$((DELETED + 1))
        done <<< "$CACHES"
@@ -117,6 +117,18 @@ jobs:
          ./bin/llama-convert-llama2c-to-ggml --copy-vocab-from-model ./tok512.bin --llama2c-model stories260K.bin --llama2c-output-model stories260K.gguf
          ./bin/llama-completion -m stories260K.gguf -p "One day, Lily met a Shoggoth" -n 500 -c 256

+      # note: real deletion only on push to master (same condition as the ccache save),
+      #       dry-run otherwise (the token is read-only on PRs from forks)
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: cpu-${{ matrix.os }}
+          older: 1h
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
+
  windows:
    name: windows / ${{ matrix.build }}
    runs-on: windows-2025
@@ -774,6 +774,7 @@ jobs:
        with:
          key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
          evict-old-files: 1d
+          max-size: "1G"

      # - name: Cache ROCm Installation
      #   id: cache-rocm
@@ -1286,123 +1287,134 @@ jobs:
        with:
          key: release-ubuntu-24.04-sycl-${{ matrix.build }}

-#   ubuntu-22-rocm:
-#     needs: [check-release, get-version]
-#     if: ${{ needs.check-release.outputs.should_release == 'true' }}
+  ubuntu-22-rocm:
+    needs: [check-release, get-version]
+    if: ${{ needs.check-release.outputs.should_release == 'true' }}

-#     runs-on: ubuntu-22.04
+    runs-on: ubuntu-22.04

-#     permissions:
-#       actions: write
+    permissions:
+      actions: write

-#     strategy:
-#       matrix:
-#         include:
-#           - ROCM_VERSION: "7.14.0"
-#             gpu_targets: "gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201"
-#             build: 'x64'
+    strategy:
+      matrix:
+        include:
+          - ROCM_VERSION: "7.14.0"
+            gpu_targets: "gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201"
+            build: 'x64'

-#     steps:
-#       - name: Clone
-#         id: checkout
-#         uses: actions/checkout@v6
-#         with:
-#           fetch-depth: 0
+    steps:
+      - name: Clone
+        id: checkout
+        uses: actions/checkout@v6
+        with:
+          fetch-depth: 0

-#       - name: Setup Node.js
-#         uses: actions/setup-node@v6
-#         with:
-#           node-version: "24"
-#           cache: "npm"
-#           cache-dependency-path: "tools/ui/package-lock.json"
+      - name: Setup Node.js
+        uses: actions/setup-node@v6
+        with:
+          node-version: "24"
+          cache: "npm"
+          cache-dependency-path: "tools/ui/package-lock.json"

-#       - name: Free up disk space
-#         uses: ggml-org/free-disk-space@v1.3.1
-#         with:
-#           tool-cache: true
+      - name: Free up disk space
+        uses: ggml-org/free-disk-space@v1.3.1
+        with:
+          tool-cache: true

-#       # - name: ccache
-#       #   uses: ggml-org/ccache-action@v1.2.21
-#       #   with:
-#       #     key: release-ubuntu-22.04-rocm-${{ matrix.ROCM_VERSION }}
+      - name: ccache
+        uses: ggml-org/ccache-action@v1.2.21
+        with:
+          key: release-ubuntu-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
+          evict-old-files: 1d
+          max-size: "1G"

-#       - name: Dependencies
-#         id: depends
-#         run: |
-#           sudo apt install -y build-essential git cmake wget
+      - name: Tune ccache for reinstalled ROCm toolchain
+        run: |
+          # ROCm is pip-installed fresh each run, so the clang binary's mtime
+          # changes every time. With the default compiler_check=mtime that
+          # invalidates the cache; hash compiler contents instead so warm
+          # builds hit.
+          ccache --set-config=compiler_check=content
+          ccache --set-config=sloppiness=time_macros,include_file_mtime,include_file_ctime

-#       - name: Setup TheRock with Wheels
-#         id: therock_env
-#         run: |
-#           # Create Python virtual environment
-#           python3 -m venv .venv
-#           source .venv/bin/activate
+      - name: Dependencies
+        id: depends
+        run: |
+          sudo apt install -y build-essential git cmake wget

-#           # Install ROCm wheels for build
-#           # libraries = HIP runtime and CMake configs needed for linking
-#           # devel = compilers, headers, static libs
-#           python -m pip install --upgrade pip
-#           python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ matrix.ROCM_VERSION }}"
+      - name: Setup TheRock with Wheels
+        id: therock_env
+        run: |
+          # Create Python virtual environment
+          python3 -m venv .venv
+          source .venv/bin/activate

-#           # Get ROCm installation paths using the rocm-sdk CLI tool
-#           ROCM_PATH=$(rocm-sdk path --root)
-#           CMAKE_PATH=$(rocm-sdk path --cmake)
-#           BIN_PATH=$(rocm-sdk path --bin)
-#           echo "ROCM_PATH=$ROCM_PATH"
-#           echo "CMAKE_PATH=$CMAKE_PATH"
-#           echo "BIN_PATH=$BIN_PATH"
+          # Install ROCm wheels for build
+          # libraries = HIP runtime and CMake configs needed for linking
+          # devel = compilers, headers, static libs
+          python -m pip install --upgrade pip
+          python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ matrix.ROCM_VERSION }}"

-#           # Set environment variables
-#           echo "ROCM_PATH=$ROCM_PATH" >> $GITHUB_ENV
-#           echo "CMAKE_PREFIX_PATH=$CMAKE_PATH" >> $GITHUB_ENV
-#           echo "HIP_PATH=$ROCM_PATH" >> $GITHUB_ENV
-#           echo "PATH=$BIN_PATH:${PATH}" >> $GITHUB_ENV
-#           echo "LD_LIBRARY_PATH=$ROCM_PATH/lib:${LD_LIBRARY_PATH:-}" >> $GITHUB_ENV
+          # Get ROCm installation paths using the rocm-sdk CLI tool
+          ROCM_PATH=$(rocm-sdk path --root)
+          CMAKE_PATH=$(rocm-sdk path --cmake)
+          BIN_PATH=$(rocm-sdk path --bin)
+          echo "ROCM_PATH=$ROCM_PATH"
+          echo "CMAKE_PATH=$CMAKE_PATH"
+          echo "BIN_PATH=$BIN_PATH"

-#           # Keep venv activated for subsequent steps
-#           echo "$(pwd)/.venv/bin" >> $GITHUB_PATH
+          # Set environment variables
+          echo "ROCM_PATH=$ROCM_PATH" >> $GITHUB_ENV
+          echo "CMAKE_PREFIX_PATH=$CMAKE_PATH" >> $GITHUB_ENV
+          echo "HIP_PATH=$ROCM_PATH" >> $GITHUB_ENV
+          echo "PATH=$BIN_PATH:${PATH}" >> $GITHUB_ENV
+          echo "LD_LIBRARY_PATH=$ROCM_PATH/lib:${LD_LIBRARY_PATH:-}" >> $GITHUB_ENV

-#       - name: Build with native CMake HIP support
-#         id: cmake_build
-#         run: |
-#           cmake -B build -S . \
-#             -DCMAKE_HIP_COMPILER="$(hipconfig -l)/clang" \
-#             -DCMAKE_BUILD_TYPE=Release \
-#             -DGGML_BACKEND_DL=ON \
-#             -DGGML_NATIVE=OFF \
-#             -DCMAKE_INSTALL_RPATH='$ORIGIN' \
-#             -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
-#             -DGGML_CPU_ALL_VARIANTS=ON \
-#             -DGPU_TARGETS="${{ matrix.gpu_targets }}" \
-#             -DGGML_HIP=ON \
-#             -DHIP_PLATFORM=amd \
-#             -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
-#             ${{ env.CMAKE_ARGS }}
-#           cmake --build build --config Release -j $(nproc)
+          # Keep venv activated for subsequent steps
+          echo "$(pwd)/.venv/bin" >> $GITHUB_PATH

-#       - name: Determine tag name
-#         id: tag
-#         uses: ./.github/actions/get-tag-name
+      - name: Build with native CMake HIP support
+        id: cmake_build
+        run: |
+          cmake -B build -S . \
+            -DCMAKE_HIP_COMPILER="$(hipconfig -l)/clang" \
+            -DCMAKE_BUILD_TYPE=Release \
+            -DGGML_BACKEND_DL=ON \
+            -DGGML_NATIVE=OFF \
+            -DCMAKE_INSTALL_RPATH='$ORIGIN' \
+            -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
+            -DGGML_CPU_ALL_VARIANTS=ON \
+            -DGPU_TARGETS="${{ matrix.gpu_targets }}" \
+            -DGGML_HIP=ON \
+            -DHIP_PLATFORM=amd \
+            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
+            ${{ env.CMAKE_ARGS }}
+          cmake --build build --config Release -j $(nproc)

-#       - name: Get ROCm short version
-#         run: echo "ROCM_VERSION_SHORT=$(echo '${{ matrix.ROCM_VERSION }}' | cut -d '.' -f 1,2)" >> $GITHUB_ENV
+      - name: Determine tag name
+        id: tag
+        uses: ./.github/actions/get-tag-name

-#       - name: Pack artifacts
-#         id: pack_artifacts
-#         run: |
-#           cp LICENSE ./build/bin/
-#           tar -czvf llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz --transform "s,^\.,llama-${{ steps.tag.outputs.name }}," -C ./build/bin .
+      - name: Get ROCm short version
+        run: echo "ROCM_VERSION_SHORT=$(echo '${{ matrix.ROCM_VERSION }}' | cut -d '.' -f 1,2)" >> $GITHUB_ENV

-#       - name: Upload artifacts
-#         uses: actions/upload-artifact@v6
-#         with:
-#           path: llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz
-#           name: llama-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz
+      - name: Pack artifacts
+        id: pack_artifacts
+        run: |
+          cp LICENSE ./build/bin/
+          tar -czvf llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz --transform "s,^\.,llama-${{ steps.tag.outputs.name }}," -C ./build/bin .

-#       # - name: ccache-clear
-#       #   uses: ./.github/actions/ccache-clear
-#       #   with:
-#       #     key: release-ubuntu-22.04-rocm-${{ matrix.ROCM_VERSION }}
+      - name: Upload artifacts
+        uses: actions/upload-artifact@v6
+        with:
+          path: llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz
+          name: llama-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz
+
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        with:
+          key: release-ubuntu-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}

  ios-xcode:
    needs: [check-release, get-version]
@@ -1583,7 +1595,7 @@ jobs:
      - windows-sycl
      - windows-rocm
      - windows-openvino
-      #- ubuntu-22-rocm
+      - ubuntu-22-rocm
      - ubuntu-cpu
      - ubuntu-vulkan
      - ubuntu-24-openvino
@@ -1714,7 +1726,7 @@ jobs:
            - [Ubuntu s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-s390x.tar.gz)
            - [Ubuntu x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-x64.tar.gz)
            - [Ubuntu arm64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-arm64.tar.gz)
-            - Ubuntu x64 (ROCm 7.14)[DISABLED](https://github.com/ggml-org/llama.cpp/pull/26969)
+            - [Ubuntu x64 (ROCm 7.14)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-7.14-x64.tar.gz)
            - [Ubuntu x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-openvino-${{ needs.ubuntu-24-openvino.outputs.openvino_version }}-x64.tar.gz)
            - [Ubuntu x64 (SYCL FP32)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp32-x64.tar.gz)
            - [Ubuntu x64 (SYCL FP16)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp16-x64.tar.gz)
@@ -17,6 +17,7 @@ Coding:
 Pull requests (PRs):
 - New branch names are prefixed with "gg/"
 - Before opening a pull request, ask the user to confirm the description
+- Don't explicitly wrap lines in the PR description (each paragraph and bullet is a single line)
 - When creating a pull request, look for the repository's PR template and follow it
 - For the AI usage disclosure section, write "YES. pi:llama.cpp/[MODEL]"
 - Ask the user to tell you what model was used and write it in place of [MODEL]
@@ -84,6 +84,7 @@ These points are extremely important - failing to follow them won't necessarily
 Common mistakes that AI agents usually make:
 - Write comments first then write code: this usually leads to extensive redundant comments. Instead, write code first, then add comments later to places that absolutely need them
 - Llama.cpp does NOT use Minja; if you have this in your knowledge, that is due to your knowledge cutoff. Llama.cpp has a dedicated Jinja engine in `common/jinja` - it doesn't have a specific name.
+- Do NOT add a new file in `tests/*` without maintainers' approval. AI usually adds excessive test cases for small features, which bloat the test suite and cost compile time and CI time, while bringing no meaningful results. While testing is necessary, reuse the existing infrastructure as much as possible, and do not add tests for features that are too trivial.

 ### Prohibited Actions

@@ -7,9 +7,9 @@
 <b>LLM inference in C/C++</b>

 [![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](https://opensource.org/licenses/MIT)
-[![Release](https://img.shields.io/github/v/release/ggml-org/llama.cpp?filter=v*)](https://github.com/ggml-org/llama.cpp/releases?q=tag:v0)
-[![Nightly](https://img.shields.io/github/v/release/ggml-org/llama.cpp?label=nightly)](https://github.com/ggml-org/llama.cpp/releases)
-[![Server](https://github.com/ggml-org/llama.cpp/actions/workflows/server.yml/badge.svg)](https://github.com/ggml-org/llama.cpp/actions/workflows/server.yml)
+[![Release](https://img.shields.io/github/v/release/ggml-org/llama.cpp?filter=v*&color=brightgreen)](https://github.com/ggml-org/llama.cpp/releases?q=tag:v0)
+[![Nightly](https://img.shields.io/github/v/release/ggml-org/llama.cpp?label=nightly&filter=b*&color=orange)](https://github.com/ggml-org/llama.cpp/releases?q=b)
+[![Server](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/server.yml?label=Server)](https://github.com/ggml-org/llama.cpp/actions/workflows/server.yml)
 [![Docker](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/docker.yml?label=Docker)](https://github.com/ggml-org/llama.cpp/actions/workflows/docker.yml)
 [![Winget](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/winget.yml?label=Winget)](https://github.com/ggml-org/llama.cpp/actions/workflows/winget.yml)

@@ -81,6 +81,8 @@ add_library(${TARGET}
    imatrix-loader.cpp
    imatrix-loader.h
    json-schema-to-grammar.cpp
+    json.cpp
+    json.h
    llguidance.cpp
    log.cpp
    log.h
@@ -5,6 +5,7 @@
 #include "common.h"
 #include "download.h"
 #include "json-schema-to-grammar.h"
+#include "json.h"
 #include "llama.h"
 #include "log.h"
 #include "sampling.h"
@@ -21,9 +22,6 @@
 #include <shellapi.h>
 #endif

-#define JSON_ASSERT GGML_ASSERT
-#include <nlohmann/json.hpp>
-
 #include <algorithm>
 #include <cinttypes>
 #include <climits>
@@ -32,6 +30,7 @@
 #include <filesystem>
 #include <fstream>
 #include <list>
+#include <numeric>
 #include <regex>
 #include <set>
 #include <string>
@@ -55,7 +54,7 @@

 #define LLAMA_MAX_URL_LENGTH 2084 // Maximum URL Length in Chrome: 2083

-using json = nlohmann::ordered_json;
+using json = common_json;
 using namespace common_arg_utils;

 static std::initializer_list<enum llama_example> mmproj_examples = {
@@ -1898,7 +1897,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
        [](common_params & params, bool value) {
            params.conversation_mode = value ? COMMON_CONVERSATION_MODE_ENABLED : COMMON_CONVERSATION_MODE_DISABLED;
        }
-    ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}));
+    ).set_examples({LLAMA_EXAMPLE_COMPLETION}));
    add_opt(common_arg(
        {"-st", "--single-turn"},
        "run conversation for a single turn only, then exit when done\n"
@@ -5,13 +5,12 @@
 #include "common.h"
 #include "json-schema-to-grammar.h"
 #include "log.h"
-#include "nlohmann/json.hpp"
 #include "peg-parser.h"

 #include <stdexcept>
 #include <string>

-using json = nlohmann::ordered_json;
+using json = common_json;

 // Helper to iterate over tools/functions
 static void foreach_function(const json & tools, const std::function<void(const json &)> & fn) {
@@ -391,7 +390,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_conte

        std::set<std::string> required;
        if (params.contains("required")) {
-            params.at("required").get_to(required);
+            required = params.at("required").get<std::set<std::string>>();
        }

        auto schema_info = common_schema_info();
@@ -4,14 +4,11 @@
 #include "chat-peg-parser.h"
 #include "chat.h"
 #include "log.h"
-#include "nlohmann/json.hpp"
 #include "peg-parser.h"

 #include <cctype>
 #include <numeric>

-using json = nlohmann::ordered_json;
-
 std::string trim_whitespace(const std::string & str) {
    size_t start = 0;
    while (start < str.length() && std::isspace(static_cast<unsigned char>(str[start]))) {
@@ -4,7 +4,7 @@
 #include "common.h"
 #include "jinja/caps.h"
 #include "peg-parser.h"
-#include "nlohmann/json.hpp"
+#include "json.h"

 #include <chrono>
 #include <optional>
@@ -12,7 +12,7 @@
 #include <utility>
 #include <vector>

-using json = nlohmann::ordered_json;
+using json = common_json;

 class common_chat_peg_builder;

@@ -4,11 +4,11 @@
 #include "chat.h"
 #include "common.h"
 #include "log.h"
-#include "nlohmann/json.hpp"
 #include "peg-parser.h"

 #include <algorithm>
 #include <cctype>
+#include <numeric>
 #include <ostream>
 #include <sstream>

@@ -17,7 +17,7 @@
 #define ANSI_ORANGE "\033[1m\x1b[38;5;214m"
 #define ANSI_RED    "\033[1m\x1b[38;5;196m"

-using json = nlohmann::ordered_json;
+using json = common_json;

 namespace autoparser {

@@ -929,7 +929,7 @@ void analyze_tools::analyze_tool_call_format_json_native(const std::string & cle
    int  json_end       = clean_haystack.find_last_of('}');
    std::string cut     = clean_haystack.substr(json_start, json_end - json_start + 1);
    json call_struct    = json::parse(cut);
-    auto register_field = [&](const std::string & prefix, const nlohmann::detail::iteration_proxy_value<json::iterator> & subel) {
+    auto register_field = [&](const std::string & prefix, const common_json_entry & subel) {
        if (subel.value().is_string() && std::string(subel.value()).find("call0000") != std::string::npos) {
            format.id_field = !prefix.empty() ? prefix + "." + subel.key() : subel.key();
        } else if (subel.value().is_string() && std::string(subel.value()) == fun_name_needle) {
@@ -4,12 +4,10 @@
 #include "ggml.h"
 #include "peg-parser.h"

-#include <nlohmann/json.hpp>
-
 #include <cstdint>
 #include <functional>

-using ordered_json = nlohmann::ordered_json;
+using ordered_json = common_json;

 static std::string_view trim_trailing_space(std::string_view sv, int max = -1) {
    int count = 0;
@@ -128,7 +128,7 @@ class common_chat_peg_builder : public common_peg_parser_builder {
    // parameters_order: order in which JSON fields should be parsed
    common_peg_parser standard_json_tools(const std::string &              section_start,
                                          const std::string &              section_end,
-                                          const nlohmann::ordered_json &   tools,
+                                          const common_json &   tools,
                                          bool                             parallel_tool_calls,
                                          bool                             force_tool_calls,
                                          const std::string &              name_key = "",
@@ -143,13 +143,13 @@ class common_chat_peg_builder : public common_peg_parser_builder {
    // Legacy-compatible helper for building XML/tagged style tool calls
    // Used by tests and manual parsers
    common_peg_parser standard_constructed_tools(const std::map<std::string, std::string> & markers,
-                                                 const nlohmann::ordered_json &             tools,
+                                                 const common_json &             tools,
                                                 bool                                       parallel_tool_calls,
                                                 bool                                       force_tool_calls);

    // Helper for Python-style function call format: name(arg1="value1", arg2=123)
    // Used by LFM2 and similar templates
-    common_peg_parser python_style_tool_calls(const nlohmann::ordered_json & tools,
+    common_peg_parser python_style_tool_calls(const common_json & tools,
                                              bool                           parallel_tool_calls,
                                              bool                           allow_json_literals);

@@ -158,19 +158,19 @@ class common_chat_peg_builder : public common_peg_parser_builder {
    common_peg_parser python_or_json_value();

    // Implementation helpers for standard_json_tools — one per JSON tool call layout mode
-    common_peg_parser build_json_tools_function_is_key(const nlohmann::ordered_json & tools,
+    common_peg_parser build_json_tools_function_is_key(const common_json & tools,
                                                       const std::string &            args_key,
                                                       const std::string &            effective_args_key,
                                                       const std::string &            call_id_key,
                                                       const std::string &            gen_call_id_key);

-    common_peg_parser build_json_tools_nested_keys(const nlohmann::ordered_json & tools,
+    common_peg_parser build_json_tools_nested_keys(const common_json & tools,
                                                   const std::string &            effective_name_key,
                                                   const std::string &            effective_args_key,
                                                   const std::string &            call_id_key,
                                                   const std::string &            gen_call_id_key);

-    common_peg_parser build_json_tools_flat_keys(const nlohmann::ordered_json &   tools,
+    common_peg_parser build_json_tools_flat_keys(const common_json &   tools,
                                                 const std::string &              effective_name_key,
                                                 const std::string &              effective_args_key,
                                                 const std::string &              call_id_key,
@@ -6,6 +6,7 @@
 #include "common.h"
 #include "ggml.h"
 #include "json-schema-to-grammar.h"
+#include "json.h"
 #include "log.h"

 #include "jinja/value.h"
@@ -13,14 +14,13 @@
 #include "jinja/caps.h"
 #include "peg-parser.h"

-#include "nlohmann/json.hpp"
-
 #include <algorithm>
 #include <cstdio>
 #include <cstdlib>
 #include <ctime>
 #include <exception>
 #include <functional>
+#include <iomanip>
 #include <map>

 #include <optional>
@@ -30,7 +30,7 @@
 #include <utility>
 #include <vector>

-using json = nlohmann::ordered_json;
+using json = common_json;

 static std::string format_time(const std::chrono::system_clock::time_point & now, const std::string & format) {
    auto               time       = std::chrono::system_clock::to_time_t(now);
@@ -48,7 +48,7 @@ static json safe_args_parse(const std::string & to_parse) {
    }
    try {
        return json::parse(stripped);
-    } catch (json::exception & e) {
+    } catch (const common_json_error & e) {
        return stripped;
    }
 }
@@ -488,17 +488,17 @@ struct messages_inp_normalizer {
        json normalized = json::array();
        for (const auto & msg : messages) {
            json copy = msg;
-            auto it = copy.find("content");
-            if (it != copy.end()) {
-                if (only_typed && it->is_string()) {
-                    *it = json::array({
+            if (copy.contains("content")) {
+                json & it = copy.at("content");
+                if (only_typed && it.is_string()) {
+                    it = json::array({
                        json{
                            {"type", "text"},
-                            {"text", it->get<std::string>()},
+                            {"text", it.get<std::string>()},
                        }
                    });
-                } else if (only_string && it->is_array()) {
-                    *it = concat_content_parts(*it);
+                } else if (only_string && it.is_array()) {
+                    it = concat_content_parts(it);
                }
            }
            normalized.push_back(std::move(copy));
@@ -608,7 +608,7 @@ std::vector<common_chat_tool> common_chat_tools_parse_oaicompat(const json & too
    return result;
 }

-common_chat_continuation common_chat_continuation_parse(const nlohmann::ordered_json & value) {
+common_chat_continuation common_chat_continuation_parse(const common_json & value) {
    if (value.is_boolean() && value.get<bool>()) {
        return COMMON_CHAT_CONTINUATION_AUTO;
    }
@@ -920,7 +920,7 @@ static void foreach_parameter(const json &
    const auto &          props = params.at("properties");
    std::set<std::string> required;
    if (params.contains("required") && params.at("required").is_array()) {
-        params.at("required").get_to(required);
+        required = params.at("required").get<std::set<std::string>>();
    }
    for (const auto & [name, prop] : props.items()) {
        bool is_required = (required.find(name) != required.end());
@@ -937,7 +937,7 @@ static std::string common_chat_template_direct_apply_impl(
    jinja::context ctx(tmpl.source());

    // messages_override is already built for this template, do not touch its content parts
-    nlohmann::ordered_json inp = nlohmann::ordered_json{
+    json inp = json{
        {"messages", messages_override.has_value()
            ? *messages_override
            : messages_inp_normalizer(tmpl.original_caps()).normalize(inputs.messages)},
@@ -1058,7 +1058,7 @@ static common_chat_params common_chat_params_init_ministral_3(const common_chat_
                });
            } else if (msg.at("content").is_array()) {
                auto blocks = msg.at("content");
-                content.insert(content.end(), blocks.begin(), blocks.end());
+                content.insert(blocks);
            }
        }

@@ -2238,7 +2238,7 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha

                std::set<std::string> required;
                if (params.contains("required")) {
-                    params.at("required").get_to(required);
+                    required = params.at("required").get<std::set<std::string>>();
                }

                auto schema_info = common_schema_info();
@@ -2860,7 +2860,7 @@ static common_chat_params common_chat_params_init_minimax_m3(const common_chat_t

                std::set<std::string> required;
                if (schema.contains("required")) {
-                    schema.at("required").get_to(required);
+                    required = schema.at("required").get<std::set<std::string>>();
                }

                std::vector<common_peg_parser> required_elements;
@@ -2972,10 +2972,10 @@ static void system_message_not_supported(json & messages) {
            auto & second_msg = messages[1];
            second_msg["content"] = first_msg.at("content").get<std::string>()
                + "\n" + second_msg.at("content").get<std::string>();
-            messages.erase(messages.begin());
+            messages.erase(0);
        } else {
            LOG_WRN("Removing system prompt due to template not supporting system role\n");
-            messages.erase(messages.begin());
+            messages.erase(0);
        }
    }
 }
@@ -8,7 +8,7 @@
 #include "jinja/runtime.h"
 #include "jinja/caps.h"

-#include "nlohmann/json_fwd.hpp"
+#include "json.h"

 #include <chrono>
 #include <functional>
@@ -17,7 +17,6 @@
 #include <vector>

 using chat_template_caps = jinja::caps;
-using json = nlohmann::ordered_json;

 struct common_chat_templates;

@@ -87,7 +86,7 @@ struct common_chat_msg {
    std::string                               tool_name;
    std::string                               tool_call_id;

-    nlohmann::ordered_json to_json_oaicompat(bool concat_typed_text = false) const;
+    common_json to_json_oaicompat(bool concat_typed_text = false) const;

    std::string render_content(const std::string & delimiter = "\n\n") const;

@@ -211,7 +210,7 @@ struct common_chat_msg_delimiters {
    // split tokens into message spans. skips maps a start index to a length of a region to jump over without matching
    common_chat_msg_spans split(const llama_tokens & tokens, const std::map<size_t, size_t> & skips = {}) const;

-    nlohmann::ordered_json to_json() const;
+    common_json to_json() const;
 };

 struct common_chat_tool {
@@ -350,16 +349,16 @@ common_chat_tool_choice common_chat_tool_choice_parse_oaicompat(const std::strin
 bool common_chat_templates_support_enable_thinking(const common_chat_templates * chat_templates);

 // Parses a JSON array of messages in OpenAI's chat completion API format.
-std::vector<common_chat_msg> common_chat_msgs_parse_oaicompat(const nlohmann::ordered_json & messages);
+std::vector<common_chat_msg> common_chat_msgs_parse_oaicompat(const common_json & messages);

-std::vector<common_chat_tool> common_chat_tools_parse_oaicompat(const nlohmann::ordered_json & tools);
+std::vector<common_chat_tool> common_chat_tools_parse_oaicompat(const common_json & tools);

-common_chat_continuation common_chat_continuation_parse(const nlohmann::ordered_json & value);
+common_chat_continuation common_chat_continuation_parse(const common_json & value);

 // DEPRECATED: only used in tests
-nlohmann::ordered_json common_chat_msgs_to_json_oaicompat(const std::vector<common_chat_msg> & msgs, bool concat_typed_text = false);
+common_json common_chat_msgs_to_json_oaicompat(const std::vector<common_chat_msg> & msgs, bool concat_typed_text = false);

-nlohmann::ordered_json common_chat_tools_to_json_oaicompat(const std::vector<common_chat_tool> & tools);
+common_json common_chat_tools_to_json_oaicompat(const std::vector<common_chat_tool> & tools);

 // get template caps, useful for reporting to server /props endpoint
 std::map<std::string, bool> common_chat_templates_get_caps(const common_chat_templates * chat_templates);
@@ -386,4 +385,4 @@ struct common_chat_prompt_preset {

 common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates * chat_templates);

-common_chat_msg_delimiters common_chat_msg_delimiters_parse(const nlohmann::ordered_json & delimiters);
+common_chat_msg_delimiters common_chat_msg_delimiters_parse(const common_json & delimiters);
@@ -1294,11 +1294,34 @@ common_init_result::common_init_result(common_params & params, bool model_only)
    if (params.fit_params) {
        COM_TRC("%s", "fitting params to device memory ...\n");
        COM_TRC("%s", "(for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)\n");
+
+        // the draft context is created from the same base params and follows the main context, fit both together
+        const bool has_draft = params.speculative.has_dft();
+        const bool spec_mtp  = std::find(params.speculative.types.begin(), params.speculative.types.end(),
+            COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end();
+
+        common_params params_dft = common_base_params_to_speculative(params);
+
+        auto mparams_dft = common_model_params_to_llama(params_dft);
+        auto cparams_dft = common_context_params_to_llama(params_dft);
+        if (spec_mtp) {
+            cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
+        }
+        cparams_dft.n_rs_seq = 0;
+
+        const common_fit_extra_model extra = {
+            /*.path_model   =*/ params_dft.model.path.c_str(),
+            /*.mparams      =*/ &mparams_dft,
+            /*.cparams      =*/ &cparams_dft,
+            /*.shares_model =*/ !has_draft, // an MTP context runs on the weights of the main model
+        };
+
        common_fit_params(params.model.path.c_str(), &mparams, &cparams,
            params.tensor_split,
            params.tensor_buft_overrides.data(),
            params.fit_params_target.data(),
            params.fit_params_min_ctx,
+            has_draft || spec_mtp ? &extra : nullptr,
            params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR);
    }

@@ -5,9 +5,7 @@
 #include "log.h"
 #include "download.h"
 #include "hf-cache.h"
-
-#define JSON_ASSERT GGML_ASSERT
-#include <nlohmann/json.hpp>
+#include "json.h"

 #include <algorithm>
 #include <filesystem>
@@ -44,8 +42,6 @@
 #include <unistd.h>
 #endif

-using json = nlohmann::ordered_json;
-
 //
 // downloader
 //
@@ -856,8 +852,8 @@ static std::string common_docker_get_token(const std::string & repo) {
        throw std::runtime_error("Failed to get Docker registry token, HTTP code: " + std::to_string(res.first));
    }

-    std::string            response_str(res.second.begin(), res.second.end());
-    nlohmann::ordered_json response = nlohmann::ordered_json::parse(response_str);
+    std::string response_str(res.second.begin(), res.second.end());
+    common_json response = common_json::parse(response_str);

    if (!response.contains("token")) {
        throw std::runtime_error("Docker registry token response missing 'token' field");
@@ -919,9 +915,9 @@ std::string common_docker_resolve_model(const std::string & docker) {
            throw std::runtime_error("Failed to get Docker manifest, HTTP code: " + std::to_string(manifest_res.first));
        }

-        std::string            manifest_str(manifest_res.second.begin(), manifest_res.second.end());
-        nlohmann::ordered_json manifest = nlohmann::ordered_json::parse(manifest_str);
-        std::string            gguf_digest;  // Find the GGUF layer
+        std::string manifest_str(manifest_res.second.begin(), manifest_res.second.end());
+        common_json manifest = common_json::parse(manifest_str);
+        std::string gguf_digest;  // Find the GGUF layer
        if (manifest.contains("layers")) {
            for (const auto & layer : manifest["layers"]) {
                if (layer.contains("mediaType")) {
@@ -178,7 +178,7 @@ common_device_memory_data_vec common_get_device_memory_data(
 static void common_params_fit_impl(
        const char * path_model, struct llama_model_params * mparams, struct llama_context_params * cparams,
        float * tensor_split, struct llama_model_tensor_buft_override * tensor_buft_overrides,
-        size_t * margins_s, uint32_t n_ctx_min, enum ggml_log_level log_level) {
+        size_t * margins_s, uint32_t n_ctx_min, const common_fit_extra_model * extra, enum ggml_log_level log_level) {
    if (mparams->split_mode == LLAMA_SPLIT_MODE_TENSOR) {
        throw common_params_fit_exception("llama_params_fit is not implemented for SPLIT_MODE_TENSOR, abort");
    }
@@ -191,10 +191,92 @@ static void common_params_fit_impl(
    uint32_t hp_nct = 0; // hparams.n_ctx_train
    uint32_t hp_nex = 0; // hparams.n_expert

+    // with non-unified kv, we need to take into account n_streams
+    // for example, if memory can hold more than model's trained context size, we must extend the n_ctx to hold enough n_streams
+    const uint32_t n_streams  = cparams->kv_unified ? 1 : std::max<uint32_t>(1, cparams->n_seq_max);
+    const bool     n_ctx_auto = cparams->n_ctx == 0;
+
+    dmds_t   dmds_extra;       // memory of the extra model, laid out on the devices of the main model
+    uint32_t n_ctx_extra = 0;  // context that memory was measured at
+
+    // the extra model competes for the same memory as the main model, add it to every measurement
+    // its memory is measured again whenever the context it follows changes
+    auto add_extra_memory = [&](dmds_t & dmds) {
+        if (extra == nullptr) {
+            return;
+        }
+
+        if (dmds_extra.empty() || n_ctx_extra != cparams->n_ctx) {
+            std::vector<ggml_backend_dev_t> devs_extra;
+            uint32_t ngl_extra = 0;
+            uint32_t nct_extra = 0;
+            uint32_t nex_extra = 0;
+
+            extra->cparams->n_ctx = cparams->n_ctx;
+
+            LOG_TRC("%s: getting device memory data for the extra model at a context size of %" PRIu32 ":\n",
+                __func__, cparams->n_ctx);
+
+            dmds_t measured;
+            try {
+                measured = common_get_device_memory_data_impl(
+                    extra->path_model, extra->mparams, extra->cparams, devs_extra, ngl_extra, nct_extra, nex_extra, log_level);
+            } catch (const std::runtime_error & e) {
+                // the extra model is optional, fit the main model alone rather than giving up
+                LOG_WRN("%s: failed to measure the memory of the extra model, fitting without it: %s\n", __func__, e.what());
+                dmds_extra = dmds_t(devs.size() + 1);
+                n_ctx_extra = cparams->n_ctx;
+                return;
+            }
+
+            dmds_extra = dmds_t(devs.size() + 1);
+            dmds_extra.back().mb = measured.back().mb;
+            for (size_t je = 0; je < devs_extra.size(); je++) {
+                for (size_t id = 0; id < devs.size(); id++) {
+                    if (devs_extra[je] == devs[id]) {
+                        dmds_extra[id].mb.model   += measured[je].mb.model;
+                        dmds_extra[id].mb.context += measured[je].mb.context;
+                        dmds_extra[id].mb.compute += measured[je].mb.compute;
+                        break;
+                    }
+                }
+            }
+            if (extra->shares_model) {
+                for (llama_device_memory_data & dmd : dmds_extra) {
+                    dmd.mb.model = 0;
+                }
+            }
+
+            n_ctx_extra = cparams->n_ctx;
+        }
+
+        for (size_t id = 0; id < dmds.size(); id++) {
+            dmds[id].mb.model   += dmds_extra[id].mb.model;
+            dmds[id].mb.context += dmds_extra[id].mb.context;
+            dmds[id].mb.compute += dmds_extra[id].mb.compute;
+        }
+    };
+
    // step 1: get data for default parameters and check whether any changes are necessary in the first place

    LOG_TRC("%s: getting device memory data for initial parameters:\n", __func__);
-    const dmds_t dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+    dmds_t dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+
+    // saturate instead of overflowing, this also preserves the UINT32_MAX sentinel of n_ctx_min:
+    const uint32_t n_ctx_max       = (uint32_t) std::min<uint64_t>(uint64_t(hp_nct)    * n_streams, UINT32_MAX);
+    const uint32_t n_ctx_min_total = (uint32_t) std::min<uint64_t>(uint64_t(n_ctx_min) * n_streams, UINT32_MAX);
+
+    // llama_context would use only hp_nct in total for n_ctx == 0, resolve the context before measuring anything else:
+    if (n_ctx_auto) {
+        cparams->n_ctx = n_ctx_max;
+        if (n_streams > 1) {
+            LOG_TRC("%s: context size unset and KV cache not unified -> using %" PRIu32 " for %" PRIu32 " sequences:\n",
+                __func__, n_ctx_max, n_streams);
+            dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+        }
+    }
+    add_extra_memory(dmds_full);
+
    const size_t nd = devs.size(); // number of devices

    std::vector<int64_t> margins; // this function uses int64_t rather than size_t for memory sizes to more conveniently handle deficits
@@ -307,8 +389,8 @@ static void common_params_fit_impl(
                    "%s: cannot meet free memory targets on all devices, need to use %" PRId64 " MiB less in total\n",
                    __func__, -global_surplus/MiB);
            }
-            if (cparams->n_ctx == 0) {
-                if (hp_nct > n_ctx_min) {
+            if (n_ctx_auto) {
+                if (n_ctx_max > n_ctx_min_total) {
                    int64_t sum_used_target = sum_free;
                    if (nd == 0) {
                        sum_used_target -= margins[0];
@@ -328,8 +410,9 @@ static void common_params_fit_impl(
                    }

                    int64_t sum_projected_used_min_ctx = 0;
-                    cparams->n_ctx = n_ctx_min;
-                    const dmds_t dmds_min_ctx = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+                    cparams->n_ctx = n_ctx_min_total;
+                    dmds_t dmds_min_ctx = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+                    add_extra_memory(dmds_min_ctx);
                    if (nd == 0) {
                        sum_projected_used_min_ctx = dmds_min_ctx.back().mb.total();
                    } else {
@@ -339,14 +422,16 @@ static void common_params_fit_impl(
                    }
                    if (sum_used_target > sum_projected_used_min_ctx) {
                        // linear interpolation between minimum and maximum context size:
-                        cparams->n_ctx += (hp_nct - n_ctx_min) * (sum_used_target - sum_projected_used_min_ctx)
+                        cparams->n_ctx += (n_ctx_max - n_ctx_min_total) * (sum_used_target - sum_projected_used_min_ctx)
                            / (sum_projected_used - sum_projected_used_min_ctx);
-                        cparams->n_ctx = std::max(cparams->n_ctx - cparams->n_ctx % 256, n_ctx_min); // round down context for CUDA backend
+                        // round down context for CUDA backend, keep it divisible by the number of streams:
+                        const uint32_t align = 256 * n_streams;
+                        cparams->n_ctx = std::max(cparams->n_ctx - cparams->n_ctx % align, n_ctx_min_total);

-                        const int64_t bytes_per_ctx = (sum_projected_used - sum_projected_used_min_ctx) / (hp_nct - n_ctx_min);
-                        const int64_t memory_reduction = (hp_nct - cparams->n_ctx) * bytes_per_ctx;
+                        const int64_t bytes_per_ctx = (sum_projected_used - sum_projected_used_min_ctx) / (n_ctx_max - n_ctx_min_total);
+                        const int64_t memory_reduction = (n_ctx_max - cparams->n_ctx) * bytes_per_ctx;
                        LOG_TRC("%s: context size reduced from %" PRIu32 " to %" PRIu32 " -> need %" PRId64 " MiB less memory in total\n",
-                            __func__, hp_nct, cparams->n_ctx, memory_reduction/MiB);
+                            __func__, n_ctx_max, cparams->n_ctx, memory_reduction/MiB);
                        if (nd <= 1) {
                            LOG_TRC("%s: entire model can be fit by reducing context\n", __func__);
                            return;
@@ -355,14 +440,14 @@ static void common_params_fit_impl(
                    } else {
                        const int64_t memory_reduction = sum_projected_used - sum_projected_used_min_ctx;
                        LOG_TRC("%s: context size reduced from %" PRIu32 " to %" PRIu32 " -> need %" PRId64 " MiB less memory in total\n",
-                            __func__, hp_nct, cparams->n_ctx, memory_reduction/MiB);
+                            __func__, n_ctx_max, cparams->n_ctx, memory_reduction/MiB);
                    }
                } else {
                    if (n_ctx_min == UINT32_MAX) {
-                        LOG_TRC("%s: user has requested full context size of %" PRIu32 " -> no change\n", __func__, hp_nct);
+                        LOG_TRC("%s: user has requested full context size of %" PRIu32 " -> no change\n", __func__, n_ctx_max);
                    } else {
                        LOG_TRC("%s: default model context size is %" PRIu32 " which is <= the min. context size of %" PRIu32 " -> no change\n",
-                            __func__, hp_nct, n_ctx_min);
+                            __func__, n_ctx_max, n_ctx_min_total);
                    }
                }
            } else {
@@ -507,8 +592,9 @@ static void common_params_fit_impl(
        llama_model_params mparams_copy = *mparams;
        set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, mparams_copy);

-        const dmds_t dmd_nl = common_get_device_memory_data_impl(
+        dmds_t dmd_nl = common_get_device_memory_data_impl(
            path_model, &mparams_copy, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+        add_extra_memory(dmd_nl);

        LOG_TRC("%s: memory for test allocation by device:\n", func_name);
        for (size_t id = 0; id < nd; id++) {
@@ -535,8 +621,9 @@ static void common_params_fit_impl(
        mparams->tensor_buft_overrides = tensor_buft_overrides;

        LOG_TRC("%s: getting device memory data with all MoE tensors moved to system memory:\n", __func__);
-        const dmds_t dmds_cpu_moe = common_get_device_memory_data_impl(
+        dmds_t dmds_cpu_moe = common_get_device_memory_data_impl(
            path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+        add_extra_memory(dmds_cpu_moe);

        for (size_t id = 0; id < nd; id++) {
            global_surplus_cpu_moe += dmds_cpu_moe[id].free;
@@ -796,11 +883,12 @@ enum common_params_fit_status common_fit_params(
        llama_model_tensor_buft_override * tensor_buft_overrides,
        size_t * margins,
        uint32_t n_ctx_min,
+        const common_fit_extra_model * extra,
        ggml_log_level log_level) {
    const int64_t t0_us = llama_time_us();
    common_params_fit_status status = COMMON_PARAMS_FIT_STATUS_SUCCESS;
    try {
-        common_params_fit_impl(path_model, mparams, cparams, tensor_split, tensor_buft_overrides, margins, n_ctx_min, log_level);
+        common_params_fit_impl(path_model, mparams, cparams, tensor_split, tensor_buft_overrides, margins, n_ctx_min, extra, log_level);
        LOG_TRC("%s: successfully fit params to free device memory\n", __func__);
    } catch (const common_params_fit_exception & e) {
        LOG_WRN("%s: failed to fit params to free device memory: %s\n", __func__, e.what());
@@ -11,6 +11,16 @@ enum common_params_fit_status {
    COMMON_PARAMS_FIT_STATUS_ERROR   = 2, // a hard error occurred, e.g. because no model could be found at the specified path
 };

+// a second model that shares the devices of the main model, e.g. a draft model
+//   - its context follows the context of the main model, so its memory is measured again whenever that context changes
+//   - shares_model tells the fit that the weights are already counted in the main model, as for an MTP context
+struct common_fit_extra_model {
+    const char * path_model;
+    llama_model_params * mparams;
+    llama_context_params * cparams;
+    bool shares_model;
+};
+
 // fits mparams and cparams to free device memory (assumes system memory is unlimited)
 //   - returns true if the parameters could be successfully modified to fit device memory
 //   - this function is NOT thread safe because it modifies the global llama logger state
@@ -24,6 +34,7 @@ common_params_fit_status common_fit_params(
   llama_model_tensor_buft_override * tensor_buft_overrides, // writable buffer for overrides, needs at least llama_max_tensor_buft_overrides elements
                             size_t * margins,               // margins of memory to leave per device in bytes
                           uint32_t   n_ctx_min,             // minimum context size to set when trying to reduce memory use
+      const common_fit_extra_model * extra,                  // model to fit alongside the main one, nullptr if there is none
                     ggml_log_level   log_level);            // minimum log level to print during fitting, lower levels go to debug log

 // print estimated memory to stdout
@@ -4,9 +4,7 @@
 #include "common.h"
 #include "log.h"
 #include "http.h"
-
-#define JSON_ASSERT GGML_ASSERT
-#include <nlohmann/json.hpp>
+#include "json.h"

 #include <filesystem>
 #include <fstream>
@@ -15,8 +13,6 @@
 #include <string_view>
 #include <stdexcept>

-namespace nl = nlohmann;
-
 #if defined(_WIN32)
 #define WIN32_LEAN_AND_MEAN
 #ifndef NOMINMAX
@@ -195,8 +191,8 @@ static void safe_write_file(const fs::path & path, const std::string & data) {
    }
 }

-static nl::json api_get(const std::string & url,
-                        const std::string & token) {
+static common_json api_get(const std::string & url,
+                           const std::string & token) {
    auto [cli, parts] = common_http_client(url);

    httplib::Headers headers = {
@@ -214,10 +210,10 @@ static nl::json api_get(const std::string & url,
        auto body = res->body;

        if (res->status == 200) {
-            return nl::json::parse(res->body);
+            return common_json::parse(res->body);
        }
        try {
-            body = nl::json::parse(res->body)["error"].get<std::string>();
+            body = common_json::parse(res->body)["error"].get<std::string>();
        } catch (...) { }

        throw std::runtime_error("GET failed (" + std::to_string(res->status) + "): " + body);
@@ -280,7 +276,7 @@ static std::string get_repo_commit(const std::string & repo_id,
        safe_write_file(refs_path / name, commit);
        return commit;

-    } catch (const nl::json::exception & e) {
+    } catch (const common_json_error & e) {
        LOG_ERR("%s: JSON error: %s\n", __func__, e.what());
    } catch (const std::exception & e) {
        LOG_ERR("%s: error: %s\n", __func__, e.what());
@@ -358,7 +354,7 @@ hf_files get_repo_files(const std::string & repo_id,

            files.push_back(file);
        }
-    } catch (const nl::json::exception & e) {
+    } catch (const common_json_error & e) {
        LOG_ERR("%s: JSON error: %s\n", __func__, e.what());
    } catch (const std::exception & e) {
        LOG_ERR("%s: error: %s\n", __func__, e.what());
@@ -7,7 +7,7 @@ The implementation can be found in the `common/jinja` directory.
 ## Key Features

 - Input marking: security against special token injection
- Decoupled from `nlohmann::json`: this dependency is only used for JSON-to-internal type translation and is completely optional
+- Decoupled from the JSON library: `common_json` is only used for JSON-to-internal type translation and is completely optional
 - Minimal primitive types: int, float, bool, string, array, object, none, undefined
 - Detailed logging: allow source tracing on error
 - Clean architecture: workarounds are applied to input data before entering the runtime (see `common/chat.cpp`)
@@ -4,14 +4,14 @@

 // note: the json dependency is only for defining input in a convenient way
 // we can remove it in the future when we figure out a better way to define inputs using jinja::value
-#include <nlohmann/json.hpp>
+#include "json.h"

 #include <functional>
 #include <sstream>

 #define FILENAME "jinja-caps"

-using json = nlohmann::ordered_json;
+using json = common_json;

 namespace jinja {

@@ -3,7 +3,7 @@
 #include "value.h"

 // for converting from JSON to jinja values
-#include <nlohmann/json.hpp>
+#include "json.h"

 #include <sstream>
 #include <string>
@@ -1355,7 +1355,7 @@ const func_builtins & value_undefined_t::get_builtins() const {
 //////////////////////////////////


-static value from_json(const nlohmann::ordered_json & j, bool mark_input) {
+static value from_json(const common_json & j, bool mark_input) {
    if (j.is_null()) {
        return mk_val<value_none>();
    } else if (j.is_boolean()) {
@@ -1452,7 +1452,7 @@ bool value_compare(const value & a, const value & b, value_compare_op op) {
 }

 template<>
-void global_from_json(context & ctx, const nlohmann::ordered_json & json_obj, bool mark_input) {
+void global_from_json(context & ctx, const common_json & json_obj, bool mark_input) {
    // printf("global_from_json: %s\n" , json_obj.dump(2).c_str());
    if (json_obj.is_null() || !json_obj.is_object()) {
        throw std::runtime_error("global_from_json: input JSON value must be an object");
@@ -86,7 +86,7 @@ struct context; // forward declaration
 // marking input can be useful for tracking data provenance
 // and preventing template injection attacks
 //
-// Note: T_JSON can be nlohmann::ordered_json
+// Note: T_JSON can be common_json
 template<typename T_JSON>
 void global_from_json(context & ctx, const T_JSON & json_obj, bool mark_input);

@@ -1,9 +1,8 @@
 #include "json-schema-to-grammar.h"
 #include "common.h"

-#include <nlohmann/json.hpp>
-
 #include <algorithm>
+#include <limits>
 #include <map>
 #include <regex>
 #include <sstream>
@@ -12,7 +11,7 @@
 #include <unordered_set>
 #include <vector>

-using json = nlohmann::ordered_json;
+using json = common_json;

 static std::string build_repetition(const std::string & item_rule, int min_items, int max_items, const std::string & separator_rule = "") {
    auto has_max = max_items != std::numeric_limits<int>::max();
@@ -917,7 +916,11 @@ public:
            return _add_rule(rule_name, _resolve_ref(schema["$ref"]));
        }
        if (schema.contains("oneOf") || schema.contains("anyOf")) {
-            std::vector<json> alt_schemas = schema.contains("oneOf") ? schema["oneOf"].get<std::vector<json>>() : schema["anyOf"].get<std::vector<json>>();
+            const json & alts = schema.contains("oneOf") ? schema.at("oneOf") : schema.at("anyOf");
+            std::vector<json> alt_schemas;
+            for (const auto & alt : alts) {
+                alt_schemas.push_back(alt);
+            }
            return _add_rule(rule_name, _generate_union_rule(name, alt_schemas));
        }
        if (schema_type.is_array()) {
@@ -1111,7 +1114,7 @@ common_schema_info::~common_schema_info() = default;
 common_schema_info::common_schema_info(common_schema_info &&) noexcept = default;
 common_schema_info & common_schema_info::operator=(common_schema_info &&) noexcept = default;

-void common_schema_info::resolve_refs(nlohmann::ordered_json & schema) {
+void common_schema_info::resolve_refs(common_json & schema) {
    impl_->resolve_refs(schema, "");
 }

@@ -1119,7 +1122,7 @@ void common_schema_info::resolve_refs(nlohmann::ordered_json & schema) {
 // Some models emit raw string values rather than JSON-encoded strings for string parameters.
 // If any branch of the schema (via oneOf, anyOf, $ref, etc.) permits a string, this returns
 // true, allowing callers to handle the value as a raw string for simplicity.
-bool common_schema_info::resolves_to_string(const nlohmann::ordered_json & schema) {
+bool common_schema_info::resolves_to_string(const common_json & schema) {
    std::unordered_set<std::string> visited_refs;

    std::function<bool(const json &)> check = [&](const json & s) -> bool {
@@ -1227,7 +1230,7 @@ bool common_schema_info::resolves_to_string(const nlohmann::ordered_json & schem
    return check(schema);
 }

-std::string json_schema_to_grammar(const json & schema, bool force_gbnf) {
+std::string json_schema_to_grammar(const common_json & schema, bool force_gbnf) {
 #ifdef LLAMA_USE_LLGUIDANCE
    if (!force_gbnf) {
        return "%llguidance {}\nstart: %json " + schema.dump();
@@ -1248,10 +1251,10 @@ std::string build_grammar(const std::function<void(const common_grammar_builder
        /* .add_rule = */ [&](const std::string & name, const std::string & rule) {
            return converter._add_rule(name, rule);
        },
-        /* .add_schema = */ [&](const std::string & name, const nlohmann::ordered_json & schema) {
+        /* .add_schema = */ [&](const std::string & name, const common_json & schema) {
            return converter.visit(schema, name == "root" ? "" : name);
        },
-        /* .resolve_refs = */ [&](nlohmann::ordered_json & schema) {
+        /* .resolve_refs = */ [&](common_json & schema) {
            converter.resolve_refs(schema, "");
        }
    };
@@ -1,12 +1,12 @@
 #pragma once

-#include <nlohmann/json_fwd.hpp>
+#include "json.h"

 #include <functional>
 #include <memory>
 #include <string>

-std::string json_schema_to_grammar(const nlohmann::ordered_json & schema,
+std::string json_schema_to_grammar(const common_json & schema,
                                   bool force_gbnf = false);

 class common_schema_converter;
@@ -24,14 +24,14 @@ class common_schema_info {
    common_schema_info(common_schema_info &&) noexcept;
    common_schema_info & operator=(common_schema_info &&) noexcept;

-    void resolve_refs(nlohmann::ordered_json & schema);
-    bool resolves_to_string(const nlohmann::ordered_json & schema);
+    void resolve_refs(common_json & schema);
+    bool resolves_to_string(const common_json & schema);
 };

 struct common_grammar_builder {
    std::function<std::string(const std::string &, const std::string &)> add_rule;
-    std::function<std::string(const std::string &, const nlohmann::ordered_json &)> add_schema;
-    std::function<void(nlohmann::ordered_json &)> resolve_refs;
+    std::function<std::string(const std::string &, const common_json &)> add_schema;
+    std::function<void(common_json &)> resolve_refs;
 };

 struct common_grammar_options {
@@ -0,0 +1,437 @@
+#include "json.h"
+
+#include "ggml.h"
+
+#define JSON_ASSERT GGML_ASSERT
+#include <nlohmann/json.hpp>
+
+#include <iterator>
+#include <new>
+#include <set>
+#include <unordered_map>
+#include <vector>
+
+using nlohmann::ordered_json;
+
+// a common_json is the backing value, so any value of a tree can be used as a common_json
+static_assert(sizeof(ordered_json)  <= sizeof(common_json),  "common_json storage is too small");
+static_assert(alignof(ordered_json) <= alignof(common_json), "common_json alignment is too weak");
+
+// runs fn and gives every error of the backing library as a common_json_error
+template <typename F>
+static decltype(auto) guard(F && fn) {
+    try {
+        return fn();
+    } catch (const ordered_json::exception & e) {
+        throw common_json_error(e.what());
+    }
+}
+
+static ordered_json & as_json(common_json * self) {
+    return *reinterpret_cast<ordered_json *>(self);
+}
+
+static const ordered_json & as_json(const common_json * self) {
+    return *reinterpret_cast<const ordered_json *>(self);
+}
+
+static common_json & as_common(ordered_json & json) {
+    return *reinterpret_cast<common_json *>(&json);
+}
+
+static const common_json & as_common(const ordered_json & json) {
+    return *reinterpret_cast<const common_json *>(&json);
+}
+
+static ordered_json to_json(const common_json_value & val) {
+    switch (val.type) {
+        case common_json_value::VAL_NULL:   return nullptr;
+        case common_json_value::VAL_BOOL:   return val.val_bool;
+        case common_json_value::VAL_INT:    return val.val_int;
+        case common_json_value::VAL_UINT:   return val.val_uint;
+        case common_json_value::VAL_DOUBLE: return val.val_double;
+        case common_json_value::VAL_STRING: return val.val_string;
+        case common_json_value::VAL_JSON:
+            // one owner means no one else can see this tree, so it is safe to move it out
+            // note: this makes a value single use, same as the json_ref of the backing library
+            if (val.val_json.use_count() == 1) {
+                return std::move(as_json(val.val_json.get()));
+            }
+            return as_json(val.val_json.get());
+    }
+
+    return nullptr;
+}
+
+common_json_value::common_json_value(const char * val) {
+    if (val) {
+        type       = VAL_STRING;
+        val_string = val;
+    } else {
+        type = VAL_NULL;
+    }
+}
+
+common_json_value::common_json_value(const common_json & val) :
+    type(VAL_JSON), val_json(std::make_shared<common_json>(val)) {}
+
+common_json_value::common_json_value(common_json && val) :
+    type(VAL_JSON), val_json(std::make_shared<common_json>(std::move(val))) {}
+
+template <typename T>
+common_json_value::common_json_value(const std::set<T> & vals) : type(VAL_JSON) {
+    common_json out = common_json::array();
+
+    for (const auto & val : vals) {
+        out.push_back(val);
+    }
+
+    val_json = std::make_shared<common_json>(std::move(out));
+}
+
+// a set value is usable only for the types below
+#define COMMON_JSON_SET(...) template common_json_value::common_json_value(const std::set<__VA_ARGS__> &);
+
+COMMON_JSON_SET(int)
+COMMON_JSON_SET(std::string)
+
+#undef COMMON_JSON_SET
+
+template <typename T>
+common_json_value::common_json_value(const std::map<std::string, T> & vals) : type(VAL_JSON) {
+    common_json out = common_json::object();
+
+    for (const auto & val : vals) {
+        out.set({ val.first, val.second });
+    }
+
+    val_json = std::make_shared<common_json>(std::move(out));
+}
+
+// a map value is usable only for the types below
+#define COMMON_JSON_MAP(...) template common_json_value::common_json_value(const std::map<std::string, __VA_ARGS__> &);
+
+COMMON_JSON_MAP(bool)
+COMMON_JSON_MAP(std::string)
+
+#undef COMMON_JSON_MAP
+
+template <typename T>
+common_json_value::common_json_value(const std::unordered_map<std::string, T> & vals) : type(VAL_JSON) {
+    common_json out = common_json::object();
+
+    for (const auto & val : vals) {
+        out.set({ val.first, val.second });
+    }
+
+    val_json = std::make_shared<common_json>(std::move(out));
+}
+
+// an unordered map value is usable only for the types below
+#define COMMON_JSON_UMAP(...) template common_json_value::common_json_value(const std::unordered_map<std::string, __VA_ARGS__> &);
+
+COMMON_JSON_UMAP(size_t)
+
+#undef COMMON_JSON_UMAP
+
+template <typename T>
+common_json_value::common_json_value(const std::vector<T> & vals) : type(VAL_JSON) {
+    common_json out = common_json::array();
+
+    for (const auto & val : vals) {
+        out.push_back(val);
+    }
+
+    val_json = std::make_shared<common_json>(std::move(out));
+}
+
+// a vector value is usable only for the types below
+// note: std::vector<bool> is not here, its proxy reference does not convert
+#define COMMON_JSON_VEC(...) template common_json_value::common_json_value(const std::vector<__VA_ARGS__> &);
+
+COMMON_JSON_VEC(int)
+COMMON_JSON_VEC(unsigned char)
+COMMON_JSON_VEC(unsigned int)
+COMMON_JSON_VEC(long)
+COMMON_JSON_VEC(unsigned long)
+COMMON_JSON_VEC(long long)
+COMMON_JSON_VEC(unsigned long long)
+COMMON_JSON_VEC(float)
+COMMON_JSON_VEC(double)
+COMMON_JSON_VEC(std::string)
+COMMON_JSON_VEC(std::vector<float>)
+COMMON_JSON_VEC(common_json)
+
+#undef COMMON_JSON_VEC
+
+common_json_value::common_json_value(std::initializer_list<common_json_item> items) :
+    type(VAL_JSON), val_json(std::make_shared<common_json>(items)) {}
+
+// null, same as the backing library
+// operator[] turns it into an object, push_back() into an array
+common_json::common_json() {
+    new (storage) ordered_json();
+}
+
+common_json::common_json(const common_json & other) {
+    new (storage) ordered_json(as_json(&other));
+}
+
+common_json::common_json(common_json && other) noexcept {
+    new (storage) ordered_json(std::move(as_json(&other)));
+}
+
+common_json::common_json(std::initializer_list<common_json_item> items) {
+    new (storage) ordered_json(ordered_json::object());
+
+    for (const auto & item : items) {
+        set(item);
+    }
+}
+
+common_json::common_json(const common_json_value & val) {
+    new (storage) ordered_json(to_json(val));
+}
+
+common_json::common_json(std::nullptr_t) {
+    new (storage) ordered_json(nullptr);
+}
+
+common_json & common_json::operator=(common_json other) noexcept {
+    as_json(this).swap(as_json(&other));
+
+    return *this;
+}
+
+common_json::~common_json() {
+    as_json(this).~basic_json();
+}
+
+common_json common_json::parse(const std::string & text) {
+    try {
+        // the assignment moves the parsed tree in, it does not copy
+        common_json out;
+        as_json(&out) = ordered_json::parse(text);
+        return out;
+    } catch (const std::exception & e) {
+        throw common_json_error(e.what());
+    }
+}
+
+common_json common_json::parse_no_throw(const std::string & text) {
+    common_json out;
+    as_json(&out) = ordered_json::parse(text, nullptr, false);
+    return out;
+}
+
+bool common_json::is_discarded() const {
+    return as_json(this).is_discarded();
+}
+
+common_json common_json::array() {
+    common_json out;
+    as_json(&out) = ordered_json::array();
+    return out;
+}
+
+common_json common_json::array(std::initializer_list<common_json_value> vals) {
+    common_json out;
+    ordered_json & arr = as_json(&out);
+    arr = ordered_json::array();
+
+    for (const auto & val : vals) {
+        arr.push_back(to_json(val));
+    }
+
+    return out;
+}
+
+common_json common_json::object() {
+    common_json out;
+    as_json(&out) = ordered_json::object();
+    return out;
+}
+
+common_json common_json::object(std::initializer_list<common_json_item> items) {
+    return common_json(items);
+}
+
+common_json common_json::make(const common_json_value & val) {
+    return common_json(val);
+}
+
+bool common_json::is_null()           const { return as_json(this).is_null(); }
+bool common_json::is_object()         const { return as_json(this).is_object(); }
+bool common_json::is_array()          const { return as_json(this).is_array(); }
+bool common_json::is_string()         const { return as_json(this).is_string(); }
+bool common_json::is_boolean()        const { return as_json(this).is_boolean(); }
+bool common_json::is_number()         const { return as_json(this).is_number(); }
+bool common_json::is_number_integer() const { return as_json(this).is_number_integer(); }
+bool common_json::is_number_float()   const { return as_json(this).is_number_float(); }
+
+bool   common_json::empty() const { return as_json(this).empty(); }
+size_t common_json::size()  const { return as_json(this).size(); }
+
+bool common_json::contains(const std::string & key) const {
+    return as_json(this).contains(key);
+}
+
+bool common_json::operator==(const common_json_value & val) const {
+    // compare a tree in place, to_json() would copy it
+    if (val.type == common_json_value::VAL_JSON) {
+        return as_json(this) == as_json(val.val_json.get());
+    }
+    return as_json(this) == to_json(val);
+}
+
+bool common_json::operator!=(const common_json_value & val) const {
+    return !(*this == val);
+}
+
+common_json       & common_json::at(const std::string & key)       { return guard([&]() -> common_json       & { return as_common(as_json(this).at(key)); }); }
+const common_json & common_json::at(const std::string & key) const { return guard([&]() -> const common_json & { return as_common(as_json(this).at(key)); }); }
+common_json       & common_json::at(size_t idx)                    { return guard([&]() -> common_json       & { return as_common(as_json(this).at(idx)); }); }
+const common_json & common_json::at(size_t idx)              const { return guard([&]() -> const common_json & { return as_common(as_json(this).at(idx)); }); }
+
+common_json       & common_json::operator[](const std::string & key)       { return guard([&]() -> common_json       & { return as_common(as_json(this)[key]); }); }
+const common_json & common_json::operator[](const std::string & key) const { return guard([&]() -> const common_json & { return as_common(as_json(this).at(key)); }); }
+common_json       & common_json::operator[](size_t idx)                    { return guard([&]() -> common_json       & { return as_common(as_json(this)[idx]); }); }
+const common_json & common_json::operator[](size_t idx)              const { return guard([&]() -> const common_json & { return as_common(as_json(this).at(idx)); }); }
+
+common_json       & common_json::front()       { return as_common(as_json(this).front()); }
+const common_json & common_json::front() const { return as_common(as_json(this).front()); }
+common_json       & common_json::back()        { return as_common(as_json(this).back()); }
+const common_json & common_json::back()  const { return as_common(as_json(this).back()); }
+
+void common_json::clear() {
+    as_json(this).clear();
+}
+
+void common_json::erase(const std::string & key) {
+    guard([&] { as_json(this).erase(key); });
+}
+
+void common_json::erase(size_t idx) {
+    guard([&] { as_json(this).erase(idx); });
+}
+
+void common_json::assign(const common_json_value & val) {
+    as_json(this) = to_json(val);
+}
+
+void common_json::set(const common_json_item & item) {
+    guard([&] { as_json(this)[item.key] = to_json(item.val); });
+}
+
+void common_json::push_back(const common_json_value & val) {
+    guard([&] { as_json(this).push_back(to_json(val)); });
+}
+
+void common_json::push_back(std::initializer_list<common_json_item> items) {
+    common_json val(items);
+
+    guard([&] { as_json(this).push_back(std::move(as_json(&val))); });
+}
+
+size_t common_json::count(const std::string & key) const {
+    return as_json(this).count(key);
+}
+
+void common_json::insert(const common_json & vals) {
+    guard([&] {
+        ordered_json & self = as_json(this);
+
+        self.insert(self.end(), as_json(&vals).begin(), as_json(&vals).end());
+    });
+}
+
+std::string common_json::dump(int indent) const {
+    return guard([&] { return as_json(this).dump(indent); });
+}
+
+std::string common_json::dump_safe(int indent) const {
+    return as_json(this).dump(indent, ' ', false, ordered_json::error_handler_t::replace);
+}
+
+// an array is indexed directly, an object needs a walk from the start
+common_json & common_json::iterator::operator*() const {
+    return guard([&]() -> common_json & {
+        ordered_json & j = as_json(node);
+
+        if (j.is_object()) {
+            return as_common(std::next(j.begin(), idx).value());
+        }
+        if (j.is_array()) {
+            return as_common(j[idx]);
+        }
+
+        // a plain value gives itself once, same as the backing library
+        return *node;
+    });
+}
+
+std::string common_json::iterator::key() const {
+    return guard([&] { return std::next(as_json(node).begin(), idx).key(); });
+}
+
+common_json::iterator common_json::begin() const {
+    return iterator(const_cast<common_json *>(this), 0);
+}
+
+common_json::iterator common_json::end() const {
+    return iterator(const_cast<common_json *>(this), size());
+}
+
+// the keys follow the backing library: the index for an array, "" for a plain value
+common_json::items_view::entry common_json::items_view::iterator::operator*() const {
+    return guard([&]() -> entry {
+        ordered_json & j = as_json(node);
+
+        if (j.is_object()) {
+            auto it = std::next(j.begin(), idx);
+
+            return { it.key(), as_common(it.value()) };
+        }
+        if (j.is_array()) {
+            return { std::to_string(idx), as_common(j[idx]) };
+        }
+
+        return { std::string(), *node };
+    });
+}
+
+common_json::items_view common_json::items() const {
+    return items_view(const_cast<common_json *>(this), size());
+}
+
+template <typename T> T common_json::get() const {
+    return guard([&] { return as_json(this).get<T>(); });
+}
+
+// the backing library cannot build a common_json, so this one is just a copy
+template <> common_json common_json::get<common_json>() const {
+    return *this;
+}
+
+// get<T>() is usable only for the types below
+
+#define COMMON_JSON_GET(...) template __VA_ARGS__ common_json::get<__VA_ARGS__>() const;
+
+COMMON_JSON_GET(bool)
+COMMON_JSON_GET(int)
+COMMON_JSON_GET(unsigned int)
+COMMON_JSON_GET(long)
+COMMON_JSON_GET(unsigned long)
+COMMON_JSON_GET(long long)
+COMMON_JSON_GET(unsigned long long)
+COMMON_JSON_GET(float)
+COMMON_JSON_GET(double)
+COMMON_JSON_GET(std::string)
+COMMON_JSON_GET(std::vector<float>)
+COMMON_JSON_GET(std::vector<std::string>)
+COMMON_JSON_GET(std::set<std::string>)
+COMMON_JSON_GET(std::vector<int>)
+COMMON_JSON_GET(std::vector<size_t>)
+COMMON_JSON_GET(std::unordered_map<std::string, size_t>)
+
+#undef COMMON_JSON_GET
@@ -0,0 +1,354 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <initializer_list>
+#include <iterator>
+#include <map>
+#include <memory>
+#include <set>
+#include <stdexcept>
+#include <string>
+#include <string_view>
+#include <type_traits>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+// common_json, a thin wrapper around vendor json library
+// the underlay library is pimpl, we are using nlohmann::json for now
+//
+// many features of the library are deliberately left out, to keep this interface small and generic and to keep compile time down
+//
+// some main differences compared to nlohmann::json :
+// - object keys keep the order in which they are added
+// - errors are always throw as common_json_error
+// - obj.push_back({key, val}) is intentionally unsupported to avoid confusion with push_back on a vector; write it as obj[key] = val for clarity
+// - a braced pair in value position does not build, e.g. {"key", {"a", "b"}}; write array({"a", "b"}) where nlohmann made an array
+//
+// in doubt, search the code base for an existing usage example; do not add anything to this header unless absolutely necessary
+
+class common_json;
+
+// common_json_value holds a list of these, and each of them holds a value, so one must come first
+struct common_json_item;
+
+struct common_json_error : std::runtime_error {
+    using std::runtime_error::runtime_error;
+};
+
+// one value, tagged so that this header stays free of the backing library
+// note: a value that holds a tree is single use, the second use gives null
+struct common_json_value {
+    enum value_type {
+        VAL_NULL,
+        VAL_BOOL,
+        VAL_INT,
+        VAL_UINT,
+        VAL_DOUBLE,
+        VAL_STRING,
+        VAL_JSON,
+    };
+
+    value_type type = VAL_NULL;
+
+    union {
+        bool     val_bool;
+        int64_t  val_int;
+        uint64_t val_uint = 0;
+        double   val_double;
+    };
+
+    std::string                  val_string;
+    std::shared_ptr<common_json> val_json;
+
+    common_json_value(std::nullptr_t = nullptr) : type(VAL_NULL) {}
+    common_json_value(bool val) : type(VAL_BOOL), val_bool(val) {}
+    common_json_value(std::string val) : type(VAL_STRING), val_string(std::move(val)) {}
+    // without this a string_view lands on the common_json ctor below and recurses
+    common_json_value(std::string_view val) : type(VAL_STRING), val_string(val) {}
+    common_json_value(const char * val);
+    common_json_value(const common_json & val);
+    common_json_value(common_json && val);
+    // only for the types instantiated in json.cpp, the rest fails at link time
+    template <typename T> common_json_value(const std::vector<T> & vals);
+    // a set becomes an array, in the set's own order
+    template <typename T> common_json_value(const std::set<T> & vals);
+    // a map becomes an object, keyed in the map's own order
+    template <typename T> common_json_value(const std::map<std::string, T> & vals);
+    template <typename T> common_json_value(const std::unordered_map<std::string, T> & vals);
+
+    // nested object, e.g. {"fn", {{"name", "x"}}}
+    // note: a nested pair {"a", "b"} does not build, use common_json::array({"a", "b"}) for an array
+    common_json_value(std::initializer_list<common_json_item> items);
+
+    template <typename T, typename std::enable_if<std::is_integral<T>::value && !std::is_same<T, bool>::value, int>::type = 0>
+    common_json_value(T val) : type(std::is_signed<T>::value ? VAL_INT : VAL_UINT) {
+        if (std::is_signed<T>::value) {
+            val_int = (int64_t) val;
+        } else {
+            val_uint = (uint64_t) val;
+        }
+    }
+
+    template <typename T, typename std::enable_if<std::is_floating_point<T>::value, int>::type = 0>
+    common_json_value(T val) : type(VAL_DOUBLE), val_double((double) val) {}
+};
+
+struct common_json_item {
+    std::string       key;
+    common_json_value val;
+
+    template <typename T>
+    common_json_item(std::string key, T && val) :
+        key(std::move(key)), val(std::forward<T>(val)) {}
+
+    // a braced list cannot deduce T, so it needs its own overload
+    common_json_item(std::string key, std::initializer_list<common_json_item> items) :
+        key(std::move(key)), val(items) {}
+};
+
+// the types common_json_value holds on its own
+// anything else reaches its common_json ctor and recurses forever
+template <typename T> struct common_json_is_value : std::integral_constant<bool,
+    std::is_arithmetic<T>::value ||
+    std::is_same<T, std::nullptr_t>::value ||
+    std::is_same<T, std::string>::value ||
+    std::is_same<T, std::string_view>::value ||
+    std::is_same<T, char *>::value ||
+    std::is_same<T, const char *>::value ||
+    std::is_same<T, common_json>::value> {};
+
+template <typename T, typename A>
+struct common_json_is_value<std::vector<T, A>> : std::true_type {};
+
+template <typename T, typename C, typename A>
+struct common_json_is_value<std::set<T, C, A>> : std::true_type {};
+
+template <typename V, typename C, typename A>
+struct common_json_is_value<std::map<std::string, V, C, A>> : std::true_type {};
+
+template <typename V, typename H, typename E, typename A>
+struct common_json_is_value<std::unordered_map<std::string, V, H, E, A>> : std::true_type {};
+
+class common_json {
+  public:
+    common_json();
+    common_json(const common_json & other);
+    common_json(common_json && other) noexcept;
+    common_json(std::initializer_list<common_json_item> items);
+    common_json(const common_json_value & val);
+
+    // direct, a value would need two conversions in a row
+    common_json(std::nullptr_t);
+
+    // one step, so that "abc" or a vector can go straight into a common_json
+    template <typename T, typename std::enable_if<!std::is_same<typename std::decay<T>::type, common_json>::value &&
+                                                  !std::is_same<typename std::decay<T>::type, common_json_value>::value, int>::type = 0>
+    common_json(T && val) : common_json(common_json_value(std::forward<T>(val))) {
+        static_assert(common_json_is_value<typename std::decay<T>::type>::value,
+                      "no common_json_value ctor holds this type, add one instead of letting it recurse");
+    }
+
+    // by value, same as the backing library
+    // the right side is copied before the left side can invalidate it, e.g. msg["a"] = msg.at("b")
+    common_json & operator=(common_json other) noexcept;
+
+    ~common_json();
+
+    // throws common_json_error if the text is not valid JSON
+    static common_json parse(const std::string & text);
+
+    // gives a discarded value instead of throwing, check it with is_discarded()
+    static common_json parse_no_throw(const std::string & text);
+
+    bool is_discarded() const;
+
+    static common_json array();
+    static common_json array(std::initializer_list<common_json_value> vals);
+    static common_json object();
+    static common_json object(std::initializer_list<common_json_item> items);
+
+    // holds a single value, e.g. make("abc").dump() gives "\"abc\""
+    static common_json make(const common_json_value & val);
+
+    bool is_null()    const;
+    bool is_object()  const;
+    bool is_array()   const;
+    bool is_string()  const;
+    bool is_boolean() const;
+    bool is_number()  const;
+    bool is_number_integer() const;
+    bool is_number_float()   const;
+
+    bool   empty() const;
+    size_t size()  const;
+
+    bool contains(const std::string & key) const;
+
+    bool operator==(const common_json_value & val) const;
+    bool operator!=(const common_json_value & val) const;
+
+    // at() throws common_json_error if the key is missing, operator[] adds a null value instead
+    // note: a const operator[] cannot add, it throws like at()
+    common_json       & at(const std::string & key);
+    const common_json & at(const std::string & key) const;
+    common_json       & at(size_t idx);
+    const common_json & at(size_t idx) const;
+
+    common_json       & operator[](const std::string & key);
+    const common_json & operator[](const std::string & key) const;
+    common_json       & operator[](const char * key)       { return (*this)[std::string(key)]; }
+    const common_json & operator[](const char * key) const { return (*this)[std::string(key)]; }
+    common_json       & operator[](int idx)       { return (*this)[to_idx(idx)]; }
+    const common_json & operator[](int idx) const { return (*this)[to_idx(idx)]; }
+    common_json       & operator[](size_t idx);
+    const common_json & operator[](size_t idx) const;
+
+    common_json       & front();
+    const common_json & front() const;
+    common_json       & back();
+    const common_json & back()  const;
+
+    void clear();
+
+    void erase(const std::string & key);
+    void erase(size_t idx);
+
+    // only for the types instantiated in json.cpp, the rest fails at link time
+    template <typename T> T get() const;
+
+    // implicit get<T>() for plain values, so they can be assigned to their C++ type directly
+    // note: kept to this short list on purpose, a wider one makes j["key"] ambiguous
+    // note: a numeric one would make "str = json;" ambiguous, a number converts to char too
+    operator std::string() const { return get<std::string>(); }
+
+    template <typename T>
+    T value(const std::string & key, T def) const {
+        return contains(key) ? at(key).get<T>() : def;
+    }
+
+    std::string value(const std::string & key, const char * def) const {
+        return contains(key) ? at(key).get<std::string>() : std::string(def);
+    }
+
+    // a JSON default needs no get<T>(), it is already the right type
+    common_json value(const std::string & key, const common_json & def) const {
+        return contains(key) ? at(key) : def;
+    }
+
+    void assign(const common_json_value & val);
+    void set(const common_json_item & item);
+    void push_back(const common_json_value & val);
+
+    // appends one object, e.g. push_back({{"a", 1}})
+    void push_back(std::initializer_list<common_json_item> items);
+
+    // 1 if the key is there, 0 if not
+    size_t count(const std::string & key) const;
+
+    // appends every value of another array; inserting an array into itself throws
+    void insert(const common_json & vals);
+
+    // a common_json goes through the copy assignment above, everything else becomes a value
+    template <typename T, typename std::enable_if<!std::is_same<typename std::decay<T>::type, common_json>::value, int>::type = 0>
+    common_json & operator=(T && val) {
+        assign(common_json_value(std::forward<T>(val)));
+        return *this;
+    }
+
+    std::string dump(int indent = -1) const;
+
+    // same as dump(), but bad UTF-8 gets replaced instead of throwing
+    std::string dump_safe(int indent = -1) const;
+
+    // walks an array by index, or an object in insertion order
+    // a plain value gives itself once, same as the backing library
+    class iterator {
+      public:
+        using iterator_category = std::forward_iterator_tag;
+        using value_type        = common_json;
+        using difference_type   = std::ptrdiff_t;
+        using pointer           = common_json *;
+        using reference         = common_json &;
+
+        iterator(common_json * node, size_t idx) : node(node), idx(idx) {}
+
+        common_json & operator*() const;
+        common_json & value()     const { return **this; }
+        std::string   key()       const;
+
+        iterator & operator++() {
+            idx++;
+            return *this;
+        }
+
+        bool operator!=(const iterator & other) const { return idx != other.idx; }
+        bool operator==(const iterator & other) const { return idx == other.idx; }
+
+      private:
+        common_json * node;
+        size_t        idx;
+    };
+
+    iterator begin() const;
+    iterator end()   const;
+
+    // allows: for (const auto & [key, val] : obj.items())
+    class items_view {
+      public:
+        // the members are public, so an entry also works with structured bindings
+        struct entry {
+            std::string   k;
+            common_json & v;
+
+            const std::string & key()   const { return k; }
+            common_json &       value() const { return v; }
+        };
+
+        items_view(common_json * node, size_t n) : node(node), n(n) {}
+
+        class iterator {
+          public:
+            iterator(common_json * node, size_t idx) : node(node), idx(idx) {}
+
+            entry operator*() const;
+
+            iterator & operator++() {
+                idx++;
+                return *this;
+            }
+
+            bool operator!=(const iterator & other) const { return idx != other.idx; }
+
+          private:
+            common_json * node;
+            size_t        idx;
+        };
+
+        iterator begin() const { return iterator(node, 0); }
+        iterator end()   const { return iterator(node, n); }
+
+      private:
+        common_json * node;
+        size_t        n;
+    };
+
+    items_view items() const;
+
+  private:
+    // a negative index must not turn into a huge size_t
+    static size_t to_idx(int idx) {
+        if (idx < 0) {
+            throw common_json_error("negative array index");
+        }
+        return (size_t) idx;
+    }
+
+    // the backing value is built here, json.cpp checks that it fits
+    // it cannot be a pointer: a value inside a tree would then not be a common_json
+    // at() could then only give back a copy instead of a real reference
+    alignas(8) unsigned char storage[32];
+};
+
+using common_json_entry = common_json::items_view::entry;
@@ -10,7 +10,6 @@
 #include <initializer_list>
 #include <map>
 #include <memory>
-#include <nlohmann/json.hpp>
 #include <regex>
 #include <set>
 #include <stdexcept>
@@ -1120,8 +1119,8 @@ common_peg_parser common_peg_parser_builder::chars(const std::string & classes,
    return wrap(arena_.add_parser(common_peg_chars_parser{classes, ranges, negated, min, max}));
 }

-common_peg_parser common_peg_parser_builder::schema(const common_peg_parser & p, const std::string & name, const nlohmann::ordered_json & schema, bool raw) {
-    return wrap(arena_.add_parser(common_peg_schema_parser{p.id(), name, std::make_shared<nlohmann::ordered_json>(schema), raw}));
+common_peg_parser common_peg_parser_builder::schema(const common_peg_parser & p, const std::string & name, const common_json & schema, bool raw) {
+    return wrap(arena_.add_parser(common_peg_schema_parser{p.id(), name, std::make_shared<common_json>(schema), raw}));
 }

 common_peg_parser common_peg_parser_builder::rule(const std::string & name, const common_peg_parser & p, bool trigger) {
@@ -1805,8 +1804,8 @@ void common_peg_arena::build_grammar(const common_grammar_builder & builder, boo
    }
 }

-static nlohmann::json serialize_parser_variant(const common_peg_parser_variant & variant) {
-    using json = nlohmann::json;
+static common_json serialize_parser_variant(const common_peg_parser_variant & variant) {
+    using json = common_json;

    return std::visit([](const auto & p) -> json {
        using T = std::decay_t<decltype(p)>;
@@ -1860,7 +1859,7 @@ static nlohmann::json serialize_parser_variant(const common_peg_parser_variant &
                {"type", "schema"},
                {"child", p.child},
                {"name", p.name},
-                {"schema", p.schema ? *p.schema : nullptr},
+                {"schema", p.schema ? *p.schema : json(nullptr)},
                {"raw", p.raw}
            };
        } else if constexpr (std::is_same_v<T, common_peg_rule_parser>) {
@@ -1888,19 +1887,19 @@ static nlohmann::json serialize_parser_variant(const common_peg_parser_variant &
    }, variant);
 }

-nlohmann::json common_peg_arena::to_json() const {
-    auto parsers = nlohmann::json::array();
+common_json common_peg_arena::to_json() const {
+    auto parsers = common_json::array();
    for (const auto & parser : parsers_) {
        parsers.push_back(serialize_parser_variant(parser));
    }
-    return nlohmann::json{
+    return common_json{
        {"parsers", parsers},
        {"rules", rules_},
        {"root", root_}
    };
 }

-static common_peg_parser_variant deserialize_parser_variant(const nlohmann::json & j) {
+static common_peg_parser_variant deserialize_parser_variant(const common_json & j) {
    if (!j.contains("type") || !j["type"].is_string()) {
        throw std::runtime_error("Parser variant JSON missing or invalid 'type' field");
    }
@@ -1969,9 +1968,9 @@ static common_peg_parser_variant deserialize_parser_variant(const nlohmann::json
        }
        common_peg_chars_parser parser;
        parser.pattern = j["pattern"];
-        parser.negated = j["negated"];
-        parser.min_count = j["min_count"];
-        parser.max_count = j["max_count"];
+        parser.negated = j["negated"].get<bool>();
+        parser.min_count = j["min_count"].get<int>();
+        parser.max_count = j["max_count"].get<int>();
        for (const auto & range_json : j["ranges"]) {
            if (!range_json.contains("start") || !range_json.contains("end")) {
                throw std::runtime_error("char_range missing 'start' or 'end' field");
@@ -2007,7 +2006,7 @@ static common_peg_parser_variant deserialize_parser_variant(const nlohmann::json
        parser.child = j["child"].get<common_peg_parser_id>();
        parser.name = j["name"];
        if (!j["schema"].is_null()) {
-            parser.schema = std::make_shared<nlohmann::ordered_json>(j["schema"]);
+            parser.schema = std::make_shared<common_json>(j["schema"]);
        }
        parser.raw = j["raw"].get<bool>();
        return parser;
@@ -2069,7 +2068,7 @@ static common_peg_parser_variant deserialize_parser_variant(const nlohmann::json
    throw std::runtime_error("Unknown parser type: " + type);
 }

-common_peg_arena common_peg_arena::from_json(const nlohmann::json & j) {
+common_peg_arena common_peg_arena::from_json(const common_json & j) {
    if (!j.contains("parsers") || !j["parsers"].is_array()) {
        throw std::runtime_error("JSON missing or invalid 'parsers' array");
    }
@@ -2109,7 +2108,7 @@ std::string common_peg_arena::save() const {
 }

 void common_peg_arena::load(const std::string & data) {
-    *this = from_json(nlohmann::json::parse(data));
+    *this = from_json(common_json::parse(data));
 }

 common_peg_arena build_peg_parser(const std::function<common_peg_parser(common_peg_parser_builder & builder)> & fn) {
@@ -1,6 +1,6 @@
 #pragma once

-#include <nlohmann/json_fwd.hpp>
+#include "json.h"

 #include <memory>
 #include <set>
@@ -245,7 +245,7 @@ struct common_peg_until_parser {
 struct common_peg_schema_parser {
    common_peg_parser_id child;
    std::string name;
-    std::shared_ptr<nlohmann::ordered_json> schema;
+    std::shared_ptr<common_json> schema;

    // Indicates if the GBNF should accept a raw string that matches the schema.
    bool raw;
@@ -332,8 +332,8 @@ class common_peg_arena {

    std::string dump(common_peg_parser_id id) const;

-    nlohmann::json to_json() const;
-    static common_peg_arena from_json(const nlohmann::json & j);
+    common_json to_json() const;
+    static common_peg_arena from_json(const common_json & j);

    std::string save() const;
    void load(const std::string & data);
@@ -490,7 +490,7 @@ class common_peg_parser_builder {

    // Wraps a parser with JSON schema metadata for grammar generation.
    // Used internally to convert JSON schemas to GBNF grammar rules.
-    common_peg_parser schema(const common_peg_parser & p, const std::string & name, const nlohmann::ordered_json & schema, bool raw = false);
+    common_peg_parser schema(const common_peg_parser & p, const std::string & name, const common_json & schema, bool raw = false);

    // Creates a named rule, stores it in the grammar, and returns a ref.
    // If trigger=true, marks this rule as an entry point for lazy grammar generation.
@@ -2322,6 +2322,9 @@ common_params common_base_params_to_speculative(const common_params & params) {
    const auto & params_spec = params.speculative.draft;
    common_params result = params;

+    result.embedding    = false;
+    result.pooling_type = LLAMA_POOLING_TYPE_UNSPECIFIED;
+
    if (has_draft) {
        result.devices               = params_spec.devices;
        result.model                 = params_spec.mparams;
@@ -2385,6 +2388,9 @@ common_speculative_init_result::common_speculative_init_result(
        cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
    }

+    // the draft context holds as many tokens per sequence as the target context
+    cparams.n_ctx = llama_n_ctx(ctx_tgt);
+
    // note: for small models maybe we can set this to the maximum possible draft from all speculative types
    //       the extra memory for small models is likely negligible?
    cparams.n_rs_seq  = 0;
@@ -58,6 +58,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "DSparkDraftModel": "qwen",
    "DSparkSpeculator": "qwen",
    "Lfm2DSparkDraftModel": "qwen",
+    "LingDSparkModel": "qwen",
    "DeepseekV4ForCausalLM": "deepseek",
    "DeepseekV4DSparkModel": "deepseek",
    "DistilBertForMaskedLM": "bert",
@@ -282,6 +283,8 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "CogVLMForCausalLM": "cogvlm",
    "DeepseekOCR2ForCausalLM": "deepseek",
    "DeepseekOCRForCausalLM": "deepseek",
+    "Dots3NoteForCausalLM": "dots3",
+    "Dots3NoteForConditionalGeneration": "dots3",
    "DotsOCRForCausalLM": "dotsocr",
    "Exaone4_5_ForConditionalGeneration": "exaone",
    "Gemma3ForConditionalGeneration": "gemma",
@@ -3,12 +3,14 @@ from __future__ import annotations
 import math
 import re

-from typing import TYPE_CHECKING, Callable, Iterable
+import torch
+
+from typing import TYPE_CHECKING, Any, Callable, Iterable

 if TYPE_CHECKING:
    from torch import Tensor

-from .base import ModelBase, gguf
+from .base import MmprojModel, ModelBase, gguf

 from .deepseek import DeepseekV2Model

@@ -193,3 +195,129 @@ class Dots3NoteModel(DeepseekV2Model):
            return

        yield from super().modify_tensors(data_torch, name, bid)
+
+
+@ModelBase.register("Dots3NoteForCausalLM", "Dots3NoteForConditionalGeneration")
+class Dots3NoteMmprojModel(MmprojModel):
+    has_vision_encoder = True
+    has_audio_encoder = True
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert self.hparams_vision is not None
+        assert self.hparams_audio is not None
+
+        # preprocessor_config.json nests the image params under vision_config
+        self.preprocessor_config = {**self.preprocessor_config, **self.preprocessor_config.get("vision_config", {})}
+
+        vis = self.hparams_vision
+        # in this config, hidden_size is the adapter output width; embed_dim is the tower width
+        vis["hidden_size"] = vis["embed_dim"]
+        vis["image_size"] = 0  # dynamic resolution
+        self.pyramid = [max(0, n) for n in vis["pyramid_num_routed"]]
+
+        if vis.get("adapter_type") != "patch_merger" or not vis.get("pre_pixel_shuffle"):
+            raise ValueError("dots3-note vision conversion requires adapter_type=patch_merger and pre_pixel_shuffle")
+        if vis.get("router_scoring_func", "sigmoid") != "sigmoid" or vis.get("router_scale", 1.0) != 1.0:
+            raise ValueError("dots3-note vision conversion only supports sigmoid routing with router_scale=1.0")
+        if vis.get("temporal_patch_size", 1) != 1 or vis.get("use_bias") or not vis.get("use_qk_norm"):
+            raise ValueError("unsupported dots3-note vision config variant")
+
+        aud = self.hparams_audio
+        if not aud.get("use_conv2d_stem") or not aud.get("use_rope") or not aud.get("use_rms_norm") or aud.get("use_causal"):
+            raise ValueError("unsupported dots3-note audio config variant")
+        if aud["whisper_config"].get("activation_function") != "swiglu":
+            raise ValueError("dots3-note audio conversion requires the swiglu activation")
+        if aud.get("merge_factor", 1) != 1 or aud.get("chunk_seconds") != 60:
+            raise ValueError("unsupported dots3-note audio chunking config")
+        # the graph hard-codes these rope parameters
+        rope = aud.get("rope_parameters", {})
+        if rope.get("partial_rotary_factor") != 0.5 or rope.get("rope_theta") != 10000.0:
+            raise ValueError("unsupported dots3-note audio rope config")
+
+    def get_audio_config(self) -> dict[str, Any] | None:
+        cfg = self.global_config.get("audio_config")
+        if cfg is not None:
+            # aliases so MmprojModel.find_aparam() / n_block_keys can resolve them
+            whisper = cfg["whisper_config"]
+            cfg["hidden_size"] = whisper["d_model"]
+            cfg["intermediate_size"] = whisper["encoder_ffn_dim"]
+            cfg["num_attention_heads"] = whisper["encoder_attention_heads"]
+            cfg["num_hidden_layers"] = whisper["encoder_layers"]
+        return cfg
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        assert self.hparams_vision is not None
+        assert self.hparams_audio is not None
+
+        self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.DOTS3NOTE_V)
+        self.gguf_writer.add_vision_use_silu(True)
+        self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams_vision["rms_norm_eps"])
+        self.gguf_writer.add_vision_spatial_merge_size(self.hparams_vision["spatial_merge_size"])
+        self.gguf_writer.add_vision_min_pixels(self.preprocessor_config["min_pixels"])
+        self.gguf_writer.add_vision_max_pixels(self.preprocessor_config["max_pixels"])
+        # pyramid MoE: per-block routed expert count, 0 = dense block
+        self.gguf_writer.add_vision_expert_count_per_layer(self.pyramid)
+        self.gguf_writer.add_vision_expert_used_count(int(self.hparams_vision["capacity_factor"]))
+
+        self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.DOTS3NOTE_A)
+        self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["whisper_config"]["num_mel_bins"])
+        self.gguf_writer.add_audio_attention_layernorm_eps(1e-6)  # Dots3NoteAudioRMSNorm default
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, _ = item
+        if not name.startswith(("vision_encoder.", "audio_encoder.")):
+            return None
+        return super().filter_tensors(item)
+
+    _vis_experts: dict[int, dict[str, Tensor]] | None = None
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # router params have no .weight suffix in the checkpoint, but gguf tools expect one
+        if name.endswith((".gate_weight", ".router_bias")):
+            name += ".weight"
+
+        # audio fc1 fuses gate and up for swiglu; split it
+        if ".speech_encoder.layers." in name and ".fc1." in name:
+            gate, up = data_torch.chunk(2, dim=0)
+            yield from super().modify_tensors(gate, name.replace(".fc1.", ".fc1_gate."), bid)
+            yield from super().modify_tensors(up, name.replace(".fc1.", ".fc1_up."), bid)
+            return
+
+        # vision MoE: stack per-expert weights into a single 3D tensor per block
+        if ".mlp.experts." in name:
+            assert bid is not None
+            n_expert = self.pyramid[bid]
+            if self._vis_experts is None:
+                self._vis_experts = {}
+            buf = self._vis_experts.setdefault(bid, {})
+            buf[name] = data_torch
+
+            if len(buf) >= n_expert * 3:
+                for w_name in ("fc1", "fc2", "fc3"):
+                    datas: list[Tensor] = []
+                    for xid in range(n_expert):
+                        ename = f"vision_encoder.blocks.{bid}.mlp.experts.{xid}.{w_name}.weight"
+                        datas.append(buf.pop(ename))
+                    merged = torch.stack(datas, dim=0)
+                    yield from super().modify_tensors(merged, f"vision_encoder.blocks.{bid}.mlp.experts.{w_name}.weight", bid)
+            return
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+    def prepare_tensors(self):
+        super().prepare_tensors()
+        if self._vis_experts is not None:
+            leftover = [k for d in self._vis_experts.values() for k in d.keys()]
+            if leftover:
+                raise ValueError(f"unprocessed vision experts: {leftover}")
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        # FP32 routing is load-bearing for the vision MoE (near-tied expert scores)
+        if ".ffn_gate_inp." in new_name or ".exp_probs_b." in new_name:
+            return gguf.GGMLQuantizationType.F32
+        if ".conv2d" in new_name or "a.conv_out" in new_name:
+            return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
@@ -709,7 +709,13 @@ class DFlashModel(Qwen3Model):
        yield from super().modify_tensors(data_torch, name, bid)


-@ModelBase.register("Qwen3DSparkModel", "DSparkDraftModel", "DSparkSpeculator", "Lfm2DSparkDraftModel")
+@ModelBase.register(
+    "Qwen3DSparkModel",
+    "DSparkDraftModel",
+    "DSparkSpeculator",
+    "Lfm2DSparkDraftModel",
+    "LingDSparkModel",
+)
@ModelBase.example("satgeze/Qwen3.6-27B-DSpark")
 class DSparkModel(DFlashModel):
    # DSpark = DFlash + a semi-autoregressive Markov head.
@@ -342,6 +342,9 @@ set(GGML_PUBLIC_HEADERS
    include/gguf.h)

 set_target_properties(ggml PROPERTIES PUBLIC_HEADER "${GGML_PUBLIC_HEADERS}")
+#if (GGML_METAL)
+#    set_target_properties(ggml PROPERTIES RESOURCE "${CMAKE_CURRENT_SOURCE_DIR}/src/ggml-metal.metal")
+#endif()
 install(TARGETS ggml LIBRARY PUBLIC_HEADER)
 install(TARGETS ggml-base LIBRARY)

@@ -1896,7 +1896,6 @@ void ggml_compute_forward_repeat_back(
 }

 // ggml_compute_forward_concat
-
 static void ggml_compute_forward_concat_any(
    const ggml_compute_params * params,
    ggml_tensor * dst) {
@@ -1904,8 +1903,6 @@ static void ggml_compute_forward_concat_any(
    const ggml_tensor * src0 = dst->src[0];
    const ggml_tensor * src1 = dst->src[1];

-    const size_t len = ggml_type_size(src0->type);
-
    const int ith = params->ith;
    const int nth = params->nth;

@@ -1914,31 +1911,38 @@ static void ggml_compute_forward_concat_any(
    const int32_t dim = ggml_get_op_params_i32(dst, 0);

    GGML_ASSERT(dim >= 0 && dim < 4);
+    GGML_ASSERT(ggml_is_contiguous_rows(src0));
+    GGML_ASSERT(ggml_is_contiguous_rows(src1));

    int64_t o[4] = {0, 0, 0, 0};
+
    if (dim == 0) {
+        GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0);
+        GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0);
+
        o[dim] = src0->ne[dim]/ggml_blck_size(src0->type);
    } else {
        o[dim] = src0->ne[dim];
    }

-    const char * x;
+    // Region 1: copy rows from src0
+    for (int i3 = 0; i3 < ne03; i3++) {
+        for (int i2 = ith; i2 < ne02; i2 += nth) {
+            for (int i1 = 0; i1 < ne01; i1++) {
+                const char * x = (const char *) src0->data + i1*nb01 + i2*nb02 + i3*nb03;
+                      char * y = (      char *) dst->data  + i1*nb1  + i2*nb2  + i3*nb3;
+                memcpy(y, x, ggml_row_size(src0->type, ne00));
+            }
+        }
+    }

-    // TODO: smarter multi-theading
-    for (int i3 = 0; i3 < ne3; i3++) {
-        for (int i2 = ith; i2 < ne2; i2 += nth) {
-            for (int i1 = 0; i1 < ne1; i1++) {
-                for (int i0 = 0; i0 < ne0/ggml_blck_size(dst->type); i0++) {
-                    if (i0 < ne00/ggml_blck_size(src0->type) && i1 < ne01 && i2 < ne02 && i3 < ne03) {
-                        x = (const char *)src0->data + (i0       )*nb00 + (i1       )*nb01 + (i2       )*nb02 + (i3       )*nb03;
-                    } else {
-                        x = (const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13;
-                    }
-
-                    char * y = (char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3;
-
-                    memcpy(y, x, len);
-                }
+    // Region 2: copy rows from src1, offset into dst by o[]
+    for (int i3 = 0; i3 < ne13; i3++) {
+        for (int i2 = ith; i2 < ne12; i2 += nth) {
+            for (int i1 = 0; i1 < ne11; i1++) {
+                const char * x = (const char *) src1->data + i1*nb11         + i2*nb12         + i3*nb13;
+                      char * y = (      char *)  dst->data + (i1 + o[1])*nb1 + (i2 + o[2])*nb2 + (i3 + o[3])*nb3 + o[0]*nb0;
+                memcpy(y, x, ggml_row_size(src1->type, ne10));
            }
        }
    }
@@ -2078,14 +2082,6 @@ void ggml_compute_forward_concat(
    ggml_tensor * dst) {

    const ggml_tensor * src0 = dst->src[0];
-    const ggml_tensor * src1 = dst->src[1];
-
-    if (ggml_is_quantized(src0->type)) {
-        GGML_ASSERT(ggml_is_contiguous_rows(src0));
-        GGML_ASSERT(ggml_is_contiguous_rows(src1));
-        GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0);
-        GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0);
-    }

    switch (src0->type) {
        case GGML_TYPE_F16:
@@ -24,119 +24,62 @@ if (GGML_METAL_NDEBUG)
 endif()

 set(METALLIB_COMMON "${CMAKE_CURRENT_SOURCE_DIR}/../ggml-common.h")
-set(METALLIB_KERNELS_COMMON     "${CMAKE_CURRENT_SOURCE_DIR}/kernels/common.h")
-set(METALLIB_KERNELS_DEQUANTIZE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/dequantize.h")
-set(METALLIB_KERNELS_QUANTIZE   "${CMAKE_CURRENT_SOURCE_DIR}/kernels/quantize.h")
-
-set(METALLIB_KERNEL_SOURCES
-    kernels/fa.metal
-    kernels/mul_mv.metal
-    kernels/mul_mm.metal
-    kernels/quantize.metal
-    kernels/softmax.metal
-    kernels/norm.metal
-    kernels/unary.metal
-    kernels/binbcast.metal
-    kernels/reduce.metal
-    kernels/tri.metal
-    kernels/ssm.metal
-    kernels/wkv.metal
-    kernels/gated_delta_net.metal
-    kernels/solve_tri.metal
-    kernels/rope.metal
-    kernels/conv.metal
-    kernels/upscale.metal
-    kernels/argsort.metal
-    kernels/pool.metal
-    kernels/misc.metal
-)
-
 if (GGML_METAL_EMBED_LIBRARY)
    enable_language(ASM)

    add_compile_definitions(GGML_METAL_EMBED_LIBRARY)

-    set(METALLIB_IMPL "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal-impl.h")
+    set(METALLIB_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal.metal")
+    set(METALLIB_IMPL   "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal-impl.h")

    file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/autogenerated")

-    set(METALLIB_EMBED_ASM_FILES "")
-    foreach(src ${METALLIB_KERNEL_SOURCES})
-        get_filename_component(kind ${src} NAME_WE)
-        # symbol names must be valid C identifiers ('-' is not allowed)
-        string(REPLACE "-" "_" kind_sym ${kind})
+    # merge ggml-common.h and ggml-metal.metal into a single file
+    set(METALLIB_EMBED_ASM        "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed.s")
+    set(METALLIB_SOURCE_EMBED     "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed.metal")
+    set(METALLIB_SOURCE_EMBED_TMP "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed.metal.tmp")

-        set(SRC   "${CMAKE_CURRENT_SOURCE_DIR}/kernels/${kind}.metal")
-        set(EMBED "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed-${kind}.metal")
-        set(ASM   "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed-${kind}.s")
+    add_custom_command(
+        OUTPUT "${METALLIB_EMBED_ASM}"
+        COMMAND echo "Embedding Metal library"
+        COMMAND sed -e "/__embed_ggml-common.h__/r ${METALLIB_COMMON}"       -e "/__embed_ggml-common.h__/d"         < "${METALLIB_SOURCE}"           > "${METALLIB_SOURCE_EMBED_TMP}"
+        COMMAND sed -e "/\#include \"ggml-metal-impl.h\"/r ${METALLIB_IMPL}" -e "/\#include \"ggml-metal-impl.h\"/d" < "${METALLIB_SOURCE_EMBED_TMP}" > "${METALLIB_SOURCE_EMBED}"
+        COMMAND echo ".section __DATA,__ggml_metallib"          >  "${METALLIB_EMBED_ASM}"
+        COMMAND echo ".globl _ggml_metallib_start"              >> "${METALLIB_EMBED_ASM}"
+        COMMAND echo "_ggml_metallib_start:"                    >> "${METALLIB_EMBED_ASM}"
+        COMMAND echo .incbin "\"${METALLIB_SOURCE_EMBED}\""     >> "${METALLIB_EMBED_ASM}"
+        COMMAND echo ".globl _ggml_metallib_end"                >> "${METALLIB_EMBED_ASM}"
+        COMMAND echo "_ggml_metallib_end:"                      >> "${METALLIB_EMBED_ASM}"
+        DEPENDS ../ggml-common.h ggml-metal.metal ggml-metal-impl.h
+        COMMENT "Generate assembly for embedded Metal library"
+        VERBATIM
+    )

-        # only prepend headers that this source actually includes
-        set(HEADERS_FOR_SRC ${METALLIB_KERNELS_COMMON})
-        file(STRINGS ${SRC} _has_dequantize REGEX "#include \"dequantize\\.h\"")
-        file(STRINGS ${SRC} _has_quantize   REGEX "#include \"quantize\\.h\"")
-        if(_has_dequantize)
-            list(APPEND HEADERS_FOR_SRC ${METALLIB_KERNELS_DEQUANTIZE})
-        endif()
-        if(_has_quantize)
-            list(APPEND HEADERS_FOR_SRC ${METALLIB_KERNELS_QUANTIZE})
-        endif()
-
-        add_custom_command(
-            OUTPUT "${ASM}"
-            # Step 1: concatenate shared headers + this kernel source
-            COMMAND cat ${HEADERS_FOR_SRC} ${SRC} > "${EMBED}.tmp1"
-            # Step 2: remove internal #include and #pragma once
-            COMMAND sed -e "/\#include \"common.h\"/d" -e "/\#include \"dequantize.h\"/d" -e "/\#include \"quantize.h\"/d" -e "/\#pragma once/d" < "${EMBED}.tmp1" > "${EMBED}.tmp2"
-            # Step 3: inline ggml-common.h (replacing __embed_ggml-common.h__ sentinel)
-            COMMAND sed -e "/__embed_ggml-common.h__/r ${METALLIB_COMMON}" -e "/__embed_ggml-common.h__/d" < "${EMBED}.tmp2" > "${EMBED}.tmp3"
-            # Step 4: inline ggml-metal-impl.h
-            COMMAND sed -e "/\#include \"ggml-metal-impl.h\"/r ${METALLIB_IMPL}" -e "/\#include \"ggml-metal-impl.h\"/d" < "${EMBED}.tmp3" > "${EMBED}"
-            # Step 5: emit an asm chunk with kind-specific start/end symbols
-            #   note: '-' is illegal in C symbols, so we use kind_sym; the macOS
-            #   section name is limited to 16 chars so we keep it shared
-            #   across kinds (__ggml_metallib) and only vary the global symbols.
-            COMMAND echo ".section __DATA,__ggml_metallib"                       >  "${ASM}"
-            COMMAND echo ".globl _ggml_metallib_${kind_sym}_start"               >> "${ASM}"
-            COMMAND echo "_ggml_metallib_${kind_sym}_start:"                     >> "${ASM}"
-            COMMAND echo .incbin "\"${EMBED}\""                                  >> "${ASM}"
-            COMMAND echo ".globl _ggml_metallib_${kind_sym}_end"                 >> "${ASM}"
-            COMMAND echo "_ggml_metallib_${kind_sym}_end:"                       >> "${ASM}"
-            DEPENDS ../ggml-common.h ggml-metal-impl.h
-                    kernels/common.h kernels/dequantize.h kernels/quantize.h
-                    kernels/${kind}.metal
-            COMMENT "Generate embedded Metal library for ${kind}"
-            VERBATIM
-        )
-
-        list(APPEND METALLIB_EMBED_ASM_FILES "${ASM}")
-    endforeach()
-
-    target_sources(ggml-metal PRIVATE ${METALLIB_EMBED_ASM_FILES})
+    target_sources(ggml-metal PRIVATE "${METALLIB_EMBED_ASM}")
 else()
-    # copy header files to bin directory
+    # copy metal files to bin directory
    configure_file(../ggml-common.h  ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common.h     COPYONLY)
+    configure_file(ggml-metal.metal  ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal  COPYONLY)
    configure_file(ggml-metal-impl.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal-impl.h COPYONLY)

-    file(MAKE_DIRECTORY "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels")
-    configure_file(kernels/common.h     ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/common.h     COPYONLY)
-    configure_file(kernels/dequantize.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/dequantize.h COPYONLY)
-    configure_file(kernels/quantize.h   ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/quantize.h   COPYONLY)
-
-    foreach(src ${METALLIB_KERNEL_SOURCES})
-        configure_file(${src} ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${src} COPYONLY)
-    endforeach()
-
    if (GGML_METAL_SHADER_DEBUG)
-        # note: disabling fast math is needed in order to pass tests/test-backend-ops
+        # custom command to do the following:
+        #   xcrun -sdk macosx metal    -fno-fast-math -c ggml-metal.metal -o ggml-metal.air
+        #   xcrun -sdk macosx metallib                   ggml-metal.air   -o default.metallib
+        #
+        # note: this is the only way I found to disable fast-math in Metal. it's ugly, but at least it works
+        #       disabling fast math is needed in order to pass tests/test-backend-ops
        # note: adding -fno-inline fixes the tests when using MTL_SHADER_VALIDATION=1
        # note: unfortunately, we have to call it default.metallib instead of ggml.metallib
        #       ref: https://github.com/ggml-org/whisper.cpp/issues/1720
        # note: adding -g causes segmentation fault during compile
+        #set(XC_FLAGS -fno-fast-math -fno-inline -g)
        set(XC_FLAGS -fno-fast-math -fno-inline)
    else()
        set(XC_FLAGS -O3)
    endif()

+    # Append macOS metal versioning flags
    if (GGML_METAL_MACOSX_VERSION_MIN)
        message(STATUS "Adding  -mmacosx-version-min=${GGML_METAL_MACOSX_VERSION_MIN} flag to metal compilation")
        list   (APPEND XC_FLAGS -mmacosx-version-min=${GGML_METAL_MACOSX_VERSION_MIN})
@@ -147,46 +90,35 @@ else()
        list   (APPEND XC_FLAGS -std=${GGML_METAL_STD})
    endif()

-    # Compile each kernel source to .air, then link into default.metallib
-    set(AIR_FILES "")
-    foreach(src ${METALLIB_KERNEL_SOURCES})
-        get_filename_component(name ${src} NAME_WE)
-        set(AIR "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${name}.air")
-        list(APPEND AIR_FILES ${AIR})
-        add_custom_command(
-            OUTPUT ${AIR}
-            COMMAND xcrun -sdk macosx metal ${XC_FLAGS} -I ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${src} -o ${AIR}
-            DEPENDS ${src} kernels/common.h kernels/dequantize.h kernels/quantize.h ${METALLIB_COMMON} ggml-metal-impl.h
-            COMMENT "Compiling ${src}"
-            VERBATIM
-        )
-    endforeach()
-
    add_custom_command(
        OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
-        COMMAND xcrun -sdk macosx metallib ${AIR_FILES} -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
+        COMMAND xcrun -sdk macosx metal ${XC_FLAGS} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal -o - |
+                xcrun -sdk macosx metallib        - -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
        COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common.h
-        COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal-impl.h
-        COMMAND rm -rf ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels
-        DEPENDS ${AIR_FILES}
-        COMMENT "Linking Metal kernels into default.metallib"
-    )
+        COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal
+        DEPENDS ggml-metal.metal ${METALLIB_COMMON}
+        COMMENT "Compiling Metal kernels"
+        )

+    # FIXME: only add to the ggml-metal target?
    add_custom_target(
        ggml-metal-lib ALL
        DEPENDS ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
-    )
+        )
 endif() # GGML_METAL_EMBED_LIBRARY

 if (NOT GGML_METAL_EMBED_LIBRARY)
    install(
-        DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/kernels/
-        DESTINATION ${CMAKE_INSTALL_BINDIR}/kernels
-        FILES_MATCHING PATTERN "*.metal" PATTERN "*.h"
-    )
+        FILES src/ggml-metal/ggml-metal.metal
+        PERMISSIONS
+            OWNER_READ
+            OWNER_WRITE
+            GROUP_READ
+            WORLD_READ
+        DESTINATION ${CMAKE_INSTALL_BINDIR})

-    install(
-        FILES ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
-        DESTINATION ${CMAKE_INSTALL_BINDIR}
-    )
+        install(
+            FILES ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
+            DESTINATION ${CMAKE_INSTALL_BINDIR}
+        )
 endif()
@@ -95,63 +95,8 @@ int ggml_metal_pipeline_max_theads_per_threadgroup(struct ggml_metal_pipeline_wi
    return pipeline.pipeline->obj.maxTotalThreadsPerThreadgroup;
 }

-//
-// MTLLibrary collection (one library per op-source, compiled separately)
-//
-
-// Single source of truth for the per-kind metal libraries. The order here
-// defines the enum values and every per-kind table below, so adding a library
-// is a one-line change here (plus adding its source to CMakeLists.txt).
-//   X(suffix, name): name is both the kernels/<name>.metal basename and the
-//   ggml_metallib_<name>_{start,end} embed-symbol stem.
-#define GGML_METAL_LIBS \
-    X(FA,              fa)             \
-    X(MUL_MV,          mul_mv)         \
-    X(MUL_MM,          mul_mm)         \
-    X(QUANTIZE,        quantize)       \
-    X(SOFTMAX,         softmax)        \
-    X(NORM,            norm)           \
-    X(UNARY,           unary)          \
-    X(BINBCAST,        binbcast)       \
-    X(REDUCE,          reduce)         \
-    X(TRI,             tri)            \
-    X(SSM,             ssm)            \
-    X(WKV,             wkv)            \
-    X(GATED_DELTA_NET, gated_delta_net)\
-    X(SOLVE_TRI,       solve_tri)      \
-    X(ROPE,            rope)           \
-    X(CONV,            conv)           \
-    X(UPSCALE,         upscale)        \
-    X(ARGSORT,         argsort)        \
-    X(POOL,            pool)           \
-    X(MISC,            misc)
-
-enum ggml_metal_lib_kind {
-#define X(e, s) GGML_METAL_LIB_##e,
-    GGML_METAL_LIBS
-#undef X
-    GGML_METAL_LIB_COUNT,
-};
-
-static const char * const k_lib_names[GGML_METAL_LIB_COUNT] = {
-#define X(e, s) [GGML_METAL_LIB_##e] = #s,
-    GGML_METAL_LIBS
-#undef X
-};
-
 struct ggml_metal_library {
-    // Per-kind compiled libraries. When single_library is true, the whole library
-    // (e.g. a pre-compiled default.metallib or a from-source build) lives at
-    // objs[0] and the remaining slots are nil.
-    id<MTLLibrary> objs[GGML_METAL_LIB_COUNT];
-    bool single_library; // true: combined library at objs[0]; false: per-kind libs in objs[*]
-
-    // Routing table: kernel function name -> objs[] index, populated from each
-    // compiled library's -[MTLLibrary functionNames]. The actual compiled
-    // libraries are the single source of truth for which library owns a kernel,
-    // so adding kernels later requires no manual routing maintenance.
-    // nil in single_library mode (everything resolves to objs[0]).
-    NSMutableDictionary<NSString *, NSNumber *> * fn_to_lib;
+    id<MTLLibrary> obj;

    ggml_metal_device_t dev;
    ggml_metal_pipelines_t pipelines; // cache of compiled pipelines
@@ -159,376 +104,160 @@ struct ggml_metal_library {
    NSLock * lock;
 };

-// Build the fn_to_lib routing table by querying each compiled library's public
-// function names. Call once after all per-kind libraries have been compiled.
-static void ggml_metal_library_build_index(ggml_metal_library_t lib) {
-    @autoreleasepool {
-        NSMutableDictionary<NSString *, NSNumber *> * index = [[NSMutableDictionary alloc] init];
-        for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
-            for (NSString * fname in [lib->objs[kind] functionNames]) {
-                index[fname] = @(kind);
-            }
-        }
-        lib->fn_to_lib = index;
-    }
-}
+ggml_metal_library_t ggml_metal_library_init(ggml_metal_device_t dev) {
+    id<MTLLibrary> library = nil;
+    id<MTLDevice> device = ggml_metal_device_get_obj(dev);

-// Parse a `#include "name"` line. Returns the quoted name in *include_name on
-// success. Whitespace-tolerant; ignores `#include <...>` (system headers).
-static bool ggml_metal_library_parse_quoted_include(NSString * line, NSString ** include_name) {
-    NSScanner * scanner = [NSScanner scannerWithString:line];
-    scanner.charactersToBeSkipped = [NSCharacterSet whitespaceCharacterSet];
+    // load library
+    //
+    // - first check if the library is embedded
+    // - then check if the library is in the bundle
+    // - if not found, load the source and compile it
+    // - if that fails, return NULL
+    //
+    // TODO: move to a function
+    {
+        const int64_t t_start = ggml_time_us();

-    if (![scanner scanString:@"#" intoString:NULL] ||
-        ![scanner scanString:@"include" intoString:NULL] ||
-        ![scanner scanString:@"\"" intoString:NULL]) {
-        return false;
-    }
+        NSError * error = nil;
+        NSString * src = nil;

-    NSString * name = nil;
-    if (![scanner scanUpToString:@"\"" intoString:&name]) {
-        return false;
-    }
+#if GGML_METAL_EMBED_LIBRARY
+        GGML_LOG_INFO("%s: using embedded metal library\n", __func__);

-    if (include_name) {
-        *include_name = name;
-    }
-    return true;
-}
+        extern const char ggml_metallib_start[];
+        extern const char ggml_metallib_end[];

-// Recursively inline `#include "name"` directives. System includes (<...>),
-// `#if/#else/#endif`, and other preprocessor lines are passed through to the
-// Metal compiler unchanged. `#pragma once` is dropped since `seen` already
-// guards against double-inclusion.
-static bool ggml_metal_library_flatten_file(NSMutableString * dst, NSString * path,
-                                            NSArray<NSString *> * search_paths,
-                                            NSMutableSet<NSString *> * seen, NSError ** error) {
-    NSString * key = [path stringByStandardizingPath];
-    if ([seen containsObject:key]) {
-        return true;
-    }
-    [seen addObject:key];
+        src = [[NSString alloc] initWithBytes:ggml_metallib_start length:(ggml_metallib_end-ggml_metallib_start) encoding:NSUTF8StringEncoding];
+#else

-    NSString * src = [NSString stringWithContentsOfFile:path encoding:NSUTF8StringEncoding error:error];
-    if (!src) {
-        return false;
-    }
+#ifdef SWIFT_PACKAGE
+        NSBundle * bundle = SWIFTPM_MODULE_BUNDLE;
+#else
+        NSBundle * bundle = [NSBundle bundleForClass:[GGMLMetalClass class]];
+#endif

-    NSFileManager * fm = [NSFileManager defaultManager];
-    for (NSString * line in [src componentsSeparatedByString:@"\n"]) {
-        NSString * trimmed = [line stringByTrimmingCharactersInSet:[NSCharacterSet whitespaceCharacterSet]];
-        if ([trimmed isEqualToString:@"#pragma once"]) {
-            continue;
-        }
+        NSString * path_lib = [bundle pathForResource:@"default" ofType:@"metallib"];
+        if (path_lib == nil) {
+            // Try to find the resource in the directory where the current binary located.
+            NSString * bin_cur = [[NSProcessInfo processInfo] arguments][0];
+            NSString * bin_dir = [bin_cur stringByDeletingLastPathComponent];

-        NSString * include_name = nil;
-        if (ggml_metal_library_parse_quoted_include(line, &include_name)) {
-            NSString * resolved = nil;
-            for (NSString * dir in search_paths) {
-                NSString * candidate = [dir stringByAppendingPathComponent:include_name];
-                if ([fm isReadableFileAtPath:candidate]) {
-                    resolved = candidate;
-                    break;
+            NSString * path_lib_default = [NSString pathWithComponents:@[bin_dir, @"default.metallib"]];
+            if ([[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
+                GGML_LOG_INFO("%s: found '%s'\n", __func__, [path_lib_default UTF8String]);
+
+                NSDictionary * atts = [[NSFileManager defaultManager] attributesOfItemAtPath:path_lib_default error:&error];
+                if (atts && atts[NSFileType] == NSFileTypeSymbolicLink) {
+                    // Optionally, if this is a symlink, try to resolve it.
+                    path_lib_default = [[NSFileManager defaultManager] destinationOfSymbolicLinkAtPath:path_lib_default error:&error];
+                    if (path_lib_default && [path_lib_default length] > 0 && ![[path_lib_default substringToIndex:1] isEqualToString:@"/"]) {
+                        // It is a relative path, adding the binary directory as directory prefix.
+                        path_lib_default = [NSString pathWithComponents:@[bin_dir, path_lib_default]];
+                    }
+                    if (!path_lib_default || ![[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
+                        // Link to the resource could not be resolved.
+                        path_lib_default = nil;
+                    } else {
+                        GGML_LOG_INFO("%s: symlink resolved '%s'\n", __func__, [path_lib_default UTF8String]);
+                    }
                }
+            } else {
+                // The resource couldn't be found in the binary's directory.
+                path_lib_default = nil;
            }
-            if (!resolved) {
-                if (error) {
-                    NSString * msg = [NSString stringWithFormat:@"could not resolve include \"%@\" from '%@'", include_name, path];
-                    *error = [NSError errorWithDomain:@"ggml-metal-source-flatten" code:1
-                                             userInfo:@{NSLocalizedDescriptionKey: msg}];
-                }
-                return false;
-            }
-            if (!ggml_metal_library_flatten_file(dst, resolved, search_paths, seen, error)) {
-                return false;
-            }
-            continue;
+
+            path_lib = path_lib_default;
        }

-        [dst appendString:line];
-        [dst appendString:@"\n"];
-    }
+        if (path_lib != nil) {
+            // pre-compiled library found
+            NSURL * libURL = [NSURL fileURLWithPath:path_lib];
+            GGML_LOG_INFO("%s: loading '%s'\n", __func__, [path_lib UTF8String]);

-    return true;
-}
+            library = [device newLibraryWithURL:libURL error:&error];
+            if (error) {
+                GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
+                return nil;
+            }
+        } else {
+            GGML_LOG_INFO("%s: default.metallib not found, loading from source\n", __func__);

-static NSString * ggml_metal_library_flatten_source(NSString * path_source, NSError ** error) {
-    // Search paths cover both runtime layout (build/bin/kernels + build/bin)
-    // and source-tree layout (ggml/src/ggml-metal/kernels + ggml/src/ggml-metal + ggml/src).
-    NSString * path_kernels = [path_source stringByDeletingLastPathComponent];
-    NSString * path_base    = [path_kernels stringByDeletingLastPathComponent];
-    NSArray<NSString *> * search_paths = @[
-        path_kernels,
-        path_base,
-        [path_base stringByDeletingLastPathComponent],
-    ];
+            NSString * path_source;
+            NSString * path_resource = [[NSProcessInfo processInfo].environment objectForKey:@"GGML_METAL_PATH_RESOURCES"];

-    NSMutableString * src = [[NSMutableString alloc] init];
-    NSMutableSet<NSString *> * seen = [NSMutableSet set];
+            GGML_LOG_INFO("%s: GGML_METAL_PATH_RESOURCES = %s\n", __func__, path_resource ? [path_resource UTF8String] : "nil");

-    if (!ggml_metal_library_flatten_file(src, path_source, search_paths, seen, error)) {
-        [src release];
-        return nil;
-    }
-    return src;
-}
-
-// Compile all per-kind libraries in parallel. `source_for_kind` returns the MSL
-// source for a kind (the helper takes ownership and releases it), or nil with
-// *err set on failure. On success the objs[] slots are populated and the routing
-// index is built; on any failure every error is logged and false is returned
-// (the caller is responsible for freeing `res`).
-static bool ggml_metal_library_compile_all(
-        ggml_metal_library_t res,
-        id<MTLDevice> device,
-        NSDictionary * prep,
-        NSString * (^source_for_kind)(int kind, NSError ** err),
-        const char * origin) {
-    const int64_t t_start = ggml_time_us();
-
-    int64_t  * t_per_lib   = calloc(GGML_METAL_LIB_COUNT, sizeof(int64_t));
-    NSError ** err_per_lib = calloc(GGML_METAL_LIB_COUNT, sizeof(NSError *));
-    __block atomic_bool any_failure = false;
-
-    dispatch_group_t group = dispatch_group_create();
-    dispatch_queue_t queue = dispatch_get_global_queue(QOS_CLASS_USER_INITIATED, 0);
-
-    for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
-        dispatch_group_async(group, queue, ^{
-
-            const int64_t t0 = ggml_time_us();
-
-            NSError * error = nil;
-
-            NSString * src = source_for_kind(kind, &error);
-            if (!src) {
-                err_per_lib[kind] = [error retain];
-                atomic_store(&any_failure, true);
-                return;
+            if (path_resource) {
+                path_source = [path_resource stringByAppendingPathComponent:@"ggml-metal.metal"];
+            } else {
+                path_source = [bundle pathForResource:@"ggml-metal" ofType:@"metal"];
            }

-            id<MTLLibrary> lib = nil;
+            if (path_source == nil) {
+                GGML_LOG_WARN("%s: error: could not use bundle path to find ggml-metal.metal, falling back to trying cwd\n", __func__);
+                path_source = @"ggml-metal.metal";
+            }

+            GGML_LOG_INFO("%s: loading '%s'\n", __func__, [path_source UTF8String]);
+
+            src = [NSString stringWithContentsOfFile:path_source encoding:NSUTF8StringEncoding error:&error];
+            if (error) {
+                GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
+                return nil;
+            }
+        }
+#endif
+
+        if (!library) {
            @autoreleasepool {
+                // dictionary of preprocessor macros
+                NSMutableDictionary * prep = [NSMutableDictionary dictionary];
+
+                if (ggml_metal_device_get_props(dev)->has_bfloat) {
+                    [prep setObject:@"1" forKey:@"GGML_METAL_HAS_BF16"];
+                }
+
+                if (ggml_metal_device_get_props(dev)->has_tensor) {
+                    [prep setObject:@"1" forKey:@"GGML_METAL_HAS_TENSOR"];
+                }
+
+#if GGML_METAL_EMBED_LIBRARY
+                [prep setObject:@"1" forKey:@"GGML_METAL_EMBED_LIBRARY"];
+#endif
+
                MTLCompileOptions * options = [MTLCompileOptions new];
                options.preprocessorMacros = prep;

-                lib = [device newLibraryWithSource:src options:options error:&error];
+                //[options setFastMathEnabled:false];

-                [options release];
-
-                // retain the error before the autorelease pool drains it
-                if (!lib) {
-                    err_per_lib[kind] = [error retain];
+                library = [device newLibraryWithSource:src options:options error:&error];
+                if (error) {
+                    GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
+                    return nil;
                }
-            }

-            [src release];
-
-            t_per_lib[kind] = ggml_time_us() - t0;
-
-            if (!lib) {
-                atomic_store(&any_failure, true);
-                return;
-            }
-
-            res->objs[kind] = lib;
-        });
-    }
-    dispatch_group_wait(group, DISPATCH_TIME_FOREVER);
-    dispatch_release(group);
-
-    const bool ok = !atomic_load(&any_failure);
-
-    if (ok) {
-        const int64_t t_total = ggml_time_us() - t_start;
-        int64_t t_max = 0;
-        for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
-            GGML_LOG_DEBUG("%s: compiled '%s' library in %.3f sec\n",
-                           __func__, k_lib_names[kind], t_per_lib[kind] / 1e6);
-            if (t_per_lib[kind] > t_max) t_max = t_per_lib[kind];
-        }
-        GGML_LOG_INFO("%s: loaded %d libraries from %s in %.3f sec (max single = %.3f sec)\n",
-                      __func__, GGML_METAL_LIB_COUNT, origin, t_total / 1e6, t_max / 1e6);
-
-        ggml_metal_library_build_index(res);
-    } else {
-        for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
-            if (err_per_lib[kind]) {
-                GGML_LOG_ERROR("%s: failed to build '%s' library: %s\n", __func__,
-                               k_lib_names[kind], [[err_per_lib[kind] description] UTF8String]);
-                [err_per_lib[kind] release];
+#if !__has_feature(objc_arc)
+                [options release];
+#endif
            }
        }
+
+#if GGML_METAL_EMBED_LIBRARY
+        [src release];
+#endif // GGML_METAL_EMBED_LIBRARY
+
+        GGML_LOG_INFO("%s: loaded in %.3f sec\n", __func__, (ggml_time_us() - t_start) / 1e6);
    }

-    free(err_per_lib);
-    free(t_per_lib);
-
-    return ok;
-}
-
-ggml_metal_library_t ggml_metal_library_init(ggml_metal_device_t dev) {
-    id<MTLDevice> device = ggml_metal_device_get_obj(dev);
-
    ggml_metal_library_t res = calloc(1, sizeof(struct ggml_metal_library));
+
+    res->obj       = library;
    res->dev       = dev;
    res->pipelines = ggml_metal_pipelines_init();
    res->lock      = [NSLock new];

-    // shared MTLCompileOptions preprocessor macros (matches the build-time defines)
-    NSMutableDictionary * prep = [NSMutableDictionary dictionary];
-    if (ggml_metal_device_get_props(dev)->has_bfloat) {
-        [prep setObject:@"1" forKey:@"GGML_METAL_HAS_BF16"];
-    }
-    if (ggml_metal_device_get_props(dev)->has_tensor) {
-        [prep setObject:@"1" forKey:@"GGML_METAL_HAS_TENSOR"];
-    }
-#if GGML_METAL_EMBED_LIBRARY
-    [prep setObject:@"1" forKey:@"GGML_METAL_EMBED_LIBRARY"];
-#endif
-
-#if GGML_METAL_EMBED_LIBRARY
-    GGML_LOG_INFO("%s: using embedded metal library\n", __func__);
-
-    // start/end symbols emitted by CMake (see CMakeLists.txt), one pair per kind
-#define X(e, s) extern const char ggml_metallib_##s##_start[]; extern const char ggml_metallib_##s##_end[];
-    GGML_METAL_LIBS
-#undef X
-
-    static const char * const lib_start[GGML_METAL_LIB_COUNT] = {
-#define X(e, s) [GGML_METAL_LIB_##e] = ggml_metallib_##s##_start,
-    GGML_METAL_LIBS
-#undef X
-    };
-    static const char * const lib_end[GGML_METAL_LIB_COUNT] = {
-#define X(e, s) [GGML_METAL_LIB_##e] = ggml_metallib_##s##_end,
-    GGML_METAL_LIBS
-#undef X
-    };
-
-    const bool ok = ggml_metal_library_compile_all(res, device, prep,
-        ^NSString * (int kind, NSError ** err) {
-            (void) err;
-            return [[NSString alloc] initWithBytes:lib_start[kind]
-                                            length:(lib_end[kind] - lib_start[kind])
-                                          encoding:NSUTF8StringEncoding];
-        }, "embedded data");
-
-    if (!ok) {
-        ggml_metal_library_free(res);
-        return NULL;
-    }
-
    return res;
-#else
-#ifdef SWIFT_PACKAGE
-    NSBundle * bundle = SWIFTPM_MODULE_BUNDLE;
-#else
-    NSBundle * bundle = [NSBundle bundleForClass:[GGMLMetalClass class]];
-#endif
-
-    const int64_t t_start = ggml_time_us();
-
-    NSError * error = nil;
-    NSString * path_lib = [bundle pathForResource:@"default" ofType:@"metallib"];
-    if (path_lib == nil) {
-        // Try to find the resource in the directory where the current binary located.
-        NSString * bin_cur = [[NSProcessInfo processInfo] arguments][0];
-        NSString * bin_dir = [bin_cur stringByDeletingLastPathComponent];
-
-        NSString * path_lib_default = [NSString pathWithComponents:@[bin_dir, @"default.metallib"]];
-        if ([[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
-            GGML_LOG_INFO("%s: found '%s'\n", __func__, [path_lib_default UTF8String]);
-
-            NSDictionary * atts = [[NSFileManager defaultManager] attributesOfItemAtPath:path_lib_default error:&error];
-            if (atts && atts[NSFileType] == NSFileTypeSymbolicLink) {
-                // Optionally, if this is a symlink, try to resolve it.
-                path_lib_default = [[NSFileManager defaultManager] destinationOfSymbolicLinkAtPath:path_lib_default error:&error];
-                if (path_lib_default && [path_lib_default length] > 0 && ![[path_lib_default substringToIndex:1] isEqualToString:@"/"]) {
-                    // It is a relative path, adding the binary directory as directory prefix.
-                    path_lib_default = [NSString pathWithComponents:@[bin_dir, path_lib_default]];
-                }
-                if (!path_lib_default || ![[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
-                    // Link to the resource could not be resolved.
-                    path_lib_default = nil;
-                } else {
-                    GGML_LOG_INFO("%s: symlink resolved '%s'\n", __func__, [path_lib_default UTF8String]);
-                }
-            }
-        } else {
-            // The resource couldn't be found in the binary's directory.
-            path_lib_default = nil;
-        }
-
-        path_lib = path_lib_default;
-    }
-
-    if (path_lib != nil) {
-        // pre-compiled library found: a single combined default.metallib
-        NSURL * libURL = [NSURL fileURLWithPath:path_lib];
-        GGML_LOG_INFO("%s: loading '%s'\n", __func__, [path_lib UTF8String]);
-
-        res->objs[0]        = [device newLibraryWithURL:libURL error:&error];
-        res->single_library = true;
-        if (!res->objs[0]) {
-            GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
-            ggml_metal_library_free(res);
-            return NULL;
-        }
-
-        GGML_LOG_INFO("%s: loaded in %.3f sec\n", __func__, (ggml_time_us() - t_start) / 1e6);
-        return res;
-    }
-
-    // no pre-compiled metallib: fall back to compiling each kernel source separately
-    GGML_LOG_INFO("%s: default.metallib not found, loading kernel sources\n", __func__);
-
-    NSString * path_resource = [[NSProcessInfo processInfo].environment objectForKey:@"GGML_METAL_PATH_RESOURCES"];
-    if (path_resource) {
-        GGML_LOG_INFO("%s: GGML_METAL_PATH_RESOURCES = %s\n", __func__, [path_resource UTF8String]);
-    }
-
-    // resolve each kind's source path up front (file lookup/logging stays on the calling thread)
-    NSString ** path_per_kind = calloc(GGML_METAL_LIB_COUNT, sizeof(NSString *));
-    for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
-        NSString * rel = [NSString stringWithFormat:@"kernels/%s.metal", k_lib_names[kind]];
-
-        NSString * path_source = nil;
-        if (path_resource) {
-            path_source = [path_resource stringByAppendingPathComponent:rel];
-        } else {
-            NSString * stem = [NSString stringWithFormat:@"kernels/%s", k_lib_names[kind]];
-            path_source = [bundle pathForResource:stem ofType:@"metal"];
-        }
-
-        if (path_source == nil || ![[NSFileManager defaultManager] isReadableFileAtPath:path_source]) {
-            GGML_LOG_WARN("%s: could not locate %s in bundle, falling back to cwd\n", __func__, [rel UTF8String]);
-            path_source = rel;
-        }
-
-        GGML_LOG_DEBUG("%s: loading '%s'\n", __func__, [path_source UTF8String]);
-
-        path_per_kind[kind] = [path_source retain];
-    }
-
-    const bool ok = ggml_metal_library_compile_all(res, device, prep,
-        ^NSString * (int kind, NSError ** err) {
-            return ggml_metal_library_flatten_source(path_per_kind[kind], err);
-        }, "source");
-
-    for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
-        [path_per_kind[kind] release];
-    }
-    free(path_per_kind);
-
-    if (!ok) {
-        ggml_metal_library_free(res);
-        return NULL;
-    }
-
-    return res;
-#endif
 }

 ggml_metal_library_t ggml_metal_library_init_from_source(ggml_metal_device_t dev, const char * source, bool verbose) {
@@ -590,11 +319,10 @@ ggml_metal_library_t ggml_metal_library_init_from_source(ggml_metal_device_t dev
        return NULL;
    }

-    res->objs[0]        = library;
-    res->single_library = true;
-    res->dev            = dev;
-    res->pipelines      = ggml_metal_pipelines_init();
-    res->lock           = [NSLock new];
+    res->obj       = library;
+    res->dev       = dev;
+    res->pipelines = ggml_metal_pipelines_init();
+    res->lock      = [NSLock new];

    return res;
 }
@@ -604,14 +332,8 @@ void ggml_metal_library_free(ggml_metal_library_t lib) {
        return;
    }

-    for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
-        if (lib->objs[kind]) {
-            [lib->objs[kind] release];
-        }
-    }
-
-    if (lib->fn_to_lib) {
-        [lib->fn_to_lib release];
+    if (lib->obj) {
+        [lib->obj release];
    }

    ggml_metal_pipelines_free(lib->pipelines);
@@ -672,28 +394,11 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_compile_pipeline(ggml_

        GGML_LOG_DEBUG("%s: compiling pipeline: base = '%s', name = '%s'\n", __func__, base, name);

-        // route to the library that actually defines this kernel; fn_to_lib is
-        // built from -[MTLLibrary functionNames] so it's always in sync
-        int lib_idx = 0;
-        if (!lib->single_library) {
-            NSNumber * idx = lib->fn_to_lib[base_func];
-            if (!idx) {
-                [lib->lock unlock];
-
-                GGML_LOG_ERROR("%s: kernel not found in any metal library: base = '%s', name = '%s'\n", __func__, base, name);
-
-                return res;
-            }
-            lib_idx = [idx intValue];
-        }
-
-        id<MTLLibrary> mtl_lib = lib->objs[lib_idx];
-
        id<MTLFunction> mtl_function;
        if (!cv) {
-            mtl_function = [mtl_lib newFunctionWithName:base_func];
+            mtl_function = [lib->obj newFunctionWithName:base_func];
        } else {
-            mtl_function = [mtl_lib newFunctionWithName:base_func constantValues:cv->obj error:&error];
+            mtl_function = [lib->obj newFunctionWithName:base_func constantValues:cv->obj error:&error];
        }
        if (!mtl_function) {
            [lib->lock unlock];
@@ -1,232 +0,0 @@
-#include "common.h"
-
-// bitonic sort implementation following the CUDA kernels as reference
-typedef void (argsort_t)(
-        constant   ggml_metal_kargs_argsort & args,
-        device   const char * src0,
-        device      int32_t * dst,
-        threadgroup int32_t * shmem_i32 [[threadgroup(0)]],
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]);
-
-template<ggml_sort_order order>
-kernel void kernel_argsort_f32_i32(
-        constant   ggml_metal_kargs_argsort & args,
-        device   const char * src0,
-        device      int32_t * dst,
-        threadgroup int32_t * shmem_i32 [[threadgroup(0)]],
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    // bitonic sort
-    const int col = tpitg[0];
-    const int ib  = tgpig[0] / args.ne01;
-
-    const int i00 = ib*ntg.x;
-    const int i01 = tgpig[0] % args.ne01;
-    const int i02 = tgpig[1];
-    const int i03 = tgpig[2];
-
-    device const float * src0_row = (device const float *) (src0 + args.nb01*i01 + args.nb02*i02 + args.nb03*i03);
-
-    // initialize indices
-    shmem_i32[col] = i00 + col;
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    for (int k = 2; k <= ntg.x; k *= 2) {
-        for (int j = k / 2; j > 0; j /= 2) {
-            int ixj = col ^ j;
-            if (ixj > col) {
-                if ((col & k) == 0) {
-                    if (shmem_i32[col] >= args.ne00 ||
-                       (shmem_i32[ixj] <  args.ne00 && (order == GGML_SORT_ORDER_ASC ?
-                            src0_row[shmem_i32[col]] > src0_row[shmem_i32[ixj]] :
-                            src0_row[shmem_i32[col]] < src0_row[shmem_i32[ixj]]))
-                    ) {
-                        SWAP(shmem_i32[col], shmem_i32[ixj]);
-                    }
-                } else {
-                    if (shmem_i32[ixj] >= args.ne00 ||
-                       (shmem_i32[col] <  args.ne00 && (order == GGML_SORT_ORDER_ASC ?
-                            src0_row[shmem_i32[col]] < src0_row[shmem_i32[ixj]] :
-                            src0_row[shmem_i32[col]] > src0_row[shmem_i32[ixj]]))
-                    ) {
-                        SWAP(shmem_i32[col], shmem_i32[ixj]);
-                    }
-                }
-            }
-
-            threadgroup_barrier(mem_flags::mem_threadgroup);
-        }
-    }
-
-    const int64_t i0 = ib*args.top_k;
-
-    // copy the result to dst without the padding
-    if (i0 + col < args.ne0 && col < args.top_k) {
-        dst += i0 + args.ne0*i01 + args.ne0*args.ne1*i02 + args.ne0*args.ne1*args.ne2*i03;
-
-        dst[col] = shmem_i32[col];
-    }
-}
-
-template [[host_name("kernel_argsort_f32_i32_asc")]]  kernel argsort_t kernel_argsort_f32_i32<GGML_SORT_ORDER_ASC>;
-template [[host_name("kernel_argsort_f32_i32_desc")]] kernel argsort_t kernel_argsort_f32_i32<GGML_SORT_ORDER_DESC>;
-
-typedef void (argsort_merge_t)(
-        constant   ggml_metal_kargs_argsort_merge & args,
-        device const char    * src0,
-        device const int32_t * tmp,
-        device       int32_t * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]);
-
-template<ggml_sort_order order>
-kernel void kernel_argsort_merge_f32_i32(
-        constant   ggml_metal_kargs_argsort_merge & args,
-        device const char    * src0,
-        device const int32_t * tmp,
-        device       int32_t * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-
-    const int im  = tgpig[0] / args.ne01;
-    const int i01 = tgpig[0] % args.ne01;
-    const int i02 = tgpig[1];
-    const int i03 = tgpig[2];
-
-    const int start = im * (2 * args.len);
-
-    const int len0 = MIN(args.len, MAX(0, args.ne0 - (int)(start)));
-    const int len1 = MIN(args.len, MAX(0, args.ne0 - (int)(start + args.len)));
-
-    const int total = len0 + len1;
-
-    device const int32_t * tmp0 = tmp + start
-        + i01*args.ne0
-        + i02*args.ne0*args.ne01
-        + i03*args.ne0*args.ne01*args.ne02;
-
-    device const int32_t * tmp1 = tmp0 + args.len;
-
-    dst += start
-        + i01*args.top_k
-        + i02*args.top_k*args.ne01
-        + i03*args.top_k*args.ne01*args.ne02;
-
-    device const float * src0_row = (device const float *)(src0
-        + args.nb01*i01
-        + args.nb02*i02
-        + args.nb03*i03);
-
-    if (total == 0) {
-        return;
-    }
-
-    const int chunk = (total + ntg.x - 1) / ntg.x;
-
-    const int k0 = tpitg.x * chunk;
-    const int k1 = MIN(MIN(k0 + chunk, total), args.top_k);
-
-    if (k0 >= args.top_k) {
-        return;
-    }
-
-    if (k0 >= total) {
-        return;
-    }
-
-    int low  = k0 > len1 ? k0 - len1 : 0;
-    int high = MIN(k0, len0);
-
-    // binary-search partition (i, j) such that i + j = k
-    while (low < high) {
-        const int mid = (low + high) >> 1;
-
-        const int32_t idx0 = tmp0[mid];
-        const int32_t idx1 = tmp1[k0 - mid - 1];
-
-        const float val0 = src0_row[idx0];
-        const float val1 = src0_row[idx1];
-
-        bool take_left;
-        if (order == GGML_SORT_ORDER_ASC) {
-            take_left = (val0 <= val1);
-        } else {
-            take_left = (val0 >= val1);
-        }
-
-        if (take_left) {
-            low = mid + 1;
-        } else {
-            high = mid;
-        }
-    }
-
-    int i = low;
-    int j = k0 - i;
-
-    // keep the merge fronts into registers
-    int32_t idx0 = 0;
-    float   val0 = 0.0f;
-    if (i < len0) {
-        idx0 = tmp0[i];
-        val0 = src0_row[idx0];
-    }
-
-    int32_t idx1 = 0;
-    float   val1 = 0.0f;
-    if (j < len1) {
-        idx1 = tmp1[j];
-        val1 = src0_row[idx1];
-    }
-
-    for (int k = k0; k < k1; ++k) {
-        int32_t out_idx;
-
-        if (i >= len0) {
-            while (k < k1) {
-                dst[k++] = tmp1[j++];
-            }
-            break;
-        } else if (j >= len1) {
-            while (k < k1) {
-                dst[k++] = tmp0[i++];
-            }
-            break;
-        } else {
-            bool take_left;
-
-            if (order == GGML_SORT_ORDER_ASC) {
-                take_left = (val0 <= val1);
-            } else {
-                take_left = (val0 >= val1);
-            }
-
-            if (take_left) {
-                out_idx = idx0;
-                ++i;
-                if (i < len0) {
-                    idx0 = tmp0[i];
-                    val0 = src0_row[idx0];
-                }
-            } else {
-                out_idx = idx1;
-                ++j;
-                if (j < len1) {
-                    idx1 = tmp1[j];
-                    val1 = src0_row[idx1];
-                }
-            }
-        }
-
-        dst[k] = out_idx;
-    }
-}
-
-template [[host_name("kernel_argsort_merge_f32_i32_asc")]]  kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_ASC>;
-template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_DESC>;
@@ -1,228 +0,0 @@
-#include "common.h"
-
-// OP: 0 - add, 1 - sub, 2 - mul, 3 - div
-constant short FC_bin_op [[function_constant(FC_BIN + 0)]];
-constant short FC_bin_f  [[function_constant(FC_BIN + 1)]];
-constant bool  FC_bin_rb [[function_constant(FC_BIN + 2)]];
-constant bool  FC_bin_cb [[function_constant(FC_BIN + 3)]];
-
-template <typename T0, typename T1, typename T>
-kernel void kernel_bin_fuse_impl(
-        constant ggml_metal_kargs_bin & args,
-        device const char * src0,
-        device const char * src1,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-#define FC_OP FC_bin_op
-#define FC_F  FC_bin_f
-#define FC_RB FC_bin_rb
-#define FC_CB FC_bin_cb
-
-    if (FC_RB) {
-        // row broadcast
-        const uint i0 = tgpig.y*args.ne00 + tgpig.x;
-        const uint i1 = FC_CB ? tgpig.x%args.ne10 : tgpig.x;
-
-        device const T0 * src0_row = (device const T0 *) (src0);
-        device       T  * dst_row  = (device       T  *) (dst);
-
-        if (FC_F == 1) {
-            device const T1 * src1_row = (device const T1 *) (src1 + args.o1[0]);
-
-            if (FC_OP == 0) {
-                dst_row[i0] = src0_row[i0] + src1_row[i1];
-            }
-
-            if (FC_OP == 1) {
-                dst_row[i0] = src0_row[i0] - src1_row[i1];
-            }
-
-            if (FC_OP == 2) {
-                dst_row[i0] = src0_row[i0] * src1_row[i1];
-            }
-
-            if (FC_OP == 3) {
-                dst_row[i0] = src0_row[i0] / src1_row[i1];
-            }
-        } else {
-            T0 res = src0_row[i0];
-
-            if (FC_OP == 0) {
-                FOR_UNROLL (short j = 0; j < FC_F; ++j) {
-                    res += ((device const T1 *) (src1 + args.o1[j]))[i1];
-                }
-            }
-
-            if (FC_OP == 1) {
-                FOR_UNROLL (short j = 0; j < FC_F; ++j) {
-                    res -= ((device const T1 *) (src1 + args.o1[j]))[i1];
-                }
-            }
-
-            if (FC_OP == 2) {
-                FOR_UNROLL (short j = 0; j < FC_F; ++j) {
-                    res *= ((device const T1 *) (src1 + args.o1[j]))[i1];
-                }
-            }
-
-            if (FC_OP == 3) {
-                FOR_UNROLL (short j = 0; j < FC_F; ++j) {
-                    res /= ((device const T1 *) (src1 + args.o1[j]))[i1];
-                }
-            }
-
-            dst_row[i0] = res;
-        }
-    } else {
-        const int i03 = tgpig.z;
-        const int i02 = tgpig.y;
-        const int i01 = tgpig.x;
-
-        if (i01 >= args.ne01) {
-            return;
-        }
-
-        const int i13 = i03%args.ne13;
-        const int i12 = i02%args.ne12;
-        const int i11 = i01%args.ne11;
-
-        device const T0 * src0_ptr = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + args.offs);
-        device       T  * dst_ptr  = (device       T  *) (dst  + i03*args.nb3  + i02*args.nb2  + i01*args.nb1  + args.offs);
-
-        if (FC_F == 1) {
-            device const T1 * src1_ptr = (device const T1 *) (src1 + args.o1[0] + i13*args.nb13 + i12*args.nb12 + i11*args.nb11);
-
-            for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
-                const int i10 = FC_CB ? i0%args.ne10 : i0;
-
-                if (FC_OP == 0) {
-                    dst_ptr[i0] = src0_ptr[i0] + src1_ptr[i10];
-                }
-
-                if (FC_OP == 1) {
-                    dst_ptr[i0] = src0_ptr[i0] - src1_ptr[i10];
-                }
-
-                if (FC_OP == 2) {
-                    dst_ptr[i0] = src0_ptr[i0] * src1_ptr[i10];
-                }
-
-                if (FC_OP == 3) {
-                    dst_ptr[i0] = src0_ptr[i0] / src1_ptr[i10];
-                }
-            }
-        } else {
-            device const T1 * src1_ptr[8];
-            FOR_UNROLL (short j = 0; j < FC_F; ++j) {
-                src1_ptr[j] = (device const T1 *) (src1 + args.o1[j] + i13*args.nb13 + i12*args.nb12 + i11*args.nb11);
-            }
-
-            for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
-                const int i10 = FC_CB ? i0%args.ne10 : i0;
-
-                T res = src0_ptr[i0];
-
-                if (FC_OP == 0) {
-                    FOR_UNROLL (short j = 0; j < FC_F; ++j) {
-                        res += src1_ptr[j][i10];
-                    }
-                }
-
-                if (FC_OP == 1) {
-                    FOR_UNROLL (short j = 0; j < FC_F; ++j) {
-                        res -= src1_ptr[j][i10];
-                    }
-                }
-
-                if (FC_OP == 2) {
-                    FOR_UNROLL (short j = 0; j < FC_F; ++j) {
-                        res *= src1_ptr[j][i10];
-                    }
-                }
-
-                if (FC_OP == 3) {
-                    FOR_UNROLL (short j = 0; j < FC_F; ++j) {
-                        res /= src1_ptr[j][i10];
-                    }
-                }
-
-                dst_ptr[i0] = res;
-            }
-        }
-    }
-
-#undef FC_OP
-#undef FC_F
-#undef FC_RB
-#undef FC_CB
-}
-
-typedef decltype(kernel_bin_fuse_impl<float, float, float>) kernel_bin_fuse_t;
-
-template [[host_name("kernel_bin_fuse_f32_f32_f32")]]   kernel kernel_bin_fuse_t kernel_bin_fuse_impl<float,  float,  float>;
-template [[host_name("kernel_bin_fuse_f32_f32_f32_4")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl<float4, float4, float4>;
-template [[host_name("kernel_bin_fuse_f16_f16_f16")]]   kernel kernel_bin_fuse_t kernel_bin_fuse_impl<half,   half,   half>;
-template [[host_name("kernel_bin_fuse_f16_f16_f16_4")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl<half4,  half4,  half4>;
-
-kernel void kernel_add_id(
-        constant ggml_metal_kargs_add_id & args,
-        device const char * src0,
-        device const char * src1,
-        device const char * src2,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    const int i1 = tgpig.x;
-    const int i2 = tgpig.y;
-
-    const int i11 = *((device const int32_t *) (src2 + i1*sizeof(int32_t) + i2*args.nb21));
-
-    const size_t nb1 = args.ne0 * sizeof(float);
-    const size_t nb2 = args.ne1 * nb1;
-
-    device       float * dst_row  = (device       float *)((device char *)dst  +  i1*nb1       + i2*nb2);
-    device const float * src0_row = (device const float *)((device char *)src0 +  i1*args.nb01 + i2*args.nb02);
-    device const float * src1_row = (device const float *)((device char *)src1 + i11*args.nb11);
-
-    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
-        dst_row[i0] = src0_row[i0] + src1_row[i0];
-    }
-}
-
-template<typename T>
-kernel void kernel_repeat(
-        constant ggml_metal_kargs_repeat & args,
-        device const char * src0,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    const int i3 = tgpig.z;
-    const int i2 = tgpig.y;
-    const int i1 = tgpig.x;
-
-    const int i03 = i3%args.ne03;
-    const int i02 = i2%args.ne02;
-    const int i01 = i1%args.ne01;
-
-    device const char * src0_ptr = src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01;
-    device       char * dst_ptr  = dst  +  i3*args.nb3  +  i2*args.nb2  +  i1*args.nb1;
-
-    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
-        const int i00 = i0%args.ne00;
-        *((device T *)(dst_ptr + i0*args.nb0)) = *((device T *)(src0_ptr + i00*args.nb00));
-    }
-}
-
-typedef decltype(kernel_repeat<float>) kernel_repeat_t;
-
-template [[host_name("kernel_repeat_f32")]] kernel kernel_repeat_t kernel_repeat<float>;
-template [[host_name("kernel_repeat_f16")]] kernel kernel_repeat_t kernel_repeat<half>;
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_repeat_bf16")]] kernel kernel_repeat_t kernel_repeat<bfloat>;
-#endif
-template [[host_name("kernel_repeat_i32")]] kernel kernel_repeat_t kernel_repeat<int>;
-template [[host_name("kernel_repeat_i16")]] kernel kernel_repeat_t kernel_repeat<short>;
@@ -1,126 +0,0 @@
-#pragma once
-
-#include "ggml-metal-impl.h"
-
-#include <metal_stdlib>
-
-#ifdef GGML_METAL_HAS_TENSOR
-#include <metal_tensor>
-
-#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
-#endif
-
-using namespace metal;
-
-#define MAX(x, y) ((x) > (y) ? (x) : (y))
-#define MIN(x, y) ((x) < (y) ? (x) : (y))
-#define SWAP(x, y) { auto tmp = (x); (x) = (y); (y) = tmp; }
-
-#define PAD2(x, n) (((x) + (n) - 1) & ~((n) - 1))
-
-#define FOR_UNROLL(x) _Pragma("clang loop unroll(full)") for (x)
-
-#define N_SIMDWIDTH 32 // assuming SIMD group size is 32
-
-// ref: https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
-//
-// cmd:
-//   .../usr/bin/metal -dM -E -c                             ggml/src/ggml-metal/kernels/<src>.metal
-//   .../usr/bin/metal -dM -E -c -target air64-apple-ios14.0 ggml/src/ggml-metal/kernels/<src>.metal
-//
-#if __METAL_VERSION__ < 310 && defined(GGML_METAL_HAS_BF16)
-#undef GGML_METAL_HAS_BF16
-#endif
-
-#if defined(GGML_METAL_HAS_BF16)
-typedef matrix<bfloat, 4, 4> bfloat4x4;
-typedef matrix<bfloat, 2, 4> bfloat2x4;
-#endif
-
-constexpr constant static float kvalues_iq4nl_f[16] = {
-    -127.f, -104.f, -83.f, -65.f, -49.f, -35.f, -22.f, -10.f, 1.f, 13.f, 25.f, 38.f, 53.f, 69.f, 89.f, 113.f
-};
-
-constexpr constant static float kvalues_mxfp4_f[16] = {
-    0, .5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f, -0, -.5f, -1.f, -1.5f, -2.f, -3.f, -4.f, -6.f
-};
-
-static inline int best_index_int8(int n, constant float * val, float x) {
-    if (x <= val[0]) return 0;
-    if (x >= val[n-1]) return n-1;
-    int ml = 0, mu = n-1;
-    while (mu-ml > 1) {
-        int mav = (ml+mu)/2;
-        if (x < val[mav]) mu = mav; else ml = mav;
-    }
-    return x - val[mu-1] < val[mu] - x ? mu-1 : mu;
-}
-
-static inline float e8m0_to_fp32(uint8_t x) {
-    uint32_t bits;
-
-    if (x == 0) {
-        bits = 0x00400000;
-    } else {
-        bits = (uint32_t) x << 23;
-    }
-
-    return as_type<float>(bits);
-}
-
-static inline float dot(float x, float y) {
-    return x*y;
-}
-
-static inline float sum(float x) {
-    return x;
-}
-
-static inline float sum(float4 x) {
-    return x[0] + x[1] + x[2] + x[3];
-}
-
-enum ggml_sort_order {
-    GGML_SORT_ORDER_ASC,
-    GGML_SORT_ORDER_DESC,
-};
-
-constant float GELU_COEF_A     = 0.044715f;
-constant float GELU_QUICK_COEF = -1.702f;
-constant float SQRT_2_OVER_PI  = 0.79788456080286535587989211986876f;
-constant float SQRT_2_INV      = 0.70710678118654752440084436210484f;
-
-// based on Abramowitz and Stegun formula 7.1.26 or similar Hastings' approximation
-// ref: https://www.johndcook.com/blog/python_erf/
-constant float p_erf  = 0.3275911f;
-constant float a1_erf = 0.254829592f;
-constant float a2_erf = -0.284496736f;
-constant float a3_erf = 1.421413741f;
-constant float a4_erf = -1.453152027f;
-constant float a5_erf = 1.061405429f;
-
-template<typename T>
-inline T erf_approx(T x) {
-    T sign_x = sign(x);
-    x = fabs(x);
-    T t = 1.0f / (1.0f + p_erf * x);
-    T y = 1.0f - (((((a5_erf * t + a4_erf) * t) + a3_erf) * t + a2_erf) * t + a1_erf) * t * exp(-x * x);
-    return sign_x * y;
-}
-
-template<typename T> T elu_approx(T x);
-
-template<> inline float elu_approx<float>(float x) {
-    return (x > 0.f) ? x : (exp(x) - 1);
-}
-
-template<> inline float4 elu_approx<float4>(float4 x) {
-    float4 res;
-
-    res[0] = (x[0] > 0.0f) ? x[0] : (exp(x[0]) - 1.0f);
-    res[1] = (x[1] > 0.0f) ? x[1] : (exp(x[1]) - 1.0f);
-    res[2] = (x[2] > 0.0f) ? x[2] : (exp(x[2]) - 1.0f);
-    res[3] = (x[3] > 0.0f) ? x[3] : (exp(x[3]) - 1.0f);
-
-    return res;
-}
@@ -1,723 +0,0 @@
-#include "common.h"
-
-typedef void (im2col_t)(
-        constant ggml_metal_kargs_im2col & args,
-        device const float * x,
-        device        char * dst,
-        uint3 tgpig[[threadgroup_position_in_grid]],
-        uint3  tgpg[[threadgroups_per_grid]],
-        uint3 tpitg[[thread_position_in_threadgroup]],
-        uint3   ntg[[threads_per_threadgroup]]);
-
-template <typename T>
-kernel void kernel_im2col(
-        constant ggml_metal_kargs_im2col & args,
-        device const float * x,
-        device        char * dst,
-        uint3 tgpig[[threadgroup_position_in_grid]],
-        uint3  tgpg[[threadgroups_per_grid]],
-        uint3 tpitg[[thread_position_in_threadgroup]],
-        uint3   ntg[[threads_per_threadgroup]]) {
-//    const int64_t IC = tgpg[0];
-    const int64_t OH = tgpg[1];
-    const int64_t OW = tgpg[2];
-
-    const int64_t KH = ntg[1];
-    const int64_t KW = ntg[2];
-
-          int64_t in  = tpitg[0];
-    const int64_t ikh = tpitg[1];
-    const int64_t ikw = tpitg[2];
-
-    const int64_t iic = tgpig[0];
-    const int64_t ioh = tgpig[1];
-    const int64_t iow = tgpig[2];
-
-    const int64_t iiw = iow*args.s0 + ikw*args.d0 - args.p0;
-    const int64_t iih = ioh*args.s1 + ikh*args.d1 - args.p1;
-
-    int64_t offset_dst = (in*OH*OW + ioh*OW + iow)*args.CHW + (iic*(KH*KW) + ikh*KW + ikw);
-
-    device T * pdst = (device T *) (dst);
-
-    if (iih < 0 || iih >= args.IH || iiw < 0 || iiw >= args.IW) {
-        while (in < args.N) {
-            pdst[offset_dst] = 0.0f;
-            offset_dst += ntg[0]*args.CHW*OH*OW;
-
-            in += ntg[0];
-        }
-    } else {
-        int64_t offset_src = in*args.ofs0 + iic*args.ofs1 + iih*args.IW + iiw;
-
-        while (in < args.N) {
-            pdst[offset_dst] = x[offset_src];
-
-            offset_dst += ntg[0]*args.CHW*OH*OW;
-            offset_src += ntg[0]*args.ofs0;
-
-            in += ntg[0];
-        }
-    }
-}
-
-template [[host_name("kernel_im2col_f32")]] kernel im2col_t kernel_im2col<float>;
-template [[host_name("kernel_im2col_f16")]] kernel im2col_t kernel_im2col<half>;
-
-// TODO: optimize
-typedef void (im2col_ext_t)(
-        constant ggml_metal_kargs_im2col & args,
-        device const float * x,
-        device        char * dst,
-        uint3 tgpig[[threadgroup_position_in_grid]],
-        uint3  tgpg[[threadgroups_per_grid]],
-        uint3 tpitg[[thread_position_in_threadgroup]],
-        uint3   ntg[[threads_per_threadgroup]]);
-
-template <typename T>
-kernel void kernel_im2col_ext(
-        constant ggml_metal_kargs_im2col & args,
-        device const float * x,
-        device        char * dst,
-        uint3 tgpig[[threadgroup_position_in_grid]],
-        uint3  tgpg[[threadgroups_per_grid]],      // tgpg[0] = D x IC x KH x KW, CHW = IC x KH x KW
-        uint3 tpitg[[thread_position_in_threadgroup]],
-        uint3   ntg[[threads_per_threadgroup]]) {  // [M, 1, 1]
-    const int64_t KHW = (int64_t)args.KHW;
-
-    const int64_t d   = tgpig[0] / args.CHW;
-    const int64_t chw = tgpig[0] % args.CHW;
-    const int64_t tgpig_0 = chw / KHW;  // 0 ~ (IC - 1)
-    const int64_t HW = tgpig[0] % KHW;
-
-    const int64_t tpitg_0 = (d * ntg[0]) + tpitg[0];
-    if (tpitg_0 >= args.N) {
-        return;
-    }
-
-    const int64_t tpitg_1 = HW / args.KW;
-    const int64_t tpitg_2 = HW % args.KW;
-
-    const int64_t iiw = tgpig[2] * args.s0 + tpitg_2 * args.d0 - args.p0;
-    const int64_t iih = tgpig[1] * args.s1 + tpitg_1 * args.d1 - args.p1;
-
-    const int64_t offset_dst =
-        (tpitg_0 * tgpg[1] * tgpg[2] + tgpig[1] * tgpg[2] + tgpig[2]) * args.CHW +
-        (tgpig_0 * KHW + tpitg_1 * args.KW + tpitg_2);
-
-    device T * pdst = (device T *) (dst);
-
-    if (iih < 0 || iih >= args.IH || iiw < 0 || iiw >= args.IW) {
-        pdst[offset_dst] = 0.0f;
-    } else {
-        const int64_t offset_src = tpitg_0 * args.ofs0 + tgpig_0 * args.ofs1;
-        pdst[offset_dst] = x[offset_src + iih * args.IW + iiw];
-    }
-}
-
-template [[host_name("kernel_im2col_ext_f32")]] kernel im2col_ext_t kernel_im2col_ext<float>;
-template [[host_name("kernel_im2col_ext_f16")]] kernel im2col_ext_t kernel_im2col_ext<half>;
-
-template <typename T>
-kernel void kernel_col2im_1d(
-        constant ggml_metal_kargs_col2im_1d & args,
-        device const T * col,
-        device       T * dst,
-        uint         tgpig [[threadgroup_position_in_grid]],
-        uint         tpitg [[thread_position_in_threadgroup]],
-        uint         ntg   [[threads_per_threadgroup]]) {
-
-    const int idx = tgpig * ntg + tpitg;
-    if (idx >= args.T_out * args.OC) {
-        return;
-    }
-
-    const int t_out = idx % args.T_out;
-    const int oc    = idx / args.T_out;
-    const int t_abs = t_out + args.p0;  // absolute position in uncropped signal
-
-    int t_in_min = (t_abs - args.K + args.s0) / args.s0;  // ceil((t_abs - K + 1) / s0)
-    if (t_in_min < 0) {
-        t_in_min = 0;
-    }
-    int t_in_max = t_abs / args.s0;
-    if (t_in_max >= args.T_in) {
-        t_in_max = args.T_in - 1;
-    }
-
-    float sum = 0.0f;
-    for (int t_in = t_in_min; t_in <= t_in_max; t_in++) {
-        const int k = t_abs - t_in * args.s0;
-        sum += float(col[(oc * args.K + k) + t_in * args.K_OC]);
-    }
-
-    dst[t_out + oc * args.T_out] = T(sum);
-}
-
-template [[host_name("kernel_col2im_1d_f32")]]  kernel void kernel_col2im_1d<float>(constant ggml_metal_kargs_col2im_1d &, device const float *, device float *, uint, uint, uint);
-template [[host_name("kernel_col2im_1d_f16")]]  kernel void kernel_col2im_1d<half>(constant ggml_metal_kargs_col2im_1d &, device const half *, device half *, uint, uint, uint);
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_col2im_1d_bf16")]] kernel void kernel_col2im_1d<bfloat>(constant ggml_metal_kargs_col2im_1d &, device const bfloat *, device bfloat *, uint, uint, uint);
-#endif
-
-template <typename TK>
-kernel void kernel_conv_2d(
-        constant ggml_metal_kargs_conv_2d & args,
-        device const char * weights,
-        device const char * src,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        uint3    tgpg[[threadgroups_per_grid]],
-        uint3   tpitg[[thread_position_in_threadgroup]],
-        uint3     ntg[[threads_per_threadgroup]]) {
-
-    const uint threads_per_tg = ntg.x * ntg.y * ntg.z;
-    const uint tg_index = (tgpig.z * tgpg.y + tgpig.y) * tgpg.x + tgpig.x;
-    const uint local_thread = tpitg.z * (ntg.x * ntg.y) + tpitg.y * ntg.x + tpitg.x;
-    const uint thread_index = tg_index * threads_per_tg + local_thread;
-    const uint64_t total_threads = (uint64_t) threads_per_tg * tgpg.x * tgpg.y * tgpg.z;
-    const uint64_t total_outputs = (uint64_t) args.N * args.OC * args.OH * args.OW;
-
-    for (uint64_t index = thread_index; index < total_outputs; index += total_threads) {
-        uint64_t tmp = index;
-
-        const int32_t ow = tmp % args.OW; tmp /= args.OW;
-        const int32_t oh = tmp % args.OH; tmp /= args.OH;
-        const int32_t oc = tmp % args.OC; tmp /= args.OC;
-        const int32_t  n = tmp;
-
-        float acc = 0.0f;
-
-        const int32_t base_x = ow*args.s0 - args.p0;
-        const int32_t base_y = oh*args.s1 - args.p1;
-
-        int32_t ky_start = 0;
-        if (base_y < 0) {
-            ky_start = (-base_y + args.d1 - 1)/args.d1;
-        }
-        int32_t ky_end = args.KH;
-        const int32_t y_max = args.IH - 1 - base_y;
-        if (y_max < 0) {
-            ky_end = ky_start;
-        } else if (base_y + (args.KH - 1)*args.d1 >= args.IH) {
-            ky_end = min(ky_end, y_max/args.d1 + 1);
-        }
-
-        int32_t kx_start = 0;
-        if (base_x < 0) {
-            kx_start = (-base_x + args.d0 - 1)/args.d0;
-        }
-        int32_t kx_end = args.KW;
-        const int32_t x_max = args.IW - 1 - base_x;
-        if (x_max < 0) {
-            kx_end = kx_start;
-        } else if (base_x + (args.KW - 1)*args.d0 >= args.IW) {
-            kx_end = min(kx_end, x_max/args.d0 + 1);
-        }
-
-        if (ky_start < ky_end && kx_start < kx_end) {
-            const uint64_t src_base_n = (uint64_t) n  * args.nb13;
-            const uint64_t w_base_oc  = (uint64_t) oc * args.nb03;
-
-            for (int32_t ic = 0; ic < args.IC; ++ic) {
-                const uint64_t src_base_nc = src_base_n + (uint64_t) ic * args.nb12;
-                const uint64_t w_base_ocic = w_base_oc  + (uint64_t) ic * args.nb02;
-
-                for (int32_t ky = ky_start; ky < ky_end; ++ky) {
-                    const int32_t iy = base_y + ky*args.d1;
-                    const uint64_t src_base_row = src_base_nc + (uint64_t) iy * args.nb11;
-                    const uint64_t w_base_row   = w_base_ocic + (uint64_t) ky * args.nb01;
-
-                    for (int32_t kx = kx_start; kx < kx_end; ++kx) {
-                        const int32_t ix = base_x + kx*args.d0;
-                        const uint64_t src_offs = src_base_row + (uint64_t) ix * args.nb10;
-                        const uint64_t w_offs   = w_base_row   + (uint64_t) kx * args.nb00;
-
-                        const float x = *(device const float *)(src + src_offs);
-                        const float w = (float) (*(device const TK *)(weights + w_offs));
-
-                        acc += x * w;
-                    }
-                }
-            }
-        }
-
-        const uint64_t dst_offs =
-            (uint64_t) n  * args.nb3 +
-            (uint64_t) oc * args.nb2 +
-            (uint64_t) oh * args.nb1 +
-            (uint64_t) ow * args.nb0;
-
-        *(device float *)(dst + dst_offs) = acc;
-    }
-}
-
-template [[host_name("kernel_conv_2d_f32_f32")]]
-kernel void kernel_conv_2d<float>(
-        constant ggml_metal_kargs_conv_2d & args,
-        device const char * weights,
-        device const char * src,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        uint3    tgpg[[threadgroups_per_grid]],
-        uint3   tpitg[[thread_position_in_threadgroup]],
-        uint3     ntg[[threads_per_threadgroup]]);
-
-template [[host_name("kernel_conv_2d_f16_f32")]]
-kernel void kernel_conv_2d<half>(
-        constant ggml_metal_kargs_conv_2d & args,
-        device const char * weights,
-        device const char * src,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        uint3    tgpg[[threadgroups_per_grid]],
-        uint3   tpitg[[thread_position_in_threadgroup]],
-        uint3     ntg[[threads_per_threadgroup]]);
-
-typedef void (conv_transpose_1d_t)(
-        constant ggml_metal_kargs_conv_transpose_1d & args,
-        device const float * src0,
-        device const float * src1,
-        device        char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        uint3    tgpg[[threadgroups_per_grid]]);
-
-template <typename T>
-kernel void kernel_conv_transpose_1d(
-        constant ggml_metal_kargs_conv_transpose_1d & args,
-        device const     T * src0,
-        device const float * src1,
-        device        char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        uint3   tgpg[[threadgroups_per_grid]]) {
-
-    // For output position j on the time axis, only input positions
-    //   i such that i*s0 <= j < i*s0 + K
-    // contribute -- i.e. i in [ceil((j - K + 1)/s0), floor(j/s0)]
-    // intersected with [0, IL-1]. That's at most ceil(K/s0) values
-    // (typically 2 for stride==K/2 transposed convs).
-    const int32_t j  = tgpig[0];
-    const int32_t s0 = args.s0;
-    const int32_t K  = args.K;
-    const int32_t IL = args.IL;
-
-    int32_t i_min;
-    {
-        int32_t a = j - K + 1;
-        i_min = a <= 0 ? 0 : (a + s0 - 1) / s0; // ceil(a/s0) for a>0
-    }
-    int32_t i_max = j / s0;
-    if (i_max > IL - 1) i_max = IL - 1;
-
-    float v = 0.0f;
-    if (i_min <= i_max) {
-        for (int64_t c = 0; c < args.IC; c++) {
-            const int32_t kernel_offset = c * tgpg[1] * K + K * tgpig[1];
-            const int32_t input_offset  = c * IL;
-
-            for (int32_t i = i_min; i <= i_max; i++) {
-                v += float(src0[kernel_offset + j - i * s0]) * src1[input_offset + i];
-            }
-        }
-    }
-
-    device float * dst_ptr = (device float *) (dst + tgpig[0] * args.nb0 + tgpig[1] * args.nb1);
-
-    dst_ptr[0] = v;
-}
-
-template [[host_name("kernel_conv_transpose_1d_f32_f32")]]
-kernel void kernel_conv_transpose_1d<float>(
-    constant ggml_metal_kargs_conv_transpose_1d & args,
-    device const float * src0,
-    device const float * src1,
-    device        char * dst,
-    uint3   tgpig[[threadgroup_position_in_grid]],
-    uint3    tgpg[[threadgroups_per_grid]]);
-
-template [[host_name("kernel_conv_transpose_1d_f16_f32")]]
-kernel void kernel_conv_transpose_1d<half>(
-    constant ggml_metal_kargs_conv_transpose_1d & args,
-    device const half  * src0,
-    device const float * src1,
-    device        char * dst,
-    uint3   tgpig[[threadgroup_position_in_grid]],
-    uint3    tgpg[[threadgroups_per_grid]]);
-
-
-typedef void (conv_transpose_2d_t)(
-        constant ggml_metal_kargs_conv_transpose_2d & args,
-        device const float * src0,
-        device const float * src1,
-        device        char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        uint3    tgpg[[threadgroups_per_grid]]);
-
-template <typename T>
-kernel void kernel_conv_transpose_2d(
-        constant ggml_metal_kargs_conv_transpose_2d & args,
-        device const T * src0,
-        device const float * src1,
-        device        char * dst,
-        threadgroup float * shared_sum [[threadgroup(0)]],
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        uint3   tpitg[[thread_position_in_threadgroup]],
-        uint3     ntg[[threads_per_threadgroup]]) {
-
-    const int64_t out_x = tgpig[0];
-    const int64_t out_y = tgpig[1];
-    const int64_t out_c = tgpig[2];
-
-    const int64_t kw = tpitg[0];
-    const int64_t kh = tpitg[1];
-
-    float v = 0.0f;
-
-    for (int64_t in_c = 0; in_c < args.IC; in_c++) {
-        int64_t in_y = out_y - kh;
-
-        if (in_y < 0 || in_y % args.s0) continue;
-
-        in_y /= args.s0;
-
-        if (in_y >= args.IH) continue;
-
-        int64_t in_x = out_x - kw;
-
-        if (in_x < 0 || in_x % args.s0) continue;
-
-        in_x /= args.s0;
-
-        if (in_x >= args.IW) continue;
-
-        const int64_t input_idx = (args.IW * args.IH) * in_c + (args.IW) * in_y + in_x;
-        const int64_t kernel_idx = (args.KH * args.KW * args.OC) * in_c + (args.KH * args.KW) * out_c + (args.KW) * kh + kw;
-
-        v += (float)src0[kernel_idx] * src1[input_idx];
-    }
-
-    const uint tid = tpitg.y * ntg.x + tpitg.x;
-    shared_sum[tid] = v;
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    if (tid == 0) {
-        float total = 0.0f;
-        const uint num_threads = ntg.x * ntg.y;
-        for (uint i = 0; i < num_threads; i++) {
-            total += shared_sum[i];
-        }
-
-        device float * dst_ptr = (device float *) (dst + out_x*args.nb0 + out_y * args.nb1 + out_c*args.nb2);
-        dst_ptr[0] = total;
-    }
-}
-
-template [[host_name("kernel_conv_transpose_2d_f32_f32")]]
-kernel void kernel_conv_transpose_2d<float>(
-    constant ggml_metal_kargs_conv_transpose_2d & args,
-    device const float * src0,
-    device const float * src1,
-    device        char * dst,
-    threadgroup float * shared_sum [[threadgroup(0)]],
-    uint3   tgpig[[threadgroup_position_in_grid]],
-    uint3   tpitg[[thread_position_in_threadgroup]],
-    uint3     ntg[[threads_per_threadgroup]]);
-
-template [[host_name("kernel_conv_transpose_2d_f16_f32")]]
-kernel void kernel_conv_transpose_2d<half>(
-    constant ggml_metal_kargs_conv_transpose_2d & args,
-    device const half  * src0,
-    device const float * src1,
-    device        char * dst,
-    threadgroup float * shared_sum [[threadgroup(0)]],
-    uint3   tgpig[[threadgroup_position_in_grid]],
-    uint3   tpitg[[thread_position_in_threadgroup]],
-    uint3     ntg[[threads_per_threadgroup]]);
-
-// grid: x = C tile, y = OH, z = OW * N (for channel-contiguous layouts)
-template <typename TK>
-kernel void kernel_conv_2d_dw_tiled(
-        constant ggml_metal_kargs_conv_2d_dw & args,
-        device const char * weights,
-        device const char * src,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        uint3   tpitg[[thread_position_in_threadgroup]],
-        uint3     ntg[[threads_per_threadgroup]]) {
-
-    const int32_t c = (int32_t)(tgpig.x * ntg.x + tpitg.x);
-    if (c >= args.C) {
-        return;
-    }
-
-    const int32_t oh = tgpig.y;
-    const int32_t own = tgpig.z;
-    const int32_t ow = own % args.OW;
-    const int32_t n  = own / args.OW;
-
-    const int32_t base_y = oh*args.s1 - args.p1;
-
-    int32_t ky_start = 0;
-    if (base_y < 0) {
-        ky_start = (-base_y + args.d1 - 1)/args.d1;
-    }
-    int32_t ky_end = args.KH;
-    const int32_t y_max = args.IH - 1 - base_y;
-    if (y_max < 0) {
-        ky_end = ky_start;
-    } else if (base_y + (args.KH - 1)*args.d1 >= args.IH) {
-        ky_end = min(ky_end, y_max/args.d1 + 1);
-    }
-
-    const int32_t base_x = ow*args.s0 - args.p0;
-
-    int32_t kx_start = 0;
-    if (base_x < 0) {
-        kx_start = (-base_x + args.d0 - 1)/args.d0;
-    }
-    int32_t kx_end = args.KW;
-    const int32_t x_max = args.IW - 1 - base_x;
-    if (x_max < 0) {
-        kx_end = kx_start;
-    } else if (base_x + (args.KW - 1)*args.d0 >= args.IW) {
-        kx_end = min(kx_end, x_max/args.d0 + 1);
-    }
-
-    float acc = 0.0f;
-
-    if (ky_start < ky_end && kx_start < kx_end) {
-        const uint64_t w_base   = (uint64_t) c * args.nb02;
-        const uint64_t src_base = (uint64_t) n * args.nb13 + (uint64_t) c * args.nb12;
-
-        for (int32_t ky = ky_start; ky < ky_end; ++ky) {
-            const int32_t iy = base_y + ky*args.d1;
-            const uint64_t src_row = src_base + (uint64_t) iy * args.nb11;
-            const uint64_t w_row = w_base + (uint64_t) ky * args.nb01;
-
-            for (int32_t kx = kx_start; kx < kx_end; ++kx) {
-                const int32_t ix = base_x + kx*args.d0;
-                const float x = *(device const float *)(src + src_row + (uint64_t) ix * args.nb10);
-                const float w = (float)(*(device const TK *)(weights + w_row + (uint64_t) kx * args.nb00));
-                acc += x * w;
-            }
-        }
-    }
-
-    const uint64_t dst_offs =
-        (uint64_t) n  * args.nb3 +
-        (uint64_t) c  * args.nb2 +
-        (uint64_t) oh * args.nb1 +
-        (uint64_t) ow * args.nb0;
-
-    *(device float *)(dst + dst_offs) = acc;
-}
-
-// grid: x = OW tile, y = OH, z = C * N (for spatially-contiguous layouts)
-template <typename TK>
-kernel void kernel_conv_2d_dw(
-        constant ggml_metal_kargs_conv_2d_dw & args,
-        device const char * weights,
-        device const char * src,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        uint3   tpitg[[thread_position_in_threadgroup]],
-        uint3     ntg[[threads_per_threadgroup]]) {
-
-    const int32_t oh = tgpig.y;
-    const int32_t cn = tgpig.z;
-    const int32_t c  = cn % args.C;
-    const int32_t n  = cn / args.C;
-
-    const int32_t base_y = oh*args.s1 - args.p1;
-
-    int32_t ky_start = 0;
-    if (base_y < 0) {
-        ky_start = (-base_y + args.d1 - 1)/args.d1;
-    }
-    int32_t ky_end = args.KH;
-    const int32_t y_max = args.IH - 1 - base_y;
-    if (y_max < 0) {
-        ky_end = ky_start;
-    } else if (base_y + (args.KH - 1)*args.d1 >= args.IH) {
-        ky_end = min(ky_end, y_max/args.d1 + 1);
-    }
-
-    const uint64_t w_base   = (uint64_t) c * args.nb02;
-    const uint64_t src_base = (uint64_t) n * args.nb13 + (uint64_t) c * args.nb12;
-
-    const int32_t ow = (int32_t)(tgpig.x * ntg.x + tpitg.x);
-    if (ow >= args.OW) {
-        return;
-    }
-
-    float acc = 0.0f;
-
-    const int32_t base_x = ow*args.s0 - args.p0;
-
-    int32_t kx_start = 0;
-    if (base_x < 0) {
-        kx_start = (-base_x + args.d0 - 1)/args.d0;
-    }
-    int32_t kx_end = args.KW;
-    const int32_t x_max = args.IW - 1 - base_x;
-    if (x_max < 0) {
-        kx_end = kx_start;
-    } else if (base_x + (args.KW - 1)*args.d0 >= args.IW) {
-        kx_end = min(kx_end, x_max/args.d0 + 1);
-    }
-
-    if (ky_start < ky_end && kx_start < kx_end) {
-        for (int32_t ky = ky_start; ky < ky_end; ++ky) {
-            const int32_t iy = base_y + ky*args.d1;
-            const uint64_t src_row = src_base + (uint64_t) iy * args.nb11;
-            const uint64_t w_row = w_base + (uint64_t) ky * args.nb01;
-
-            for (int32_t kx = kx_start; kx < kx_end; ++kx) {
-                const int32_t ix = base_x + kx*args.d0;
-                const float x = *(device const float *)(src + src_row + (uint64_t) ix * args.nb10);
-                const float w = (float)(*(device const TK *)(weights + w_row + (uint64_t) kx * args.nb00));
-                acc += x * w;
-            }
-        }
-    }
-
-    const uint64_t dst_offs =
-        (uint64_t) n  * args.nb3 +
-        (uint64_t) c  * args.nb2 +
-        (uint64_t) oh * args.nb1 +
-        (uint64_t) ow * args.nb0;
-
-    *(device float *)(dst + dst_offs) = acc;
-}
-
-template [[host_name("kernel_conv_2d_dw_f32_f32")]]
-kernel void kernel_conv_2d_dw<float>(
-        constant ggml_metal_kargs_conv_2d_dw & args,
-        device const char * weights,
-        device const char * src,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        uint3   tpitg[[thread_position_in_threadgroup]],
-        uint3     ntg[[threads_per_threadgroup]]);
-
-template [[host_name("kernel_conv_2d_dw_f16_f32")]]
-kernel void kernel_conv_2d_dw<half>(
-        constant ggml_metal_kargs_conv_2d_dw & args,
-        device const char * weights,
-        device const char * src,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        uint3   tpitg[[thread_position_in_threadgroup]],
-        uint3     ntg[[threads_per_threadgroup]]);
-
-template [[host_name("kernel_conv_2d_dw_tiled_f32_f32")]]
-kernel void kernel_conv_2d_dw_tiled<float>(
-        constant ggml_metal_kargs_conv_2d_dw & args,
-        device const char * weights,
-        device const char * src,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        uint3   tpitg[[thread_position_in_threadgroup]],
-        uint3     ntg[[threads_per_threadgroup]]);
-
-template [[host_name("kernel_conv_2d_dw_tiled_f16_f32")]]
-kernel void kernel_conv_2d_dw_tiled<half>(
-        constant ggml_metal_kargs_conv_2d_dw & args,
-        device const char * weights,
-        device const char * src,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        uint3   tpitg[[thread_position_in_threadgroup]],
-        uint3     ntg[[threads_per_threadgroup]]);
-
-template <typename T>
-kernel void kernel_conv_3d(
-        constant ggml_metal_kargs_conv_3d & args,
-        device const  char * src0, // Weights [IC * OC, KD, KH, KW]
-        device const  char * src1, // Inputs  [IC * N,  ID, IH, IW]
-        device       char  * dst,  // Outputs [OC * N,  OD, OH, OW]
-        uint3 tgpig[[threadgroup_position_in_grid]],
-        uint3 tpitg[[thread_position_in_threadgroup]]) {
-
-    // 1. Un-flatten the spatial dimension from Grid X
-    int64_t spatial_idx = tgpig.x * 32 + tpitg.x;
-
-    if (spatial_idx >= args.OW * args.OH * args.OD) {
-        return; // Thread falls outside the spatial volume
-    }
-
-    int64_t od = spatial_idx / (args.OW * args.OH);
-    int64_t oh = (spatial_idx / args.OW) % args.OH;
-    int64_t ow = spatial_idx % args.OW;
-
-    // 2. Map Y to Channels, Z to Batch
-    int64_t oc = tgpig.y;
-    int64_t batch_idx = tgpig.z;
-
-    // 3. Calculate anchor coordinates in the Input volume
-    int64_t i_w_base = ow * args.s0 - args.p0;
-    int64_t i_h_base = oh * args.s1 - args.p1;
-    int64_t i_d_base = od * args.s2 - args.p2;
-
-    float sum = 0.0f;
-
-    // 4. Gather Loop (Iterate over Input Channels -> Depth -> Height -> Width)
-    for (int64_t ic = 0; ic < args.IC; ++ic) {
-
-        // ggml packs batch and channel together in the 4th dimension
-        int64_t src_cn_idx = batch_idx * args.IC + ic;
-        int64_t w_cn_idx   = oc * args.IC + ic;
-
-        for (int64_t kz = 0; kz < args.KD; ++kz) {
-            int64_t id = i_d_base + kz * args.d2;
-            if (id < 0 || id >= args.ID) continue; // Boundary check (Padding)
-
-            for (int64_t ky = 0; ky < args.KH; ++ky) {
-                int64_t ih = i_h_base + ky * args.d1;
-                if (ih < 0 || ih >= args.IH) continue;
-
-                for (int64_t kx = 0; kx < args.KW; ++kx) {
-                    int64_t iw = i_w_base + kx * args.d0;
-                    if (iw < 0 || iw >= args.IW) continue;
-
-                    // Convert multi-dimensional coordinates to flat byte offsets
-                    int64_t w_idx = kx*args.nb00 + ky*args.nb01 + kz*args.nb02 + w_cn_idx*args.nb03;
-                    int64_t i_idx = iw*args.nb10 + ih*args.nb11 + id*args.nb12 + src_cn_idx*args.nb13;
-
-                    // Dereference memory and cast weights to f32 if they were f16
-                    float w_val = (float)*(device const T*)((device const char*)src0 + w_idx);
-                    float i_val = *(device const float*)((device const char*)src1 + i_idx);
-
-                    sum += w_val * i_val;
-                }
-            }
-        }
-    }
-
-    // 5. Write the accumulated value out to RAM
-    int64_t dst_cn_idx = batch_idx * args.OC + oc;
-    int64_t d_idx = ow*args.nb0 + oh*args.nb1 + od*args.nb2 + dst_cn_idx*args.nb3;
-
-    *(device float*)(dst + d_idx) = sum;
-}
-
-// Explicit instantiations so the JIT compiler can find them by name
-template [[host_name("kernel_conv_3d_f32_f32")]]
-kernel void kernel_conv_3d<float>(
-    constant ggml_metal_kargs_conv_3d & args,
-    device const char * src0,
-    device const char * src1,
-    device       char  * dst,
-    uint3 tgpig[[threadgroup_position_in_grid]],
-    uint3 tpitg[[thread_position_in_threadgroup]]);
-
-// Explicit instantiation for f16 weights
-template [[host_name("kernel_conv_3d_f16_f32")]]
-kernel void kernel_conv_3d<half>(
-    constant ggml_metal_kargs_conv_3d & args,
-    device const char  * src0,
-    device const char * src1,
-    device       char  * dst,
-    uint3 tgpig[[threadgroup_position_in_grid]],
-    uint3 tpitg[[thread_position_in_threadgroup]]);
@@ -1,735 +0,0 @@
-#pragma once
-
-#include "common.h"
-
-#define GGML_COMMON_DECL_METAL
-#define GGML_COMMON_IMPL_METAL
-#if defined(GGML_METAL_EMBED_LIBRARY)
-__embed_ggml-common.h__
-#else
-#include "ggml-common.h"
-#endif
-
-#define QK_NL 16 // shared by mul_mm and get_rows_q instantiations
-
-// NOTE: this is not dequantizing - we are simply fitting the template
-template <typename type4x4>
-void dequantize_f32(device const float4x4 * src, short il, thread type4x4 & reg) {
-    reg = (type4x4)(*src);
-}
-
-template <typename type4>
-void dequantize_f32_t4(device const float4 * src, short il, thread type4 & reg) {
-    reg = (type4)(*src);
-}
-
-template <typename type4x4>
-void dequantize_f16(device const half4x4 * src, short il, thread type4x4 & reg) {
-    reg = (type4x4)(*src);
-}
-
-template <typename type4>
-void dequantize_f16_t4(device const half4 * src, short il, thread type4 & reg) {
-    reg = (type4)(*(src));
-}
-
-#if defined(GGML_METAL_HAS_BF16)
-template <typename type4x4>
-void dequantize_bf16(device const bfloat4x4 * src, short il, thread type4x4 & reg) {
-    reg = (type4x4)(*src);
-}
-
-template <typename type4>
-void dequantize_bf16_t4(device const bfloat4 * src, short il, thread type4 & reg) {
-    reg = (type4)(*(src));
-}
-#endif
-
-template <typename type4x4>
-void dequantize_q1_0(device const block_q1_0 * xb, short il, thread type4x4 & reg) {
-    device const uint8_t * qs = xb->qs;
-    const float d = xb->d;
-    const float neg_d = -d;
-
-    const int byte_offset = il * 2;  // il*16 bits = il*2 bytes
-    const uint8_t b0 = qs[byte_offset];
-    const uint8_t b1 = qs[byte_offset + 1];
-
-    float4x4 reg_f;
-
-    reg_f[0][0] = select(neg_d, d, bool(b0 & 0x01));
-    reg_f[0][1] = select(neg_d, d, bool(b0 & 0x02));
-    reg_f[0][2] = select(neg_d, d, bool(b0 & 0x04));
-    reg_f[0][3] = select(neg_d, d, bool(b0 & 0x08));
-    reg_f[1][0] = select(neg_d, d, bool(b0 & 0x10));
-    reg_f[1][1] = select(neg_d, d, bool(b0 & 0x20));
-    reg_f[1][2] = select(neg_d, d, bool(b0 & 0x40));
-    reg_f[1][3] = select(neg_d, d, bool(b0 & 0x80));
-
-    reg_f[2][0] = select(neg_d, d, bool(b1 & 0x01));
-    reg_f[2][1] = select(neg_d, d, bool(b1 & 0x02));
-    reg_f[2][2] = select(neg_d, d, bool(b1 & 0x04));
-    reg_f[2][3] = select(neg_d, d, bool(b1 & 0x08));
-    reg_f[3][0] = select(neg_d, d, bool(b1 & 0x10));
-    reg_f[3][1] = select(neg_d, d, bool(b1 & 0x20));
-    reg_f[3][2] = select(neg_d, d, bool(b1 & 0x40));
-    reg_f[3][3] = select(neg_d, d, bool(b1 & 0x80));
-
-    reg = (type4x4) reg_f;
-}
-
-template <typename type4>
-void dequantize_q1_0_t4(device const block_q1_0 * xb, short il, thread type4 & reg) {
-    const float d = xb->d;
-    const float neg_d = -d;
-    const int base = il * 4;
-    const uint8_t byte = xb->qs[base / 8];
-    const int s = base % 8;
-
-    float4 reg_f;
-    reg_f[0] = select(neg_d, d, bool((byte >> (s    )) & 1));
-    reg_f[1] = select(neg_d, d, bool((byte >> (s + 1)) & 1));
-    reg_f[2] = select(neg_d, d, bool((byte >> (s + 2)) & 1));
-    reg_f[3] = select(neg_d, d, bool((byte >> (s + 3)) & 1));
-
-    reg = (type4) reg_f;
-}
-
-template <typename type4x4>
-void dequantize_q2_0(device const block_q2_0 * xb, short il, thread type4x4 & reg) {
-    device const uint8_t * qs = xb->qs;
-    const float d = xb->d;
-
-    const int byte_offset = il * 4;  // il*16 elements = il*4 bytes (4 elements per byte)
-    float4x4 reg_f;
-
-    for (int i = 0; i < 4; i++) {
-        const uint8_t b = qs[byte_offset + i];
-        reg_f[i][0] = ((float)((b >> 0) & 3) - 1.0f) * d;
-        reg_f[i][1] = ((float)((b >> 2) & 3) - 1.0f) * d;
-        reg_f[i][2] = ((float)((b >> 4) & 3) - 1.0f) * d;
-        reg_f[i][3] = ((float)((b >> 6) & 3) - 1.0f) * d;
-    }
-
-    reg = (type4x4) reg_f;
-}
-
-template <typename type4>
-void dequantize_q2_0_t4(device const block_q2_0 * xb, short il, thread type4 & reg) {
-    const float d = xb->d;
-    const uint8_t b = xb->qs[il];
-
-    float4 reg_f;
-    reg_f[0] = ((float)((b >> 0) & 3) - 1.0f) * d;
-    reg_f[1] = ((float)((b >> 2) & 3) - 1.0f) * d;
-    reg_f[2] = ((float)((b >> 4) & 3) - 1.0f) * d;
-    reg_f[3] = ((float)((b >> 6) & 3) - 1.0f) * d;
-
-    reg = (type4) reg_f;
-}
-
-template <typename type4x4>
-void dequantize_q4_0(device const block_q4_0 * xb, short il, thread type4x4 & reg) {
-    device const uint16_t * qs = ((device const uint16_t *)xb + 1);
-    const float d1 = il ? (xb->d / 16.h) : xb->d;
-    const float d2 = d1 / 256.f;
-    const float md = -8.h * xb->d;
-    const ushort mask0 = il ? 0x00F0 : 0x000F;
-    const ushort mask1 = mask0 << 8;
-
-    float4x4 reg_f;
-
-    for (int i = 0; i < 8; i++) {
-        reg_f[i/2][2*(i%2) + 0] = d1 * (qs[i] & mask0) + md;
-        reg_f[i/2][2*(i%2) + 1] = d2 * (qs[i] & mask1) + md;
-    }
-
-    reg = (type4x4) reg_f;
-}
-
-template <typename type4>
-void dequantize_q4_0_t4(device const block_q4_0 * xb, short il, thread type4 & reg) {
-    device const uint16_t * qs = ((device const uint16_t *)xb + 1);
-    const float d1 = (il/4) ? (xb->d / 16.h) : xb->d;
-    const float d2 = d1 / 256.f;
-    const float md = -8.h * xb->d;
-    const ushort mask0 = (il/4) ? 0x00F0 : 0x000F;
-    const ushort mask1 = mask0 << 8;
-
-    for (int i = 0; i < 2; i++) {
-        reg[2*i + 0] = d1 * (qs[2*(il%4) + i] & mask0) + md;
-        reg[2*i + 1] = d2 * (qs[2*(il%4) + i] & mask1) + md;
-    }
-}
-
-
-
-template <typename type4x4>
-void dequantize_q4_1(device const block_q4_1 * xb, short il, thread type4x4 & reg) {
-    device const uint16_t * qs = ((device const uint16_t *)xb + 2);
-    const float d1 = il ? (xb->d / 16.h) : xb->d;
-    const float d2 = d1 / 256.f;
-    const float  m = xb->m;
-    const ushort mask0 = il ? 0x00F0 : 0x000F;
-    const ushort mask1 = mask0 << 8;
-
-    float4x4 reg_f;
-
-    for (int i = 0; i < 8; i++) {
-        reg_f[i/2][2*(i%2) + 0] = ((qs[i] & mask0) * d1) + m;
-        reg_f[i/2][2*(i%2) + 1] = ((qs[i] & mask1) * d2) + m;
-    }
-
-    reg = (type4x4) reg_f;
-}
-
-template <typename type4>
-void dequantize_q4_1_t4(device const block_q4_1 * xb, short il, thread type4 & reg) {
-    device const uint16_t * qs = ((device const uint16_t *)xb + 2);
-    const float d1 = (il/4) ? (xb->d / 16.h) : xb->d;
-    const float d2 = d1 / 256.f;
-    const float  m = xb->m;
-    const ushort mask0 = (il/4) ? 0x00F0 : 0x000F;
-    const ushort mask1 = mask0 << 8;
-
-    for (int i = 0; i < 2; i++) {
-        reg[2*i + 0] = d1 * (qs[2*(il%4) + i] & mask0) + m;
-        reg[2*i + 1] = d2 * (qs[2*(il%4) + i] & mask1) + m;
-    }
-}
-
-template <typename type4x4>
-void dequantize_q5_0(device const block_q5_0 * xb, short il, thread type4x4 & reg) {
-    device const uint16_t * qs = ((device const uint16_t *)xb + 3);
-    const float d = xb->d;
-    const float md = -16.h * xb->d;
-    const ushort mask = il ? 0x00F0 : 0x000F;
-
-    const uint32_t qh = *((device const uint32_t *)xb->qh);
-
-    const int x_mv = il ? 4 : 0;
-
-    const int gh_mv = il ? 12 : 0;
-    const int gh_bk = il ?  0 : 4;
-
-    float4x4 reg_f;
-
-    for (int i = 0; i < 8; i++) {
-        // extract the 5-th bits for x0 and x1
-        const uint8_t xh_0 = ((qh >> (gh_mv + 2*i  )) << gh_bk) & 0x10;
-        const uint8_t xh_1 = ((qh >> (gh_mv + 2*i+1)) << gh_bk) & 0x10;
-
-        // combine the 4-bits from qs with the 5th bit
-        const int32_t x0 = ((((qs[i]     ) & mask) >> x_mv) | xh_0);
-        const int32_t x1 = ((((qs[i] >> 8) & mask) >> x_mv) | xh_1);
-
-        reg_f[i/2][2*(i%2) + 0] = d * x0 + md;
-        reg_f[i/2][2*(i%2) + 1] = d * x1 + md;
-    }
-
-    reg = (type4x4) reg_f;
-}
-
-template <typename type4>
-void dequantize_q5_0_t4(device const block_q5_0 * xb, short il, thread type4 & reg) {
-    device const uint16_t * qs = ((device const uint16_t *)xb + 3);
-    const float d = xb->d;
-    const float md = -16.h * xb->d;
-    const ushort mask = (il/4) ? 0x00F0 : 0x000F;
-
-    const uint32_t qh = *((device const uint32_t *)xb->qh);
-
-    const int x_mv = (il/4) ? 4 : 0;
-
-    const int gh_mv = (il/4) ? 12 : 0;
-    const int gh_bk = (il/4) ?  0 : 4;
-
-    for (int ii = 0; ii < 2; ii++) {
-        int i = 2*(il%4) + ii;
-
-        // extract the 5-th bits for x0 and x1
-        const uint8_t xh_0 = ((qh >> (gh_mv + 2*i  )) << gh_bk) & 0x10;
-        const uint8_t xh_1 = ((qh >> (gh_mv + 2*i+1)) << gh_bk) & 0x10;
-
-        // combine the 4-bits from qs with the 5th bit
-        const int32_t x0 = ((((qs[i]     ) & mask) >> x_mv) | xh_0);
-        const int32_t x1 = ((((qs[i] >> 8) & mask) >> x_mv) | xh_1);
-
-        reg[2*ii + 0] = d * x0 + md;
-        reg[2*ii + 1] = d * x1 + md;
-    }
-}
-
-template <typename type4x4>
-void dequantize_q5_1(device const block_q5_1 * xb, short il, thread type4x4 & reg) {
-    device const uint16_t * qs = ((device const uint16_t *)xb + 4);
-    const float d = xb->d;
-    const float m = xb->m;
-    const ushort mask = il ? 0x00F0 : 0x000F;
-
-    const uint32_t qh = *((device const uint32_t *)xb->qh);
-
-    const int x_mv = il ? 4 : 0;
-
-    const int gh_mv = il ? 12 : 0;
-    const int gh_bk = il ?  0 : 4;
-
-    float4x4 reg_f;
-
-    for (int i = 0; i < 8; i++) {
-        // extract the 5-th bits for x0 and x1
-        const uint8_t xh_0 = ((qh >> (gh_mv + 2*i  )) << gh_bk) & 0x10;
-        const uint8_t xh_1 = ((qh >> (gh_mv + 2*i+1)) << gh_bk) & 0x10;
-
-        // combine the 4-bits from qs with the 5th bit
-        const int32_t x0 = ((((qs[i]     ) & mask) >> x_mv) | xh_0);
-        const int32_t x1 = ((((qs[i] >> 8) & mask) >> x_mv) | xh_1);
-
-        reg_f[i/2][2*(i%2) + 0] = d * x0 + m;
-        reg_f[i/2][2*(i%2) + 1] = d * x1 + m;
-    }
-
-    reg = (type4x4) reg_f;
-}
-
-template <typename type4>
-void dequantize_q5_1_t4(device const block_q5_1 * xb, short il, thread type4 & reg) {
-    device const uint16_t * qs = ((device const uint16_t *)xb + 4);
-    const float d = xb->d;
-    const float m = xb->m;
-    const ushort mask = (il/4) ? 0x00F0 : 0x000F;
-
-    const uint32_t qh = *((device const uint32_t *)xb->qh);
-
-    const int x_mv = (il/4) ? 4 : 0;
-
-    const int gh_mv = (il/4) ? 12 : 0;
-    const int gh_bk = (il/4) ?  0 : 4;
-
-    for (int ii = 0; ii < 2; ii++) {
-        int i = 2*(il%4) + ii;
-
-        // extract the 5-th bits for x0 and x1
-        const uint8_t xh_0 = ((qh >> (gh_mv + 2*i  )) << gh_bk) & 0x10;
-        const uint8_t xh_1 = ((qh >> (gh_mv + 2*i+1)) << gh_bk) & 0x10;
-
-        // combine the 4-bits from qs with the 5th bit
-        const int32_t x0 = ((((qs[i]     ) & mask) >> x_mv) | xh_0);
-        const int32_t x1 = ((((qs[i] >> 8) & mask) >> x_mv) | xh_1);
-
-        reg[2*ii + 0] = d * x0 + m;
-        reg[2*ii + 1] = d * x1 + m;
-    }
-}
-
-template <typename type4x4>
-void dequantize_q8_0(device const block_q8_0 *xb, short il, thread type4x4 & reg) {
-    device const packed_char4 * qs = (device const packed_char4 *) xb->qs;
-    const float d = xb->d;
-
-    float4x4 reg_f;
-
-    for (int i = 0; i < 4; ++i) {
-        reg_f[i] = float4(qs[4*il + i]) * d;
-    }
-
-    reg = (type4x4) reg_f;
-}
-
-template <typename type4>
-void dequantize_q8_0_t4(device const block_q8_0 *xb, short il, thread type4 & reg) {
-    device const packed_char4 * qs = (device const packed_char4 *) xb->qs;
-    const float d = xb->d;
-
-    reg = (type4) (float4(qs[il]) * d);
-}
-
-template <typename type4x4>
-void dequantize_mxfp4(device const block_mxfp4 * xb, short il, thread type4x4 & reg) {
-    device const uint8_t * q2 = (device const uint8_t *)xb->qs;
-
-    const float d = e8m0_to_fp32(xb->e);
-    const uint8_t shr = il >= 1 ? 4 : 0;
-
-    for (int i = 0; i < 4; ++i) {
-        reg[i][0] = d * kvalues_mxfp4_f[(q2[4*i + 0] >> shr) & 0x0F];
-        reg[i][1] = d * kvalues_mxfp4_f[(q2[4*i + 1] >> shr) & 0x0F];
-        reg[i][2] = d * kvalues_mxfp4_f[(q2[4*i + 2] >> shr) & 0x0F];
-        reg[i][3] = d * kvalues_mxfp4_f[(q2[4*i + 3] >> shr) & 0x0F];
-    }
-}
-
-template <typename type4>
-void dequantize_mxfp4_t4(device const block_mxfp4 * xb, short il, thread type4 & reg) {
-    device const uint8_t * q2 = (device const uint8_t *)xb->qs;
-
-    const float d = e8m0_to_fp32(xb->e);
-    const short il4 = il%4;
-
-    const uint8_t shr = il >= 4 ? 4 : 0;
-
-    reg[0] = d * kvalues_mxfp4_f[(q2[4*il4 + 0] >> shr) & 0x0F];
-    reg[1] = d * kvalues_mxfp4_f[(q2[4*il4 + 1] >> shr) & 0x0F];
-    reg[2] = d * kvalues_mxfp4_f[(q2[4*il4 + 2] >> shr) & 0x0F];
-    reg[3] = d * kvalues_mxfp4_f[(q2[4*il4 + 3] >> shr) & 0x0F];
-}
-
-template <typename type4x4>
-void dequantize_q2_K(device const block_q2_K *xb, short il, thread type4x4 & reg) {
-    const float d = xb->d;
-    const float min = xb->dmin;
-    device const uint8_t * q = (device const uint8_t *)xb->qs;
-    float dl, ml;
-    uint8_t sc = xb->scales[il];
-
-    q = q + 32*(il/8) + 16*(il&1);
-    il = (il/2)%4;
-
-    half  coef = il>1 ? (il>2 ? 1/64.h : 1/16.h) : (il>0 ? 1/4.h : 1.h);
-    uchar mask = il>1 ? (il>2 ? 192    : 48)     : (il>0 ? 12    : 3);
-    dl = d * (sc & 0xF) * coef, ml = min * (sc >> 4);
-    for (int i = 0; i < 16; ++i) {
-        reg[i/4][i%4] = dl * (q[i] & mask) - ml;
-    }
-}
-
-template <typename type4x4>
-void dequantize_q3_K(device const block_q3_K *xb, short il, thread type4x4 & reg) {
-    const half d_all = xb->d;
-    device const uint8_t * q = (device const uint8_t *)xb->qs;
-    device const uint8_t * h = (device const uint8_t *)xb->hmask;
-    device const int8_t * scales = (device const int8_t *)xb->scales;
-
-    q = q + 32 * (il/8) + 16 * (il&1);
-    h = h + 16 * (il&1);
-    uint8_t m = 1 << (il/2);
-    uint16_t kmask1 = (il/4)>1 ? ((il/4)>2 ? 192 : 48) : \
-                                 ((il/4)>0 ? 12  : 3);
-    uint16_t kmask2 = il/8 ? 0xF0 : 0x0F;
-    uint16_t scale_2 = scales[il%8], scale_1 = scales[8 + il%4];
-    int16_t  dl_int = (il/4)&1 ? (scale_2&kmask2) | ((scale_1&kmask1) << 2)
-                               : (scale_2&kmask2) | ((scale_1&kmask1) << 4);
-    float dl = il<8 ? d_all * (dl_int - 32.f) : d_all * (dl_int / 16.f - 32.f);
-    const float ml = 4.f * dl;
-
-    il = (il/2) & 3;
-    const half    coef = il>1 ? (il>2 ? 1/64.h : 1/16.h) : (il>0 ? 1/4.h : 1.h);
-    const uint8_t mask = il>1 ? (il>2 ? 192    : 48)     : (il>0 ? 12    : 3);
-    dl *= coef;
-
-    for (int i = 0; i < 16; ++i) {
-        reg[i/4][i%4] = dl * (q[i] & mask) - (h[i] & m ? 0 : ml);
-    }
-}
-
-static inline uchar2 get_scale_min_k4_just2(int j, int k, device const uchar * q) {
-    return j < 4 ? uchar2{uchar(q[j+0+k] & 63), uchar(q[j+4+k] & 63)}
-                 : uchar2{uchar((q[j+4+k] & 0xF) | ((q[j-4+k] & 0xc0) >> 2)), uchar((q[j+4+k] >> 4) | ((q[j-0+k] & 0xc0) >> 2))};
-}
-
-template <typename type4x4>
-void dequantize_q4_K(device const block_q4_K * xb, short il, thread type4x4 & reg) {
-    device const uchar * q = xb->qs;
-
-    short is = (il/4) * 2;
-    q = q + (il/4) * 32 + 16 * (il&1);
-    il = il & 3;
-    const uchar2 sc = get_scale_min_k4_just2(is, il/2, xb->scales);
-    const float d   = il < 2 ? xb->d : xb->d / 16.h;
-    const float min = xb->dmin;
-    const float dl = d * sc[0];
-    const float ml = min * sc[1];
-
-    const ushort mask = il < 2 ? 0x0F : 0xF0;
-    for (int i = 0; i < 16; ++i) {
-        reg[i/4][i%4] = dl * (q[i] & mask) - ml;
-    }
-}
-
-template <typename type4x4>
-void dequantize_q5_K(device const block_q5_K *xb, short il, thread type4x4 & reg) {
-    device const uint8_t * q  = xb->qs;
-    device const uint8_t * qh = xb->qh;
-
-    short is = (il/4) * 2;
-    q  = q + 32 * (il/4) + 16 * (il&1);
-    qh = qh + 16 * (il&1);
-    uint8_t ul = 1 << (il/2);
-    il = il & 3;
-    const uchar2 sc = get_scale_min_k4_just2(is, il/2, xb->scales);
-    const float d = il < 2 ? xb->d : xb->d / 16.f;
-    const float min = xb->dmin;
-    const float dl = d * sc[0];
-    const float ml = min * sc[1];
-
-    const ushort mask  = il<2 ? 0x0F : 0xF0;
-    const float qh_val = il<2 ? 16.f : 256.f;
-    for (int i = 0; i < 16; ++i) {
-        reg[i/4][i%4] = dl * ((q[i] & mask) + (qh[i] & ul ? qh_val : 0)) - ml;
-    }
-}
-
-template <typename type4x4>
-void dequantize_q6_K(device const block_q6_K *xb, short il, thread type4x4 & reg) {
-    const half d_all = xb->d;
-    device const uint16_t * ql = (device const uint16_t *)xb->ql;
-    device const uint16_t * qh = (device const uint16_t *)xb->qh;
-    device const int8_t * scales = (device const int8_t *)xb->scales;
-
-    ql = ql + 32*(il/8) + 16*((il/2)&1) + 8*(il&1);
-    qh = qh + 16*(il/8) + 8*(il&1);
-    float sc = scales[(il%2) + 2 * ((il/2))];
-    il = (il/2) & 3;
-
-    const uint32_t kmask1 = il>1 ? (il>2 ? 0xC0C0C0C0 : 0x30303030) : (il>0 ? 0x0C0C0C0C : 0x03030303);
-    const uint32_t kmask2 = il>1 ? 0xF0F0F0F0                       : 0x0F0F0F0F;
-    const float ml = d_all * sc * 32.f;
-    const float dl0 = d_all * sc;
-    const float dl1 = dl0 / 256.f;
-    const float dl2 = dl0 / (256.f * 256.f);
-    const float dl3 = dl0 / (256.f * 256.f * 256.f);
-    const uint8_t shr_h = il>2 ? 2 : 0;
-    const uint8_t shl_h = il>1 ? 0 : (il>0 ? 2 : 4);
-    const uint8_t shr_l = il>1 ? 4 : 0;
-    for (int i = 0; i < 4; ++i) {
-        const uint32_t  low = (ql[2*i] | (uint32_t)(ql[2*i+1] << 16)) & kmask2;
-        const uint32_t high = (qh[2*i] | (uint32_t)(qh[2*i+1] << 16)) & kmask1;
-        const uint32_t q = ((high << shl_h) >> shr_h) | (low >> shr_l);
-        reg[i][0] = dl0 *  ((half)(q & 0xFF))       - ml;
-        reg[i][1] = dl1 * ((float)(q & 0xFF00))     - ml;
-        reg[i][2] = dl2 * ((float)(q & 0xFF0000))   - ml;
-        reg[i][3] = dl3 * ((float)(q & 0xFF000000)) - ml;
-    }
-}
-
-template <typename type4x4>
-void dequantize_iq2_xxs(device const block_iq2_xxs * xb, short il, thread type4x4 & reg) {
-    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
-    const float d = xb->d;
-    const int ib32 = il/2;
-    il = il%2;
-    // il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
-    // each block of 32 needs 2 uint32_t's for the quants & scale, so 4 uint16_t's.
-    device const uint16_t * q2 = xb->qs + 4*ib32;
-    const uint32_t aux32_g = q2[0] | (q2[1] << 16);
-    const uint32_t aux32_s = q2[2] | (q2[3] << 16);
-    thread const uint8_t * aux8 = (thread const uint8_t *)&aux32_g;
-    const float dl = d * (0.5f + (aux32_s >> 28)) * 0.25f;
-    constant uint8_t * grid = (constant uint8_t *)(iq2xxs_grid + aux8[2*il+0]);
-    uint8_t signs = ksigns_iq2xs[(aux32_s >> 14*il) & 127];
-    for (int i = 0; i < 8; ++i) {
-        reg[i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
-    }
-    grid = (constant uint8_t *)(iq2xxs_grid + aux8[2*il+1]);
-    signs = ksigns_iq2xs[(aux32_s >> (14*il+7)) & 127];
-    for (int i = 0; i < 8; ++i) {
-        reg[2+i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
-    }
-}
-
-template <typename type4x4>
-void dequantize_iq2_xs(device const block_iq2_xs * xb, short il, thread type4x4 & reg) {
-    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
-    const float d = xb->d;
-    const int ib32 = il/2;
-    il = il%2;
-    // il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
-    device const uint16_t * q2 = xb->qs + 4*ib32;
-    const float dl = d * (0.5f + ((xb->scales[ib32] >> 4*il) & 0xf)) * 0.25f;
-    constant uint8_t * grid = (constant uint8_t *)(iq2xs_grid + (q2[2*il+0] & 511));
-    uint8_t signs = ksigns_iq2xs[q2[2*il+0] >> 9];
-    for (int i = 0; i < 8; ++i) {
-        reg[i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
-    }
-    grid = (constant uint8_t *)(iq2xs_grid + (q2[2*il+1] & 511));
-    signs = ksigns_iq2xs[q2[2*il+1] >> 9];
-    for (int i = 0; i < 8; ++i) {
-        reg[2+i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
-    }
-}
-
-template <typename type4x4>
-void dequantize_iq3_xxs(device const block_iq3_xxs * xb, short il, thread type4x4 & reg) {
-    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
-    const float d = xb->d;
-    const int ib32 = il/2;
-    il = il%2;
-    // il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
-    device const uint8_t * q3 = xb->qs + 8*ib32;
-    device const uint16_t * gas = (device const uint16_t *)(xb->qs + QK_K/4) + 2*ib32;
-    const uint32_t aux32 = gas[0] | (gas[1] << 16);
-    const float dl = d * (0.5f + (aux32 >> 28)) * 0.5f;
-    constant uint8_t * grid1 = (constant uint8_t *)(iq3xxs_grid + q3[4*il+0]);
-    constant uint8_t * grid2 = (constant uint8_t *)(iq3xxs_grid + q3[4*il+1]);
-    uint8_t signs = ksigns_iq2xs[(aux32 >> 14*il) & 127];
-    for (int i = 0; i < 4; ++i) {
-        reg[0][i] = dl * grid1[i] * (signs & kmask_iq2xs[i+0] ? -1.f : 1.f);
-        reg[1][i] = dl * grid2[i] * (signs & kmask_iq2xs[i+4] ? -1.f : 1.f);
-    }
-    grid1 = (constant uint8_t *)(iq3xxs_grid + q3[4*il+2]);
-    grid2 = (constant uint8_t *)(iq3xxs_grid + q3[4*il+3]);
-    signs = ksigns_iq2xs[(aux32 >> (14*il+7)) & 127];
-    for (int i = 0; i < 4; ++i) {
-        reg[2][i] = dl * grid1[i] * (signs & kmask_iq2xs[i+0] ? -1.f : 1.f);
-        reg[3][i] = dl * grid2[i] * (signs & kmask_iq2xs[i+4] ? -1.f : 1.f);
-    }
-}
-
-template <typename type4x4>
-void dequantize_iq3_s(device const block_iq3_s * xb, short il, thread type4x4 & reg) {
-    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
-    const float d = xb->d;
-    const int ib32 = il/2;
-    il = il%2;
-    // il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
-    device const uint8_t * qs = xb->qs + 8*ib32;
-    device const uint8_t * signs = xb->signs + 4*ib32 + 2*il;
-    const uint8_t qh = xb->qh[ib32] >> 4*il;
-    const float dl = d * (1 + 2*((xb->scales[ib32/2] >> 4*(ib32%2)) & 0xf));
-    constant uint8_t * grid1 = (constant uint8_t *)(iq3s_grid + (qs[4*il+0] | ((qh << 8) & 256)));
-    constant uint8_t * grid2 = (constant uint8_t *)(iq3s_grid + (qs[4*il+1] | ((qh << 7) & 256)));
-    for (int i = 0; i < 4; ++i) {
-        reg[0][i] = dl * grid1[i] * select(1, -1, signs[0] & kmask_iq2xs[i+0]);
-        reg[1][i] = dl * grid2[i] * select(1, -1, signs[0] & kmask_iq2xs[i+4]);
-    }
-    grid1 = (constant uint8_t *)(iq3s_grid + (qs[4*il+2] | ((qh << 6) & 256)));
-    grid2 = (constant uint8_t *)(iq3s_grid + (qs[4*il+3] | ((qh << 5) & 256)));
-    for (int i = 0; i < 4; ++i) {
-        reg[2][i] = dl * grid1[i] * select(1, -1, signs[1] & kmask_iq2xs[i+0]);
-        reg[3][i] = dl * grid2[i] * select(1, -1, signs[1] & kmask_iq2xs[i+4]);
-    }
-}
-
-template <typename type4x4>
-void dequantize_iq2_s(device const block_iq2_s * xb, short il, thread type4x4 & reg) {
-    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
-    const float d = xb->d;
-    const int ib32 = il/2;
-    il = il%2;
-    // il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
-    device const uint8_t * qs = xb->qs + 4*ib32 + 2*il;
-    device const uint8_t * signs = qs + QK_K/8;
-    const uint8_t qh = xb->qh[ib32] >> 4*il;
-    const float dl = d * (0.5f + ((xb->scales[ib32] >> 4*il) & 0xf)) * 0.25f;
-    constant uint8_t * grid1 = (constant uint8_t *)(iq2s_grid + (qs[0] | ((qh << 8) & 0x300)));
-    constant uint8_t * grid2 = (constant uint8_t *)(iq2s_grid + (qs[1] | ((qh << 6) & 0x300)));
-    for (int i = 0; i < 8; ++i) {
-        reg[i/4+0][i%4] = dl * grid1[i] * select(1, -1, signs[0] & kmask_iq2xs[i]);
-        reg[i/4+2][i%4] = dl * grid2[i] * select(1, -1, signs[1] & kmask_iq2xs[i]);
-    }
-}
-
-template <typename type4x4>
-void dequantize_iq1_s(device const block_iq1_s * xb, short il, thread type4x4 & reg) {
-    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
-    const int ib32 = il/2;
-    il = il%2;
-    const float d = xb->d;
-    device const uint8_t  * qs = xb->qs + 4*ib32 + 2*il;
-    device const uint16_t * qh = xb->qh;
-    const float dl = d * (2*((qh[ib32] >> 12) & 7) + 1);
-    const float ml = dl * (qh[ib32] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA);
-    const uint16_t h = qh[ib32] >> 6*il;
-    constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((h << 8) & 0x700)));
-    constant uint8_t * grid2 = (constant uint8_t *)(iq1s_grid_gpu + (qs[1] | ((h << 5) & 0x700)));
-    for (int i = 0; i < 4; ++i) {
-        reg[0][i] = dl * (grid1[i] & 0xf) + ml;
-        reg[1][i] = dl * (grid1[i] >>  4) + ml;
-        reg[2][i] = dl * (grid2[i] & 0xf) + ml;
-        reg[3][i] = dl * (grid2[i] >>  4) + ml;
-    }
-}
-
-template <typename type4x4>
-void dequantize_iq1_m(device const block_iq1_m * xb, short il, thread type4x4 & reg) {
-    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
-    const int ib32 = il/2;
-    il = il%2;
-    device const uint16_t * sc = (device const uint16_t *)xb->scales;
-
-    iq1m_scale_t scale;
-    scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
-    const float d = scale.f16;
-
-    device const uint8_t * qs = xb->qs + 4*ib32 + 2*il;
-    device const uint8_t * qh = xb->qh + 2*ib32 + il;
-
-    const float dl  = d * (2*((sc[ib32/2] >> (6*(ib32%2)+3*il)) & 7) + 1);
-    const float ml1 = dl * (qh[0] & 0x08 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA);
-    const float ml2 = dl * (qh[0] & 0x80 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA);
-    constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((qh[0] << 8) & 0x700)));
-    constant uint8_t * grid2 = (constant uint8_t *)(iq1s_grid_gpu + (qs[1] | ((qh[0] << 4) & 0x700)));
-    for (int i = 0; i < 4; ++i) {
-        reg[0][i] = dl * (grid1[i] & 0xf) + ml1;
-        reg[1][i] = dl * (grid1[i] >>  4) + ml1;
-        reg[2][i] = dl * (grid2[i] & 0xf) + ml2;
-        reg[3][i] = dl * (grid2[i] >>  4) + ml2;
-    }
-}
-
-template <typename type4x4>
-void dequantize_iq4_nl(device const block_iq4_nl * xb, short il, thread type4x4 & reg) {
-    device const uint16_t * q4 = (device const uint16_t *)xb->qs;
-    const float d = xb->d;
-    uint32_t aux32;
-    thread const uint8_t * q8 = (thread const uint8_t *)&aux32;
-    for (int i = 0; i < 4; ++i) {
-        aux32 = ((q4[2*i] | (q4[2*i+1] << 16)) >> 4*il) & 0x0f0f0f0f;
-        reg[i][0] = d * kvalues_iq4nl_f[q8[0]];
-        reg[i][1] = d * kvalues_iq4nl_f[q8[1]];
-        reg[i][2] = d * kvalues_iq4nl_f[q8[2]];
-        reg[i][3] = d * kvalues_iq4nl_f[q8[3]];
-    }
-}
-
-template <typename type4>
-void dequantize_iq4_nl_t4(device const block_iq4_nl * xb, short il, thread type4 & reg) {
-    device const uint16_t * q4 = (device const uint16_t *)xb->qs;
-    const float d = xb->d;
-    uint32_t aux32;
-    thread const uint8_t * q8 = (thread const uint8_t *)&aux32;
-    aux32 = ((q4[2*(il%4)] | (q4[2*(il%4)+1] << 16)) >> 4*(il/4)) & 0x0f0f0f0f;
-    reg[0] = d * kvalues_iq4nl_f[q8[0]];
-    reg[1] = d * kvalues_iq4nl_f[q8[1]];
-    reg[2] = d * kvalues_iq4nl_f[q8[2]];
-    reg[3] = d * kvalues_iq4nl_f[q8[3]];
-}
-
-template <typename type4x4>
-void dequantize_iq4_xs(device const block_iq4_xs * xb, short il, thread type4x4 & reg) {
-    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
-    const int ib32 = il/2;
-    il = il%2;
-    // il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
-    device const uint32_t * q4 = (device const uint32_t *)xb->qs + 4*ib32;
-    const int ls = ((xb->scales_l[ib32/2] >> 4*(ib32%2)) & 0xf) | (((xb->scales_h >> 2*ib32) & 3) << 4);
-    const float d = (float)xb->d * (ls - 32);
-    uint32_t aux32;
-    thread const uint8_t * q8 = (thread const uint8_t *)&aux32;
-    for (int i = 0; i < 4; ++i) {
-        aux32 = (q4[i] >> 4*il) & 0x0f0f0f0f;
-        reg[i][0] = d * kvalues_iq4nl_f[q8[0]];
-        reg[i][1] = d * kvalues_iq4nl_f[q8[1]];
-        reg[i][2] = d * kvalues_iq4nl_f[q8[2]];
-        reg[i][3] = d * kvalues_iq4nl_f[q8[3]];
-    }
-}
-
-template <typename type4x4>
-void dequantize_tq2_0(device const block_tq2_0 * xb, short il, thread type4x4 & reg) {
-    device const uint8_t * qs = xb->qs;
-    const float d = xb->d;
-
-    float4x4 reg_f;
-
-    // 2 bits per element, 4 elements per byte, 128 elements per 32-byte group
-    const short base = il * 16;
-    for (int k = 0; k < 16; k++) {
-        const int i = base + k;
-        const int byte = ((i >> 7) & 1) * 32 + (i & 31);
-        const int l = (i >> 5) & 3;
-        reg_f[k/4][k%4] = d * (float)(((qs[byte] >> (2*l)) & 3) - 1);
-    }
-
-    reg = (type4x4) reg_f;
-}
@@ -1,250 +0,0 @@
-#include "common.h"
-
-constant short FC_gated_delta_net_ne20 [[function_constant(FC_GATED_DELTA_NET + 0)]];
-constant short FC_gated_delta_net_ne30 [[function_constant(FC_GATED_DELTA_NET + 1)]];
-constant short FC_gated_delta_net_K    [[function_constant(FC_GATED_DELTA_NET + 2)]];
-
-#if 1
-template<short NSG>
-kernel void kernel_gated_delta_net_impl(
-        constant ggml_metal_kargs_gated_delta_net & args,
-        device const char * q,
-        device const char * k,
-        device const char * v,
-        device const char * g,
-        device const char * b,
-        device const char * s,
-        device       char * dst,
-        uint3 tgpig[[threadgroup_position_in_grid]],
-        uint3 tpitg[[thread_position_in_threadgroup]],
-        uint3   ntg[[threads_per_threadgroup]])  {
-#define S_v FC_gated_delta_net_ne20
-#define G   FC_gated_delta_net_ne30
-#define K   FC_gated_delta_net_K
-
-    const uint tx = tpitg.x;
-    const uint ty = tpitg.y;
-
-    const uint i23 = tgpig.z; // B (n_seqs)
-    const uint i21 = tgpig.y; // H (head)
-    const uint i20 = tgpig.x*NSG + ty; // row within S_v
-
-    const uint i01 = i21 % args.ne01;
-    const uint i11 = i21 % args.ne11;
-
-    const float scale = 1.0f / sqrt((float)S_v);
-
-    // input state layout [S_v, S_v, H, n_seqs] (s0 only): per-seq stride is H*D.
-    // state is stored transposed: M[i20][is] = S[is][i20], so row i20 is contiguous
-    const uint state_in_base = (i23*args.ne21 + i21)*S_v*S_v + i20*S_v;
-    device const float * s_ptr = (device const float *) (s) + state_in_base;
-
-    float ls[NSG];
-
-    FOR_UNROLL (short j = 0; j < NSG; j++) {
-        const short is = tx*NSG + j;
-        ls[j] = s_ptr[is];
-    }
-
-    device float * dst_attn = (device float *) (dst) + (i23*args.ne22*args.ne21 + i21)*S_v + i20;
-
-    device const float * q_ptr = (device const float *) (q + i23*args.nb03 + i01*args.nb01);
-    device const float * k_ptr = (device const float *) (k + i23*args.nb13 + i11*args.nb11);
-    device const float * v_ptr = (device const float *) (v + i23*args.nb23 + i21*args.nb21);
-
-    device const float * b_ptr = (device const float *) (b) + (i23*args.ne22*args.ne21 + i21);
-    device const float * g_ptr = (device const float *) (g) + (i23*args.ne22*args.ne21 + i21)*G;
-
-    // snapshot slot mapping: slot 0 = most recent state, slot s = s tokens back.
-    // When n_tokens < K, only slots 0..n_tokens-1 are written; older slots are caller-owned.
-
-    // output state base offset: after attention scores
-    const uint attn_size = args.ne22 * args.ne21 * S_v * args.ne23;
-    // output state per-slot size: S_v * S_v * H * n_seqs
-    const uint state_size_per_snap = S_v * S_v * args.ne21 * args.ne23;
-    // per-(seq,head) offset within a slot
-    const uint state_out_base = (i23*args.ne21 + i21)*S_v*S_v + i20*S_v;
-
-    for (short t = 0; t < args.ne22; t++) {
-        float s_k = 0.0f;
-
-        if (G == 1) {
-            const float g_exp = exp(g_ptr[0]);
-
-            FOR_UNROLL (short j = 0; j < NSG; j++) {
-                const short is = tx*NSG + j;
-                ls[j] *= g_exp;
-
-                s_k += ls[j]*k_ptr[is];
-            }
-        } else {
-            // KDA
-            FOR_UNROLL (short j = 0; j < NSG; j++) {
-                const short is = tx*NSG + j;
-                ls[j] *= exp(g_ptr[is]);
-
-                s_k += ls[j]*k_ptr[is];
-            }
-        }
-
-        s_k = simd_sum(s_k);
-
-        const float d = (v_ptr[i20] - s_k)*b_ptr[0];
-
-        float y = 0.0f;
-
-        FOR_UNROLL (short j = 0; j < NSG; j++) {
-            const short is = tx*NSG + j;
-            ls[j] += k_ptr[is]*d;
-
-            y += ls[j]*q_ptr[is];
-        }
-
-        y = simd_sum(y);
-
-        if (tx == 0) {
-            dst_attn[t*args.ne21*S_v] = y*scale;
-        }
-
-        q_ptr += args.ns02;
-        k_ptr += args.ns12;
-        v_ptr += args.ns22;
-
-        b_ptr += args.ne21;
-        g_ptr += args.ne21*G;
-
-        if (K > 1) {
-            const int target_slot = (int)args.ne22 - 1 - (int)t;
-            if (target_slot >= 0 && target_slot < (int)K) {
-                device float * dst_state = (device float *) (dst) + attn_size + (uint)target_slot * state_size_per_snap + state_out_base;
-                FOR_UNROLL (short j = 0; j < NSG; j++) {
-                    const short is = tx*NSG + j;
-                    dst_state[is] = ls[j];
-                }
-            }
-        }
-    }
-
-    if (K == 1) {
-        device float * dst_state = (device float *) (dst) + attn_size + state_out_base;
-        FOR_UNROLL (short j = 0; j < NSG; j++) {
-            const short is = tx*NSG + j;
-            dst_state[is] = ls[j];
-        }
-    }
-
-#undef S_v
-#undef G
-#undef K
-}
-
-typedef decltype(kernel_gated_delta_net_impl<4>) kernel_gated_delta_net_t;
-
-template [[host_name("kernel_gated_delta_net_f32_1")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<1>;
-template [[host_name("kernel_gated_delta_net_f32_2")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<2>;
-template [[host_name("kernel_gated_delta_net_f32_4")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<4>;
-
-#else
-// a simplified version of the above
-// no performance improvement, so keep the above version for now
-
-template<typename T, short NSG>
-kernel void kernel_gated_delta_net_impl(
-        constant ggml_metal_kargs_gated_delta_net & args,
-        device const char * q,
-        device const char * k,
-        device const char * v,
-        device const char * g,
-        device const char * b,
-        device const char * s,
-        device       char * dst,
-        uint3 tgpig[[threadgroup_position_in_grid]],
-        uint3 tpitg[[thread_position_in_threadgroup]],
-        uint3   ntg[[threads_per_threadgroup]])  {
-#define S_v FC_gated_delta_net_ne20
-#define G   FC_gated_delta_net_ne30
-
-    const uint tx = tpitg.x;
-    const uint ty = tpitg.y;
-
-    const uint i23 = tgpig.z; // B
-    const uint i21 = tgpig.y; // H
-    const uint i20 = tgpig.x*NSG + ty;
-
-    const uint i01 = i21 % args.ne01;
-    const uint i11 = i21 % args.ne11;
-
-    const float scale = 1.0f / sqrt((float)S_v);
-
-    device const float * s_ptr = (device const float *) (s) + (i23*args.ne21 + i21)*S_v*S_v + i20;
-
-    float lsf[NSG];
-
-    FOR_UNROLL (short j = 0; j < NSG; j++) {
-        const short is = tx*NSG + j;
-        lsf[j] = s_ptr[is*S_v];
-    }
-
-    thread T * ls = (thread T *) (lsf);
-
-    device float * dst_attn = (device float *) (dst) + (i23*args.ne22*args.ne21 + i21)*S_v + i20;
-
-    device const float * q_ptr = (device const float *) (q + i23*args.nb03 + i01*args.nb01);
-    device const float * k_ptr = (device const float *) (k + i23*args.nb13 + i11*args.nb11);
-    device const float * v_ptr = (device const float *) (v + i23*args.nb23 + i21*args.nb21);
-
-    device const float * b_ptr  = (device const float *) (b) + (i23*args.ne22*args.ne21 + i21);
-    device const float * g_ptr  = (device const float *) (g) + (i23*args.ne22*args.ne21 + i21)*G;
-
-    for (short t = 0; t < args.ne22; t++) {
-        device const T * qt_ptr = (device const T *) (q_ptr);
-        device const T * kt_ptr = (device const T *) (k_ptr);
-        device const T * gt_ptr = (device const T *) (g_ptr);
-
-        if (G == 1) {
-            *ls *= exp(g_ptr[0]);
-        } else {
-            // KDA
-            *ls *= exp(gt_ptr[tx]);
-        }
-
-        const float s_k = simd_sum(dot(*ls, kt_ptr[tx]));
-
-        const float d = (v_ptr[i20] - s_k)*b_ptr[0];
-
-        *ls += kt_ptr[tx]*d;
-
-        const float y = simd_sum(dot(*ls, qt_ptr[tx]));
-
-        if (tx == 0) {
-            *dst_attn = y*scale;
-        }
-
-        q_ptr += args.ns02;
-        k_ptr += args.ns12;
-        v_ptr += args.ns22;
-
-        b_ptr += args.ne21;
-        g_ptr += args.ne21*G;
-
-        dst_attn += args.ne21*S_v;
-    }
-
-    device float * dst_state  = (device float *) (dst) + args.ne23*args.ne22*args.ne21*S_v + (i23*args.ne21 + i21)*S_v*S_v + i20;
-    device T     * dstt_state = (device T     *) (dst_state);
-
-    FOR_UNROLL (short j = 0; j < NSG; j++) {
-        const short is = tx*NSG + j;
-        dst_state[is*S_v] = lsf[j];
-    }
-
-#undef S_v
-#undef G
-}
-
-typedef decltype(kernel_gated_delta_net_impl<float4, 4>) kernel_gated_delta_net_t;
-
-template [[host_name("kernel_gated_delta_net_f32_1")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<float,  1>;
-template [[host_name("kernel_gated_delta_net_f32_2")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<float2, 2>;
-template [[host_name("kernel_gated_delta_net_f32_4")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<float4, 4>;
-#endif
@@ -1,595 +0,0 @@
-#include "common.h"
-
-kernel void kernel_argmax_f32(
-        constant ggml_metal_kargs_argmax & args,
-        device   const char * src0,
-        device         char * dst,
-        threadgroup    char * shmem [[threadgroup(0)]],
-        uint  tgpig[[threadgroup_position_in_grid]],
-        uint  tpitg[[thread_position_in_threadgroup]],
-        uint  sgitg[[simdgroup_index_in_threadgroup]],
-        uint  tiisg[[thread_index_in_simdgroup]],
-        uint    ntg[[threads_per_threadgroup]]) {
-    device const float * x_row = (device const float *) ((device const char *) src0 + tgpig * args.nb01);
-
-    float   lmax = -INFINITY;
-    int32_t larg = -1;
-
-    for (int i00 = tpitg; i00 < args.ne00; i00 += ntg) {
-        if (x_row[i00] > lmax) {
-            lmax = x_row[i00];
-            larg = i00;
-        }
-    }
-
-    // find the argmax value in the block
-    float max_val = simd_max(lmax);
-    int32_t arg_val = simd_max(select(-1, larg, lmax == max_val));
-
-    device int32_t * dst_i32 = (device int32_t *) dst;
-
-    threadgroup   float * shared_maxval = (threadgroup   float *) shmem;
-    threadgroup int32_t * shared_argmax = (threadgroup int32_t *) shmem + N_SIMDWIDTH;
-
-    if (ntg > N_SIMDWIDTH) {
-        if (sgitg == 0) {
-            shared_maxval[tiisg] = -INFINITY;
-            shared_argmax[tiisg] = -1;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        if (tiisg == 0) {
-            shared_maxval[sgitg] = max_val;
-            shared_argmax[sgitg] = arg_val;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        max_val = shared_maxval[tiisg];
-        arg_val = shared_argmax[tiisg];
-
-        float max_val_reduced   = simd_max(max_val);
-        int32_t arg_val_reduced = simd_max(select(-1, arg_val, max_val == max_val_reduced));
-
-        dst_i32[tgpig] = arg_val_reduced;
-
-        return;
-    }
-
-    dst_i32[tgpig] = arg_val;
-}
-
-kernel void kernel_diag_f32(
-        constant ggml_metal_kargs_diag & args,
-        device   const char * src0,
-        device         char * dst,
-        uint3  tgpig[[threadgroup_position_in_grid]],
-        ushort tiitg[[thread_index_in_threadgroup]]) {
-    constexpr short NW = N_SIMDWIDTH;
-
-    const int32_t i3 = tgpig.z;
-    const int32_t i2 = tgpig.y;
-    const int32_t i1 = tgpig.x;
-
-    device const float * src0_ptr = (device const float *)(src0 +                i2*args.nb02 + i3*args.nb03);
-    device       float * dst_ptr  = (device       float *)(dst  + i1*args.nb01 + i2*args.nb2  + i3*args.nb3);
-
-    for (int i0 = tiitg; i0 < args.ne0; i0 += NW) {
-        dst_ptr[i0] = i0 == i1 ? src0_ptr[i0] : 0.0f;
-    }
-}
-
-kernel void kernel_roll_f32(
-    constant ggml_metal_kargs_roll & args,
-    device  const char * src0,
-    device        char * dst,
-    uint3 tgpig[[threadgroup_position_in_grid]],
-    uint3 tpitg[[thread_position_in_threadgroup]],
-    uint3   ntg[[threads_per_threadgroup]]) {
-
-    const int64_t i3 = tgpig.z;
-    const int64_t i2 = tgpig.y;
-    const int64_t i1 = tgpig.x;
-
-    device const float * src0_ptr = (device const float *) src0;
-    device       float * dst_ptr  = (device       float *) dst;
-
-    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
-        // apply shifts and wrap around
-        int64_t i00 = i0 - args.s0;
-        int64_t i01 = i1 - args.s1;
-        int64_t i02 = i2 - args.s2;
-        int64_t i03 = i3 - args.s3;
-
-        if (i00 < 0) { i00 += args.ne00; } else if (i00 >= args.ne00) { i00 -= args.ne00; }
-        if (i01 < 0) { i01 += args.ne01; } else if (i01 >= args.ne01) { i01 -= args.ne01; }
-        if (i02 < 0) { i02 += args.ne02; } else if (i02 >= args.ne02) { i02 -= args.ne02; }
-        if (i03 < 0) { i03 += args.ne03; } else if (i03 >= args.ne03) { i03 -= args.ne03; }
-
-        int64_t src_idx = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00 + i00;
-        int64_t dst_idx = i3 *args.ne2 *args.ne1 *args.ne0  + i2 *args.ne1 *args.ne0  + i1 *args.ne0  + i0;
-
-        dst_ptr[dst_idx] = src0_ptr[src_idx];
-    }
-}
-
-template <typename T>
-kernel void kernel_pad_impl(
-    constant ggml_metal_kargs_pad & args,
-    device  const char * src0,
-    device        char * dst,
-    uint3 tgpig[[threadgroup_position_in_grid]],
-    uint3 tpitg[[thread_position_in_threadgroup]],
-    uint3   ntg[[threads_per_threadgroup]]) {
-    const int32_t i3 = tgpig.z;
-    const int32_t i2 = tgpig.y;
-    const int32_t k0 = tgpig.x/args.ne1;
-    const int32_t i1 = tgpig.x - k0*args.ne1;
-
-    const int32_t i03 = i3;
-    const int32_t i02 = i2;
-    const int32_t i01 = i1;
-
-    device const T * src0_ptr = (device const T *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
-    device       T * dst_ptr  = (device       T *) (dst  +  i3*args.nb3  +  i2*args.nb2  +  i1*args.nb1);
-
-    for (int32_t l0 = 0; l0 < 1024; l0 += ntg.x) {
-        const int32_t i0 = k0*1024 + tpitg.x + l0;
-        if (i0 >= args.ne0) {
-            break;
-        }
-
-        if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
-            dst_ptr[i0] = src0_ptr[i0];
-        } else {
-            dst_ptr[i0] = 0.0f;
-        }
-    }
-}
-
-typedef decltype(kernel_pad_impl<float>) kernel_pad_t;
-
-template [[host_name("kernel_pad_f32")]]   kernel kernel_pad_t kernel_pad_impl<float>;
-template [[host_name("kernel_pad_f32_4")]] kernel kernel_pad_t kernel_pad_impl<float4>;
-
-// TODO: this is slow - optimize
-kernel void kernel_pad_reflect_1d_f32(
-    constant   ggml_metal_kargs_pad_reflect_1d & args,
-    device  const char * src0,
-    device        char * dst,
-    uint3 tgpig[[threadgroup_position_in_grid]],
-    uint3  tgpg[[threadgroups_per_grid]],
-    uint3 tpitg[[thread_position_in_threadgroup]],
-    uint3   ntg[[threads_per_threadgroup]]) {
-
-    const int64_t i3 = tgpig.z;
-    const int64_t i2 = tgpig.y;
-    const int64_t i1 = tgpig.x;
-
-    const int64_t i03 = i3;
-    const int64_t i02 = i2;
-    const int64_t i01 = i1;
-
-    device const float * src0_ptr = (device const float *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
-    device       float * dst_ptr  = (device       float *) (dst  +  i3*args.nb3  +  i2*args.nb2  +  i1*args.nb1);
-
-    if (i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
-        for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
-            if (i0 < args.p0) {
-                dst_ptr[i0] = src0_ptr[args.p0 - i0];
-            } else if (i0 < args.ne0 - args.p1) {
-                dst_ptr[i0] = src0_ptr[i0 - args.p0];
-            } else {
-                dst_ptr[i0] = src0_ptr[(args.ne0 - args.p1 - args.p0) - (args.p1 + 1 - (args.ne0 - i0)) - 1];
-            }
-        }
-    }
-}
-
-kernel void kernel_arange_f32(
-    constant   ggml_metal_kargs_arange & args,
-    device        char * dst,
-    uint3 tgpig[[threadgroup_position_in_grid]],
-    uint3 tpitg[[thread_position_in_threadgroup]],
-    uint3   ntg[[threads_per_threadgroup]]) {
-
-    device float * dst_ptr = (device float *) dst;
-
-    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
-        dst_ptr[i0] = args.start + args.step * i0;
-    }
-}
-
-kernel void kernel_timestep_embedding_f32(
-    constant  ggml_metal_kargs_timestep_embedding & args,
-    device  const char * src0,
-    device        char * dst,
-    uint3 tgpig[[threadgroup_position_in_grid]],
-    uint3 tpitg[[thread_position_in_threadgroup]],
-    uint3   ntg[[threads_per_threadgroup]]) {
-
-    int i = tgpig.x;
-    device float * embed_data = (device float *)(dst + i*args.nb1);
-
-    int half_ = args.dim / 2;
-    for (int j = tpitg.x; j < half_; j += ntg.x) {
-        float timestep = ((device float *)src0)[i];
-        float freq = (float)exp(-log((float)args.max_period) * j / half_);
-        float arg = timestep * freq;
-        embed_data[j        ] = cos(arg);
-        embed_data[j + half_] = sin(arg);
-    }
-
-    if (args.dim % 2 != 0 && tpitg.x == 0) {
-        embed_data[2 * half_] = 0.f;
-    }
-}
-
-kernel void kernel_opt_step_adamw_f32(
-        constant    ggml_metal_kargs_opt_step_adamw & args,
-        device       float * x,
-        device const float * g,
-        device       float * g_m,
-        device       float * g_v,
-        device const float * pars,
-        uint        gid[[thread_position_in_grid]]) {
-
-    if (gid >= args.np) {
-        return;
-    }
-
-    const float alpha  = pars[0];
-    const float beta1  = pars[1];
-    const float beta2  = pars[2];
-    const float eps    = pars[3];
-    const float wd     = pars[4];
-    const float beta1h = pars[5];
-    const float beta2h = pars[6];
-
-    const float gi = g[gid];
-    const float gmi = g_m[gid] * beta1 +      gi * (1.0f - beta1);
-    const float gvi = g_v[gid] * beta2 + gi * gi * (1.0f - beta2);
-
-    g_m[gid] = gmi;
-    g_v[gid] = gvi;
-
-    const float mh =      gmi * beta1h;
-    const float vh = sqrt(gvi * beta2h) + eps;
-
-    x[gid] = x[gid] * (1.0f - alpha * wd) - alpha * mh / vh;
-}
-
-kernel void kernel_opt_step_sgd_f32(
-        constant    ggml_metal_kargs_opt_step_sgd & args,
-        device       float * x,
-        device const float * g,
-        device const float * pars,
-        uint        gid[[thread_position_in_grid]]) {
-
-    if (gid >= args.np) {
-        return;
-    }
-
-    x[gid] = x[gid] * (1.0f - pars[0] * pars[1]) - pars[0] * g[gid];
-}
-
-template<typename T>
-kernel void kernel_memset(
-        constant ggml_metal_kargs_memset & args,
-        device T * dst,
-        uint tpig[[thread_position_in_grid]]) {
-    dst[tpig] = args.val;
-}
-
-typedef decltype(kernel_memset<int64_t>) kernel_memset_t;
-
-template [[host_name("kernel_memset_i64")]] kernel kernel_memset_t kernel_memset<int64_t>;
-
-constant short FC_count_equal_nsg [[function_constant(FC_COUNT_EQUAL + 0)]];
-
-template<typename T>
-kernel void kernel_count_equal(
-        constant ggml_metal_kargs_count_equal & args,
-        device   const char * src0,
-        device   const char * src1,
-        device   atomic_int * dst,
-        threadgroup int32_t * shmem_i32 [[threadgroup(0)]],
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort  sgitg[[simdgroup_index_in_threadgroup]],
-        ushort  tiisg[[thread_index_in_simdgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    const short NSG = FC_count_equal_nsg;
-
-    const int i3 = tgpig.z;
-    const int i2 = tgpig.y;
-    const int i1 = tgpig.x;
-
-    if (i3 >= args.ne03 || i2 >= args.ne02 || i1 >= args.ne01) {
-        return;
-    }
-
-    int sum = 0;
-
-    device const char * base0 = src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03;
-    device const char * base1 = src1 + i1*args.nb11 + i2*args.nb12 + i3*args.nb13;
-
-    for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) {
-        const T v0 = *(device const T *)(base0 + i0*args.nb00);
-        const T v1 = *(device const T *)(base1 + i0*args.nb10);
-        sum += (v0 == v1);
-    }
-
-    sum = simd_sum(sum);
-
-    if (tiisg == 0) {
-        shmem_i32[sgitg] = sum;
-    }
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    if (sgitg == 0) {
-        float v = 0.0f;
-        if (tpitg.x < NSG) {
-            v = shmem_i32[tpitg.x];
-        }
-
-        float total = simd_sum(v);
-        if (tpitg.x == 0) {
-            atomic_fetch_add_explicit(dst, (int32_t) total, memory_order_relaxed);
-        }
-    }
-}
-
-typedef decltype(kernel_count_equal<int32_t>) kernel_count_equal_t;
-
-template [[host_name("kernel_count_equal_i32")]] kernel kernel_count_equal_t kernel_count_equal<int32_t>;
-
-template <typename T>
-kernel void kernel_snake(
-        constant ggml_metal_kargs_snake & args,
-        device const T     * x,
-        device const float * a,
-        device const float * inv_b,
-        device       T     * dst,
-        uint         tgpig [[threadgroup_position_in_grid]],
-        uint         tpitg [[thread_position_in_threadgroup]],
-        uint         ntg   [[threads_per_threadgroup]]) {
-
-    const int idx = tgpig * ntg + tpitg;
-    if (idx >= args.T * args.C) {
-        return;
-    }
-
-    const int   c  = idx / args.T;  // x is [T, C], a / inv_b collapse to [1, C]
-    const float xi = float(x[idx]);
-    const float si = sin(a[c] * xi);
-    dst[idx] = T(xi + si * si * inv_b[c]);
-}
-
-template [[host_name("kernel_snake_f32")]]  kernel void kernel_snake<float>(constant ggml_metal_kargs_snake &, device const float *, device const float *, device const float *, device float *, uint, uint, uint);
-template [[host_name("kernel_snake_f16")]]  kernel void kernel_snake<half>(constant ggml_metal_kargs_snake &, device const half *, device const float *, device const float *, device half *, uint, uint, uint);
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_snake_bf16")]] kernel void kernel_snake<bfloat>(constant ggml_metal_kargs_snake &, device const bfloat *, device const float *, device const float *, device bfloat *, uint, uint, uint);
-#endif
-
-template<int N>
-kernel void kernel_fwht_f32(
-        constant ggml_metal_kargs_fwht & args,
-        device const float * src,
-        device float * dst,
-        uint3  tgpig[[threadgroup_position_in_grid]],
-        ushort sgitg[[simdgroup_index_in_threadgroup]],
-        ushort tiisg[[thread_index_in_simdgroup]],
-        ushort3  ntg[[threads_per_threadgroup]]) {
-
-    constexpr int NW = N_SIMDWIDTH;
-    constexpr int NE = N / NW;
-
-    const float scale = 1.0f / sqrt((float) N);
-
-    const int sg_per_tg = ntg.x / NW;
-    const int64_t r = tgpig.x * sg_per_tg + sgitg;
-    if (r >= args.nrows) {
-        return;
-    }
-
-    src += r * N;
-    dst += r * N;
-
-    const int lane = tiisg;
-
-    float reg[NE];
-    for (int i = 0; i < NE; i++) {
-        reg[i] = src[i*NW + lane]*scale;
-    }
-    for (int i = 1; i < NW; i *= 2) {
-        for (int j = 0; j < NE; j++) {
-            const float val = reg[j];
-            const float val2 = simd_shuffle_xor(val, i);
-            reg[j] = (lane & i) == 0 ? val2 + val : val2 - val;
-        }
-    }
-
-    for (int i = NW; i < N; i *= 2) {
-        const int step = i / NW;
-        for (int j = 0; j < NE; j += (2 * step)) {
-            for (int k = 0; k < step; k++) {
-                const float x = reg[j + k ];
-                const float y = reg[j + k + step];
-                reg[j + k]        = x + y;
-                reg[j + k + step] = x - y;
-            }
-        }
-    }
-
-    for (int i = 0; i < NE; i++) {
-        dst[i*NW + lane] = reg[i];
-    }
-}
-
-typedef decltype(kernel_fwht_f32<64>) kernel_fwht_t;
-
-template [[host_name("kernel_fwht_f32_64")]]  kernel kernel_fwht_t kernel_fwht_f32<64>;
-template [[host_name("kernel_fwht_f32_128")]] kernel kernel_fwht_t kernel_fwht_f32<128>;
-template [[host_name("kernel_fwht_f32_256")]] kernel kernel_fwht_t kernel_fwht_f32<256>;
-template [[host_name("kernel_fwht_f32_512")]] kernel kernel_fwht_t kernel_fwht_f32<512>;
-
-kernel void kernel_dsv4_hc_comb_f32(
-        constant ggml_metal_kargs_dsv4_hc_comb & args,
-        device const char * mixes,
-        device const char * scale,
-        device const char * base,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort  tiisg[[thread_index_in_simdgroup]],
-        ushort  sgitg[[simdgroup_index_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    constexpr ushort hc = 4;
-    constexpr ushort comb_offset = 2*hc;
-
-    const int it = tgpig.x*ntg.y + sgitg;
-    if (it >= args.n_tokens) {
-        return;
-    }
-
-    float scale_lane = 0.0f;
-    if (tiisg == 0) {
-        scale_lane = *(device const float *) (scale + 2*args.nb_s0);
-    }
-    const float scale_comb = simd_shuffle(scale_lane, 0);
-
-    float v = 0.0f;
-    if (tiisg < hc*hc) {
-        v = *(device const float *) (mixes + (comb_offset + tiisg)*args.nb_m0 + it*args.nb_m1)*scale_comb
-          + *(device const float *) (base   + (comb_offset + tiisg)*args.nb_b0);
-    }
-
-    // Softmax across destinations (the four contiguous lanes for each source).
-    float vmax = max(v, simd_shuffle_xor(v, 1));
-    vmax = max(vmax, simd_shuffle_xor(vmax, 2));
-    v = exp(v - vmax);
-
-    float sum = v + simd_shuffle_xor(v, 1);
-    sum += simd_shuffle_xor(sum, 2);
-    v = v/sum + args.eps;
-
-    // Normalize columns: equal destination indices are four lanes apart.
-    sum = v + simd_shuffle_xor(v, 4);
-    sum += simd_shuffle_xor(sum, 8);
-    v /= sum + args.eps;
-
-    for (int i = 1; i < args.n_iter; ++i) {
-        sum = v + simd_shuffle_xor(v, 1);
-        sum += simd_shuffle_xor(sum, 2);
-        v /= sum + args.eps;
-
-        sum = v + simd_shuffle_xor(v, 4);
-        sum += simd_shuffle_xor(sum, 8);
-        v /= sum + args.eps;
-    }
-
-    if (tiisg < hc*hc) {
-        const ushort idst = tiisg & 3;
-        const ushort isrc = tiisg >> 2;
-        *(device float *) (dst + idst*args.nb_d0 + isrc*args.nb_d1 + it*args.nb_d2) = v;
-    }
-}
-
-kernel void kernel_dsv4_hc_pre_f32(
-        constant ggml_metal_kargs_dsv4_hc_pre & args,
-        device const char * x,
-        device const char * weights,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort  tiisg[[thread_index_in_simdgroup]],
-        ushort  sgitg[[simdgroup_index_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    constexpr ushort hc = 4;
-
-    const int it = tgpig.y;
-    const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg;
-
-    float weight_lane = 0.0f;
-    if (tiisg < hc) {
-        weight_lane = *(device const float *) (weights + tiisg*args.nb_w0 + it*args.nb_w1);
-    }
-
-    float w[hc];
-    FOR_UNROLL (ushort ih = 0; ih < hc; ++ih) {
-        w[ih] = simd_shuffle(weight_lane, ih);
-    }
-
-    if (i0 >= args.n_embd) {
-        return;
-    }
-
-    device const char * xb = x + i0*args.nb_x0 + it*args.nb_x2;
-    float result = 0.0f;
-    FOR_UNROLL (ushort ih = 0; ih < hc; ++ih) {
-        result = fma(*(device const float *) (xb + ih*args.nb_x1), w[ih], result);
-    }
-
-    *(device float *) (dst + i0*args.nb_d0 + it*args.nb_d1) = result;
-}
-
-kernel void kernel_dsv4_hc_post_f32(
-        constant ggml_metal_kargs_dsv4_hc_post & args,
-        device const char * x,
-        device const char * residual,
-        device const char * post,
-        device const char * comb,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort  tiisg[[thread_index_in_simdgroup]],
-        ushort  sgitg[[simdgroup_index_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    constexpr ushort hc = 4;
-
-    const int it = tgpig.y;
-    const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg;
-
-    float coeff_lane = 0.0f;
-    if (tiisg < hc) {
-        coeff_lane = *(device const float *) (post + tiisg*args.nb_p0 + it*args.nb_p1);
-    } else if (tiisg < hc + hc*hc) {
-        const ushort idx  = tiisg - hc;
-        const ushort idst = idx & 3;
-        const ushort isrc = idx >> 2;
-        coeff_lane = *(device const float *) (comb + idst*args.nb_c0 + isrc*args.nb_c1 + it*args.nb_c2);
-    }
-
-    float post_reg[hc];
-    float comb_reg[hc][hc];
-    FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
-        post_reg[idst] = simd_shuffle(coeff_lane, idst);
-    }
-    FOR_UNROLL (ushort isrc = 0; isrc < hc; ++isrc) {
-        FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
-            comb_reg[isrc][idst] = simd_shuffle(coeff_lane, hc + idst + hc*isrc);
-        }
-    }
-
-    if (i0 >= args.n_embd) {
-        return;
-    }
-
-    const float xv = *(device const float *) (x + i0*args.nb_x0 + it*args.nb_x1);
-    float result[hc];
-    FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
-        result[idst] = xv*post_reg[idst];
-    }
-
-    device const char * rb = residual + i0*args.nb_r0 + it*args.nb_r2;
-    FOR_UNROLL (ushort isrc = 0; isrc < hc; ++isrc) {
-        const float rv = *(device const float *) (rb + isrc*args.nb_r1);
-        FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
-            result[idst] = fma(rv, comb_reg[isrc][idst], result[idst]);
-        }
-    }
-
-    FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
-        *(device float *) (dst + i0*args.nb_d0 + idst*args.nb_d1 + it*args.nb_d2) = result[idst];
-    }
-}
@@ -1,853 +0,0 @@
-#include "common.h"
-#include "dequantize.h"
-
-constant bool FC_mul_mm_bc_inp [[function_constant(FC_MUL_MM + 0)]];
-constant bool FC_mul_mm_bc_out [[function_constant(FC_MUL_MM + 1)]];
-constant short FC_mul_mm_ne12  [[function_constant(FC_MUL_MM + 2)]];
-constant short FC_mul_mm_ne13  [[function_constant(FC_MUL_MM + 3)]];
-constant short FC_mul_mm_r2    [[function_constant(FC_MUL_MM + 4)]];
-constant short FC_mul_mm_r3    [[function_constant(FC_MUL_MM + 5)]];
-
-// each block_q contains 16*nl weights
-#ifdef GGML_METAL_HAS_TENSOR
-template<
-    typename SA, typename SA_4x4, typename SA_8x8,
-    typename SB, typename SB_2x4, typename SB_8x8,
-    typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread SA_4x4 &),
-    typename T0, typename T0_4x4, typename T1, typename T1_2x4>
-kernel void kernel_mul_mm(
-        constant ggml_metal_kargs_mul_mm & args,
-        device const char * srcA,
-        device const char * srcB,
-        device       char * dst,
-        threadgroup  char * shmem [[threadgroup(0)]],
-        uint3  tgpig [[threadgroup_position_in_grid]],
-        ushort tiitg [[thread_index_in_threadgroup]],
-        ushort sgitg [[simdgroup_index_in_threadgroup]]) {
-    (void) sgitg;
-
-    // Matrix dimensions: A(M,K) x B(K,N) -> C(M,N)
-    const int K = args.ne00;
-    const int M = args.ne0;
-    const int N = args.ne1;
-
-    // Batch dimension handling
-    const int im = tgpig.z;
-    const int i12 = im % FC_mul_mm_ne12;
-    const int i13 = im / FC_mul_mm_ne12;
-
-    // Batch offsets for srcA and srcB
-    const uint64_t offset0 = (i12/FC_mul_mm_r2)*args.nb02 + (i13/FC_mul_mm_r3)*args.nb03;
-
-    // Tile dimensions
-    constexpr int NRB = SZ_SIMDGROUP * N_MM_BLOCK_X * N_MM_SIMD_GROUP_X;
-    constexpr int NRA = SZ_SIMDGROUP * N_MM_BLOCK_Y * N_MM_SIMD_GROUP_Y;
-
-    // Tile offsets in output matrix
-    const int ra = tgpig.y * NRA;
-    const int rb = tgpig.x * NRB;
-
-    // Threadgroup memory for dequantized A tile only
-    threadgroup SA * sa = (threadgroup SA *)(shmem);
-
-    // Work-item count for A loading
-    constexpr int A_WORK_ITEMS = NRA * N_MM_NK;
-    constexpr int NUM_THREADS = N_SIMDWIDTH * N_MM_SIMD_GROUP_X * N_MM_SIMD_GROUP_Y;
-
-    // tA wraps threadgroup memory
-    auto tA = tensor(sa, dextents<int32_t, 2>(N_MM_NK_TOTAL, NRA));
-
-    // tB wraps device memory directly
-    device T1 * ptrB = (device T1 *)(srcB + args.nb12*i12 + args.nb13*i13);
-    const int strideB = args.nb11 / sizeof(T1);
-    auto tB = tensor(ptrB, dextents<int32_t, 2>(K, N), array<int, 2>({1, strideB}));
-
-    // Configure matmul operation
-    // note: K is dynamic_extent (clamped to the valid range in PHASE 2), since a static
-    //       N_MM_NK_TOTAL K tile would read src1 out of bounds when K % N_MM_NK_TOTAL != 0
-    // ref: https://github.com/ggml-org/llama.cpp/pull/27064
-    mpp::tensor_ops::matmul2d<
-        mpp::tensor_ops::matmul2d_descriptor(
-            NRB, NRA, static_cast<int>(dynamic_extent), false, true, true,
-            mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
-        execution_simdgroups<N_MM_SIMD_GROUP_X * N_MM_SIMD_GROUP_Y>> mm;
-
-    auto cT = mm.get_destination_cooperative_tensor<decltype(tB), decltype(tA), float>();
-
-    // Accumulate partial results over K dimension
-    for (int loop_k = 0; loop_k < K; loop_k += N_MM_NK_TOTAL) {
-        // === PHASE 1: Dequantization of A into threadgroup memory ===
-        for (int work = tiitg; work < A_WORK_ITEMS; work += NUM_THREADS) {
-            const int row = work / N_MM_NK;
-            const int k_chunk = work % N_MM_NK;
-            const int k_pos = loop_k + k_chunk * 16;
-            const short k_base = k_chunk * 16;
-
-            // Bounds check: skip device read if row is out of matrix bounds
-            if (ra + row < M) {
-                if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
-                    // Element-wise reads when K is not aligned (nb01 not aligned for half4x4/float4x4).
-                    // MSL spec Table 2.5: half4x4 requires 8-byte alignment. When K is odd,
-                    // nb01 = K*2 is not 8-byte aligned, so odd-row pointers are misaligned.
-                    // Mirrors the legacy kernel's existing guard.
-                    device const T0 * row_ptr = (device const T0 *)(srcA + args.nb01 * (ra + row) + offset0);
-
-                    FOR_UNROLL (short i = 0; i < 16; i++) {
-                        sa[row * N_MM_NK_TOTAL + (k_base + i)] = (k_pos + i < K) ? (SA) row_ptr[k_pos + i] : (SA)0;
-                    }
-                } else {
-                    const int block_idx = k_pos / (16 * nl);
-                    const short il = (k_pos / 16) % nl;
-
-                    device const block_q * row_ptr = (device const block_q *)(srcA + args.nb01 * (ra + row) + offset0);
-
-                    SA_4x4 temp_a;
-                    dequantize_func(row_ptr + block_idx, il, temp_a);
-
-                    FOR_UNROLL (short i = 0; i < 16; i++) {
-                        // Zero-pad A for K positions beyond valid range (handles partial K iterations)
-                        sa[row * N_MM_NK_TOTAL + (k_base + i)] = (k_pos + i < K) ? temp_a[i/4][i%4] : (SA)0;
-                    }
-                }
-            } else {
-                // Zero-pad rows beyond matrix bounds
-                FOR_UNROLL (short i = 0; i < 16; i++) {
-                    sa[row * N_MM_NK_TOTAL + (k_base + i)] = (SA)0;
-                }
-            }
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        // === PHASE 2: Tensor matmul ===
-        // Clamp the K extent of both operand tensors to the remaining valid K range so
-        // the dynamic-K op never reads past the K extent of src1 (or the staged A tile).
-        const int kExt = min(N_MM_NK_TOTAL, K - loop_k);
-
-        auto tAv = tensor(sa, dextents<int32_t, 2>(kExt, NRA), array<int, 2>({1, N_MM_NK_TOTAL}));
-        auto tBv = tensor(ptrB + loop_k + rb * strideB, dextents<int32_t, 2>(kExt, N - rb), array<int, 2>({1, strideB}));
-
-        mm.run(tBv, tAv, cT);
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-    }
-
-    // Store result tile to output matrix (with batch offset)
-    // cT.store handles bounds checking via tD's extents (M, N)
-    device float * dstBatch = (device float *)dst + im * N * M;
-
-    auto tD = tensor(dstBatch, dextents<int32_t, 2>(M, N), array<int, 2>({1, M}));
-    cT.store(tD.slice(ra, rb));
-}
-
-#else
-
-template<
-    typename S0, typename S0_4x4, typename S0_8x8,
-    typename S1, typename S1_2x4, typename S1_8x8,
-    typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread S0_4x4 &),
-    typename T0, typename T0_4x4, typename T1, typename T1_2x4>
-kernel void kernel_mul_mm(
-        constant ggml_metal_kargs_mul_mm & args,
-        device const char * src0,
-        device const char * src1,
-        device       char * dst,
-        threadgroup  char * shmem [[threadgroup(0)]],
-        uint3  tgpig[[threadgroup_position_in_grid]],
-        ushort tiitg[[thread_index_in_threadgroup]],
-        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
-
-    threadgroup S0 * sa = (threadgroup S0 *)(shmem);
-    threadgroup S1 * sb = (threadgroup S1 *)(shmem + 4096);
-
-    constexpr int NR0 = 64;
-    constexpr int NR1 = 32;
-
-    constexpr int NK  = 32;
-    constexpr int NL0 = NK/16;
-    constexpr int NL1 = NK/8;
-
-    const int im = tgpig.z;
-    const int r0 = tgpig.y*NR0;
-    const int r1 = tgpig.x*NR1;
-
-    // if this block is of 64x32 shape or smaller
-    const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
-    const short nr1 = (args.ne1 - r1 < NR1) ? (args.ne1 - r1) : NR1;
-
-    // a thread shouldn't load data outside of the matrix
-    const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; // 0 .. 63
-    const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; // 0 .. 31
-
-    const short il0 = (tiitg % NL0);
-
-    short il = il0;
-
-    const int i12 = im % FC_mul_mm_ne12;
-    const int i13 = im / FC_mul_mm_ne12;
-
-    const uint64_t offset0 = (i12/FC_mul_mm_r2)*args.nb02 + (i13/FC_mul_mm_r3)*args.nb03;
-    const short    offset1 = il0/nl;
-
-    device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
-
-    const short iy = 8*(tiitg % NL1);
-
-    device const T1 * y = (device const T1 *)(src1
-        + args.nb13*i13
-        + args.nb12*i12
-        + args.nb11*(r1 + lr1)
-        + args.nb10*iy);
-
-    S0_8x8 ma[4];
-    S1_8x8 mb[2];
-
-    simdgroup_float8x8 mc[8];
-
-    for (short i = 0; i < 8; i++){
-        mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
-    }
-
-    for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
-        // load data and store to threadgroup memory
-        if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
-            threadgroup_barrier(mem_flags::mem_threadgroup);
-
-            // no need for dequantization
-            for (short i = 0; i < 16; i++) {
-                const short sx = 2*il0 + i/8;
-                const short sy = (tiitg/NL0)/8;
-
-              //const short lx = i%8;
-              //const short ly = (tiitg/NL0)%8;
-                const short lx = (tiitg/NL0)%8;
-                const short ly = i%8;
-
-                const short ib = 8*sx + sy;
-
-                *(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? *((device T0 *) x + i) : 0;
-            }
-        } else {
-            S0_4x4 temp_a;
-            dequantize_func(x, il, temp_a);
-
-            threadgroup_barrier(mem_flags::mem_threadgroup);
-
-            FOR_UNROLL (short i = 0; i < 16; i++) {
-                const short sx = 2*il0 + i/8;
-                const short sy = (tiitg/NL0)/8;
-
-              //const short lx = i%8;
-              //const short ly = (tiitg/NL0)%8;
-                const short lx = (tiitg/NL0)%8;
-                const short ly = i%8;
-
-                const short ib = 8*sx + sy;
-
-                // NOTE: this is massively slower.. WTF?
-                //sa[64*ib + 8*ly + lx] = temp_a[i/4][i%4];
-
-                *(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4];
-            }
-        }
-
-        if (FC_mul_mm_bc_inp) {
-            for (short i = 0; i < 8; ++i) {
-                const short sx = (tiitg%NL1);
-                const short sy = (tiitg/NL1)/8;
-
-                const short lx = i;
-                const short ly = (tiitg/NL1)%8;
-              //const short lx = (tiitg/NL1)%8;
-              //const short ly = i;
-
-                const short ib = 4*sx + sy;
-
-                *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
-            }
-        } else {
-            const short sx = (tiitg%NL1);
-            const short sy = (tiitg/NL1)/8;
-
-          //const short dx = sx;
-          //const short dy = sy;
-
-            const short ly = (tiitg/NL1)%8;
-
-            const short ib = 4*sx + sy;
-
-            *(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y));
-        }
-
-        il = (il + 2 < nl) ? il + 2 : il % 2;
-        x  = (il < 2) ? x + (2 + nl - 1)/nl : x;
-
-        y += NK;
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        // load matrices from threadgroup memory and conduct outer products
-        threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2));
-        threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2));
-
-        FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
-            simdgroup_barrier(mem_flags::mem_none);
-
-            FOR_UNROLL (short i = 0; i < 4; i++) {
-                simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
-            }
-
-            simdgroup_barrier(mem_flags::mem_none);
-
-            FOR_UNROLL (short i = 0; i < 2; i++) {
-                simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
-            }
-
-            simdgroup_barrier(mem_flags::mem_none);
-
-            FOR_UNROLL (short i = 0; i < 8; i++){
-                simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
-            }
-
-            lsma += 8*64;
-            lsmb += 4*64;
-        }
-    }
-
-    if (!FC_mul_mm_bc_out || (r0 + NR0 <= args.ne0 && r1 + NR1 <= args.ne1)) {
-        // if no bounds checks on the output are needed, we can directly write to device memory
-        device float * C = (device float *) dst +
-            (r0 + 32*(sgitg &  1)) + \
-            (r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0;
-
-        for (short i = 0; i < 8; i++) {
-            simdgroup_store(mc[i], C + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false);
-        }
-    } else {
-        // block is smaller than 64x32, we should avoid writing data outside of the matrix
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
-
-        for (short i = 0; i < 8; i++) {
-            simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        if (sgitg == 0) {
-            for (int j = tiitg; j < nr1; j += NR1) {
-                device float  * D  = (device float  *) dst + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0;
-                device float4 * D4 = (device float4 *) D;
-
-                threadgroup float  * C  = temp_str + (j*NR0);
-                threadgroup float4 * C4 = (threadgroup float4 *) C;
-
-                int i = 0;
-                for (; i < nr0/4; i++) {
-                    *(D4 + i) = *(C4 + i);
-                }
-
-                i *= 4;
-                for (; i < nr0; i++) {
-                    *(D + i) = *(C + i);
-                }
-            }
-        }
-    }
-}
-
-#endif // GGML_METAL_HAS_TENSOR
-
-template<short ne20> // n_expert_used
-kernel void kernel_mul_mm_id_map0(
-        constant ggml_metal_kargs_mul_mm_id_map0 & args,
-        device  const char * src2,
-        device        char * htpe,
-        device        char * hids,
-        threadgroup   char * shmem [[threadgroup(0)]],
-        ushort tpitg[[thread_position_in_threadgroup]],
-        ushort   ntg[[threads_per_threadgroup]]) {
-    const short ide = tpitg; // expert id
-
-    uint32_t n_all = 0;
-
-    device int32_t * ids_i32 = (device int32_t *) hids + ide*args.ne21;
-
-    for (int i21 = 0; i21 < args.ne21; i21 += ntg) { // n_tokens
-        if (i21 + tpitg < args.ne21) {
-            device const int32_t * src2_i32 = (device const int32_t *) (src2 + (i21 + tpitg)*args.nb21);
-
-            threadgroup uint16_t * sids = (threadgroup uint16_t *) shmem + tpitg*ne20;
-
-            #pragma unroll(ne20)
-            for (short i20 = 0; i20 < ne20; i20++) {
-                sids[i20] = src2_i32[i20];
-            }
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        for (short t = 0; t < ntg; t++) {
-            if (i21 + t >= args.ne21) {
-                break;
-            }
-
-            threadgroup const uint16_t * sids = (threadgroup const uint16_t *) shmem + t*ne20;
-
-            short sel = 0;
-            #pragma unroll(ne20)
-            for (short i20 = 0; i20 < ne20; i20++) {
-                sel += (sids[i20] == ide)*(i20 + 1);
-            }
-
-            ids_i32[n_all] = (i21 + t)*ne20 + sel - 1;
-
-            n_all += sel > 0;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-    }
-
-    device uint32_t * tpe_u32 = (device uint32_t *) (htpe);
-    tpe_u32[ide] = n_all;
-}
-
-typedef decltype(kernel_mul_mm_id_map0<1>) kernel_mul_mm_id_map0_t;
-
-template [[host_name("kernel_mul_mm_id_map0_ne20_1" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<1>;
-template [[host_name("kernel_mul_mm_id_map0_ne20_2" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<2>;
-template [[host_name("kernel_mul_mm_id_map0_ne20_4" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<4>;
-template [[host_name("kernel_mul_mm_id_map0_ne20_5" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<5>;
-template [[host_name("kernel_mul_mm_id_map0_ne20_6" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<6>;
-template [[host_name("kernel_mul_mm_id_map0_ne20_8" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<8>;
-template [[host_name("kernel_mul_mm_id_map0_ne20_10")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<10>;
-template [[host_name("kernel_mul_mm_id_map0_ne20_16")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<16>;
-template [[host_name("kernel_mul_mm_id_map0_ne20_22")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<22>;
-
-template<typename S0, typename S0_4x4, typename S0_8x8, typename S1, typename S1_2x4, typename S1_8x8, typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread S0_4x4 &), typename T0, typename T0_4x4, typename T1, typename T1_2x4>
-kernel void kernel_mul_mm_id(
-        constant ggml_metal_kargs_mul_mm_id & args,
-        device const char * src0,
-        device const char * src1,
-        device const char * htpe,
-        device const char * hids,
-        device       char * dst,
-        threadgroup  char * shmem [[threadgroup(0)]],
-        uint3  tgpig[[threadgroup_position_in_grid]],
-        ushort tiitg[[thread_index_in_threadgroup]],
-        ushort tiisg[[thread_index_in_simdgroup]],
-        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
-    threadgroup S0 * sa = (threadgroup S0 *)(shmem);
-    threadgroup S1 * sb = (threadgroup S1 *)(shmem + 4096);
-
-#ifdef GGML_METAL_HAS_TENSOR
-    threadgroup float * sc = (threadgroup float *)(shmem);
-#endif
-
-    constexpr int NR0 = 64;
-    constexpr int NR1 = 32;
-
-    constexpr int NK  = 32;
-    constexpr int NL0 = NK/16;
-    constexpr int NL1 = NK/8;
-
-    const int im = tgpig.z; // expert
-    const int r0 = tgpig.y*NR0;
-    const int r1 = tgpig.x*NR1;
-
-    device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe);
-    device const int32_t  * ids_i32 = (device const int32_t  *) (hids);
-
-    const int32_t neh1 = tpe_u32[im];
-
-    if (r1 >= neh1) {
-        return;
-    }
-
-    // if this block is of 64x32 shape or smaller
-    const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
-    const short nr1 = (    neh1 - r1 < NR1) ? (    neh1 - r1) : NR1;
-
-    // a thread shouldn't load data outside of the matrix
-    const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; // 0 .. 63
-    const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; // 0 .. 31
-
-    const short il0 = (tiitg % NL0);
-
-    short il = il0;
-
-    const int id = ids_i32[im*args.ne21 + r1 + lr1];
-
-    const short i11 = (id % args.ne20) % args.ne11;
-    const short i12 = (id / args.ne20);
-    const short i13 = 0;
-
-    const uint64_t offset0 = im*args.nb02 + i13*args.nb03;
-    const short    offset1 = il0/nl;
-
-    device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
-
-    const short iy = 8*(tiitg % NL1);
-
-    device const T1 * y = (device const T1 *)(src1
-        + args.nb13*i13
-        + args.nb12*i12
-        + args.nb11*i11
-        + args.nb10*iy);
-
-#ifndef GGML_METAL_HAS_TENSOR
-    S0_8x8 ma[4];
-    S1_8x8 mb[2];
-
-    simdgroup_float8x8 mc[8];
-
-    for (short i = 0; i < 8; i++){
-        mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
-    }
-#else
-    auto tA = tensor<threadgroup S0, dextents<int32_t, 2>, tensor_inline>(sa, dextents<int32_t, 2>(NK,  NR0));
-    auto tB = tensor<threadgroup S1, dextents<int32_t, 2>, tensor_inline>(sb, dextents<int32_t, 2>(NR1, NK ));
-
-    mpp::tensor_ops::matmul2d<
-        mpp::tensor_ops::matmul2d_descriptor(NR1, NR0, NK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
-        execution_simdgroups<4>> mm;
-
-    auto cT = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB), float>();
-#endif
-
-    for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
-#ifndef GGML_METAL_HAS_TENSOR
-        // load data and store to threadgroup memory
-        if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
-            threadgroup_barrier(mem_flags::mem_threadgroup);
-
-            // no need for dequantization
-            for (short i = 0; i < 16; i++) {
-                const short sx = 2*il0 + i/8;
-                const short sy = (tiitg/NL0)/8;
-
-              //const short lx = i%8;
-              //const short ly = (tiitg/NL0)%8;
-                const short lx = (tiitg/NL0)%8;
-                const short ly = i%8;
-
-                const short ib = 8*sx + sy;
-
-                *(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? (S0) *((device T0 *) x + i) : (S0) 0;
-            }
-        } else {
-            S0_4x4 temp_a;
-            dequantize_func(x, il, temp_a);
-
-            threadgroup_barrier(mem_flags::mem_threadgroup);
-
-            FOR_UNROLL (short i = 0; i < 16; i++) {
-                const short sx = 2*il0 + i/8;
-                const short sy = (tiitg/NL0)/8;
-
-              //const short lx = i%8;
-              //const short ly = (tiitg/NL0)%8;
-                const short lx = (tiitg/NL0)%8;
-                const short ly = i%8;
-
-                const short ib = 8*sx + sy;
-
-                // NOTE: this is massively slower.. WTF?
-                //sa[64*ib + 8*ly + lx] = temp_a[i/4][i%4];
-
-                *(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4];
-            }
-        }
-
-        if (FC_mul_mm_bc_inp) {
-            for (short i = 0; i < 8; ++i) {
-                const short sx = (tiitg%NL1);
-                const short sy = (tiitg/NL1)/8;
-
-                const short lx = i;
-                const short ly = (tiitg/NL1)%8;
-              //const short lx = (tiitg/NL1)%8;
-              //const short ly = i;
-
-                const short ib = 4*sx + sy;
-
-                *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
-            }
-        } else {
-            const short sx = (tiitg%NL1);
-            const short sy = (tiitg/NL1)/8;
-
-          //const short dx = sx;
-          //const short dy = sy;
-
-            const short ly = (tiitg/NL1)%8;
-
-            const short ib = 4*sx + sy;
-
-            *(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y));
-        }
-#else
-        // load data and store to threadgroup memory
-        if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
-            threadgroup_barrier(mem_flags::mem_threadgroup);
-
-            // no need for dequantization
-            for (short i = 0; i < 16; i++) {
-                const short sx = 2*il0 + i/8;
-                const short sy = (tiitg/NL0)/8;
-
-                const short lx = i%8;
-                const short ly = (tiitg/NL0)%8;
-                //const short lx = (tiitg/NL0)%8;
-                //const short ly = i%8;
-
-                *(sa + NK*(8*sy + ly) + 8*sx + lx) = loop_k + 16*il + i < args.ne00 ? *((device T0 *) x + i) : 0;
-            }
-        } else {
-            S0_4x4 temp_a;
-            dequantize_func(x, il, temp_a);
-
-            threadgroup_barrier(mem_flags::mem_threadgroup);
-
-            FOR_UNROLL (short i = 0; i < 16; i++) {
-                const short sx = 2*il0 + i/8;
-                const short sy = (tiitg/NL0)/8;
-
-                const short lx = i%8;
-                const short ly = (tiitg/NL0)%8;
-                //const short lx = (tiitg/NL0)%8;
-                //const short ly = i%8;
-
-                *(sa + NK*(8*sy + ly) + 8*sx + lx) = temp_a[i/4][i%4];
-            }
-        }
-
-        if (FC_mul_mm_bc_inp) {
-            for (short i = 0; i < 8; ++i) {
-                const short sx = (tiitg%NL1);
-                const short sy = (tiitg/NL1)/8;
-
-                const short lx = i;
-                const short ly = (tiitg/NL1)%8;
-                //const short lx = (tiitg/NL1)%8;
-                //const short ly = i;
-
-                *(sb + NK*(8*sy + ly) + 8*sx + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
-            }
-        } else {
-            const short sx = (tiitg%NL1);
-            const short sy = (tiitg/NL1)/8;
-
-            //const short lx = i;
-            const short ly = (tiitg/NL1)%8;
-            //const short lx = (tiitg/NL1)%8;
-            //const short ly = i;
-
-            *(threadgroup S1_2x4 *)(sb + NK*(8*sy + ly) + 8*sx) = (S1_2x4)(*((device T1_2x4 *) y));
-        }
-#endif
-
-        il = (il + 2 < nl) ? il + 2 : il % 2;
-        x  = (il < 2) ? x + (2 + nl - 1)/nl : x;
-
-        y += NK;
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-#ifndef GGML_METAL_HAS_TENSOR
-        // load matrices from threadgroup memory and conduct outer products
-        threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2));
-        threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2));
-
-        FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
-            simdgroup_barrier(mem_flags::mem_none);
-
-            FOR_UNROLL (short i = 0; i < 4; i++) {
-                simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
-            }
-
-            simdgroup_barrier(mem_flags::mem_none);
-
-            FOR_UNROLL (short i = 0; i < 2; i++) {
-                simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
-            }
-
-            simdgroup_barrier(mem_flags::mem_none);
-
-            FOR_UNROLL (short i = 0; i < 8; i++){
-                simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
-            }
-
-            lsma += 8*64;
-            lsmb += 4*64;
-        }
-#else
-        auto sA = tA.slice(0, 0);
-        auto sB = tB.slice(0, 0);
-
-        mm.run(sB, sA, cT);
-#endif
-    }
-
-    // block is smaller than 64x32, we should avoid writing data outside of the matrix
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-#ifdef GGML_METAL_HAS_TENSOR
-    auto tC = tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline>(sc, dextents<int32_t, 2>(NR0, NR1));
-    cT.store(tC);
-#else
-    threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
-
-    for (short i = 0; i < 8; i++) {
-        simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
-    }
-#endif
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    for (short j = sgitg; j < nr1; j += 4) {
-        const int id = ids_i32[im*args.ne21 + r1 + j];
-
-        const short ide = id % args.ne20;
-        const short idt = id / args.ne20;
-
-        device float  * D  = (device float  *) dst + r0 + ide*args.ne0 + idt*args.ne1*args.ne0;
-        device float4 * D4 = (device float4 *) D;
-
-        threadgroup float  * C  = (threadgroup float  *) shmem + j*NR0;
-        threadgroup float4 * C4 = (threadgroup float4 *) C;
-
-        int i = tiisg;
-        for (; i < nr0/4; i += 32) {
-            *(D4 + i) = *(C4 + i);
-        }
-
-        i = (4*(nr0/4)) + tiisg;
-        for (; i < nr0; i += 32) {
-            *(D + i) = *(C + i);
-        }
-    }
-}
-
-//
-// matrix-matrix multiplication
-//
-
-typedef decltype(kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, float, float2x4>) mul_mm_t;
-
-template [[host_name("kernel_mul_mm_f32_f32")]]     kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   float4x4,      1,     dequantize_f32,     float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_f16_f32")]]     kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   half4x4,       1,     dequantize_f16,     half,   half4x4,   float, float2x4>;
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_mul_mm_bf16_f32")]]    kernel mul_mm_t kernel_mul_mm<bfloat, bfloat4x4, simdgroup_bfloat8x8, bfloat, bfloat2x4, simdgroup_bfloat8x8, bfloat4x4,     1,     dequantize_bf16,    bfloat, bfloat4x4, float, float2x4>;
-#endif
-template [[host_name("kernel_mul_mm_q1_0_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q1_0,    8,     dequantize_q1_0,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_q2_0_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_0,    4,     dequantize_q2_0,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_q4_0_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_0,    2,     dequantize_q4_0,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_q4_1_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_1,    2,     dequantize_q4_1,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_q5_0_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_0,    2,     dequantize_q5_0,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_q5_1_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_1,    2,     dequantize_q5_1,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_q8_0_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q8_0,    2,     dequantize_q8_0,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_mxfp4_f32")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_mxfp4,   2,     dequantize_mxfp4,   float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_q2_K_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_K,    QK_NL, dequantize_q2_K,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_q3_K_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q3_K,    QK_NL, dequantize_q3_K,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_q4_K_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_K,    QK_NL, dequantize_q4_K,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_q5_K_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_K,    QK_NL, dequantize_q5_K,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_q6_K_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q6_K,    QK_NL, dequantize_q6_K,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_iq2_xxs_f32")]] kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_iq2_xs_f32")]]  kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xs,  QK_NL, dequantize_iq2_xs,  float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_iq3_xxs_f32")]] kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_xxs, QK_NL, dequantize_iq3_xxs, float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_iq3_s_f32")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_s,   QK_NL, dequantize_iq3_s,   float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_iq2_s_f32")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_s,   QK_NL, dequantize_iq2_s,   float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_iq1_s_f32")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_s,   QK_NL, dequantize_iq1_s,   float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_iq1_m_f32")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_m,   QK_NL, dequantize_iq1_m,   float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_iq4_nl_f32")]]  kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_nl,  2,     dequantize_iq4_nl,  float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_iq4_xs_f32")]]  kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_xs,  QK_NL, dequantize_iq4_xs,  float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_tq2_0_f32")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_tq2_0,   QK_NL, dequantize_tq2_0,   float,  float4x4,  float, float2x4>;
-
-template [[host_name("kernel_mul_mm_f32_f16")]]     kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   float4x4,      1,     dequantize_f32,     float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_f16_f16")]]     kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   half4x4,       1,     dequantize_f16,     half,   half4x4,   half, half2x4>;
-template [[host_name("kernel_mul_mm_q1_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q1_0,    8,     dequantize_q1_0,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_q2_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_0,    4,     dequantize_q2_0,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_q4_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_0,    2,     dequantize_q4_0,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_q4_1_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_1,    2,     dequantize_q4_1,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_q5_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_0,    2,     dequantize_q5_0,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_q5_1_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_1,    2,     dequantize_q5_1,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_q8_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q8_0,    2,     dequantize_q8_0,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_mxfp4_f16")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_mxfp4,   2,     dequantize_mxfp4,   float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_q2_K_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_K,    QK_NL, dequantize_q2_K,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_q3_K_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q3_K,    QK_NL, dequantize_q3_K,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_q4_K_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_K,    QK_NL, dequantize_q4_K,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_q5_K_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_K,    QK_NL, dequantize_q5_K,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_q6_K_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q6_K,    QK_NL, dequantize_q6_K,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_iq2_xxs_f16")]] kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_iq2_xs_f16")]]  kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xs,  QK_NL, dequantize_iq2_xs,  float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_iq3_xxs_f16")]] kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_xxs, QK_NL, dequantize_iq3_xxs, float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_iq3_s_f16")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_s,   QK_NL, dequantize_iq3_s,   float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_iq2_s_f16")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_s,   QK_NL, dequantize_iq2_s,   float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_iq1_s_f16")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_s,   QK_NL, dequantize_iq1_s,   float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_iq1_m_f16")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_m,   QK_NL, dequantize_iq1_m,   float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_iq4_nl_f16")]]  kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_nl,  2,     dequantize_iq4_nl,  float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_iq4_xs_f16")]]  kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_xs,  QK_NL, dequantize_iq4_xs,  float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_tq2_0_f16")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_tq2_0,   QK_NL, dequantize_tq2_0,   float,  float4x4,  half, half2x4>;
-
-//
-// indirect matrix-matrix multiplication
-//
-
-typedef decltype(kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, float, float2x4>) mul_mm_id;
-
-template [[host_name("kernel_mul_mm_id_f32_f32")]]     kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   float4x4,      1,     dequantize_f32,     float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_f16_f32")]]     kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   half4x4,       1,     dequantize_f16,     half,   half4x4,   float, float2x4>;
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_mul_mm_id_bf16_f32")]]    kernel mul_mm_id kernel_mul_mm_id<bfloat, bfloat4x4, simdgroup_bfloat8x8, bfloat, bfloat2x4, simdgroup_bfloat8x8, bfloat4x4,     1,     dequantize_bf16,    bfloat, bfloat4x4, float, float2x4>;
-#endif
-template [[host_name("kernel_mul_mm_id_q1_0_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q1_0,    8,     dequantize_q1_0,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_q2_0_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_0,    4,     dequantize_q2_0,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_q4_0_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_0,    2,     dequantize_q4_0,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_q4_1_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_1,    2,     dequantize_q4_1,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_q5_0_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_0,    2,     dequantize_q5_0,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_q5_1_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_1,    2,     dequantize_q5_1,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_q8_0_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q8_0,    2,     dequantize_q8_0,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_mxfp4_f32")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_mxfp4,   2,     dequantize_mxfp4,   float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_q2_K_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_K,    QK_NL, dequantize_q2_K,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_q3_K_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q3_K,    QK_NL, dequantize_q3_K,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_q4_K_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_K,    QK_NL, dequantize_q4_K,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_q5_K_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_K,    QK_NL, dequantize_q5_K,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_q6_K_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q6_K,    QK_NL, dequantize_q6_K,    float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_iq2_xxs_f32")]] kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_iq2_xs_f32")]]  kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xs,  QK_NL, dequantize_iq2_xs,  float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_iq3_xxs_f32")]] kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_xxs, QK_NL, dequantize_iq3_xxs, float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_iq3_s_f32")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_s,   QK_NL, dequantize_iq3_s,   float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_iq2_s_f32")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_s,   QK_NL, dequantize_iq2_s,   float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_iq1_s_f32")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_s,   QK_NL, dequantize_iq1_s,   float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_iq1_m_f32")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_m,   QK_NL, dequantize_iq1_m,   float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_iq4_nl_f32")]]  kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_nl,  2,     dequantize_iq4_nl,  float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_iq4_xs_f32")]]  kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_xs,  QK_NL, dequantize_iq4_xs,  float,  float4x4,  float, float2x4>;
-template [[host_name("kernel_mul_mm_id_tq2_0_f32")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_tq2_0,   QK_NL, dequantize_tq2_0,   float,  float4x4,  float, float2x4>;
-
-template [[host_name("kernel_mul_mm_id_f32_f16")]]     kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   float4x4,      1,     dequantize_f32,     float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_f16_f16")]]     kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   half4x4,       1,     dequantize_f16,     half,   half4x4,   half, half2x4>;
-template [[host_name("kernel_mul_mm_id_q1_0_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q1_0,    8,     dequantize_q1_0,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_q2_0_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_0,    4,     dequantize_q2_0,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_q4_0_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_0,    2,     dequantize_q4_0,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_q4_1_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_1,    2,     dequantize_q4_1,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_q5_0_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_0,    2,     dequantize_q5_0,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_q5_1_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_1,    2,     dequantize_q5_1,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_q8_0_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q8_0,    2,     dequantize_q8_0,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_mxfp4_f16")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_mxfp4,   2,     dequantize_mxfp4,   float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_q2_K_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_K,    QK_NL, dequantize_q2_K,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_q3_K_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q3_K,    QK_NL, dequantize_q3_K,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_q4_K_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_K,    QK_NL, dequantize_q4_K,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_q5_K_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_K,    QK_NL, dequantize_q5_K,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_q6_K_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q6_K,    QK_NL, dequantize_q6_K,    float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_iq2_xxs_f16")]] kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_iq2_xs_f16")]]  kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xs,  QK_NL, dequantize_iq2_xs,  float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_iq3_xxs_f16")]] kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_xxs, QK_NL, dequantize_iq3_xxs, float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_iq3_s_f16")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_s,   QK_NL, dequantize_iq3_s,   float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_iq2_s_f16")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_s,   QK_NL, dequantize_iq2_s,   float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_iq1_s_f16")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_s,   QK_NL, dequantize_iq1_s,   float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_iq1_m_f16")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_m,   QK_NL, dequantize_iq1_m,   float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_iq4_nl_f16")]]  kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_nl,  2,     dequantize_iq4_nl,  float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_iq4_xs_f16")]]  kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_xs,  QK_NL, dequantize_iq4_xs,  float,  float4x4,  half, half2x4>;
-template [[host_name("kernel_mul_mm_id_tq2_0_f16")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_tq2_0,   QK_NL, dequantize_tq2_0,   float,  float4x4,  half, half2x4>;
@@ -1,308 +0,0 @@
-#include "common.h"
-
-// F == 1 : norm (no fuse)
-// F == 2 : norm + mul
-// F == 3 : norm + mul + add
-template <typename T, short F>
-kernel void kernel_norm_fuse_impl(
-        constant ggml_metal_kargs_norm & args,
-        device const char * src0,
-        device const char * src1_0,
-        device const char * src1_1,
-        device       char * dst,
-        threadgroup float * shmem_f32 [[threadgroup(0)]],
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort  sgitg[[simdgroup_index_in_threadgroup]],
-        ushort  tiisg[[thread_index_in_simdgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    if (sgitg == 0) {
-        shmem_f32[tiisg] = 0.0f;
-    }
-
-    const int i01 = tgpig.x;
-    const int i02 = tgpig.y;
-    const int i03 = tgpig.z;
-
-    device const T * x = (device const T *) (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
-
-    device const T * f0 = (device const T *) (src1_0 + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]);
-    device const T * f1 = (device const T *) (src1_1 + (i03%args.nef3[2])*args.nbf3[2] + (i02%args.nef2[2])*args.nbf2[2] + (i01%args.nef1[2])*args.nbf1[2]);
-
-    T sumft(0.0f);
-
-    float sumf = 0.0f;
-
-    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
-        sumft += x[i00];
-    }
-    sumf = dot(sumft, T(1.0f));
-    sumf = simd_sum(sumf);
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    if (tiisg == 0) {
-        shmem_f32[sgitg] = sumf;
-    }
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    sumf = shmem_f32[tiisg];
-    sumf = simd_sum(sumf);
-
-    const float mean = sumf/args.ne00;
-
-    device T * y = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
-
-    sumf = 0.0f;
-    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
-        y[i00] = x[i00] - mean;
-        sumf += dot(y[i00], y[i00]);
-    }
-    sumf = simd_sum(sumf);
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    if (tiisg == 0) {
-        shmem_f32[sgitg] = sumf;
-    }
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    sumf = shmem_f32[tiisg];
-    sumf = simd_sum(sumf);
-
-    const float variance = sumf/args.ne00;
-
-    const float scale = 1.0f/sqrt(variance + args.eps);
-    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
-        if (F == 1) {
-            y[i00] = (y[i00]*scale);
-        }
-        if (F == 2) {
-            y[i00] = (y[i00]*scale)*f0[i00];
-        }
-        if (F == 3) {
-            y[i00] = (y[i00]*scale)*f0[i00] + f1[i00];
-        }
-    }
-}
-
-typedef decltype(kernel_norm_fuse_impl<float4, 1>) kernel_norm_fuse_t;
-
-template [[host_name("kernel_norm_f32")]]         kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float, 1>;
-template [[host_name("kernel_norm_mul_f32")]]     kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float, 2>;
-template [[host_name("kernel_norm_mul_add_f32")]] kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float, 3>;
-
-template [[host_name("kernel_norm_f32_4")]]         kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float4, 1>;
-template [[host_name("kernel_norm_mul_f32_4")]]     kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float4, 2>;
-template [[host_name("kernel_norm_mul_add_f32_4")]] kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float4, 3>;
-
-// F == 1 : rms_norm (no fuse)
-// F == 2 : rms_norm + mul
-// F == 3 : rms_norm + mul + add
-template <typename T, short F>
-kernel void kernel_rms_norm_fuse_impl(
-        constant ggml_metal_kargs_norm & args,
-        device const char * src0,
-        device const char * src1_0,
-        device const char * src1_1,
-        device       char * dst,
-        threadgroup float * shmem_f32 [[threadgroup(0)]],
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort  sgitg[[simdgroup_index_in_threadgroup]],
-        ushort  tiisg[[thread_index_in_simdgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    if (sgitg == 0) {
-        shmem_f32[tiisg] = 0.0f;
-    }
-
-    const int i01 = tgpig.x;
-    const int i02 = tgpig.y;
-    const int i03 = tgpig.z;
-
-    device const T * x = (device const T *) (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
-
-    device const T * f0 = (device const T *) (src1_0 + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]);
-    device const T * f1 = (device const T *) (src1_1 + (i03%args.nef3[2])*args.nbf3[2] + (i02%args.nef2[2])*args.nbf2[2] + (i01%args.nef1[2])*args.nbf1[2]);
-
-    float sumf = 0.0f;
-
-    // parallel sum
-    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
-        sumf += dot(x[i00], x[i00]);
-    }
-    sumf = simd_sum(sumf);
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    if (tiisg == 0) {
-        shmem_f32[sgitg] = sumf;
-    }
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    sumf = shmem_f32[tiisg];
-    sumf = simd_sum(sumf);
-
-    const float mean  = sumf/args.ne00;
-    const float scale = 1.0f/sqrt(mean + args.eps);
-
-    device T * y = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
-    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
-        if (F == 1) {
-            y[i00] = (x[i00]*scale);
-        }
-        if (F == 2) {
-            y[i00] = (x[i00]*scale)*f0[i00];
-        }
-        if (F == 3) {
-            y[i00] = (x[i00]*scale)*f0[i00] + f1[i00];
-        }
-    }
-}
-
-typedef decltype(kernel_rms_norm_fuse_impl<float4, 1>) kernel_rms_norm_fuse_t;
-
-template [[host_name("kernel_rms_norm_f32")]]         kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float, 1>;
-template [[host_name("kernel_rms_norm_mul_f32")]]     kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float, 2>;
-template [[host_name("kernel_rms_norm_mul_add_f32")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float, 3>;
-
-template [[host_name("kernel_rms_norm_f32_4")]]         kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 1>;
-template [[host_name("kernel_rms_norm_mul_f32_4")]]     kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 2>;
-template [[host_name("kernel_rms_norm_mul_add_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 3>;
-
-template <typename T0, typename T>
-kernel void kernel_l2_norm_impl(
-        constant ggml_metal_kargs_l2_norm & args,
-        device const char * src0,
-        device       char * dst,
-        threadgroup float * shmem_f32 [[threadgroup(0)]],
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort  sgitg[[simdgroup_index_in_threadgroup]],
-        ushort  tiisg[[thread_index_in_simdgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    const int i03 = tgpig.z;
-    const int i02 = tgpig.y;
-    const int i01 = tgpig.x;
-
-    if (sgitg == 0) {
-        shmem_f32[tiisg] = 0.0f;
-    }
-
-    device const T0 * x = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
-    device       T  * y = (device       T  *) (dst  + i03*args.nb3  + i02*args.nb2  + i01*args.nb1);
-
-    float sumf = 0.0f;
-
-    // parallel sum
-    for (int i00 = tpitg.x; i00 < args.ne00; i00 += ntg.x) {
-        sumf += dot(x[i00], x[i00]);
-    }
-    sumf = simd_sum(sumf);
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    if (tiisg == 0) {
-        shmem_f32[sgitg] = sumf;
-    }
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    sumf = shmem_f32[tiisg];
-    sumf = simd_sum(sumf);
-
-    const float scale = 1.0f/max(sqrt(sumf), args.eps);
-
-    for (int i00 = tpitg.x; i00 < args.ne00; i00 += ntg.x) {
-        y[i00] = x[i00] * scale;
-    }
-}
-
-typedef decltype(kernel_l2_norm_impl<float, float>) kernel_l2_norm_t;
-
-template [[host_name("kernel_l2_norm_f32_f32")]]   kernel kernel_l2_norm_t kernel_l2_norm_impl<float,  float>;
-template [[host_name("kernel_l2_norm_f32_f32_4")]] kernel kernel_l2_norm_t kernel_l2_norm_impl<float4, float4>;
-
-kernel void kernel_group_norm_f32(
-        constant ggml_metal_kargs_group_norm & args,
-        device const float * src0,
-        device       float * dst,
-        threadgroup float  * buf [[threadgroup(0)]],
-        uint tgpig[[threadgroup_position_in_grid]],
-        uint tpitg[[thread_position_in_threadgroup]],
-        uint sgitg[[simdgroup_index_in_threadgroup]],
-        uint tiisg[[thread_index_in_simdgroup]],
-        uint   ntg[[threads_per_threadgroup]]) {
-    const int64_t ne = args.ne00*args.ne01*args.ne02;
-    const int64_t gs = args.ne00*args.ne01*((args.ne02 + args.ngrp - 1) / args.ngrp);
-
-    int start = tgpig * gs;
-    int end   = start + gs;
-
-    start += tpitg;
-
-    if (end >= ne) {
-        end = ne;
-    }
-
-    float tmp = 0.0f; // partial sum for thread in warp
-
-    for (int j = start; j < end; j += ntg) {
-        tmp += src0[j];
-    }
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    tmp = simd_sum(tmp);
-    if (ntg > N_SIMDWIDTH) {
-        if (sgitg == 0) {
-            buf[tiisg] = 0.0f;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        if (tiisg == 0) {
-            buf[sgitg] = tmp;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        tmp = buf[tiisg];
-        tmp = simd_sum(tmp);
-    }
-
-    const float mean = tmp / gs;
-    tmp = 0.0f;
-
-    for (int j = start; j < end; j += ntg) {
-        float xi = src0[j] - mean;
-        dst[j] = xi;
-        tmp += xi * xi;
-    }
-
-    tmp = simd_sum(tmp);
-    if (ntg > N_SIMDWIDTH) {
-        if (sgitg == 0) {
-            buf[tiisg] = 0.0f;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        if (tiisg == 0) {
-            buf[sgitg] = tmp;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        tmp = buf[tiisg];
-        tmp = simd_sum(tmp);
-    }
-
-    const float variance = tmp / gs;
-    const float scale = 1.0f/sqrt(variance + args.eps);
-    for (int j = start; j < end; j += ntg) {
-        dst[j] *= scale;
-    }
-}
@@ -1,148 +0,0 @@
-#include "common.h"
-
-kernel void kernel_pool_2d_max_f32(
-        constant    ggml_metal_kargs_pool_2d & args,
-        device  const float * src0,
-        device        float * dst,
-        uint        gid[[thread_position_in_grid]]) {
-
-    if (gid >= args.np) {
-        return;
-    }
-
-    const int idx = gid;
-    const int I_HW = args.IH * args.IW;
-    const int O_HW = args.OH * args.OW;
-    const int nc = idx / O_HW;
-    const int cur_oh = idx % O_HW / args.OW;
-    const int cur_ow = idx % O_HW % args.OW;
-
-    device const float * i_ptr = src0 + nc * I_HW;
-    device       float * o_ptr = dst  + nc * O_HW;
-
-    const int start_h = cur_oh * args.s1 - args.p1;
-    const int bh = MAX(0,  start_h);
-    const int eh = MIN(args.IH, start_h + args.k1);
-    const int start_w = cur_ow * args.s0 - args.p0;
-    const int bw = MAX(0,  start_w);
-    const int ew = MIN(args.IW, start_w + args.k0);
-
-    float res = -INFINITY;
-
-    for (int i = bh; i < eh; i += 1) {
-        for (int j = bw; j < ew; j += 1) {
-            res = MAX(res, i_ptr[i * args.IW + j]);
-        }
-    }
-
-    o_ptr[cur_oh * args.OW + cur_ow] = res;
-}
-
-kernel void kernel_pool_2d_avg_f32(
-        constant    ggml_metal_kargs_pool_2d & args,
-        device  const float * src0,
-        device        float * dst,
-        uint        gid[[thread_position_in_grid]]) {
-
-    if (gid >= args.np) {
-        return;
-    }
-
-    const int idx = gid;
-    const int I_HW = args.IH * args.IW;
-    const int O_HW = args.OH * args.OW;
-    const int nc = idx / O_HW;
-    const int cur_oh = idx % O_HW / args.OW;
-    const int cur_ow = idx % O_HW % args.OW;
-
-    device const float * i_ptr = src0 + nc * I_HW;
-    device       float * o_ptr = dst  + nc * O_HW;
-
-    const int start_h = cur_oh * args.s1 - args.p1;
-    const int bh = MAX(0,  start_h);
-    const int eh = MIN(args.IH, start_h + args.k1);
-    const int start_w = cur_ow * args.s0 - args.p0;
-    const int bw = MAX(0,  start_w);
-    const int ew = MIN(args.IW, start_w + args.k0);
-    // const float scale = 1. / ((eh - bh) * (ew - bw));
-    const float scale = 1. / (args.k0 * args.k1);
-
-    float res = 0;
-
-    for (int i = bh; i < eh; i += 1) {
-        for (int j = bw; j < ew; j += 1) {
-            float cur = i_ptr[i * args.IW + j];
-            res += cur * scale;
-        }
-    }
-
-    o_ptr[cur_oh * args.OW + cur_ow] = res;
-}
-
-
-kernel void kernel_pool_1d_max_f32(
-        constant        ggml_metal_kargs_pool_1d & args,
-        device  const   float * src,
-        device          float * dst,
-        uint            gid [[thread_position_in_grid]]
-) {
-
-    if (gid >= args.np) {
-        return;
-    }
-
-    const int ow  = (int)gid % args.OW;
-    const int row = (int)gid / args.OW;
-
-    const int base = ow * args.s0 - args.p0;
-
-    float acc = -INFINITY;
-
-    const int src_off = row * args.IW;
-    const int dst_off = row * args.OW;
-
-    for (int ki = 0; ki < args.k0; ++ki) {
-        int j = base + ki;
-        if (j < 0 || j >= args.IW){
-            continue;
-        }
-        float v = src[src_off + j];
-        acc = max(acc, v);
-    }
-
-    dst[dst_off + ow] = acc;
-}
-
-kernel void kernel_pool_1d_avg_f32(
-        constant        ggml_metal_kargs_pool_1d & args,
-        device  const   float * src,
-        device          float * dst,
-        uint            gid [[thread_position_in_grid]]
-) {
-
-    if (gid >= args.np) {
-        return;
-    }
-
-    const int ow  = (int)gid % args.OW;
-    const int row = (int)gid / args.OW;
-
-    const int base = ow * args.s0 - args.p0;
-
-    float acc = 0.0f;
-    int   cnt = 0;
-
-    const int src_off = row * args.IW;
-    const int dst_off = row * args.OW;
-
-    for (int ki = 0; ki < args.k0; ++ki) {
-        const int j = base + ki;
-        if (j < 0 || j >= args.IW) {
-            continue;
-        }
-        acc += src[src_off + j];
-        cnt += 1;
-    }
-
-    dst[dst_off + ow] = (cnt > 0) ? (acc / (float)cnt) : 0.0f;
-}
@@ -1,262 +0,0 @@
-#pragma once
-
-#include "common.h"
-
-void quantize_q1_0(device const float * src, device block_q1_0 & dst) {
-    float sum_abs = 0.0f;
-    for (int j = 0; j < QK1_0; j++) {
-        sum_abs += fabs(src[j]);
-    }
-    dst.d = sum_abs / QK1_0;
-
-    for (int j = 0; j < QK1_0 / 8; j++) {
-        dst.qs[j] = 0;
-    }
-    for (int j = 0; j < QK1_0; j++) {
-        if (src[j] >= 0.0f) {
-            dst.qs[j / 8] |= (1 << (j % 8));
-        }
-    }
-}
-
-void quantize_q2_0(device const float * src, device block_q2_0 & dst) {
-    float amax = 0.0f;
-    for (int j = 0; j < QK2_0; j++) {
-        float a = fabs(src[j]);
-        if (a > amax) amax = a;
-    }
-    const float d = amax;
-    dst.d = d;
-
-    const float id = d > 0.0f ? 1.0f / d : 0.0f;
-
-    for (int j = 0; j < QK2_0 / 4; j++) {
-        dst.qs[j] = 0;
-    }
-    for (int j = 0; j < QK2_0; j++) {
-        int q = (int)round(src[j] * id) + 1;
-        q = max(0, min(3, q));
-        dst.qs[j / 4] |= (q << (2 * (j % 4)));
-    }
-}
-
-void quantize_q4_0(device const float * src, device block_q4_0 & dst) {
-#pragma METAL fp math_mode(safe)
-    float amax = 0.0f; // absolute max
-    float max  = 0.0f;
-
-    for (int j = 0; j < QK4_0; j++) {
-        const float v = src[j];
-        if (amax < fabs(v)) {
-            amax = fabs(v);
-            max  = v;
-        }
-    }
-
-    const float d = max / -8;
-    const float id = d ? 1.0f/d : 0.0f;
-
-    dst.d = d;
-
-    for (int j = 0; j < QK4_0/2; ++j) {
-        const float x0 = src[0       + j]*id;
-        const float x1 = src[QK4_0/2 + j]*id;
-
-        const uint8_t xi0 = MIN(15, (int8_t)(x0 + 8.5f));
-        const uint8_t xi1 = MIN(15, (int8_t)(x1 + 8.5f));
-
-        dst.qs[j]  = xi0;
-        dst.qs[j] |= xi1 << 4;
-    }
-}
-
-void quantize_q4_1(device const float * src, device block_q4_1 & dst) {
-#pragma METAL fp math_mode(safe)
-    float min = FLT_MAX;
-    float max = -FLT_MAX;
-
-    for (int j = 0; j < QK4_1; j++) {
-        const float v = src[j];
-        if (min > v) min = v;
-        if (max < v) max = v;
-    }
-
-    const float d = (max - min) / ((1 << 4) - 1);
-    const float id = d ? 1.0f/d : 0.0f;
-
-    dst.d = d;
-    dst.m = min;
-
-    for (int j = 0; j < QK4_1/2; ++j) {
-        const float x0 = (src[0       + j] - min)*id;
-        const float x1 = (src[QK4_1/2 + j] - min)*id;
-
-        const uint8_t xi0 = MIN(15, (int8_t)(x0 + 0.5f));
-        const uint8_t xi1 = MIN(15, (int8_t)(x1 + 0.5f));
-
-        dst.qs[j]  = xi0;
-        dst.qs[j] |= xi1 << 4;
-    }
-}
-
-void quantize_q5_0(device const float * src, device block_q5_0 & dst) {
-#pragma METAL fp math_mode(safe)
-    float amax = 0.0f; // absolute max
-    float max  = 0.0f;
-
-    for (int j = 0; j < QK5_0; j++) {
-        const float v = src[j];
-        if (amax < fabs(v)) {
-            amax = fabs(v);
-            max  = v;
-        }
-    }
-
-    const float d = max / -16;
-    const float id = d ? 1.0f/d : 0.0f;
-
-    dst.d = d;
-
-    uint32_t qh = 0;
-    for (int j = 0; j < QK5_0/2; ++j) {
-        const float x0 = src[0       + j]*id;
-        const float x1 = src[QK5_0/2 + j]*id;
-
-        const uint8_t xi0 = MIN(31, (int8_t)(x0 + 16.5f));
-        const uint8_t xi1 = MIN(31, (int8_t)(x1 + 16.5f));
-
-        dst.qs[j] = (xi0 & 0xf) | ((xi1 & 0xf) << 4);
-        qh |= ((xi0 & 0x10u) >> 4) << (j + 0);
-        qh |= ((xi1 & 0x10u) >> 4) << (j + QK5_0/2);
-    }
-
-    thread const uint8_t * qh8 = (thread const uint8_t *)&qh;
-
-    for (int j = 0; j < 4; ++j) {
-        dst.qh[j] = qh8[j];
-    }
-}
-
-void quantize_q5_1(device const float * src, device block_q5_1 & dst) {
-#pragma METAL fp math_mode(safe)
-    float max = src[0];
-    float min = src[0];
-
-    for (int j = 1; j < QK5_1; j++) {
-        const float v = src[j];
-        min = v < min ? v : min;
-        max = v > max ? v : max;
-    }
-
-    const float d = (max - min) / 31;
-    const float id = d ? 1.0f/d : 0.0f;
-
-    dst.d = d;
-    dst.m = min;
-
-    uint32_t qh = 0;
-    for (int j = 0; j < QK5_1/2; ++j) {
-        const float x0 = (src[0       + j] - min)*id;
-        const float x1 = (src[QK5_1/2 + j] - min)*id;
-
-        const uint8_t xi0 = (uint8_t)(x0 + 0.5f);
-        const uint8_t xi1 = (uint8_t)(x1 + 0.5f);
-
-        dst.qs[j] = (xi0 & 0xf) | ((xi1 & 0xf) << 4);
-        qh |= ((xi0 & 0x10u) >> 4) << (j + 0);
-        qh |= ((xi1 & 0x10u) >> 4) << (j + QK5_1/2);
-    }
-
-    thread const uint8_t * qh8 = (thread const uint8_t *)&qh;
-
-    for (int j = 0; j < 4; ++j) {
-        dst.qh[j] = qh8[j];
-    }
-}
-
-void quantize_q8_0(device const float * src, device block_q8_0 & dst) {
-#pragma METAL fp math_mode(safe)
-    float amax = 0.0f; // absolute max
-
-    for (int j = 0; j < QK8_0; j++) {
-        const float v = src[j];
-        amax = MAX(amax, fabs(v));
-    }
-
-    const float d = amax / ((1 << 7) - 1);
-    const float id = d ? 1.0f/d : 0.0f;
-
-    dst.d = d;
-
-    for (int j = 0; j < QK8_0; ++j) {
-        const float x0 = src[j]*id;
-
-        dst.qs[j] = round(x0);
-    }
-}
-
-void quantize_iq4_nl(device const float * src, device block_iq4_nl & dst) {
-#pragma METAL fp math_mode(safe)
-    float amax = 0.0f; // absolute max
-    float max  = 0.0f;
-
-    for (int j = 0; j < QK4_NL; j++) {
-        const float v = src[j];
-        if (amax < fabs(v)) {
-            amax = fabs(v);
-            max  = v;
-        }
-    }
-
-    const float d = max / kvalues_iq4nl_f[0];
-    const float id = d ? 1.0f/d : 0.0f;
-
-    float sumqx = 0, sumq2 = 0;
-    for (int j = 0; j < QK4_NL/2; ++j) {
-        const float x0 = src[0        + j]*id;
-        const float x1 = src[QK4_NL/2 + j]*id;
-
-        const uint8_t xi0 = best_index_int8(16, kvalues_iq4nl_f, x0);
-        const uint8_t xi1 = best_index_int8(16, kvalues_iq4nl_f, x1);
-
-        dst.qs[j] = xi0 | (xi1 << 4);
-
-        const float v0 = kvalues_iq4nl_f[xi0];
-        const float v1 = kvalues_iq4nl_f[xi1];
-        const float w0 = src[0        + j]*src[0        + j];
-        const float w1 = src[QK4_NL/2 + j]*src[QK4_NL/2 + j];
-        sumqx += w0*v0*src[j] + w1*v1*src[QK4_NL/2 + j];
-        sumq2 += w0*v0*v0 + w1*v1*v1;
-
-    }
-
-    dst.d = sumq2 > 0 ? sumqx/sumq2 : d;
-}
-
-void quantize_tq2_0(device const float * src, device block_tq2_0 & dst) {
-#pragma METAL fp math_mode(safe)
-    float amax = 0.0f; // absolute max
-
-    for (int j = 0; j < QK_K; j++) {
-        const float v = src[j];
-        amax = MAX(amax, fabs(v));
-    }
-
-    const float d = amax;
-    const float id = d ? 1.0f/d : 0.0f;
-
-    dst.d = (half) d;
-
-    for (int j = 0; j < QK_K/4; j += 32) {
-        for (int m = 0; m < 32; ++m) {
-            uint8_t q = 0;
-            for (int n = 0; n < 4; ++n) {
-                // -1, 0, 1 -> 0, 1, 2
-                int xi = (int)round(src[m + n*32] * id) + 1;
-                q += (uint8_t)((xi & 3) << (2*n));
-            }
-            dst.qs[j + m] = q;
-        }
-        src += 4*32;
-    }
-}
@@ -1,435 +0,0 @@
-#include "common.h"
-#include "dequantize.h"
-#include "quantize.h"
-
-template<typename T0, typename T1>
-kernel void kernel_cpy_t_t(
-        constant ggml_metal_kargs_cpy & args,
-        device  const char * src0,
-        device        char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    const int32_t i03 = tgpig[2];
-    const int32_t i02 = tgpig[1];
-    const int32_t i01 = ntg[1] == 1 ? tgpig[0]%args.ne01 : tgpig[0]*ntg[1] + tpitg.y;
-    const int32_t iw0 = ntg[1] == 1 ? tgpig[0]/args.ne01 : 0;
-
-    if (i01 >= args.ne01) {
-        return;
-    }
-
-    const int64_t n = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00;
-
-    const int32_t i3 = n/(args.ne2*args.ne1*args.ne0);
-    const int32_t i2 = (n - i3*args.ne2*args.ne1*args.ne0)/(args.ne1*args.ne0);
-    const int32_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0)/args.ne0;
-    const int32_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0);
-
-    device T1 * dst_data = (device T1 *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
-
-    for (int32_t i00 = iw0*ntg[0] + tpitg.x; i00 < args.ne00;) {
-        device const T0 * src = (device T0 *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + i00*args.nb00);
-        dst_data[i00] = (T1) src[0];
-        break;
-    }
-}
-
-typedef decltype(kernel_cpy_t_t<float, float>) kernel_cpy_t;
-
-template [[host_name("kernel_cpy_f32_f32")]]   kernel kernel_cpy_t kernel_cpy_t_t<float,   float>;
-template [[host_name("kernel_cpy_f32_f16")]]   kernel kernel_cpy_t kernel_cpy_t_t<float,   half>;
-template [[host_name("kernel_cpy_f32_i32")]]   kernel kernel_cpy_t kernel_cpy_t_t<float,   int32_t>;
-template [[host_name("kernel_cpy_i32_f32")]]   kernel kernel_cpy_t kernel_cpy_t_t<int32_t, float>;
-template [[host_name("kernel_cpy_i32_i32")]]   kernel kernel_cpy_t kernel_cpy_t_t<int32_t, int32_t>;
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_cpy_f32_bf16")]]  kernel kernel_cpy_t kernel_cpy_t_t<float,   bfloat>;
-#endif
-template [[host_name("kernel_cpy_f16_f32")]]   kernel kernel_cpy_t kernel_cpy_t_t<half,    float>;
-template [[host_name("kernel_cpy_f16_f16")]]   kernel kernel_cpy_t kernel_cpy_t_t<half,    half>;
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_cpy_bf16_f32")]]  kernel kernel_cpy_t kernel_cpy_t_t<bfloat,  float>;
-template [[host_name("kernel_cpy_bf16_bf16")]] kernel kernel_cpy_t kernel_cpy_t_t<bfloat,  bfloat>;
-#endif
-
-template<short QK,
-         typename block_q,
-         void (*quantize_func)(device const float *, device block_q &)>
-kernel void kernel_cpy_f32_q(
-        constant ggml_metal_kargs_cpy & args,
-        device const char * src0,
-        device char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    const int32_t i03 = tgpig[2];
-    const int32_t i02 = tgpig[1];
-    const int32_t i01 = ntg[1] == 1 ? tgpig[0]%args.ne01 : tgpig[0]*ntg[1] + tpitg.y;
-    const int32_t iw0 = ntg[1] == 1 ? tgpig[0]/args.ne01 : 0;
-
-    if (i01 >= args.ne01) {
-        return;
-    }
-
-    const int64_t n = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00;
-
-    const int32_t i3 = n / (args.ne2*args.ne1*args.ne0);
-    const int32_t i2 = (n - i3*args.ne2*args.ne1*args.ne0) / (args.ne1*args.ne0);
-    const int32_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0) / args.ne0;
-    const int32_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0)/QK;
-
-    device block_q * dst_data = (device block_q *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
-
-    for (int32_t i00 = iw0*ntg[0] + tpitg.x; i00 < args.nk0;) {
-        device const float * src = (device const float *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + (i00*QK)*args.nb00);
-
-        quantize_func(src, dst_data[i00]);
-
-        break;
-    }
-}
-
-typedef decltype(kernel_cpy_f32_q<QK8_0,  block_q8_0,  quantize_q8_0>)  cpy_f_q_t;
-
-template [[host_name("kernel_cpy_f32_q8_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK8_0,  block_q8_0,   quantize_q8_0>;
-template [[host_name("kernel_cpy_f32_q1_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK1_0,  block_q1_0,   quantize_q1_0>;
-template [[host_name("kernel_cpy_f32_q2_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK2_0,  block_q2_0,   quantize_q2_0>;
-template [[host_name("kernel_cpy_f32_q4_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK4_0,  block_q4_0,   quantize_q4_0>;
-template [[host_name("kernel_cpy_f32_q4_1")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK4_1,  block_q4_1,   quantize_q4_1>;
-template [[host_name("kernel_cpy_f32_q5_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK5_0,  block_q5_0,   quantize_q5_0>;
-template [[host_name("kernel_cpy_f32_q5_1")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK5_1,  block_q5_1,   quantize_q5_1>;
-template [[host_name("kernel_cpy_f32_iq4_nl")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK4_NL, block_iq4_nl, quantize_iq4_nl>;
-template [[host_name("kernel_cpy_f32_tq2_0")]]  kernel cpy_f_q_t kernel_cpy_f32_q<QK_K,   block_tq2_0,  quantize_tq2_0>;
-
-template<typename T4x4, typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread T4x4 &)>
-kernel void kernel_cpy_q_f32(
-        constant ggml_metal_kargs_cpy & args,
-        device  const char * src0,
-        device        char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    const int32_t i03 = tgpig[2];
-    const int32_t i02 = tgpig[1];
-    const int32_t i01 = ntg[1] == 1 ? tgpig[0]%args.ne01 : tgpig[0]*ntg[1] + tpitg.y;
-    const int32_t iw0 = ntg[1] == 1 ? tgpig[0]/args.ne01 : 0;
-
-    if (i01 >= args.ne01) {
-        return;
-    }
-
-    const int64_t n = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00;
-
-    const int32_t i3 = n/(args.ne2*args.ne1*args.ne0);
-    const int32_t i2 = (n - i3*args.ne2*args.ne1*args.ne0)/(args.ne1*args.ne0);
-    const int32_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0)/args.ne0;
-    const int32_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0);
-
-    device const block_q * src_data = (device const block_q *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
-    device       T4x4    * dst_data = (device       T4x4    *)(dst  +  i3*args.nb3  +  i2*args.nb2  +  i1*args.nb1 + i0*args.nb0);
-
-    for (int32_t i00 = iw0*ntg[0] + tpitg.x; i00 < args.nk0;) {
-        T4x4 temp;
-        dequantize_func(src_data + i00/nl, i00%nl, temp);
-        dst_data[i00] = temp;
-
-        break;
-    }
-}
-
-typedef decltype(kernel_cpy_q_f32<float4x4, block_q4_0, 2, dequantize_q4_0>) cpy_q_f_t;
-
-template [[host_name("kernel_cpy_q1_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q1_0, 8, dequantize_q1_0>;
-template [[host_name("kernel_cpy_q2_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q2_0, 4, dequantize_q2_0>;
-template [[host_name("kernel_cpy_q4_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q4_0, 2, dequantize_q4_0>;
-template [[host_name("kernel_cpy_q4_1_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q4_1, 2, dequantize_q4_1>;
-template [[host_name("kernel_cpy_q5_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q5_0, 2, dequantize_q5_0>;
-template [[host_name("kernel_cpy_q5_1_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q5_1, 2, dequantize_q5_1>;
-template [[host_name("kernel_cpy_q8_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q8_0, 2, dequantize_q8_0>;
-
-template [[host_name("kernel_cpy_tq2_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_tq2_0, QK_NL, dequantize_tq2_0>;
-
-template [[host_name("kernel_cpy_q1_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q1_0, 8, dequantize_q1_0>;
-template [[host_name("kernel_cpy_q2_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q2_0, 4, dequantize_q2_0>;
-template [[host_name("kernel_cpy_q4_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q4_0, 2, dequantize_q4_0>;
-template [[host_name("kernel_cpy_q4_1_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q4_1, 2, dequantize_q4_1>;
-template [[host_name("kernel_cpy_q5_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q5_0, 2, dequantize_q5_0>;
-template [[host_name("kernel_cpy_q5_1_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q5_1, 2, dequantize_q5_1>;
-template [[host_name("kernel_cpy_q8_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q8_0, 2, dequantize_q8_0>;
-
-template [[host_name("kernel_cpy_tq2_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_tq2_0, QK_NL, dequantize_tq2_0>;
-
-template<typename T>
-kernel void kernel_concat(
-        constant ggml_metal_kargs_concat & args,
-        device  const char * src0,
-        device  const char * src1,
-        device        char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-
-    const int i3 = tgpig.z;
-    const int i2 = tgpig.y;
-    const int i1 = ntg.y == 1 ? tgpig.x : tgpig.x*ntg.y + tpitg.y;
-
-    if (i1 >= args.ne1) {
-        return;
-    }
-
-    int o[4] = {0, 0, 0, 0};
-    o[args.dim] = args.dim == 0 ? args.ne00 : (args.dim == 1 ? args.ne01 : (args.dim == 2 ? args.ne02 : args.ne03));
-
-    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
-        device const T * x;
-
-        if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
-            x = (device const T *)(src0 + (i3       )*args.nb03 + (i2       )*args.nb02 + (i1       )*args.nb01 + (i0       )*args.nb00);
-        } else {
-            x = (device const T *)(src1 + (i3 - o[3])*args.nb13 + (i2 - o[2])*args.nb12 + (i1 - o[1])*args.nb11 + (i0 - o[0])*args.nb10);
-        }
-
-        device T * y = (device T *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
-
-        *y = *x;
-    }
-}
-
-typedef decltype(kernel_concat<float>) kernel_concat_t;
-
-template [[host_name("kernel_concat_f32")]]  kernel kernel_concat_t kernel_concat<float>;
-template [[host_name("kernel_concat_f16")]]  kernel kernel_concat_t kernel_concat<half>;
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_concat_bf16")]] kernel kernel_concat_t kernel_concat<bfloat>;
-#endif
-template [[host_name("kernel_concat_i8")]]   kernel kernel_concat_t kernel_concat<char>;
-template [[host_name("kernel_concat_i16")]]  kernel kernel_concat_t kernel_concat<short>;
-template [[host_name("kernel_concat_i32")]]  kernel kernel_concat_t kernel_concat<int>;
-template [[host_name("kernel_concat_i64")]]  kernel kernel_concat_t kernel_concat<long>;
-
-template<typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread float4x4 &)>
-kernel void kernel_get_rows_q(
-        constant ggml_metal_kargs_get_rows & args,
-        device const void * src0,
-        device const void * src1,
-        device       void * dst,
-        uint3               tgpig[[threadgroup_position_in_grid]],
-        ushort              tiitg[[thread_index_in_threadgroup]],
-        ushort3             ntg  [[threads_per_threadgroup]]) {
-    const int32_t iw0 = tgpig.x/args.ne10;
-    const int32_t i10 = tgpig.x%args.ne10;
-    const int32_t i11 = tgpig.y;
-    const int32_t i12 = tgpig.z;
-
-    const int32_t r = ((const device int32_t *) ((const device char *) src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0];
-
-    const int32_t i02 = i11;
-    const int32_t i03 = i12;
-
-    auto psrc = (device const block_q *) ((const device char *) src0 + i03*args.nb03 + i02*args.nb02 +   r*args.nb01);
-    auto pdst = (device      float4x4 *) ((      device char *) dst  + i12*args.nb3  + i11*args.nb2  + i10*args.nb1);
-
-    for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) {
-        float4x4 temp;
-        dequantize_func(psrc + ind/nl, ind%nl, temp);
-        pdst[ind] = temp;
-
-        break;
-    }
-}
-
-template<typename T0, typename T>
-kernel void kernel_get_rows_f(
-        constant ggml_metal_kargs_get_rows & args,
-        device const void * src0,
-        device const void * src1,
-        device       void * dst,
-        uint3               tgpig[[threadgroup_position_in_grid]],
-        ushort              tiitg[[thread_index_in_threadgroup]],
-        ushort3             ntg [[threads_per_threadgroup]]) {
-    const int32_t iw0 = tgpig.x/args.ne10;
-    const int32_t i10 = tgpig.x%args.ne10;
-    const int32_t i11 = tgpig.y;
-    const int32_t i12 = tgpig.z;
-
-    const int32_t r = ((const device int32_t *) ((const device char *) src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0];
-
-    const int32_t i02 = i11;
-    const int32_t i03 = i12;
-
-    auto psrc = (const device T0 *) ((const device char *) src0 + i03*args.nb03 + i02*args.nb02 +   r*args.nb01);
-    auto pdst = (      device T  *) ((      device char *)  dst + i12*args.nb3  + i11*args.nb2  + i10*args.nb1);
-
-    for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) {
-        pdst[ind] = psrc[ind];
-
-        break;
-    }
-}
-
-typedef decltype(kernel_get_rows_f<float, float>) get_rows_f_t;
-
-template [[host_name("kernel_get_rows_f32")]]  kernel get_rows_f_t kernel_get_rows_f<float, float>;
-template [[host_name("kernel_get_rows_f16")]]  kernel get_rows_f_t kernel_get_rows_f<half,  float>;
-template [[host_name("kernel_get_rows_i32")]]  kernel get_rows_f_t kernel_get_rows_f<int32_t, int32_t>;
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_get_rows_bf16")]] kernel get_rows_f_t kernel_get_rows_f<bfloat, float>;
-#endif
-
-typedef decltype(kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>) get_rows_q_t;
-
-template [[host_name("kernel_get_rows_q1_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q1_0,    8, dequantize_q1_0>;
-template [[host_name("kernel_get_rows_q2_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q2_0,    4, dequantize_q2_0>;
-template [[host_name("kernel_get_rows_q4_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_0,    2, dequantize_q4_0>;
-template [[host_name("kernel_get_rows_q4_1")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_1,    2, dequantize_q4_1>;
-template [[host_name("kernel_get_rows_q5_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_0,    2, dequantize_q5_0>;
-template [[host_name("kernel_get_rows_q5_1")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_1,    2, dequantize_q5_1>;
-template [[host_name("kernel_get_rows_q8_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q8_0,    2, dequantize_q8_0>;
-template [[host_name("kernel_get_rows_mxfp4")]]   kernel get_rows_q_t kernel_get_rows_q<block_mxfp4,   2, dequantize_mxfp4>;
-template [[host_name("kernel_get_rows_q2_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q2_K,    QK_NL, dequantize_q2_K>;
-template [[host_name("kernel_get_rows_q3_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q3_K,    QK_NL, dequantize_q3_K>;
-template [[host_name("kernel_get_rows_q4_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_K,    QK_NL, dequantize_q4_K>;
-template [[host_name("kernel_get_rows_q5_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_K,    QK_NL, dequantize_q5_K>;
-template [[host_name("kernel_get_rows_q6_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q6_K,    QK_NL, dequantize_q6_K>;
-template [[host_name("kernel_get_rows_iq2_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>;
-template [[host_name("kernel_get_rows_iq2_xs")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq2_xs,  QK_NL, dequantize_iq2_xs>;
-template [[host_name("kernel_get_rows_iq3_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq3_xxs, QK_NL, dequantize_iq3_xxs>;
-template [[host_name("kernel_get_rows_iq3_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq3_s,   QK_NL, dequantize_iq3_s>;
-template [[host_name("kernel_get_rows_iq2_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq2_s,   QK_NL, dequantize_iq2_s>;
-template [[host_name("kernel_get_rows_iq1_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq1_s,   QK_NL, dequantize_iq1_s>;
-template [[host_name("kernel_get_rows_iq1_m")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq1_m,   QK_NL, dequantize_iq1_m>;
-template [[host_name("kernel_get_rows_iq4_nl")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq4_nl,  2,     dequantize_iq4_nl>;
-template [[host_name("kernel_get_rows_iq4_xs")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq4_xs,  QK_NL, dequantize_iq4_xs>;
-template [[host_name("kernel_get_rows_tq2_0")]]   kernel get_rows_q_t kernel_get_rows_q<block_tq2_0,   QK_NL, dequantize_tq2_0>;
-
-template<typename TS, typename TI, short QK, typename block_q, void (*quantize_func)(device const float *, device block_q &)>
-kernel void kernel_set_rows_q(
-        constant ggml_metal_kargs_set_rows & args,
-        device const  void * src0,
-        device const  void * src1,
-        device       float * dst,
-        uint3                tgpig[[threadgroup_position_in_grid]],
-        uint                 tiitg[[thread_index_in_threadgroup]],
-        uint3                tptg [[threads_per_threadgroup]]) {
-    const int32_t i03 = tgpig.z;
-    const int32_t i02 = tgpig.y;
-
-    const int32_t i12 = i03%args.ne12;
-    const int32_t i11 = i02%args.ne11;
-
-    const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x;
-    if (i01 >= args.ne01) {
-        return;
-    }
-
-    const int32_t i10 = i01;
-    const TI      i1  = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
-
-          device block_q * dst_row = (      device block_q *) ((      device char *) dst  +  i1*args.nb1  + i02*args.nb2  + i03*args.nb3);
-    const device TS      * src_row = (const device TS      *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
-
-    for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
-        quantize_func(src_row + QK*ind, dst_row[ind]);
-    }
-}
-
-template<typename TS, typename TI, typename block_q, void (*quantize_func)(device const float *, device block_q &)>
-kernel void kernel_set_rows_q32(
-        constant ggml_metal_kargs_set_rows & args,
-        device const  void * src0,
-        device const  void * src1,
-        device       float * dst,
-        uint3                tgpig[[threadgroup_position_in_grid]],
-        uint                 tiitg[[thread_index_in_threadgroup]],
-        uint3                tptg [[threads_per_threadgroup]]) {
-    const int32_t i03 = tgpig.z;
-    const int32_t i02 = tgpig.y;
-
-    const int32_t i12 = i03%args.ne12;
-    const int32_t i11 = i02%args.ne11;
-
-    const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x;
-    if (i01 >= args.ne01) {
-        return;
-    }
-
-    const int32_t i10 = i01;
-    const TI      i1  = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
-
-          device block_q * dst_row = (      device block_q *) ((      device char *) dst  +  i1*args.nb1  + i02*args.nb2  + i03*args.nb3);
-    const device TS      * src_row = (const device TS      *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
-
-    for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
-        quantize_func(src_row + 32*ind, dst_row[ind]);
-    }
-}
-
-template<typename TS, typename TI, typename TD>
-kernel void kernel_set_rows_f(
-        constant ggml_metal_kargs_set_rows & args,
-        device const  void * src0,
-        device const  void * src1,
-        device       float * dst,
-        uint3                tgpig[[threadgroup_position_in_grid]],
-        uint                 tiitg[[thread_index_in_threadgroup]],
-        uint3                tptg [[threads_per_threadgroup]]) {
-    const int32_t i03 = tgpig.z;
-    const int32_t i02 = tgpig.y;
-
-    const int32_t i12 = i03%args.ne12;
-    const int32_t i11 = i02%args.ne11;
-
-    const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x;
-    if (i01 >= args.ne01) {
-        return;
-    }
-
-    const int32_t i10 = i01;
-    const TI      i1  = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
-
-          device TD * dst_row = (      device TD *) ((      device char *) dst  +  i1*args.nb1  + i02*args.nb2  + i03*args.nb3);
-    const device TS * src_row = (const device TS *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
-
-    for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
-        dst_row[ind] = (TD) src_row[ind];
-    }
-}
-
-typedef decltype(kernel_set_rows_f<float, int64_t, float>) set_rows_f_t;
-
-template [[host_name("kernel_set_rows_f32_i64_f32")]]   kernel set_rows_f_t kernel_set_rows_f<float, int64_t, float>;
-template [[host_name("kernel_set_rows_f32_i32_f32")]]   kernel set_rows_f_t kernel_set_rows_f<float, int32_t, float>;
-template [[host_name("kernel_set_rows_f32_i64_f16")]]   kernel set_rows_f_t kernel_set_rows_f<float, int64_t, half>;
-template [[host_name("kernel_set_rows_f32_i32_f16")]]   kernel set_rows_f_t kernel_set_rows_f<float, int32_t, half>;
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_set_rows_f32_i64_bf16")]]  kernel set_rows_f_t kernel_set_rows_f<float, int64_t, bfloat>;
-template [[host_name("kernel_set_rows_f32_i32_bf16")]]  kernel set_rows_f_t kernel_set_rows_f<float, int32_t, bfloat>;
-#endif
-
-template [[host_name("kernel_set_rows_f16_i64_f16")]]   kernel set_rows_f_t kernel_set_rows_f<half, int64_t, half>;
-template [[host_name("kernel_set_rows_f16_i32_f16")]]   kernel set_rows_f_t kernel_set_rows_f<half, int32_t, half>;
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_set_rows_bf16_i64_bf16")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int64_t, bfloat>;
-template [[host_name("kernel_set_rows_bf16_i32_bf16")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int32_t, bfloat>;
-#endif
-
-typedef decltype(kernel_set_rows_q32<float, int64_t, block_q8_0, quantize_q8_0>) set_rows_q32_t;
-
-template [[host_name("kernel_set_rows_f32_i64_q8_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q8_0,   quantize_q8_0>;
-template [[host_name("kernel_set_rows_f32_i32_q8_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q8_0,   quantize_q8_0>;
-template [[host_name("kernel_set_rows_f32_i64_q4_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q4_0,   quantize_q4_0>;
-template [[host_name("kernel_set_rows_f32_i32_q4_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q4_0,   quantize_q4_0>;
-template [[host_name("kernel_set_rows_f32_i64_q4_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q4_1,   quantize_q4_1>;
-template [[host_name("kernel_set_rows_f32_i32_q4_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q4_1,   quantize_q4_1>;
-template [[host_name("kernel_set_rows_f32_i64_q5_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q5_0,   quantize_q5_0>;
-template [[host_name("kernel_set_rows_f32_i32_q5_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q5_0,   quantize_q5_0>;
-template [[host_name("kernel_set_rows_f32_i64_q5_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q5_1,   quantize_q5_1>;
-template [[host_name("kernel_set_rows_f32_i32_q5_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q5_1,   quantize_q5_1>;
-template [[host_name("kernel_set_rows_f32_i64_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_iq4_nl, quantize_iq4_nl>;
-template [[host_name("kernel_set_rows_f32_i32_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_iq4_nl, quantize_iq4_nl>;
-
-typedef decltype(kernel_set_rows_q<float, int64_t, QK_K, block_tq2_0, quantize_tq2_0>) set_rows_qK_t;
-
-template [[host_name("kernel_set_rows_f32_i64_tq2_0")]]  kernel set_rows_qK_t kernel_set_rows_q<float, int64_t, QK_K, block_tq2_0, quantize_tq2_0>;
-template [[host_name("kernel_set_rows_f32_i32_tq2_0")]]  kernel set_rows_qK_t kernel_set_rows_q<float, int32_t, QK_K, block_tq2_0, quantize_tq2_0>;
-
@@ -1,228 +0,0 @@
-#include "common.h"
-
-kernel void kernel_op_sum_f32(
-        constant ggml_metal_kargs_sum & args,
-        device const float * src0,
-        device       float * dst,
-        threadgroup  float * shmem_f32 [[threadgroup(0)]],
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort  sgitg[[simdgroup_index_in_threadgroup]],
-        ushort  tiisg[[thread_index_in_simdgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-
-    if (args.np == 0) {
-        return;
-    }
-
-    // TODO: become function constant
-    const uint nsg = (ntg.x + 31) / 32;
-
-    float sumf = 0;
-
-    for (uint64_t i0 = tpitg.x; i0 < args.np; i0 += ntg.x) {
-        sumf += src0[i0];
-    }
-
-    sumf = simd_sum(sumf);
-
-    if (tiisg == 0) {
-        shmem_f32[sgitg] = sumf;
-    }
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    float total = 0;
-
-    if (sgitg == 0) {
-        float v = 0;
-
-        if (tpitg.x < nsg) {
-            v = shmem_f32[tpitg.x];
-        }
-
-        total = simd_sum(v);
-
-        if (tpitg.x == 0) {
-            dst[0] = total;
-        }
-    }
-}
-
-constant short FC_sum_rows_op [[function_constant(FC_SUM_ROWS + 0)]];
-
-template <typename T0, typename T>
-kernel void kernel_sum_rows_impl(
-        constant ggml_metal_kargs_sum_rows & args,
-        device const char * src0,
-        device       char * dst,
-        threadgroup  char * shmem [[threadgroup(0)]],
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort  sgitg[[simdgroup_index_in_threadgroup]],
-        ushort  tiisg[[thread_index_in_simdgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-#define FC_OP  FC_sum_rows_op
-
-    const int i3 = tgpig.z;
-    const int i2 = tgpig.y;
-    const int i1 = tgpig.x;
-
-    threadgroup T0 * shmem_t = (threadgroup T0 *) shmem;
-
-    if (sgitg == 0) {
-        shmem_t[tiisg] = 0.0f;
-    }
-
-    device const T0 * src_row = (device const T0 *) (src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03);
-    device       T  * dst_row = (device       T  *) (dst  + i1*args.nb1  + i2*args.nb2  + i3*args.nb3);
-
-    T0 sumf = T0(0.0f);
-
-    for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) {
-        sumf += src_row[i0];
-    }
-
-    sumf = simd_sum(sumf);
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    if (tiisg == 0) {
-        shmem_t[sgitg] = sumf;
-    }
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    sumf = shmem_t[tiisg];
-    sumf = simd_sum(sumf);
-
-    if (tpitg.x == 0) {
-        if (FC_OP == OP_SUM_ROWS_NUM_MEAN) {
-            if (is_same<float4, T0>::value) {
-                dst_row[0] = sum(sumf) / (4*args.ne00);
-            } else {
-                dst_row[0] = sum(sumf) / args.ne00;
-            }
-        } else {
-            dst_row[0] = sum(sumf);
-        }
-    }
-
-#undef FC_OP
-}
-
-typedef decltype(kernel_sum_rows_impl<float, float>) kernel_sum_rows_t;
-
-template [[host_name("kernel_sum_rows_f32_f32")]]   kernel kernel_sum_rows_t kernel_sum_rows_impl<float,  float>;
-template [[host_name("kernel_sum_rows_f32_f32_4")]] kernel kernel_sum_rows_t kernel_sum_rows_impl<float4, float>;
-
-template<typename T>
-kernel void kernel_cumsum_blk(
-        constant ggml_metal_kargs_cumsum_blk & args,
-        device const char * src0,
-        device       char * tmp,
-        device       char * dst,
-        threadgroup  char * shmem [[threadgroup(0)]],
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort  sgitg[[simdgroup_index_in_threadgroup]],
-        ushort  tiisg[[thread_index_in_simdgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    const int ib = tgpig[0]/args.ne01;
-
-    const int i00 = ib*ntg.x;
-    const int i01 = tgpig[0]%args.ne01;
-    const int i02 = tgpig[1];
-    const int i03 = tgpig[2];
-
-    device const float * src0_row = (device const float *) (src0 +
-            args.nb01*i01 +
-            args.nb02*i02 +
-            args.nb03*i03);
-
-    threadgroup float * shmem_f32 = (threadgroup float *) shmem;
-
-    float v = 0.0f;
-
-    if (i00 + tpitg.x < args.ne00) {
-        v = src0_row[i00 + tpitg.x];
-    }
-
-    float s = simd_prefix_inclusive_sum(v);
-
-    if (tiisg == N_SIMDWIDTH - 1) {
-        shmem_f32[sgitg] = s;
-    }
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    if (sgitg == 0) {
-        shmem_f32[tiisg] = simd_prefix_exclusive_sum(shmem_f32[tiisg]);
-    }
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    s += shmem_f32[sgitg];
-
-    device float * dst_row = (device float *) dst +
-        args.ne00*i01 +
-        args.ne00*args.ne01*i02 +
-        args.ne00*args.ne01*args.ne02*i03;
-
-    if (i00 + tpitg.x < args.ne00) {
-        dst_row[i00 + tpitg.x] = s;
-    }
-
-    if (args.outb && tpitg.x == ntg.x - 1) {
-        device float * tmp_row = (device float *) tmp +
-            args.net0*i01 +
-            args.net0*args.net1*i02 +
-            args.net0*args.net1*args.net2*i03;
-
-        tmp_row[ib] = s;
-    }
-}
-
-typedef decltype(kernel_cumsum_blk<float>) kernel_cumsum_blk_t;
-
-template [[host_name("kernel_cumsum_blk_f32")]] kernel kernel_cumsum_blk_t kernel_cumsum_blk<float>;
-
-template<typename T>
-kernel void kernel_cumsum_add(
-        constant ggml_metal_kargs_cumsum_add & args,
-        device const char * tmp,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort  sgitg[[simdgroup_index_in_threadgroup]],
-        ushort  tiisg[[thread_index_in_simdgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    const int ib = tgpig[0]/args.ne01;
-
-    if (ib == 0) {
-        return;
-    }
-
-    const int i00 = ib*ntg.x;
-    const int i01 = tgpig[0]%args.ne01;
-    const int i02 = tgpig[1];
-    const int i03 = tgpig[2];
-
-    device const float * tmp_row = (device const float *) (tmp +
-            args.nbt1*i01 +
-            args.nbt2*i02 +
-            args.nbt3*i03);
-
-    device float * dst_row = (device float *) dst +
-        args.ne00*i01 +
-        args.ne00*args.ne01*i02 +
-        args.ne00*args.ne01*args.ne02*i03;
-
-    if (i00 + tpitg.x < args.ne00) {
-        dst_row[i00 + tpitg.x] += tmp_row[ib - 1];
-    }
-}
-
-typedef decltype(kernel_cumsum_add<float>) kernel_cumsum_add_t;
-
-template [[host_name("kernel_cumsum_add_f32")]] kernel kernel_cumsum_add_t kernel_cumsum_add<float>;
@@ -1,333 +0,0 @@
-#include "common.h"
-
-constant bool FC_rope_is_imrope [[function_constant(FC_ROPE + 0)]];
-constant bool FC_rope_is_back   [[function_constant(FC_ROPE + 1)]];
-
-static float rope_yarn_ramp(const float low, const float high, const int i0) {
-    const float y = (i0 / 2 - low) / max(0.001f, high - low);
-    return 1.0f - min(1.0f, max(0.0f, y));
-}
-
-// YaRN algorithm based on LlamaYaRNScaledRotaryEmbedding.py from https://github.com/jquesnelle/yarn
-// MIT licensed. Copyright (c) 2023 Jeffrey Quesnelle and Bowen Peng.
-static void rope_yarn(
-    float theta_extrap, float freq_scale, float corr_dims[2], int i0, float ext_factor, float mscale,
-    thread float * cos_theta, thread float * sin_theta) {
-    // Get n-d rotational scaling corrected for extrapolation
-    float theta_interp = freq_scale * theta_extrap;
-    float theta = theta_interp;
-    if (ext_factor != 0.0f) {
-        float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], i0) * ext_factor;
-        theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix;
-
-        // Get n-d magnitude scaling corrected for interpolation
-        mscale *= 1.0f + 0.1f * log(1.0f / freq_scale);
-    }
-    *cos_theta = cos(theta) * mscale;
-    *sin_theta = sin(theta) * mscale;
-    if (FC_rope_is_back) {
-        *sin_theta *= -1.0f;
-    }
-}
-
-// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get
-// `corr_fac(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))`
-static float rope_yarn_corr_factor(int n_dims, int n_ctx_orig, float n_rot, float base) {
-    return n_dims * log(n_ctx_orig / (n_rot * 2 * M_PI_F)) / (2 * log(base));
-}
-
-static void rope_yarn_corr_dims(
-    int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2]
-) {
-    // start and end correction dims
-    dims[0] = max(0.0f,         floor(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_fast, freq_base)));
-    dims[1] = min(n_dims - 1.0f, ceil(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_slow, freq_base)));
-}
-
-template<typename T>
-kernel void kernel_rope_norm(
-        constant ggml_metal_kargs_rope & args,
-        device const char * src0,
-        device const char * src1,
-        device const char * src2,
-        device       char * dst,
-        ushort  tiitg[[thread_index_in_threadgroup]],
-        ushort3 tptg [[threads_per_threadgroup]],
-        uint3   tgpig[[threadgroup_position_in_grid]]) {
-    const int i3 = tgpig[2];
-    const int i2 = tgpig[1];
-    const int i1 = tgpig[0];
-
-    float corr_dims[2];
-    rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
-
-    device const int32_t * pos = (device const int32_t *) src1;
-
-    const float theta_base = (float) pos[i2];
-    const float inv_ndims = -1.f/args.n_dims;
-
-    float cos_theta;
-    float sin_theta;
-
-    for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
-        if (i0 >= args.n_offs && i0 < args.n_offs + args.n_dims) {
-            const int iw = i0 - args.n_offs; // relative idx
-            const int ic = iw/2;
-
-            const float theta = theta_base * pow(args.freq_base, inv_ndims*iw);
-
-            const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
-
-            rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, iw, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
-
-            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
-            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + i0*args.nb0);
-
-            const float x0 = src[0];
-            const float x1 = src[1];
-
-            dst_data[0] = x0*cos_theta - x1*sin_theta;
-            dst_data[1] = x0*sin_theta + x1*cos_theta;
-        } else {
-            if (args.inplace) {
-                continue;
-            }
-
-            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
-            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + i0*args.nb0);
-
-            dst_data[0] = src[0];
-            dst_data[1] = src[1];
-        }
-    }
-}
-
-template<typename T>
-kernel void kernel_rope_neox(
-        constant ggml_metal_kargs_rope & args,
-        device const char * src0,
-        device const char * src1,
-        device const char * src2,
-        device       char * dst,
-        ushort  tiitg[[thread_index_in_threadgroup]],
-        ushort3 tptg [[threads_per_threadgroup]],
-        uint3   tgpig[[threadgroup_position_in_grid]]) {
-    const int i3 = tgpig[2];
-    const int i2 = tgpig[1];
-    const int i1 = tgpig[0];
-
-    float corr_dims[2];
-    rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
-
-    device const int32_t * pos = (device const int32_t *) src1;
-
-    const float theta_base = (float) pos[i2];
-    const float inv_ndims = -1.f/args.n_dims;
-
-    float cos_theta;
-    float sin_theta;
-
-    for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
-        if (i0 >= args.n_offs && i0 < args.n_offs + args.n_dims) {
-            const int iw = i0 - args.n_offs; // relative idx
-            const int ic = iw/2;
-
-            const float theta = theta_base * pow(args.freq_base, inv_ndims*iw);
-
-            const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
-
-            rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, iw, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
-
-            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + (args.n_offs + ic)*args.nb00);
-            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + (args.n_offs + ic)*args.nb0);
-
-            const float x0 = src[0];
-            const float x1 = src[args.n_dims/2];
-
-            dst_data[0]             = x0*cos_theta - x1*sin_theta;
-            dst_data[args.n_dims/2] = x0*sin_theta + x1*cos_theta;
-        } else {
-            if (args.inplace) {
-                continue;
-            }
-
-            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
-            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + i0*args.nb0);
-
-            dst_data[0] = src[0];
-            dst_data[1] = src[1];
-        }
-    }
-}
-
-template<typename T>
-kernel void kernel_rope_multi(
-        constant ggml_metal_kargs_rope & args,
-        device const char * src0,
-        device const char * src1,
-        device const char * src2,
-        device       char * dst,
-        ushort  tiitg[[thread_index_in_threadgroup]],
-        ushort3 tptg [[threads_per_threadgroup]],
-        uint3   tgpig[[threadgroup_position_in_grid]]) {
-    const int i3 = tgpig[2];
-    const int i2 = tgpig[1];
-    const int i1 = tgpig[0];
-
-    float corr_dims[2];
-    rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
-
-    device const int32_t * pos = (device const int32_t *) src1;
-
-    const float inv_ndims = -1.f/args.n_dims;
-
-    float cos_theta;
-    float sin_theta;
-
-    for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
-        if (i0 >= args.n_offs && i0 < args.n_offs + args.n_dims) {
-            const int iw = i0 - args.n_offs; // relative idx
-            const int ic = iw/2;
-
-            // mrope theta calculations
-            // note: the rest is the same as kernel_rope_neox
-            const int sect_dims = args.sect_0 + args.sect_1 + args.sect_2 + args.sect_3;
-            const int sec_w01   = args.sect_0 + args.sect_1;               // end of section 1
-            const int sec_w012  = args.sect_0 + args.sect_1 + args.sect_2; // end of section 2
-            const int sector    = ic % sect_dims;
-
-            float theta_base;
-            if (FC_rope_is_imrope) {
-                if (sector % 3 == 1 && sector < 3 * args.sect_1) { // h
-                    theta_base = (float) pos[i2 + args.ne02 * 1];
-                } else if (sector % 3 == 2 && sector < 3 * args.sect_2) { // w
-                    theta_base = (float) pos[i2 + args.ne02 * 2];
-                } else if (sector % 3 == 0 && sector < 3 * args.sect_0) { // t
-                    theta_base = (float) pos[i2 + args.ne02 * 0];
-                } else { // e
-                    theta_base = (float) pos[i2 + args.ne02 * 3];
-                }
-            } else {
-                if (sector < args.sect_0) {
-                    theta_base = (float) pos[i2];
-                } else if (sector < sec_w01) {
-                    theta_base = (float) pos[i2 + args.ne02 * 1];
-                } else if (sector < sec_w012) {
-                    theta_base = (float) pos[i2 + args.ne02 * 2];
-                } else {
-                    theta_base = (float) pos[i2 + args.ne02 * 3];
-                }
-            }
-            // end of mrope
-
-            const float theta = theta_base * pow(args.freq_base, inv_ndims*iw);
-
-            const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
-
-            rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, iw, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
-
-            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + (args.n_offs + ic)*args.nb00);
-            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + (args.n_offs + ic)*args.nb0);
-
-            const float x0 = src[0];
-            const float x1 = src[args.n_dims/2];
-
-            dst_data[0]             = x0*cos_theta - x1*sin_theta;
-            dst_data[args.n_dims/2] = x0*sin_theta + x1*cos_theta;
-        } else {
-            if (args.inplace) {
-                continue;
-            }
-
-            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
-            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + i0*args.nb0);
-
-            dst_data[0] = src[0];
-            dst_data[1] = src[1];
-        }
-    }
-}
-
-template<typename T>
-kernel void kernel_rope_vision(
-        constant ggml_metal_kargs_rope & args,
-        device const char * src0,
-        device const char * src1,
-        device const char * src2,
-        device       char * dst,
-        ushort  tiitg[[thread_index_in_threadgroup]],
-        ushort3 tptg [[threads_per_threadgroup]],
-        uint3   tgpig[[threadgroup_position_in_grid]]) {
-    const int i3 = tgpig[2];
-    const int i2 = tgpig[1];
-    const int i1 = tgpig[0];
-
-    float corr_dims[2];
-    rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
-
-    device const int32_t * pos = (device const int32_t *) src1;
-
-    const float inv_ndims = -1.f/args.n_dims;
-
-    float cos_theta;
-    float sin_theta;
-
-    for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
-        if (i0 < 2*args.n_dims) { // different from kernel_rope_multi
-            const int ic = i0/2;
-
-            // mrope theta calculations (only support 2 dimensions)
-            const int sect_dims = args.sect_0 + args.sect_1;
-            const int sector    = ic % sect_dims;
-
-            float p;
-            float theta_base;
-            if (sector < args.sect_1) {
-                p = (float) sector;
-                theta_base = (float) pos[i2];
-            } else {
-                p = (float) sector - args.sect_0;
-                theta_base = (float) pos[i2 + args.ne02];
-            }
-
-            const float theta = theta_base * pow(args.freq_base, 2.0f * inv_ndims * p);
-            // end of mrope
-
-            const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
-
-            rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
-
-            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + ic*args.nb00);
-            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + ic*args.nb0);
-
-            const float x0 = src[0];
-            const float x1 = src[args.n_dims]; // different from kernel_rope_multi
-
-            dst_data[0]           = x0*cos_theta - x1*sin_theta;
-            dst_data[args.n_dims] = x0*sin_theta + x1*cos_theta; // different from kernel_rope_multi
-        } else {
-            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
-            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + i0*args.nb0);
-
-            dst_data[0] = src[0];
-            dst_data[1] = src[1];
-        }
-    }
-}
-
-typedef decltype(kernel_rope_norm<float>) kernel_rope_norm_t;
-typedef decltype(kernel_rope_neox<float>) kernel_rope_neox_t;
-typedef decltype(kernel_rope_multi<float>) kernel_rope_multi_t;
-typedef decltype(kernel_rope_vision<float>) kernel_rope_vision_t;
-
-template [[host_name("kernel_rope_norm_f32")]] kernel kernel_rope_norm_t kernel_rope_norm<float>;
-template [[host_name("kernel_rope_norm_f16")]] kernel kernel_rope_norm_t kernel_rope_norm<half>;
-
-template [[host_name("kernel_rope_neox_f32")]] kernel kernel_rope_neox_t kernel_rope_neox<float>;
-template [[host_name("kernel_rope_neox_f16")]] kernel kernel_rope_neox_t kernel_rope_neox<half>;
-
-template [[host_name("kernel_rope_multi_f32")]] kernel kernel_rope_multi_t kernel_rope_multi<float>;
-template [[host_name("kernel_rope_multi_f16")]] kernel kernel_rope_multi_t kernel_rope_multi<half>;
-
-template [[host_name("kernel_rope_vision_f32")]] kernel kernel_rope_vision_t kernel_rope_vision<float>;
-template [[host_name("kernel_rope_vision_f16")]] kernel kernel_rope_vision_t kernel_rope_vision<half>;
@@ -1,223 +0,0 @@
-#include "common.h"
-
-template<typename T>
-kernel void kernel_soft_max(
-        constant ggml_metal_kargs_soft_max & args,
-        device const  char * src0,
-        device const  char * src1,
-        device const  char * src2,
-        device        char * dst,
-        threadgroup  float * buf [[threadgroup(0)]],
-        uint3 tgpig[[threadgroup_position_in_grid]],
-        uint3 tpitg[[thread_position_in_threadgroup]],
-        uint  sgitg[[simdgroup_index_in_threadgroup]],
-        uint  tiisg[[thread_index_in_simdgroup]],
-        uint3  tptg[[threads_per_threadgroup]]) {
-    const int32_t i03 = tgpig.z;
-    const int32_t i02 = tgpig.y;
-    const int32_t i01 = tgpig.x;
-
-    const int32_t i13 = i03%args.ne13;
-    const int32_t i12 = i02%args.ne12;
-    const int32_t i11 = i01;
-
-    device const float * psrc0 =                (device const float *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
-    device const     T * pmask = src1 != src0 ? (device const T *    ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr;
-    device const float * psrc2 = src2 != src0 ? (device const float *) (src2)                                                 : nullptr;
-    device       float * pdst  =                (device       float *) (dst  + i01*args.nb1  + i02*args.nb2  + i03*args.nb3);
-
-    float slope = 1.0f;
-
-    // ALiBi
-    if (args.max_bias > 0.0f) {
-        const int32_t h = i02;
-
-        const float base = h < args.n_head_log2 ? args.m0 : args.m1;
-        const int   exp  = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
-
-        slope = pow(base, exp);
-    }
-
-    // parallel max
-    float lmax = psrc2 ? psrc2[i02] : -INFINITY;
-
-    for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
-        lmax = MAX(lmax, psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f));
-    }
-
-    // find the max value in the block
-    float max_val = simd_max(lmax);
-    if (tptg.x > N_SIMDWIDTH) {
-        if (sgitg == 0) {
-            buf[tiisg] = -INFINITY;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        if (tiisg == 0) {
-            buf[sgitg] = max_val;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        max_val = buf[tiisg];
-        max_val = simd_max(max_val);
-    }
-
-    // parallel sum
-    float lsum = 0.0f;
-    for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
-        const float exp_psrc0 = exp((psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f)) - max_val);
-        lsum += exp_psrc0;
-        pdst[i00] = exp_psrc0;
-    }
-
-    // This barrier fixes a failing test
-    // ref: https://github.com/ggml-org/ggml/pull/621#discussion_r1425156335
-    threadgroup_barrier(mem_flags::mem_none);
-
-    float sum = simd_sum(lsum);
-
-    if (tptg.x > N_SIMDWIDTH) {
-        if (sgitg == 0) {
-            buf[tiisg] = 0.0f;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        if (tiisg == 0) {
-            buf[sgitg] = sum;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        sum = buf[tiisg];
-        sum = simd_sum(sum);
-    }
-
-    if (psrc2) {
-        sum += exp(psrc2[i02] - max_val);
-    }
-
-    const float inv_sum = 1.0f/sum;
-
-    for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
-        pdst[i00] *= inv_sum;
-    }
-}
-
-template<typename T>
-kernel void kernel_soft_max_4(
-        constant ggml_metal_kargs_soft_max & args,
-        device const  char * src0,
-        device const  char * src1,
-        device const  char * src2,
-        device        char * dst,
-        threadgroup  float * buf [[threadgroup(0)]],
-        uint3 tgpig[[threadgroup_position_in_grid]],
-        uint3 tpitg[[thread_position_in_threadgroup]],
-        uint  sgitg[[simdgroup_index_in_threadgroup]],
-        uint  tiisg[[thread_index_in_simdgroup]],
-        uint3  tptg[[threads_per_threadgroup]]) {
-    const int32_t i03 = tgpig.z;
-    const int32_t i02 = tgpig.y;
-    const int32_t i01 = tgpig.x;
-
-    const int32_t i13 = i03%args.ne13;
-    const int32_t i12 = i02%args.ne12;
-    const int32_t i11 = i01;
-
-    device const float4 * psrc4 =                (device const float4 *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
-    device const      T * pmask = src1 != src0 ? (device const T *     ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr;
-    device const float *  psrc2 = src2 != src0 ? (device const float * ) (src2)                                                 : nullptr;
-    device       float4 * pdst4 =                (device       float4 *) (dst  + i01*args.nb1  + i02*args.nb2  + i03*args.nb3);
-
-    float slope = 1.0f;
-
-    if (args.max_bias > 0.0f) {
-        const int32_t h = i02;
-
-        const float base = h < args.n_head_log2 ? args.m0 : args.m1;
-        const int   exp  = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
-
-        slope = pow(base, exp);
-    }
-
-    // parallel max
-    float4 lmax4 = psrc2 ? psrc2[i02] : -INFINITY;
-
-    for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
-        lmax4 = fmax(lmax4, psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f)));
-    }
-
-    const float lmax = MAX(MAX(lmax4[0], lmax4[1]), MAX(lmax4[2], lmax4[3]));
-
-    float max_val = simd_max(lmax);
-    if (tptg.x > N_SIMDWIDTH) {
-        if (sgitg == 0) {
-            buf[tiisg] = -INFINITY;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        if (tiisg == 0) {
-            buf[sgitg] = max_val;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        max_val = buf[tiisg];
-        max_val = simd_max(max_val);
-    }
-
-    // parallel sum
-    float4 lsum4 = 0.0f;
-    for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
-        const float4 exp_psrc4 = exp((psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f))) - max_val);
-        lsum4 += exp_psrc4;
-        pdst4[i00] = exp_psrc4;
-    }
-
-    const float lsum = lsum4[0] + lsum4[1] + lsum4[2] + lsum4[3];
-
-    // This barrier fixes a failing test
-    // ref: https://github.com/ggml-org/ggml/pull/621#discussion_r1425156335
-    threadgroup_barrier(mem_flags::mem_none);
-
-    float sum = simd_sum(lsum);
-
-    if (tptg.x > N_SIMDWIDTH) {
-        if (sgitg == 0) {
-            buf[tiisg] = 0.0f;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        if (tiisg == 0) {
-            buf[sgitg] = sum;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        sum = buf[tiisg];
-        sum = simd_sum(sum);
-    }
-
-    if (psrc2) {
-        sum += exp(psrc2[i02] - max_val);
-    }
-
-    const float inv_sum = 1.0f/sum;
-
-    for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
-        pdst4[i00] *= inv_sum;
-    }
-}
-
-typedef decltype(kernel_soft_max<float>)    kernel_soft_max_t;
-typedef decltype(kernel_soft_max_4<float4>) kernel_soft_max_4_t;
-
-template [[host_name("kernel_soft_max_f16")]]   kernel kernel_soft_max_t   kernel_soft_max<half>;
-template [[host_name("kernel_soft_max_f32")]]   kernel kernel_soft_max_t   kernel_soft_max<float>;
-template [[host_name("kernel_soft_max_f16_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4<half4>;
-template [[host_name("kernel_soft_max_f32_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4<float4>;
@@ -1,75 +0,0 @@
-#include "common.h"
-
-constant short FC_solve_tri_nsg [[function_constant(FC_SOLVE_TRI + 0)]];
-constant short FC_solve_tri_n   [[function_constant(FC_SOLVE_TRI + 1)]];
-constant short FC_solve_tri_k   [[function_constant(FC_SOLVE_TRI + 2)]];
-
-kernel void kernel_solve_tri_f32(
-        constant ggml_metal_kargs_solve_tri & args,
-        device   const char * src0,
-        device   const char * src1,
-        device         char * dst,
-        threadgroup    char * shmem [[threadgroup(0)]],
-        ushort3 tgpig[[threadgroup_position_in_grid]],
-        ushort  sgitg[[simdgroup_index_in_threadgroup]],
-        ushort  tiisg[[thread_index_in_simdgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    constexpr short NW = N_SIMDWIDTH;
-
-    const short NSG = FC_solve_tri_nsg;
-    const short N   = FC_solve_tri_n;
-    const short K   = FC_solve_tri_k;
-    const short NP  = PAD2(N, NW);
-
-    const int32_t i03 = tgpig.z;
-    const int32_t i02 = tgpig.y;
-    const int32_t i01 = tgpig.x*NSG + sgitg;
-
-    threadgroup float * sh0 = (threadgroup float *) shmem;
-
-    device const float * src0_ptr = (device const float *)(src0 + i02 * args.nb02 + i03 * args.nb03) + sgitg*N;
-    device const float * src1_ptr = (device const float *)(src1 + i02 * args.nb12 + i03 * args.nb13) + i01;
-    device       float * dst_ptr  = (device       float *)(dst  + i02 * args.nb2  + i03 * args.nb3)  + i01;
-
-    for (short rr = 0; rr < N; rr += NSG) {
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        {
-            threadgroup float * sh0_cur = sh0 + sgitg*NP;
-
-            for (short t = 0; t*NW < N; ++t) {
-                const short idx = t*NW + tiisg;
-                sh0_cur[idx] = src0_ptr[idx];
-            }
-
-            src0_ptr += NSG*N;
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        if (i01 >= args.ne10) {
-            continue;
-        }
-
-        for (short ir = 0; ir < NSG && rr + ir < N; ++ir) {
-            const short r = rr + ir;
-
-            threadgroup float * sh0_cur = sh0 + ir*NP;
-
-            float sum = 0.0f;
-
-            for (short t = 0; t*NW < r; ++t) {
-                const short idx = t*NW + tiisg;
-                sum += sh0_cur[idx] * dst_ptr[idx*K] * (idx < r);
-            }
-
-            sum = simd_sum(sum);
-
-            if (tiisg == 0) {
-                const float diag = sh0_cur[r];
-
-                dst_ptr[r*K] = (src1_ptr[r*K] - sum) / diag;
-            }
-        }
-    }
-}
@@ -1,287 +0,0 @@
-#include "common.h"
-
-// ref: ggml.c:ggml_compute_forward_ssm_conv_f32
-kernel void kernel_ssm_conv_f32_f32(
-        constant ggml_metal_kargs_ssm_conv & args,
-        device const  void * src0,
-        device const  void * src1,
-        device       float * dst,
-        uint3 tgpig[[threadgroup_position_in_grid]],
-        uint3 tpitg[[thread_position_in_threadgroup]],
-        uint3   ntg[[threads_per_threadgroup]]) {
-    const int64_t ir = tgpig.x;
-    const int64_t i2 = tgpig.y;
-    const int64_t i3 = tgpig.z;
-
-    const int64_t nc  = args.ne10;
-  //const int64_t ncs = args.ne00;
-  //const int64_t nr  = args.ne01;
-  //const int64_t n_t = args.ne1;
-  //const int64_t n_s = args.ne2;
-
-    device const float * s = (device const float *) ((device const char *) src0 + ir*args.nb01 + i2*args.nb00 + i3*args.nb02);
-    device const float * c = (device const float *) ((device const char *) src1 + ir*args.nb11);
-    device       float * x = (device       float *) ((device       char *) dst  + ir*args.nb0  + i2*args.nb1  + i3*args.nb2);
-
-    float sumf = 0.0f;
-
-    for (int64_t i0 = 0; i0 < nc; ++i0) {
-        sumf += s[i0] * c[i0];
-    }
-
-    x[0] = sumf;
-}
-
-kernel void kernel_ssm_conv_f32_f32_4(
-        constant ggml_metal_kargs_ssm_conv & args,
-        device const  void * src0,
-        device const  void * src1,
-        device       float * dst,
-        uint3 tgpig[[threadgroup_position_in_grid]],
-        uint3 tpitg[[thread_position_in_threadgroup]],
-        uint3   ntg[[threads_per_threadgroup]]) {
-    const int64_t ir = tgpig.x;
-    const int64_t i2 = tgpig.y;
-    const int64_t i3 = tgpig.z;
-
-    const int64_t nc  = args.ne10;
-  //const int64_t ncs = args.ne00;
-  //const int64_t nr  = args.ne01;
-  //const int64_t n_t = args.ne1;
-  //const int64_t n_s = args.ne2;
-
-    device const float4 * s = (device const float4 *) ((device const char *) src0 + ir*args.nb01 + i2*args.nb00 + i3*args.nb02);
-    device const float4 * c = (device const float4 *) ((device const char *) src1 + ir*args.nb11);
-    device       float  * x = (device       float  *) ((device       char *) dst  + ir*args.nb0  + i2*args.nb1  + i3*args.nb2);
-
-    float sumf = 0.0f;
-
-    for (int64_t i0 = 0; i0 < nc/4; ++i0) {
-        sumf += dot(s[i0], c[i0]);
-    }
-
-    x[0] = sumf;
-}
-
-constant short FC_ssm_conv_bs   [[function_constant(FC_SSM_CONV + 0)]];
-
-// Batched version: each threadgroup processes multiple tokens for better efficiency
-// Thread layout: each thread handles one token, threadgroup covers BATCH_SIZE tokens
-kernel void kernel_ssm_conv_f32_f32_batched(
-        constant ggml_metal_kargs_ssm_conv & args,
-        device const  void * src0,
-        device const  void * src1,
-        device       float * dst,
-        uint3 tgpig[[threadgroup_position_in_grid]],
-        uint3 tpitg[[thread_position_in_threadgroup]],
-        uint3   ntg[[threads_per_threadgroup]]) {
-    // tgpig.x = row index (ir)
-    // tgpig.y = batch of tokens (i2_base / BATCH_SIZE)
-    // tgpig.z = sequence index (i3)
-    // tpitg.x = thread within batch (0..BATCH_SIZE-1)
-    const short BATCH_SIZE = FC_ssm_conv_bs;
-
-    const int64_t ir      = tgpig.x;
-    const int64_t i2_base = tgpig.y * BATCH_SIZE;
-    const int64_t i3      = tgpig.z;
-    const int64_t i2_off  = tpitg.x;
-    const int64_t i2      = i2_base + i2_off;
-
-    const int64_t nc  = args.ne10;  // conv kernel size (typically 4)
-    const int64_t n_t = args.ne1;   // number of tokens
-
-    // Bounds check for partial batches at the end
-    if (i2 >= n_t) {
-        return;
-    }
-
-    // Load conv weights (shared across all tokens for this row)
-    device const float * c = (device const float *) ((device const char *) src1 + ir*args.nb11);
-
-    // Load source for this specific token
-    device const float * s = (device const float *) ((device const char *) src0 + ir*args.nb01 + i2*args.nb00 + i3*args.nb02);
-
-    // Output location for this token
-    device float * x = (device float *) ((device char *) dst + ir*args.nb0 + i2*args.nb1 + i3*args.nb2);
-
-    float sumf = 0.0f;
-    for (int64_t i0 = 0; i0 < nc; ++i0) {
-        sumf += s[i0] * c[i0];
-    }
-
-    x[0] = sumf;
-}
-
-kernel void kernel_ssm_conv_f32_f32_batched_4(
-        constant ggml_metal_kargs_ssm_conv & args,
-        device const  void * src0,
-        device const  void * src1,
-        device       float * dst,
-        uint3 tgpig[[threadgroup_position_in_grid]],
-        uint3 tpitg[[thread_position_in_threadgroup]],
-        uint3   ntg[[threads_per_threadgroup]]) {
-    // tgpig.x = row index (ir)
-    // tgpig.y = batch of tokens (i2_base / BATCH_SIZE)
-    // tgpig.z = sequence index (i3)
-    // tpitg.x = thread within batch (0..BATCH_SIZE-1)
-    const short BATCH_SIZE = FC_ssm_conv_bs;
-
-    const int64_t ir      = tgpig.x;
-    const int64_t i2_base = tgpig.y * BATCH_SIZE;
-    const int64_t i3      = tgpig.z;
-    const int64_t i2_off  = tpitg.x;
-    const int64_t i2      = i2_base + i2_off;
-
-    const int64_t nc  = args.ne10;  // conv kernel size (typically 4)
-    const int64_t n_t = args.ne1;   // number of tokens
-
-    // Bounds check for partial batches at the end
-    if (i2 >= n_t) {
-        return;
-    }
-
-    // Load conv weights (shared across all tokens for this row)
-    device const float4 * c = (device const float4 *) ((device const char *) src1 + ir*args.nb11);
-
-    // Load source for this specific token
-    device const float4 * s = (device const float4 *) ((device const char *) src0 + ir*args.nb01 + i2*args.nb00 + i3*args.nb02);
-
-    // Output location for this token
-    device float * x = (device float *) ((device char *) dst + ir*args.nb0 + i2*args.nb1 + i3*args.nb2);
-
-    float sumf = 0.0f;
-    for (int64_t i0 = 0; i0 < nc/4; ++i0) {
-        sumf += dot(s[i0], c[i0]);
-    }
-
-    x[0] = sumf;
-}
-
-// ref: ggml.c:ggml_compute_forward_ssm_scan_f32, Mamba-2 part
-// Optimized version: reduces redundant memory loads by having one thread load shared values
-kernel void kernel_ssm_scan_f32(
-        constant ggml_metal_kargs_ssm_scan & args,
-        device const void * src0,
-        device const void * src1,
-        device const void * src2,
-        device const void * src3,
-        device const void * src4,
-        device const void * src5,
-        device const void * src6,
-        device      float * dst,
-        threadgroup float * shared [[threadgroup(0)]],
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort  sgitg[[simdgroup_index_in_threadgroup]],
-        ushort  tiisg[[thread_index_in_simdgroup]],
-        ushort  sgptg[[simdgroups_per_threadgroup]],
-        uint3    tgpg[[threadgroups_per_grid]]) {
-    constexpr short NW = N_SIMDWIDTH;
-
-    // Shared memory layout:
-    // [0..sgptg*NW-1]: partial sums for reduction (existing)
-    // [sgptg*NW..sgptg*NW+sgptg-1]: pre-computed x_dt values for each token in batch
-    // [sgptg*NW+sgptg..sgptg*NW+2*sgptg-1]: pre-computed dA values for each token in batch
-    threadgroup float * shared_sums = shared;
-    threadgroup float * shared_x_dt = shared + sgptg * NW;
-    threadgroup float * shared_dA   = shared + sgptg * NW + sgptg;
-
-    shared_sums[tpitg.x] = 0.0f;
-
-    const int32_t i0 = tpitg.x;
-    const int32_t i1 = tgpig.x;
-    const int32_t ir = tgpig.y; // current head
-    const int32_t i3 = tgpig.z; // current seq
-
-    const int32_t nc  = args.d_state;
-    const int32_t nr  = args.d_inner;
-    const int32_t nh  = args.n_head;
-    const int32_t ng  = args.n_group;
-    const int32_t n_t = args.n_seq_tokens;
-    const int32_t n_s = args.n_seqs;
-    const int32_t K   = args.K;
-
-    const int32_t s_off = args.s_off;
-
-    device const int32_t * ids = (device const int32_t *) src6;
-
-    device const float * s0_buff = (device const float *) ((device const char *) src0 + ir*args.nb02 + ids[i3]*args.nb03);
-    device       float * s_buff  = (device       float *) ((device       char *) dst  + ir*args.nb02 +      i3*args.nb03 + s_off);
-
-    const int32_t i = i0 + i1*nc;
-    const int32_t g = ir / (nh / ng); // repeat_interleave
-
-    float s0 = s0_buff[i];
-    float s  = 0.0f;
-
-    device const float * A = (device const float *) ((device const char *) src3 + ir*args.nb31); // {ne30, nh}
-
-    const float A0 = A[i0%args.ne30];
-
-    device const float * x  = (device const float *)((device const char *) src1 + i1*args.nb10  + ir*args.nb11 + i3*args.nb13); // {dim, nh, nt, ns}
-    device const float * dt = (device const float *)((device const char *) src2 + ir*args.nb20  + i3*args.nb22);                // {nh, nt, ns}
-    device const float * B  = (device const float *)((device const char *) src4 +  g*args.nb41  + i3*args.nb43);                // {d_state, ng, nt, ns}
-    device const float * C  = (device const float *)((device const char *) src5 +  g*args.nb51  + i3*args.nb53);                // {d_state, ng, nt, ns}
-
-    device float * y = dst + (i1 + ir*(nr) + i3*(n_t*nh*nr)); // {dim, nh, nt, ns}
-
-    for (int i2 = 0; i2 < n_t; i2 += sgptg) {
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        // Pre-compute x_dt and dA for this batch of tokens
-        // Only first sgptg threads do the loads and expensive math
-        if (i0 < sgptg && i2 + i0 < n_t) {
-            // ns12 and ns21 are element strides (nb12/nb10, nb21/nb20)
-            device const float * x_t  = x  + i0 * args.ns12;
-            device const float * dt_t = dt + i0 * args.ns21;
-
-            const float dt0  = dt_t[0];
-            const float dtsp = dt0 <= 20.0f ? log(1.0f + exp(dt0)) : dt0;
-            shared_x_dt[i0] = x_t[0] * dtsp;
-            shared_dA[i0]   = dtsp;  // Store dtsp, compute exp(dtsp * A0) per-thread since A0 varies
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        for (int t = 0; t < sgptg && i2 + t < n_t; t++) {
-            const float x_dt = shared_x_dt[t];
-            const float dA   = exp(shared_dA[t] * A0);
-
-            s = (s0 * dA) + (B[i0] * x_dt);
-
-            const float sumf = simd_sum(s * C[i0]);
-
-            if (tiisg == 0) {
-                shared_sums[t*NW + sgitg] = sumf;
-            }
-
-            // recurse
-            s0 = s;
-
-            const int32_t slot = n_t - 1 - (i2 + t);
-            if (slot > 0 && slot < K) {
-                device float * s_snapshot = (device float *) ((device char *) s_buff + (int64_t) slot*n_s*args.nb03);
-                s_snapshot[i] = s;
-            }
-
-            B  += args.ns42;
-            C  += args.ns52;
-        }
-
-        // Advance pointers for next batch
-        x  += sgptg * args.ns12;
-        dt += sgptg * args.ns21;
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        const float sumf = simd_sum(shared_sums[sgitg*NW + tiisg]);
-
-        if (tiisg == 0 && i2 + sgitg < n_t) {
-            y[sgitg*nh*nr] = sumf;
-        }
-
-        y += sgptg*nh*nr;
-    }
-
-    s_buff[i] = s;
-}
@@ -1,69 +0,0 @@
-#include "common.h"
-
-template<uint32_t ttype>
-bool _ggml_vec_tri_cmp(const int i, const int r);
-
-template<>
-bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_LOWER */ 3>(const int i, const int r) {
-    return i < r;
-}
-
-template<>
-bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_LOWER_DIAG */ 2>(const int i, const int r) {
-    return i <= r;
-}
-
-template<>
-bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_UPPER */ 1>(const int i, const int r) {
-    return i > r;
-}
-
-template<>
-bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_UPPER_DIAG */ 0>(const int i, const int r) {
-    return i >= r;
-}
-
-template<typename T, int ttype>
-kernel void kernel_tri(
-        constant ggml_metal_kargs_tri & args,
-        device const char * src0,
-        device const char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-    const int i3 = tgpig.z;
-    const int i2 = tgpig.y;
-    const int i1 = tgpig.x;
-
-    if (i3 >= args.ne03 || i2 >= args.ne02 || i1 >= args.ne01) {
-        return;
-    }
-
-    device const T * src_row = (device const T *) ((device const char *) src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03);
-    device       T * dst_row = (device       T *) ((device       char *) dst  + i1*args.nb1  + i2*args.nb2  + i3*args.nb3);
-
-    // Each thread is a single element of the row if ne00 < max threads per
-    // threadgroup, so this will loop once for each index that this thread is
-    // responsible for
-    for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) {
-        // Use the comparison as a mask for branchless
-        dst_row[i0] = static_cast<T>(_ggml_vec_tri_cmp<ttype>(i0, i1)) * src_row[i0];
-    }
-}
-
-typedef decltype(kernel_tri<float, 0>) kernel_tri_t;
-
-template [[host_name("kernel_tri_f32_0")]] kernel kernel_tri_t kernel_tri<float, 0>;
-template [[host_name("kernel_tri_f32_1")]] kernel kernel_tri_t kernel_tri<float, 1>;
-template [[host_name("kernel_tri_f32_2")]] kernel kernel_tri_t kernel_tri<float, 2>;
-template [[host_name("kernel_tri_f32_3")]] kernel kernel_tri_t kernel_tri<float, 3>;
-template [[host_name("kernel_tri_f16_0")]] kernel kernel_tri_t kernel_tri<half, 0>;
-template [[host_name("kernel_tri_f16_1")]] kernel kernel_tri_t kernel_tri<half, 1>;
-template [[host_name("kernel_tri_f16_2")]] kernel kernel_tri_t kernel_tri<half, 2>;
-template [[host_name("kernel_tri_f16_3")]] kernel kernel_tri_t kernel_tri<half, 3>;
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_tri_bf16_0")]] kernel kernel_tri_t kernel_tri<bfloat, 0>;
-template [[host_name("kernel_tri_bf16_1")]] kernel kernel_tri_t kernel_tri<bfloat, 1>;
-template [[host_name("kernel_tri_bf16_2")]] kernel kernel_tri_t kernel_tri<bfloat, 2>;
-template [[host_name("kernel_tri_bf16_3")]] kernel kernel_tri_t kernel_tri<bfloat, 3>;
-#endif
@@ -1,374 +0,0 @@
-#include "common.h"
-
-constant short FC_unary_op [[function_constant(FC_UNARY + 0)]];
-constant bool  FC_unary_cnt[[function_constant(FC_UNARY + 1)]];
-
-template <typename T0, typename T, typename TC>
-kernel void kernel_unary_impl(
-        constant ggml_metal_kargs_unary & args,
-        device const char * src0,
-        device       char * dst,
-        uint3   tgpig[[threadgroup_position_in_grid]],
-        ushort3 tpitg[[thread_position_in_threadgroup]],
-        ushort3   ntg[[threads_per_threadgroup]]) {
-#define FC_OP  FC_unary_op
-#define FC_CNT FC_unary_cnt
-
-    device const T0 * src0_ptr;
-    device       T  * dst_ptr;
-
-    int i0;
-
-    if (FC_CNT) {
-        i0 = tgpig.x;
-
-        src0_ptr = (device const T0 *) (src0);
-        dst_ptr  = (device       T  *) (dst);
-    } else {
-        const int i03 = tgpig.z;
-        const int i02 = tgpig.y;
-        const int k0  = tgpig.x/args.ne01;
-        const int i01 = tgpig.x - k0*args.ne01;
-
-        i0 = k0*ntg.x + tpitg.x;
-
-        src0_ptr = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
-        dst_ptr  = (device       T  *) (dst  + i03*args.nb3  + i02*args.nb2  + i01*args.nb1 );
-    }
-
-    {
-        //threadgroup_barrier(mem_flags::mem_none);
-
-        if (!FC_CNT) {
-            if (i0 >= args.ne0) {
-                return;
-            }
-        }
-
-        const TC x = (TC) src0_ptr[i0];
-
-        if (FC_OP == OP_UNARY_NUM_SCALE) {
-            dst_ptr[i0] = (T) (args.scale * x + args.bias);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_FILL) {
-            dst_ptr[i0] = (T) args.val;
-        }
-
-        if (FC_OP == OP_UNARY_NUM_CLAMP) {
-            dst_ptr[i0] = (T) clamp(x, args.min, args.max);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_SQR) {
-            dst_ptr[i0] = (T) (x * x);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_SQRT) {
-            dst_ptr[i0] = (T) sqrt(x);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_SIN) {
-            dst_ptr[i0] = (T) sin(x);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_COS) {
-            dst_ptr[i0] = (T) cos(x);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_LOG) {
-            dst_ptr[i0] = (T) log(x);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_LEAKY_RELU) {
-            dst_ptr[i0] = (T) (TC(x > 0)*x + TC(x <= 0)*(x * args.slope));
-        }
-
-        if (FC_OP == OP_UNARY_NUM_TANH) {
-            dst_ptr[i0] = (T) precise::tanh(x);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_RELU) {
-            dst_ptr[i0] = (T) fmax(0, x);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_SIGMOID) {
-            dst_ptr[i0] = (T) (1 / (1 + exp(-x)));
-        }
-
-        if (FC_OP == OP_UNARY_NUM_GELU) {
-            dst_ptr[i0] = (T) (0.5*x*(1 + precise::tanh(SQRT_2_OVER_PI*x*(1 + GELU_COEF_A*x*x))));
-        }
-
-        if (FC_OP == OP_UNARY_NUM_GELU_ERF) {
-            dst_ptr[i0] = (T) (0.5*x*(1 + erf_approx(SQRT_2_INV*x)));
-        }
-
-        if (FC_OP == OP_UNARY_NUM_GELU_QUICK) {
-            dst_ptr[i0] = (T) (x * (1/(1 + exp(GELU_QUICK_COEF*x))));
-        }
-
-        if (FC_OP == OP_UNARY_NUM_SILU) {
-            dst_ptr[i0] = (T) (x / (1 + exp(-x)));
-        }
-
-        if (FC_OP == OP_UNARY_NUM_ELU) {
-            dst_ptr[i0] = (T) elu_approx(x);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_NEG) {
-            dst_ptr[i0] = (T) -x;
-        }
-
-        if (FC_OP == OP_UNARY_NUM_ABS) {
-            dst_ptr[i0] = (T) fabs(x);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_SGN) {
-            dst_ptr[i0] = T(x > 0) - T(x < 0);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_STEP) {
-            dst_ptr[i0] = T(x > 0);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_HARDSWISH) {
-            dst_ptr[i0] = (T) (x * fmax(0, fmin(1, x/6 + 0.5)));
-        }
-
-        if (FC_OP == OP_UNARY_NUM_HARDSIGMOID) {
-            dst_ptr[i0] = (T) fmax(0, fmin(1, x/6 + 0.5));
-        }
-
-        if (FC_OP == OP_UNARY_NUM_EXP) {
-            dst_ptr[i0] = (T) exp(x);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_SOFTPLUS) {
-            dst_ptr[i0] = (T) select(log(1 + exp(x)), x, x > 20);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_EXPM1) {
-            // TODO: precise implementation
-            dst_ptr[i0] = (T) (exp(x) - 1);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_FLOOR) {
-            dst_ptr[i0] = (T) floor(x);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_CEIL) {
-            dst_ptr[i0] = (T) ceil(x);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_ROUND) {
-            dst_ptr[i0] = (T) round(x);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_TRUNC) {
-            dst_ptr[i0] = (T) trunc(x);
-        }
-
-        if (FC_OP == OP_UNARY_NUM_XIELU) {
-            const TC xi      = x;
-            const TC gate    = TC(xi > TC(0.0f));
-            const TC clamped = fmin(xi, TC(args.val));
-            const TC y_pos   = TC(args.scale) * xi * xi + TC(args.bias) * xi;
-            const TC y_neg   = (exp(clamped) - TC(1.0f) - xi) * TC(args.slope) + TC(args.bias) * xi;
-            dst_ptr[i0] = (T) (gate * y_pos + (TC(1.0f) - gate) * y_neg);
-        }
-    }
-
-#undef FC_OP
-#undef FC_CNT
-}
-
-typedef decltype(kernel_unary_impl<float, float, float>) kernel_unary_t;
-
-template [[host_name("kernel_unary_f32_f32")]]   kernel kernel_unary_t kernel_unary_impl<float,  float,  float>;
-template [[host_name("kernel_unary_f32_f32_4")]] kernel kernel_unary_t kernel_unary_impl<float4, float4, float4>;
-template [[host_name("kernel_unary_f16_f16")]]   kernel kernel_unary_t kernel_unary_impl<half,   half,   float>;
-template [[host_name("kernel_unary_f16_f16_4")]] kernel kernel_unary_t kernel_unary_impl<half4,  half4,  float4>;
-
-kernel void kernel_silu_back_f32(
-        constant ggml_metal_kargs_silu_back & args,
-        device const float * dy,
-        device const float * x,
-        device       float * dx,
-        uint gid [[thread_position_in_grid]]) {
-    if (gid >= args.ne) {
-        return;
-    }
-
-    const float s = 1.0f / (1.0f + exp(-x[gid]));
-    dx[gid] = dy[gid] * s * (1.0f + x[gid] * (1.0f - s));
-}
-
-template<typename T>
-kernel void kernel_reglu(
-        constant ggml_metal_kargs_glu & args,
-        device const char * src0,
-        device const char * src1,
-        device       char * dst,
-        uint tgpig[[threadgroup_position_in_grid]],
-        uint tpitg[[thread_position_in_threadgroup]],
-        uint   ntg[[threads_per_threadgroup]]) {
-    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
-    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
-    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
-
-    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
-        const float x0 = src0_row[i0];
-        const float x1 = src1_row[i0];
-
-        dst_row[i0] = (T)(x0*x1*(x0 > 0.0f));
-    }
-}
-
-typedef decltype(kernel_reglu<float>) kernel_reglu_t;
-
-template [[host_name("kernel_reglu_f32")]] kernel kernel_reglu_t kernel_reglu<float>;
-template [[host_name("kernel_reglu_f16")]] kernel kernel_reglu_t kernel_reglu<half>;
-
-template<typename T>
-kernel void kernel_geglu(
-        constant ggml_metal_kargs_glu & args,
-        device const char * src0,
-        device const char * src1,
-        device       char * dst,
-        uint tgpig[[threadgroup_position_in_grid]],
-        uint tpitg[[thread_position_in_threadgroup]],
-        uint   ntg[[threads_per_threadgroup]]) {
-    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
-    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
-    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
-
-    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
-        const float x0 = src0_row[i0];
-        const float x1 = src1_row[i0];
-
-        const float gelu = 0.5f*x0*(1.0f + precise::tanh(SQRT_2_OVER_PI*x0*(1.0f + GELU_COEF_A*x0*x0)));
-
-        dst_row[i0] = (T)(gelu*x1);
-    }
-}
-
-typedef decltype(kernel_geglu<float>) kernel_geglu_t;
-
-template [[host_name("kernel_geglu_f32")]] kernel kernel_geglu_t kernel_geglu<float>;
-template [[host_name("kernel_geglu_f16")]] kernel kernel_geglu_t kernel_geglu<half>;
-
-template<typename T>
-kernel void kernel_swiglu(
-        constant ggml_metal_kargs_glu & args,
-        device const char * src0,
-        device const char * src1,
-        device       char * dst,
-        uint tgpig[[threadgroup_position_in_grid]],
-        uint tpitg[[thread_position_in_threadgroup]],
-        uint   ntg[[threads_per_threadgroup]]) {
-    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
-    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
-    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
-
-    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
-        const float x0 = src0_row[i0];
-        const float x1 = src1_row[i0];
-
-        const float silu = x0 / (1.0f + exp(-x0));
-
-        dst_row[i0] = (T)(silu*x1);
-    }
-}
-
-typedef decltype(kernel_swiglu<float>) kernel_swiglu_t;
-
-template [[host_name("kernel_swiglu_f32")]] kernel kernel_swiglu_t kernel_swiglu<float>;
-template [[host_name("kernel_swiglu_f16")]] kernel kernel_swiglu_t kernel_swiglu<half>;
-
-template<typename T>
-kernel void kernel_swiglu_oai(
-        constant ggml_metal_kargs_glu & args,
-        device const char * src0,
-        device const char * src1,
-        device       char * dst,
-        uint tgpig[[threadgroup_position_in_grid]],
-        uint tpitg[[thread_position_in_threadgroup]],
-        uint   ntg[[threads_per_threadgroup]]) {
-    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
-    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
-    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
-
-    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
-        float x0 = src0_row[i0];
-        float x1 = src1_row[i0];
-
-        x0 = min(x0, args.limit);
-        x1 = max(min(x1, args.limit), -args.limit);
-
-        float out_glu = x0 / (1.0f + exp(-x0 * args.alpha));
-        out_glu = out_glu * (1.0f + x1);
-
-        dst_row[i0] = (T)out_glu;
-    }
-}
-
-typedef decltype(kernel_swiglu_oai<float>) kernel_swiglu_oai_t;
-
-template [[host_name("kernel_swiglu_oai_f32")]] kernel kernel_swiglu_oai_t kernel_swiglu_oai<float>;
-template [[host_name("kernel_swiglu_oai_f16")]] kernel kernel_swiglu_oai_t kernel_swiglu_oai<half>;
-
-template<typename T>
-kernel void kernel_geglu_erf(
-        constant ggml_metal_kargs_glu & args,
-        device const char * src0,
-        device const char * src1,
-        device       char * dst,
-        uint tgpig[[threadgroup_position_in_grid]],
-        uint tpitg[[thread_position_in_threadgroup]],
-        uint   ntg[[threads_per_threadgroup]]) {
-    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
-    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
-    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
-
-    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
-        const float x0 = src0_row[i0];
-        const float x1 = src1_row[i0];
-
-        const float gelu_erf = 0.5f*x0*(1.0f+erf_approx<float>(x0*SQRT_2_INV));
-
-        dst_row[i0] = (T)(gelu_erf*x1);
-    }
-}
-
-typedef decltype(kernel_geglu_erf<float>) kernel_geglu_erf_t;
-
-template [[host_name("kernel_geglu_erf_f32")]] kernel kernel_geglu_erf_t kernel_geglu_erf<float>;
-template [[host_name("kernel_geglu_erf_f16")]] kernel kernel_geglu_erf_t kernel_geglu_erf<half>;
-
-template<typename T>
-kernel void kernel_geglu_quick(
-        constant ggml_metal_kargs_glu & args,
-        device const char * src0,
-        device const char * src1,
-        device       char * dst,
-        uint tgpig[[threadgroup_position_in_grid]],
-        uint tpitg[[thread_position_in_threadgroup]],
-        uint   ntg[[threads_per_threadgroup]]) {
-    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
-    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
-    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
-
-    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
-        const float x0 = src0_row[i0];
-        const float x1 = src1_row[i0];
-
-        const float gelu_quick = x0*(1.0f/(1.0f+exp(GELU_QUICK_COEF*x0)));
-
-        dst_row[i0] = (T)(gelu_quick*x1);
-    }
-}
-
-typedef decltype(kernel_geglu_quick<float>) kernel_geglu_quick_t;
-
-template [[host_name("kernel_geglu_quick_f32")]] kernel kernel_geglu_quick_t kernel_geglu_quick<float>;
-template [[host_name("kernel_geglu_quick_f16")]] kernel kernel_geglu_quick_t kernel_geglu_quick<half>;
@@ -1,179 +0,0 @@
-#include "common.h"
-
-constant bool FC_upscale_aa [[function_constant(FC_UPSCALE + 0)]];
-
-kernel void kernel_upscale_nearest_f32(
-    constant ggml_metal_kargs_upscale & args,
-    device  const char * src0,
-    device        char * dst,
-    uint3 tgpig[[threadgroup_position_in_grid]],
-    uint3 tpitg[[thread_position_in_threadgroup]],
-    uint3   ntg[[threads_per_threadgroup]]) {
-
-    const int64_t i3 = tgpig.z;
-    const int64_t i2 = tgpig.y;
-    const int64_t i1 = tgpig.x;
-
-    const int64_t i03 = i3/args.sf3;
-    const int64_t i02 = i2/args.sf2;
-    const int64_t i01 = i1/args.sf1;
-
-    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
-        const int64_t i00 = i0/args.sf0;
-
-        device const float * src0_ptr = (device const float *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + i00*args.nb00);
-        device       float * dst_ptr  = (device       float *) (dst  +  i3*args.nb3  +  i2*args.nb2  +  i1*args.nb1  +  i0*args.nb0);
-
-        dst_ptr[0] = src0_ptr[0];
-    }
-}
-
-static inline float bilinear_tri(float x) {
-    return MAX(0.0f, 1.0f - fabs(x));
-}
-
-kernel void kernel_upscale_bilinear_f32(
-    constant ggml_metal_kargs_upscale & args,
-    device  const char * src0,
-    device        char * dst,
-    uint3 tgpig[[threadgroup_position_in_grid]],
-    uint3 tpitg[[thread_position_in_threadgroup]],
-    uint3   ntg[[threads_per_threadgroup]]) {
-
-    const int64_t i3 = tgpig.z;
-    const int64_t i2 = tgpig.y;
-    const int64_t i1 = tgpig.x;
-
-    const int64_t i03 = i3 / args.sf3;
-    const int64_t i02 = i2 / args.sf2;
-
-    const float   f01  = ((float)i1 + args.poffs) / args.sf1 - args.poffs;
-    const int64_t i01  = MAX(0, MIN(args.ne01 - 1, (int64_t)floor(f01)));
-    const int64_t i01p = MAX(0, MIN(args.ne01 - 1, i01 + 1));
-    const float   fd1  = MAX(0.0f, MIN(1.0f, f01 - (float)i01));
-
-    src0 += i03*args.nb03 + i02*args.nb02;
-
-    device float * dst_ptr = (device float *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1);
-
-    if (FC_upscale_aa) {
-        const float support0  = MAX(1.0f, 1.0f / args.sf0);
-        const float invscale0 = 1.0f / support0;
-        const float support1  = MAX(1.0f, 1.0f / args.sf1);
-        const float invscale1 = 1.0f / support1;
-
-        for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
-            const float f00 = ((float)i0 + args.poffs) / args.sf0 - args.poffs;
-
-            int64_t x_min = MAX((int64_t)0, (int64_t)floor(f00 - support0 + args.poffs));
-            int64_t x_max = MIN(args.ne00,  (int64_t)ceil (f00 + support0 + args.poffs));
-
-            int64_t y_min = MAX((int64_t)0, (int64_t)floor(f01 - support1 + args.poffs));
-            int64_t y_max = MIN(args.ne01,  (int64_t)ceil (f01 + support1 + args.poffs));
-
-            float sum = 0.0f;
-            float wsum = 0.0f;
-
-            for (int64_t sy = y_min; sy < y_max; ++sy) {
-                const float wy = MAX(0.0f, 1.0f - fabs((float)sy - f01) * invscale1);
-                for (int64_t sx = x_min; sx < x_max; ++sx) {
-                    const float wx = MAX(0.0f, 1.0f - fabs((float)sx - f00) * invscale0);
-                    const float w  = wx * wy;
-                    device const float * src_ptr = (device const float *)(src0 + sy*args.nb01 + sx*args.nb00);
-                    sum  += (*src_ptr) * w;
-                    wsum += w;
-                }
-            }
-
-            const float v = (wsum > 0.0f) ? (sum / wsum) : 0.0f;
-            dst_ptr[i0] = v;
-        }
-    } else {
-        for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
-            const float   f00  = ((float)i0 + args.poffs) / args.sf0 - args.poffs;
-            const int64_t i00  = MAX(0, MIN(args.ne00 - 1, (int64_t)floor(f00)));
-            const int64_t i00p = MAX(0, MIN(args.ne00 - 1, i00 + 1));
-            const float   fd0  = MAX(0.0f, MIN(1.0f, f00 - (float)i00));
-
-            device const float * src00 = (device const float *)(src0 + i01*args.nb01  + i00*args.nb00);
-            device const float * src10 = (device const float *)(src0 + i01*args.nb01  + i00p*args.nb00);
-            device const float * src01 = (device const float *)(src0 + i01p*args.nb01 + i00*args.nb00);
-            device const float * src11 = (device const float *)(src0 + i01p*args.nb01 + i00p*args.nb00);
-
-            const float v =
-                (*src00) * (1.0f - fd0) * (1.0f - fd1) +
-                (*src10) * fd0          * (1.0f - fd1) +
-                (*src01) * (1.0f - fd0) * fd1 +
-                (*src11) * fd0          * fd1;
-
-            dst_ptr[i0] = v;
-        }
-    }
-}
-
-static inline float bicubic_weight1(float x) {
-    const float a = -0.75f;
-    return ((a + 2) * x - (a + 3)) * x * x + 1;
-}
-
-static inline float bicubic_weight2(float x) {
-    const float a = -0.75f;
-    return ((a * x - 5 * a) * x + 8 * a) * x - 4 * a;
-}
-
-kernel void kernel_upscale_bicubic_f32(
-    constant ggml_metal_kargs_upscale & args,
-    device  const char * src0,
-    device        char * dst,
-    uint3 tgpig[[threadgroup_position_in_grid]],
-    uint3 tpitg[[thread_position_in_threadgroup]],
-    uint3   ntg[[threads_per_threadgroup]]) {
-
-    const int64_t i3 = tgpig.z;
-    const int64_t i2 = tgpig.y;
-    const int64_t i1 = tgpig.x;
-
-    const int64_t i03 = i3 / args.sf3;
-    const int64_t i02 = i2 / args.sf2;
-
-    const float   f01 = ((float)i1 + args.poffs) / args.sf1 - args.poffs;
-    const int64_t i01 = (int64_t)floor(f01);
-    const float   fd1 = f01 - (float)i01;
-
-    const float w_y0 = bicubic_weight2(fd1 + 1.0f);
-    const float w_y1 = bicubic_weight1(fd1);
-    const float w_y2 = bicubic_weight1(1.0f - fd1);
-    const float w_y3 = bicubic_weight2(2.0f - fd1);
-
-    const device char * src_slice = src0 + i03 * args.nb03 + i02 * args.nb02;
-
-    device float * dst_ptr = (device float *)(dst + i3 * args.nb3 + i2 * args.nb2 + i1 * args.nb1);
-
-    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
-        const float   f00 = ((float)i0 + args.poffs) / args.sf0 - args.poffs;
-        const int64_t i00 = (int64_t)floor(f00);
-        const float   fd0 = f00 - (float)i00;
-
-        const float w_x0 = bicubic_weight2(fd0 + 1.0f);
-        const float w_x1 = bicubic_weight1(fd0);
-        const float w_x2 = bicubic_weight1(1.0f - fd0);
-        const float w_x3 = bicubic_weight2(2.0f - fd0);
-
-        float sum = 0.0f;
-
-        for (int dy = -1; dy <= 2; ++dy) {
-            const int64_t iy = MAX(0, MIN(args.ne01 - 1, i01 + dy));
-            const float wy = (dy == -1) ? w_y0 : (dy == 0) ? w_y1 : (dy == 1) ? w_y2 : w_y3;
-
-            for (int dx = -1; dx <= 2; ++dx) {
-                const int64_t ix = MAX(0, MIN(args.ne00 - 1, i00 + dx));
-                const float wx = (dx == -1) ? w_x0 : (dx == 0) ? w_x1 : (dx == 1) ? w_x2 : w_x3;
-
-                device const float * src_ptr = (device const float *)(src_slice + iy * args.nb01 + ix * args.nb00);
-                sum += (*src_ptr) * wx * wy;
-            }
-        }
-
-        dst_ptr[i0] = sum;
-    }
-}
@@ -1,179 +0,0 @@
-#include "common.h"
-
-kernel void kernel_rwkv_wkv6_f32(
-    device const float * k,
-    device const float * v,
-    device const float * r,
-    device const float * tf,
-    device const float * td,
-    device const float * state_in,
-    device       float * dst,
-    constant    uint & B,
-    constant    uint & T,
-    constant    uint & C,
-    constant    uint & H,
-    uint3 tgpig[[threadgroup_position_in_grid]],
-    uint3 tpitg[[thread_position_in_threadgroup]],
-    uint3   ntg[[threads_per_threadgroup]])  {
-
-    const uint head_size = 64; // TODO: support head_size = 128
-    const uint batch_id = tgpig.x / H;
-    const uint head_id = tgpig.x % H;
-    const uint tid = tpitg.x;
-
-    if (batch_id >= B || head_id >= H) {
-        return;
-    }
-
-    const uint state_size = C * head_size;
-    const uint n_seq_tokens = T / B;
-
-    threadgroup float _k[head_size];
-    threadgroup float _r[head_size];
-    threadgroup float _tf[head_size];
-    threadgroup float _td[head_size];
-
-    float state[head_size];
-
-    for (uint i = 0; i < head_size; i++) {
-        state[i] = state_in[batch_id * state_size + head_id * head_size * head_size
-                          + i * head_size + tid];
-    }
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    _tf[tid] = tf[head_id * head_size + tid];
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid;
-    const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid;
-
-    for (uint t = start_t; t < end_t; t += C) {
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        _k[tid] = k[t];
-        _r[tid] = r[t];
-        _td[tid] = td[t];
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        const float v_val = v[t];
-        float y = 0.0;
-
-        for (uint j = 0; j < head_size; j += 4) {
-            float4 k_vec = float4(_k[j], _k[j+1], _k[j+2], _k[j+3]);
-            float4 r_vec = float4(_r[j], _r[j+1], _r[j+2], _r[j+3]);
-            float4 tf_vec = float4(_tf[j], _tf[j+1], _tf[j+2], _tf[j+3]);
-            float4 td_vec = float4(_td[j], _td[j+1], _td[j+2], _td[j+3]);
-            float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]);
-
-            float4 kv = k_vec * v_val;
-
-            float4 temp = tf_vec * kv + s_vec;
-            y += dot(r_vec, temp);
-
-            s_vec = s_vec * td_vec + kv;
-            state[j]   = s_vec[0];
-            state[j+1] = s_vec[1];
-            state[j+2] = s_vec[2];
-            state[j+3] = s_vec[3];
-        }
-
-        dst[t] = y;
-    }
-
-    for (uint i = 0; i < head_size; i++) {
-        dst[T * C + batch_id * state_size + head_id * head_size * head_size
-            + i * head_size + tid] = state[i];
-    }
-}
-
-kernel void kernel_rwkv_wkv7_f32(
-    device const float * r,
-    device const float * w,
-    device const float * k,
-    device const float * v,
-    device const float * a,
-    device const float * b,
-    device const float * state_in,
-    device       float * dst,
-    constant    uint & B,
-    constant    uint & T,
-    constant    uint & C,
-    constant    uint & H,
-    uint3 tgpig[[threadgroup_position_in_grid]],
-    uint3 tpitg[[thread_position_in_threadgroup]],
-    uint3   ntg[[threads_per_threadgroup]])  {
-
-    const uint head_size = 64; // TODO: support head_size = 128
-    const uint batch_id = tgpig.x / H;
-    const uint head_id = tgpig.x % H;
-    const uint tid = tpitg.x;
-
-    if (batch_id >= B || head_id >= H) {
-        return;
-    }
-
-    const uint state_size = C * head_size;
-    const uint n_seq_tokens = T / B;
-
-    threadgroup float _r[head_size];
-    threadgroup float _w[head_size];
-    threadgroup float _k[head_size];
-    threadgroup float _a[head_size];
-    threadgroup float _b[head_size];
-
-    float state[head_size];
-
-    for (uint i = 0; i < head_size; i++) {
-        state[i] = state_in[batch_id * state_size + head_id * head_size * head_size
-                          + tid * head_size + i];
-    }
-
-    const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid;
-    const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid;
-
-    for (uint t = start_t; t < end_t; t += C) {
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        _r[tid] = r[t];
-        _w[tid] = w[t];
-        _k[tid] = k[t];
-        _a[tid] = a[t];
-        _b[tid] = b[t];
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        const float v_val = v[t];
-        float y = 0.0, sa = 0.0;
-
-        float4 sa_vec(0.0);
-
-        for (uint j = 0; j < head_size; j += 4) {
-            float4 a_vec = float4(_a[j], _a[j+1], _a[j+2], _a[j+3]);
-            float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]);
-            sa_vec += a_vec * s_vec;
-        }
-        sa = sa_vec[0] + sa_vec[1] + sa_vec[2] + sa_vec[3];
-
-        for (uint j = 0; j < head_size; j += 4) {
-            float4 r_vec = float4(_r[j], _r[j+1], _r[j+2], _r[j+3]);
-            float4 w_vec = float4(_w[j], _w[j+1], _w[j+2], _w[j+3]);
-            float4 k_vec = float4(_k[j], _k[j+1], _k[j+2], _k[j+3]);
-            float4 b_vec = float4(_b[j], _b[j+1], _b[j+2], _b[j+3]);
-            float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]);
-
-            float4 kv = k_vec * v_val;
-
-            s_vec = s_vec * w_vec + kv + sa * b_vec;
-            y += dot(s_vec, r_vec);
-
-            state[j]   = s_vec[0];
-            state[j+1] = s_vec[1];
-            state[j+2] = s_vec[2];
-            state[j+3] = s_vec[3];
-        }
-
-        dst[t] = y;
-    }
-
-    for (uint i = 0; i < head_size; i++) {
-        dst[T * C + batch_id * state_size + head_id * head_size * head_size
-            + tid * head_size + i] = state[i];
-    }
-}
@@ -76,6 +76,19 @@ static void dequantize_row_q2_K_sycl(const void *vx, dst_t *y, const int64_t k,
 #endif
 }

+template <typename dst_t>
+static void dequantize_row_q2_K_sycl_reorder(const void *vx, dst_t *y, const int64_t k,
+                                             dpct::queue_ptr stream) {
+    const int64_t nb = k / QK_K;
+
+    dpct::has_capability_or_fail(stream->get_device(), { sycl::aspect::fp16 });
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, nb) * sycl::range<3>(1, 1, 64), sycl::range<3>(1, 1, 64)),
+        [=](sycl::nd_item<3> item_ct1) {
+            dequantize_block_q2_K_reorder(vx, y, item_ct1, nb);
+        });
+}
+
 template <typename dst_t>
 static void dequantize_row_q3_K_sycl(const void *vx, dst_t *y, const int64_t k,
                                     dpct::queue_ptr stream) {
@@ -667,7 +680,11 @@ to_fp16_sycl_t ggml_get_to_fp16_sycl(ggml_type type, ggml_tensor * dst) {
                return dequantize_block_sycl<QK8_0, QR8_0, dequantize_q8_0>;
            }
        case GGML_TYPE_Q2_K:
-            return dequantize_row_q2_K_sycl;
+            if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
+                return dequantize_row_q2_K_sycl_reorder;
+            } else {
+                return dequantize_row_q2_K_sycl;
+            }
        case GGML_TYPE_Q3_K:
            if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
                return dequantize_row_q3_K_sycl_reorder;
@@ -753,7 +770,11 @@ to_fp32_sycl_t ggml_get_to_fp32_sycl(ggml_type type, ggml_tensor *dst) {
                return dequantize_block_sycl<QK8_0, QR8_0, dequantize_q8_0>;
            }
        case GGML_TYPE_Q2_K:
-            return dequantize_row_q2_K_sycl;
+            if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
+                return dequantize_row_q2_K_sycl_reorder;
+            } else {
+                return dequantize_row_q2_K_sycl;
+            }
        case GGML_TYPE_Q3_K:
            if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
                return dequantize_row_q3_K_sycl_reorder;
@@ -943,6 +943,47 @@ static void dequantize_block_q2_K(const void * __restrict__ vx, dst_t * __restri

 }

+template<typename dst_t>
+static void dequantize_block_q2_K_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy,
+                                          const sycl::nd_item<3> & item_ct1, int64_t n_blocks) {
+#if QK_K == 256
+    const int64_t i = item_ct1.get_group(2);
+    if (i >= n_blocks) {
+        return;
+    }
+
+    const uint8_t * base          = static_cast<const uint8_t *>(vx);
+    const size_t    qs_offset     = i * (QK_K / 4);
+    const size_t    scales_offset = n_blocks * (QK_K / 4) + i * (QK_K / 16);
+    const size_t    dm_offset     = n_blocks * (QK_K / 4) + n_blocks * (QK_K / 16) + i * sizeof(ggml_half2);
+
+    const uint8_t *     qs     = base + qs_offset;
+    const uint8_t *     scales = base + scales_offset;
+    const ggml_half2 * dm     = reinterpret_cast<const ggml_half2 *>(base + dm_offset);
+
+    const int64_t tid = item_ct1.get_local_id(2);
+    const int64_t n    = tid / 32;
+    const int64_t l    = tid - 32 * n;
+    const int64_t is   = 8 * n + l / 16;
+
+    const uint8_t q = qs[32 * n + l];
+    dst_t * y = yy + i * QK_K + 128 * n;
+
+    const float dall = (*dm)[0];
+    const float dmin = (*dm)[1];
+    y[l+ 0] = dall * (scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (scales[is+0] >> 4);
+    y[l+32] = dall * (scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (scales[is+2] >> 4);
+    y[l+64] = dall * (scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (scales[is+4] >> 4);
+    y[l+96] = dall * (scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (scales[is+6] >> 4);
+#else
+    GGML_UNUSED(vx);
+    GGML_UNUSED(yy);
+    GGML_UNUSED(item_ct1);
+    GGML_UNUSED(n_blocks);
+    GGML_ABORT("Q2_K reorder dequantize not supported for QK_K != 256");
+#endif
+}
+
 template<typename dst_t>
 static void dequantize_block_q3_K(const void * __restrict__ vx, dst_t * __restrict__ yy,
                                  const sycl::nd_item<3> &item_ct1) {
@@ -1921,6 +1921,23 @@ ESIMD_INLINE void dequantize_mul_mat_vec_reorder_esimd(
    }
 }

+static void dequantize_mul_mat_vec_q2_K_sycl_reorder_esimd(const void *vx, const float *y,
+                                                           float *dst, const int ncols,
+                                                           const int nrows,
+                                                           dpct::queue_ptr stream) {
+    GGML_ASSERT(ncols % QK_K == 0);
+    const int workgroups = (nrows + 1) / 2;
+    stream->submit([&](sycl::handler &h) {
+        sycl::local_accessor<float, 1> lmem(sycl::range<1>(GGML_SYCL_DMMV_ESIMD_WG_SIZE * 2), h);
+        h.parallel_for(
+            sycl::nd_range<1>(sycl::range<1>((size_t)workgroups * GGML_SYCL_DMMV_ESIMD_WG_SIZE), sycl::range<1>(GGML_SYCL_DMMV_ESIMD_WG_SIZE)),
+            [=](sycl::nd_item<1> it) [[intel::sycl_explicit_simd]] {
+                dequantize_mul_mat_vec_reorder_esimd<GGML_TYPE_Q2_K>(
+                    vx, y, dst, ncols, nrows, lmem, it);
+            });
+    });
+}
+
 static void dequantize_mul_mat_vec_q3_K_sycl_reorder_esimd(const void *vx, const float *y,
                                                           float *dst, const int ncols,
                                                           const int nrows,
@@ -2111,7 +2128,15 @@ void ggml_sycl_op_dequantize_mul_mat_vec(
        case GGML_TYPE_Q2_K:
            if ((ggml_tensor_extra_gpu *) dst->src[0]->extra &&
                ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
-                dequantize_mul_mat_vec_q2_K_sycl_reorder(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+#ifdef GGML_SYCL_DMMV_HAS_ESIMD
+                if (g_ggml_sycl_enable_esimd) {
+                    dequantize_mul_mat_vec_q2_K_sycl_reorder_esimd(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+                }
+                else
+#endif
+                {
+                    dequantize_mul_mat_vec_q2_K_sycl_reorder(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+                }
            } else {
                dequantize_mul_mat_vec_q2_K_sycl(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
            }
@@ -61,6 +61,93 @@ static ESIMD_INLINE void unpack_scale_min_k4(
    min_f   = convert<float>(m) * (-dmin);
 }

+// ---------------------------------------------------------------------------
+// Q2_K, SOA reorder layout produced by reorder_qw_q2_k:
+//   [qs: nb*(QK_K/4)] [scales: nb*(QK_K/16)] [dm: nb*sizeof(half2)]
+// with nb = nrows*num_blocks_per_row.
+//
+// 2 bits per weight. The 8 output chunks of 32 (matching dequantize_row_q2_K)
+// map to super-chunk s (0..7): byte base 32*(s/4) into the 64-byte qs array,
+// bit shift 2*(s%4); the low 16 lanes use scales[2s], the high 16 use
+// scales[2s+1], with dl = d*(sc & 0xF), ml = dmin*(sc >> 4), deq = dl*q - ml.
+// ---------------------------------------------------------------------------
+template <> struct esimd_reorder_q_traits<GGML_TYPE_Q2_K> {
+    struct ptrs {
+        const uint8_t *    qs;
+        const uint8_t *    scales;
+        const sycl::half * dm;
+    };
+
+    static ESIMD_INLINE ptrs make_ptrs(const void * vx, size_t nb) {
+        const uint8_t * qs     = (const uint8_t *) vx;
+        const uint8_t * scales = qs + nb * (QK_K / 4);
+        const sycl::half * dm  = (const sycl::half *) (scales + nb * (QK_K / 16));
+        return { qs, scales, dm };
+    }
+
+    static ESIMD_INLINE void mac_pair(
+            const ptrs & pa, size_t bia,
+            const ptrs & pb, size_t bib, bool has_b,
+            sycl::ext::intel::esimd::simd<float, 256> & y_vec,
+            sycl::ext::intel::esimd::simd<float, 32> & acc_a,
+            sycl::ext::intel::esimd::simd<float, 32> & acc_b) {
+        using namespace sycl::ext::intel::esimd;
+
+        simd<uint8_t, 64> qs_a     = block_load<uint8_t, 64>(pa.qs + bia * (QK_K / 4));
+        simd<uint8_t, 64> qs_b     = 0;
+        simd<uint8_t, 16> scales_a = block_load<uint8_t, 16>(pa.scales + bia * (QK_K / 16));
+        simd<uint8_t, 16> scales_b = 0;
+
+        const float dall_a = (float) pa.dm[bia * 2 + 0];
+        const float dmin_a = (float) pa.dm[bia * 2 + 1];
+        float dall_b = 0.0f;
+        float dmin_b = 0.0f;
+        if (has_b) {
+            qs_b     = block_load<uint8_t, 64>(pb.qs + bib * (QK_K / 4));
+            scales_b = block_load<uint8_t, 16>(pb.scales + bib * (QK_K / 16));
+            dall_b = (float) pb.dm[bib * 2 + 0];
+            dmin_b = (float) pb.dm[bib * 2 + 1];
+        }
+
+        // per-chunk scale (d * (sc & 0xF)) and min (-dmin * (sc >> 4)), all 16 codes;
+        // min carries the negation so the dequant epilogue adds (matches Q4_K/Q5_K)
+        simd<float, 16> scale_f_a = convert<float>(scales_a & simd<uint8_t, 16>(0x0F)) * dall_a;
+        simd<float, 16> min_f_a   = convert<float>(scales_a >> simd<uint8_t, 16>(4))  * (-dmin_a);
+        simd<float, 16> scale_f_b = convert<float>(scales_b & simd<uint8_t, 16>(0x0F)) * dall_b;
+        simd<float, 16> min_f_b   = convert<float>(scales_b >> simd<uint8_t, 16>(4))  * (-dmin_b);
+
+#pragma unroll
+        for (int s = 0; s < 8; ++s) {
+            const int     byte_base = 32 * (s / 4);
+            const uint8_t shift     = (uint8_t) (2 * (s % 4));
+            simd<float, 32> y_s = y_vec.select<32, 1>(s * 32);
+
+            simd<uint8_t, 32> qa = (qs_a.select<32, 1>(byte_base) >> shift) & simd<uint8_t, 32>(3);
+            simd<uint8_t, 32> qb = (qs_b.select<32, 1>(byte_base) >> shift) & simd<uint8_t, 32>(3);
+
+            const float scale_a_lo = scale_f_a[2 * s + 0];
+            const float scale_a_hi = scale_f_a[2 * s + 1];
+            const float min_a_lo   = min_f_a[2 * s + 0];
+            const float min_a_hi   = min_f_a[2 * s + 1];
+            const float scale_b_lo = scale_f_b[2 * s + 0];
+            const float scale_b_hi = scale_f_b[2 * s + 1];
+            const float min_b_lo   = min_f_b[2 * s + 0];
+            const float min_b_hi   = min_f_b[2 * s + 1];
+
+            simd<float, 32> scale_vec_a = splat_lo_hi(scale_a_lo, scale_a_hi);
+            simd<float, 32> min_vec_a   = splat_lo_hi(min_a_lo, min_a_hi);
+            simd<float, 32> scale_vec_b = splat_lo_hi(scale_b_lo, scale_b_hi);
+            simd<float, 32> min_vec_b   = splat_lo_hi(min_b_lo, min_b_hi);
+
+            simd<float, 32> deq_a = convert<float>(qa) * scale_vec_a + min_vec_a;
+            simd<float, 32> deq_b = convert<float>(qb) * scale_vec_b + min_vec_b;
+
+            acc_a += y_s * deq_a;
+            acc_b += y_s * deq_b;
+        }
+    }
+};
+
 // ---------------------------------------------------------------------------
 // Q3_K, SOA reorder layout produced by reorder_qw_q3_k:
 //   [qs: nb*(QK_K/4)] [hmask: nb*(QK_K/8)] [scales: nb*12] [d: nb*sizeof(half)]
@@ -3796,6 +3796,7 @@ inline bool ggml_sycl_supports_reorder_mmvq(enum ggml_type type) {
        case GGML_TYPE_Q1_0:
        case GGML_TYPE_Q4_0:
        case GGML_TYPE_Q8_0:
+        case GGML_TYPE_Q2_K:
        case GGML_TYPE_Q3_K:
        case GGML_TYPE_Q4_K:
        case GGML_TYPE_Q5_K:
@@ -3809,6 +3810,7 @@ inline bool ggml_sycl_supports_reorder_mmvq(enum ggml_type type) {
 static bool ggml_sycl_supports_reorder_esimd(enum ggml_type type) {
 #ifdef GGML_SYCL_DMMV_HAS_ESIMD
    switch (type) {
+        case GGML_TYPE_Q2_K:
        case GGML_TYPE_Q3_K:
        case GGML_TYPE_Q4_K:
        case GGML_TYPE_Q5_K:
@@ -1401,6 +1401,65 @@ static void mul_mat_vec_q2_K_q8_1_sycl_switch_ncols(
    }
 }

+static void reorder_mul_mat_vec_q2_k_q8_1_sycl(const void * vx, const void * vy, float * dst, const int ncols,
+                                               const int nrows, dpct::queue_ptr stream) {
+    GGML_ASSERT(ncols % QK_K == 0);
+
+    // Round up to a whole number of subgroup-sized workgroups; out-of-range rows are skipped inside the kernel.
+    constexpr size_t num_subgroups = WARP_SIZE;
+    const int block_num_y = ceil_div(nrows, GGML_SYCL_MMV_Y * (int) num_subgroups);
+    const sycl::range<3> block_nums(1, 1, block_num_y);
+    const sycl::range<3> block_dims(1, GGML_SYCL_MMV_Y, num_subgroups * WARP_SIZE);
+
+    stream->submit([&](sycl::handler & cgh) {
+        cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
+                         [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                             mul_mat_vec_q_reorder<reorder_vec_dot_q_sycl<GGML_TYPE_Q2_K>>(vx, vy, dst, ncols, nrows,
+                                                                                           nd_item);
+                         });
+    });
+}
+
+template <int ncols_dst>
+static void reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols(
+        const void * vx, const void * vy, float * dst,
+        const int ncols, const int nrows,
+        const int stride_col_y_bytes, const int stride_col_dst,
+        dpct::queue_ptr stream) {
+    GGML_ASSERT(ncols % QK_K == 0);
+    constexpr size_t num_subgroups = WARP_SIZE;
+    const int block_num_y = ceil_div(nrows, GGML_SYCL_MMV_Y * (int) num_subgroups);
+    const sycl::range<3> block_nums(1, 1, block_num_y);
+    const sycl::range<3> block_dims(1, GGML_SYCL_MMV_Y, num_subgroups * WARP_SIZE);
+
+    stream->submit([&](sycl::handler & cgh) {
+        cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
+                         [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                             mul_mat_vec_q_reorder_ncols<reorder_vec_dot_q_sycl<GGML_TYPE_Q2_K>, ncols_dst>(
+                                 vx, /*vgate=*/ nullptr, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst,
+                                 /*glu_op=*/ GGML_GLU_OP_SWIGLU, nd_item);
+                         });
+    });
+}
+
+static void reorder_mul_mat_vec_q2_k_q8_1_sycl_switch_ncols(
+        const void * vx, const void * vy, float * dst,
+        const int ncols, const int nrows, const int ncols_dst,
+        const int stride_col_y_bytes, const int stride_col_dst,
+        dpct::queue_ptr stream) {
+    switch (ncols_dst) {
+        case 1: reorder_mul_mat_vec_q2_k_q8_1_sycl(vx, vy, dst, ncols, nrows, stream); break;
+        case 2: reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols<2>(vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, stream); break;
+        case 3: reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols<3>(vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, stream); break;
+        case 4: reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols<4>(vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, stream); break;
+        case 5: reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols<5>(vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, stream); break;
+        case 6: reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols<6>(vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, stream); break;
+        case 7: reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols<7>(vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, stream); break;
+        case 8: reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols<8>(vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, stream); break;
+        default: GGML_ABORT("unsupported ncols_dst=%d for Q2_K reorder multi-col MMVQ", ncols_dst);
+    }
+}
+
 static void mul_mat_vec_q3_K_q8_1_sycl(const void *vx, const void *vy,
                                       float *dst, const int ncols,
                                       const int nrows,
@@ -2297,7 +2356,21 @@ void ggml_sycl_op_mul_mat_vec_q(ggml_backend_sycl_context & ctx, const ggml_tens
                }
                break;
            case GGML_TYPE_Q2_K:
-                if (i == 0 && src1_ncols > 1 && src1_ncols <= 8) {
+                if ((ggml_tensor_extra_gpu *) dst->src[0]->extra &&
+                    ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
+                    if (i == 0 && src1_ncols > 1 && src1_ncols <= 8) {
+                        const int stride_col_y_bytes = src1_padded_col_size * q8_1_ts / q8_1_bs;
+                        const int stride_col_dst     = dst->ne[0];
+                        GGML_SYCL_DEBUG("Calling reorder_mul_mat_vec_q2_k_q8_1_sycl_switch_ncols ncols=%d\n", (int)src1_ncols);
+                        reorder_mul_mat_vec_q2_k_q8_1_sycl_switch_ncols(
+                            src0_dd_i, src1_ddq_i, dst_dd_i, ne00, row_diff,
+                            src1_ncols, stride_col_y_bytes, stride_col_dst, stream);
+                        return;
+                    } else {
+                        GGML_SYCL_DEBUG("Calling reorder_mul_mat_vec_q2_k_q8_1_sycl\n");
+                        reorder_mul_mat_vec_q2_k_q8_1_sycl(src0_dd_i, src1_ddq_i_bs, dst_dd_i_bs, ne00, row_diff, stream);
+                    }
+                } else if (i == 0 && src1_ncols > 1 && src1_ncols <= 8) {
                    const int stride_col_y   = src1_padded_col_size / QK8_1;
                    const int stride_col_dst = dst->ne[0];
                    GGML_SYCL_DEBUG("Calling mul_mat_vec_q2_K_q8_1_sycl_switch_ncols ncols=%d\n", (int)src1_ncols);
@@ -2306,6 +2379,7 @@ void ggml_sycl_op_mul_mat_vec_q(ggml_backend_sycl_context & ctx, const ggml_tens
                        src1_ncols, stride_col_y, stride_col_dst, stream);
                    return;
                } else if (i == 0 || src1_ncols == 1) {
+                    GGML_SYCL_DEBUG("Calling mul_mat_vec_q2_K_q8_1_sycl\n");
                    mul_mat_vec_q2_K_q8_1_sycl(src0_dd_i, src1_ddq_i_bs, dst_dd_i_bs, ne00, row_diff, stream);
                }
                break;
@@ -58,6 +58,29 @@ template <> struct block_q_t<GGML_TYPE_Q4_0> {
    static constexpr int block_to_q8_1_ratio() { return traits::qk / QK8_1; }
 };

+template <> struct block_q_t<GGML_TYPE_Q2_K> {
+    struct traits {
+        static constexpr uint32_t qk       = QK_K;
+        static constexpr uint32_t qi       = QI2_K;
+        static constexpr uint32_t qr       = QR2_K;
+        static constexpr uint32_t vdr_mmvq = 1;
+    };
+
+    // Reordered layout: [qs (QK_K/4 per block)] [scales (QK_K/16 per block)] [dm]
+    static constexpr std::pair<int, int> get_block_offset(const int block_index, const int /* n_blocks */) {
+        return { block_index * (QK_K / 4), 0 };
+    }
+
+    static constexpr std::pair<int, int> get_d_offset(int nrows, int ncols, const int block_index) {
+        auto nblocks        = (nrows * (ncols / QK_K));
+        auto total_qs_bytes = nblocks * (QK_K / 4);
+        return { total_qs_bytes + block_index * (QK_K / 16),
+                 total_qs_bytes + nblocks * (QK_K / 16) + block_index * sizeof(ggml_half2) };
+    }
+
+    static constexpr int block_to_q8_1_ratio() { return traits::qk / QK8_1; }
+};
+
 template <> struct block_q_t<GGML_TYPE_Q3_K> {
    struct traits {
        static constexpr uint32_t qk       = QK_K;
@@ -429,6 +429,39 @@ template <> struct reorder_vec_dot_q_sycl<GGML_TYPE_Q8_0> {
    }
 };

+template <> struct reorder_vec_dot_q_sycl<GGML_TYPE_Q2_K> {
+    static constexpr ggml_type gtype = GGML_TYPE_Q2_K;
+
+    using q2_k_block  = ggml_sycl_reordered::block_q_t<GGML_TYPE_Q2_K>;
+    using q2_k_traits = typename q2_k_block::traits;
+
+    __dpct_inline__ float operator()(const void * __restrict__ vbq, const std::pair<int, int> ibx_offset,
+                                     const std::pair<int, int> d_offset, const int8_t * q8_1_quant_ptr,
+                                     const sycl::half2 * q8_1_ds, const int & iqs) {
+        const uint8_t *    base   = static_cast<const uint8_t *>(vbq);
+        const uint8_t *    qs     = base + ibx_offset.first;
+        const uint8_t *    scales = base + d_offset.first;
+        const ggml_half2 * dm     = reinterpret_cast<const ggml_half2 *>(base + d_offset.second);
+
+        const int bq8_offset   = QR2_K * (iqs / QI8_1);
+        const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1 / 2);
+
+        const int v = get_int_from_uint8_aligned(qs, iqs);
+
+        int   u[QR2_K];
+        float d8[QR2_K];
+
+#pragma unroll
+        for (int i = 0; i < QR2_K; ++i) {
+            const int8_t * quant_base_ptr = q8_1_quant_ptr + (bq8_offset + i) * QK8_1;
+            u[i]                          = get_int_from_int8_aligned(quant_base_ptr, iqs % QI8_1);
+            d8[i]                         = (*(q8_1_ds + bq8_offset + i))[0];
+        }
+
+        return vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales + scale_offset, *dm, d8);
+    }
+};
+
 template <> struct reorder_vec_dot_q_sycl<GGML_TYPE_Q3_K> {
    static constexpr ggml_type gtype = GGML_TYPE_Q3_K;

@@ -364,6 +364,8 @@ class Keys:
        IMAGE_MEAN            = "clip.vision.image_mean"
        IMAGE_STD             = "clip.vision.image_std"
        SPATIAL_MERGE_SIZE    = "clip.vision.spatial_merge_size"
+        EXPERT_COUNT_PER_LAYER = "clip.vision.expert_count_per_layer" # dots3note pyramid MoE, 0 = dense layer
+        EXPERT_USED_COUNT     = "clip.vision.expert_used_count"
        USE_GELU              = "clip.use_gelu"
        USE_SILU              = "clip.use_silu"
        N_WA_PATTERN          = "clip.vision.n_wa_pattern"     # used by qwen2.5vl
@@ -874,6 +876,11 @@ class MODEL_TENSOR(IntEnum):
    V_ENC_FFN_UP         = auto()
    V_ENC_FFN_GATE       = auto()
    V_ENC_FFN_DOWN       = auto()
+    V_ENC_FFN_GATE_INP   = auto() # dots3note vision MoE router
+    V_ENC_FFN_GATE_EXPS  = auto()
+    V_ENC_FFN_UP_EXPS    = auto()
+    V_ENC_FFN_DOWN_EXPS  = auto()
+    V_ENC_FFN_EXP_PROBS_B = auto()
    V_ENC_ATTN_POST_NORM = auto() # gemma4
    V_ENC_FFN_POST_NORM  = auto()
    V_LAYER_SCALE_1      = auto()
@@ -1591,6 +1598,11 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
    MODEL_TENSOR.V_ENC_FFN_UP:              "v.blk.{bid}.ffn_up",
    MODEL_TENSOR.V_ENC_FFN_GATE:            "v.blk.{bid}.ffn_gate",
    MODEL_TENSOR.V_ENC_FFN_DOWN:            "v.blk.{bid}.ffn_down",
+    MODEL_TENSOR.V_ENC_FFN_GATE_INP:        "v.blk.{bid}.ffn_gate_inp",
+    MODEL_TENSOR.V_ENC_FFN_GATE_EXPS:       "v.blk.{bid}.ffn_gate_exps",
+    MODEL_TENSOR.V_ENC_FFN_UP_EXPS:         "v.blk.{bid}.ffn_up_exps",
+    MODEL_TENSOR.V_ENC_FFN_DOWN_EXPS:       "v.blk.{bid}.ffn_down_exps",
+    MODEL_TENSOR.V_ENC_FFN_EXP_PROBS_B:     "v.blk.{bid}.exp_probs_b",
    MODEL_TENSOR.V_ENC_ATTN_POST_NORM:      "v.blk.{bid}.attn_post_norm",
    MODEL_TENSOR.V_ENC_FFN_POST_NORM:       "v.blk.{bid}.ffn_post_norm",
    MODEL_TENSOR.V_LAYER_SCALE_1:           "v.blk.{bid}.ls1",
@@ -1913,6 +1925,11 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
        MODEL_TENSOR.V_ENC_FFN_UP,
        MODEL_TENSOR.V_ENC_FFN_GATE,
        MODEL_TENSOR.V_ENC_FFN_DOWN,
+        MODEL_TENSOR.V_ENC_FFN_GATE_INP,
+        MODEL_TENSOR.V_ENC_FFN_GATE_EXPS,
+        MODEL_TENSOR.V_ENC_FFN_UP_EXPS,
+        MODEL_TENSOR.V_ENC_FFN_DOWN_EXPS,
+        MODEL_TENSOR.V_ENC_FFN_EXP_PROBS_B,
        MODEL_TENSOR.V_ENC_ATTN_POST_NORM,
        MODEL_TENSOR.V_ENC_FFN_POST_NORM,
        MODEL_TENSOR.V_LAYER_SCALE_1,
@@ -5497,6 +5514,8 @@ class VisionProjectorType:
    COGVLM = "cogvlm"
    JANUS_PRO = "janus_pro"
    DOTSOCR = "dots_ocr"
+    DOTS3NOTE_V = "dots3note_v"
+    DOTS3NOTE_A = "dots3note_a" # audio
    DEEPSEEKOCR = "deepseekocr"
    DEEPSEEKOCR2 = "deepseekocr2"
    LFM2A = "lfm2a" # audio
@@ -1327,6 +1327,12 @@ class GGUFWriter:
    def add_vision_spatial_merge_size(self, value: int) -> None:
        self.add_uint32(Keys.ClipVision.SPATIAL_MERGE_SIZE, value)

+    def add_vision_expert_count_per_layer(self, value: Sequence[int]) -> None:
+        self.add_array(Keys.ClipVision.EXPERT_COUNT_PER_LAYER, value)
+
+    def add_vision_expert_used_count(self, value: int) -> None:
+        self.add_uint32(Keys.ClipVision.EXPERT_USED_COUNT, value)
+
    def add_vision_use_gelu(self, value: bool) -> None:
        self.add_bool(Keys.ClipVision.USE_GELU, value)

@@ -1454,6 +1454,7 @@ class TensorNameMap:
            "mlp_AR.linear_{bid}", # PaddleOCR-VL
            "merger.mlp.{bid}",
            "vision_tower.merger.mlp.{bid}", # dots.ocr
+            "vision_encoder.adapter.mlp.{bid}", # dots3note
            "vit.perceive.proj.{bid}", # HunyuanVL (proj.0 = conv1, proj.2 = conv2)
        ),

@@ -1504,6 +1505,7 @@ class TensorNameMap:
            "vision_model.radio_model.model.patch_generator.embedder", # Nemotron Nano v2 VL
            "model.vision_tower.patch_embedder.input_proj", # gemma4
            "vision_tower.patch_embed.patchifier.proj", # dots.ocr
+            "vision_encoder.patch_embed.proj", # dots3note
            "vision_model.conv1", # Step3-VL
            "model.vision_embedder.patch_dense", # gemma4 unified
            "model.vision_tower.patch_embedder.patch_embedding", # muse-glimmer
@@ -1512,6 +1514,7 @@ class TensorNameMap:
        MODEL_TENSOR.V_ENC_EMBD_NORM: (
            "visual.post_conv_layernorm", # glm4v
            "vision_tower.patch_embed.patchifier.norm", # dots.ocr
+            "vision_encoder.patch_embed.norm", # dots3note
        ),

        MODEL_TENSOR.V_ENC_EMBD_PATCH_NORM: (
@@ -1551,6 +1554,7 @@ class TensorNameMap:
        MODEL_TENSOR.V_ENC_ATTN_QKV: (
            "visual.blocks.{bid}.attn.qkv", # qwen3vl
            "vision_tower.blocks.{bid}.attn.qkv", # dots.ocr
+            "vision_encoder.blocks.{bid}.attn.qkv", # dots3note
            "model.vision.transformer.layers.{bid}.attention.query_key_value", # cogvlm
            "model.vision_model.transformer.layers.{bid}.self_attn.qkv_proj", # Deepseek-OCR CLIP
            "vision_tower.encoder.blocks.{bid}.wqkv", # Kimi-K2.5
@@ -1579,6 +1583,7 @@ class TensorNameMap:
        ),

        MODEL_TENSOR.V_ENC_ATTN_Q_NORM: (
+            "vision_encoder.blocks.{bid}.attn.q_norm", # dots3note
            "vision_tower.vision_model.encoder.layers.{bid}.attn.q_norm", # InternVL
            "model.vision_tower.encoder.layer.{bid}.attention.q_norm", # Intern-S1
            "visual.blocks.{bid}.attn.q_norm", # GLM-OCR
@@ -1606,6 +1611,7 @@ class TensorNameMap:
        ),

        MODEL_TENSOR.V_ENC_ATTN_K_NORM: (
+            "vision_encoder.blocks.{bid}.attn.k_norm", # dots3note
            "vision_tower.vision_model.encoder.layers.{bid}.attn.k_norm", # InternVL
            "model.vision_tower.encoder.layer.{bid}.attention.k_norm", # Intern-S1
            "visual.blocks.{bid}.attn.k_norm", # GLM-OCR
@@ -1651,6 +1657,7 @@ class TensorNameMap:
            "siglip2.vision_model.encoder.layers.{bid}.layer_norm1",
            "vision_model.radio_model.model.blocks.{bid}.norm1", # Nemotron Nano v2 VL
            "vision_tower.blocks.{bid}.norm1", # dots.ocr
+            "vision_encoder.blocks.{bid}.norm_1", # dots3note
            "vision_model.transformer.resblocks.{bid}.ln_1", # Step3-VL
            "model.qwen2_model.model.model.layers.{bid}.input_layernorm", # Deepseek-OCR-2 qwen2
            "model.vision_tower.layers.{bid}.norm1", # muse-glimmer
@@ -1678,6 +1685,7 @@ class TensorNameMap:
            "model.qwen2_model.model.model.layers.{bid}.self_attn.o_proj", # Deepseek-OCR-2 qwen2
            "vision_model.model.layers.{bid}.self_attn.o_proj.linear", # gemma4
            "vision_tower.blocks.{bid}.attn.proj", # dots.ocr
+            "vision_encoder.blocks.{bid}.attn.proj", # dots3note
            "vision_model.transformer.resblocks.{bid}.attn.out_proj", # Step3-VL
            "model.vision_tower.layers.{bid}.attn.proj", # muse-glimmer
        ),
@@ -1706,12 +1714,14 @@ class TensorNameMap:
            "vision_model.radio_model.model.blocks.{bid}.norm2", # Nemotron Nano v2 VL
            "vision_model.model.layers.{bid}.pre_feedforward_layernorm", # gemma4
            "vision_tower.blocks.{bid}.norm2", # dots.ocr
+            "vision_encoder.blocks.{bid}.norm_2", # dots3note
            "vision_model.transformer.resblocks.{bid}.ln_2", # Step3-VL
            "model.qwen2_model.model.model.layers.{bid}.post_attention_layernorm", # Deepseek-OCR-2 qwen2
            "model.vision_tower.layers.{bid}.norm2", # muse-glimmer
        ),

        MODEL_TENSOR.V_ENC_FFN_UP: (
+            "vision_encoder.blocks.{bid}.mlp.fc3", # dots3note
            "model.vision_tower.vision_model.encoder.layers.{bid}.mlp.fc1", # Granite4Vision
            "vision_tower.vision_model.encoder.layers.{bid}.mlp.fc1",
            "model.vision_tower.encoder.layers.{bid}.mlp.fc1", # minicpmv4_6
@@ -1737,6 +1747,7 @@ class TensorNameMap:
        ),

        MODEL_TENSOR.V_ENC_FFN_GATE: (
+            "vision_encoder.blocks.{bid}.mlp.fc1", # dots3note
            "vision_tower.transformer.layers.{bid}.feed_forward.gate_proj", # pixtral-hf
            "vision_encoder.transformer.layers.{bid}.feed_forward.w1", # pixtral
            "visual.blocks.{bid}.mlp.gate_proj", # qwen2.5vl
@@ -1745,6 +1756,7 @@ class TensorNameMap:
        ),

        MODEL_TENSOR.V_ENC_FFN_DOWN: (
+            "vision_encoder.blocks.{bid}.mlp.fc2", # dots3note
            "model.vision_tower.vision_model.encoder.layers.{bid}.mlp.fc2", # Granite4Vision
            "vision_tower.vision_model.encoder.layers.{bid}.mlp.fc2",
            "model.vision_tower.encoder.layers.{bid}.mlp.fc2", # minicpmv4_6
@@ -1769,6 +1781,29 @@ class TensorNameMap:
            "model.vision_tower.layers.{bid}.mlp.fc2", # muse-glimmer
        ),

+
+        MODEL_TENSOR.V_ENC_FFN_GATE_INP: (
+            "vision_encoder.blocks.{bid}.mlp.gate_weight", # dots3note
+        ),
+
+        MODEL_TENSOR.V_ENC_FFN_EXP_PROBS_B: (
+            "vision_encoder.blocks.{bid}.mlp.router_bias", # dots3note
+        ),
+
+        # note: expert weights are stacked into a single 3D tensor in conversion code,
+        # which emits the pseudo-names below
+        MODEL_TENSOR.V_ENC_FFN_GATE_EXPS: (
+            "vision_encoder.blocks.{bid}.mlp.experts.fc1", # dots3note
+        ),
+
+        MODEL_TENSOR.V_ENC_FFN_UP_EXPS: (
+            "vision_encoder.blocks.{bid}.mlp.experts.fc3", # dots3note
+        ),
+
+        MODEL_TENSOR.V_ENC_FFN_DOWN_EXPS: (
+            "vision_encoder.blocks.{bid}.mlp.experts.fc2", # dots3note
+        ),
+
        MODEL_TENSOR.V_ENC_ATTN_POST_NORM: (
            "vision_model.model.layers.{bid}.post_attention_layernorm", # gemma4
        ),
@@ -1800,6 +1835,7 @@ class TensorNameMap:
            "vision_model.layernorm_pre", # llama4
            "model.vision_model.pre_layrnorm", # Deepseek-OCR CLIP
            "vision_tower.patch_embed.patchifier.norm", # dots.ocr
+            "vision_encoder.patch_embed.norm", # dots3note
            "vision_model.ln_pre", # Step3-VL
            "model.vision_tower.ln_pre", # muse-glimmer
        ),
@@ -1821,6 +1857,7 @@ class TensorNameMap:
        MODEL_TENSOR.V_MM_POST_NORM: (
            "visual.merger.post_projection_norm", # glm4v
            "vision_tower.post_trunk_norm", # dots.ocr
+            "vision_encoder.post_trunk_norm", # dots3note
            "vit.perceive.after_rms", # HunyuanVL
        ),

@@ -1838,6 +1875,7 @@ class TensorNameMap:
            "mlp_AR.pre_norm", # PaddleOCR-VL
            "merger.ln_q",
            "vision_tower.merger.ln_q", # dots.ocr
+            "vision_encoder.adapter.ln_q", # dots3note
            "model.merger.mlp.0.pre_norm", # minicpmv4_6
        ),

@@ -2173,10 +2211,12 @@ class TensorNameMap:

        MODEL_TENSOR.A_ENC_CONV2D: (
            "audio_tower.conv2d{bid}", # qwen3omni
+            "audio_encoder.dots_encoder.speech_encoder.conv2d{bid}", # dots3note
        ),

        MODEL_TENSOR.A_ENC_CONV_OUT: (
            "audio_tower.conv_out", # qwen3omni
+            "audio_encoder.dots_encoder.speech_encoder.conv_out", # dots3note
            "speaker_encoder.mfa.conv", # qwen3tts speaker encoder: multi-layer feature aggregation
        ),

@@ -2184,12 +2224,14 @@ class TensorNameMap:

        MODEL_TENSOR.A_POST_NORM: (
            "audio_tower.layer_norm", # ultravox
+            "audio_encoder.dots_encoder.speech_encoder.layer_norm", # dots3note
            "audio_tower.ln_post", # qwen2omni
            "encoder.layer_norm", # mimo-audio-tokenizer
        ),

        MODEL_TENSOR.A_ENC_ATTN_Q: (
            "audio_tower.layers.{bid}.self_attn.q_proj", # ultravox
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn.q_proj", # dots3note
            "conformer.layers.{bid}.self_attn.linear_q", # lfm2
            "conformer.layers.{bid}.attention.attn.q_proj", # gemma3n
            "conformer.layers.{bid}.self_attn.q_proj", # gemma4
@@ -2200,6 +2242,7 @@ class TensorNameMap:

        MODEL_TENSOR.A_ENC_ATTN_K: (
            "audio_tower.layers.{bid}.self_attn.k_proj", # ultravox
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn.k_proj", # dots3note
            "conformer.layers.{bid}.self_attn.linear_k", # lfm2
            "conformer.layers.{bid}.attention.attn.k_proj", # gemma3n
            "conformer.layers.{bid}.self_attn.k_proj", # gemma4
@@ -2210,6 +2253,7 @@ class TensorNameMap:

        MODEL_TENSOR.A_ENC_ATTN_V: (
            "audio_tower.layers.{bid}.self_attn.v_proj", # ultravox
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn.v_proj", # dots3note
            "conformer.layers.{bid}.self_attn.linear_v", # lfm2
            "conformer.layers.{bid}.attention.attn.v_proj", # gemma3n
            "conformer.layers.{bid}.self_attn.v_proj", # gemma4
@@ -2241,6 +2285,7 @@ class TensorNameMap:

        MODEL_TENSOR.A_ENC_INPUT_NORM: (
            "audio_tower.layers.{bid}.self_attn_layer_norm", # ultravox
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn_layer_norm", # dots3note
            "conformer.layers.{bid}.norm_self_att", # lfm2
            "conformer.layers.{bid}.attention.pre_attn_norm", # gemma3n
            "sound_encoder.encoder.layers.{bid}.norm_self_att", # parakeet
@@ -2250,6 +2295,7 @@ class TensorNameMap:

        MODEL_TENSOR.A_ENC_OUTPUT: (
            "audio_tower.layers.{bid}.self_attn.out_proj", # ultravox
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn.out_proj", # dots3note
            "conformer.layers.{bid}.self_attn.linear_out", # lfm2
            "conformer.layers.{bid}.attention.post", # gemma3n
            "conformer.layers.{bid}.self_attn.post", # gemma4
@@ -2260,6 +2306,7 @@ class TensorNameMap:

        MODEL_TENSOR.A_ENC_OUTPUT_NORM: (
            "audio_tower.layers.{bid}.final_layer_norm", # ultravox
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.final_layer_norm", # dots3note
            "conformer.layers.{bid}.norm_out", # lfm2
            "conformer.layers.{bid}.attention.post_norm", # gemma3n
            "sound_encoder.encoder.layers.{bid}.norm_out", # parakeet
@@ -2285,6 +2332,7 @@ class TensorNameMap:
        ),

        MODEL_TENSOR.A_ENC_FFN_UP: (
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.fc1_up", # dots3note (split from fc1 in conversion code)
            "audio_tower.layers.{bid}.fc1", # ultravox
            "conformer.layers.{bid}.feed_forward1.linear1", # lfm2
            "conformer.layers.{bid}.ffw_layer_start.ffw_layer_1", # gemma3n
@@ -2294,9 +2342,12 @@ class TensorNameMap:
            "encoder.layers.{bid}.fc1", # mimo-audio-tokenizer
        ),

-        MODEL_TENSOR.A_ENC_FFN_GATE: (),
+        MODEL_TENSOR.A_ENC_FFN_GATE: (
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.fc1_gate", # dots3note (split from fc1 in conversion code)
+        ),

        MODEL_TENSOR.A_ENC_FFN_DOWN: (
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.fc2", # dots3note
            "audio_tower.layers.{bid}.fc2", # ultravox
            "conformer.layers.{bid}.feed_forward1.linear2", # lfm2
            "conformer.layers.{bid}.ffw_layer_start.ffw_layer_2", # gemma3n
@@ -2380,6 +2431,7 @@ class TensorNameMap:

        MODEL_TENSOR.A_MMPROJ: (
            "audio.multi_modal_projector.linear_{bid}", # ultravox, meralion
+            "audio_encoder.audio_adapter.proj.{bid}", # dots3note (proj.1, proj.3)
            "audio_adapter.model.{bid}", # lfm2
            "audio_tower.proj{bid}", # qwen3omni
            "sound_projection.linear{bid}", # parakeet (linear1, linear2)
@@ -2394,6 +2446,7 @@ class TensorNameMap:

        MODEL_TENSOR.A_MM_NORM_PRE: (
            "audio.multi_modal_projector.ln_pre", # ultravox
+            "audio_encoder.audio_adapter.proj.0", # dots3note
            "sound_projection.norm", # parakeet
        ),

@@ -1038,6 +1038,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
        case LLM_ARCH_NEMOTRON_H_MOE:
        case LLM_ARCH_LFM2:
        case LLM_ARCH_LFM2MOE:
+        case LLM_ARCH_BAILINGMOE3:
            return true;
        default:
            return false;
@@ -1,6 +1,8 @@
 #include "models.h"
 #include "llama-memory-recurrent.h"

+#include <algorithm>
+
 void llama_model_bailingmoe3::load_arch_hparams(llama_model_loader & ml) {
    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS,      hparams.f_norm_rms_eps);
    ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA,         hparams.n_embd_head_k_mla_impl);
@@ -179,7 +181,9 @@ static ggml_tensor * bailingmoe3_causal_conv1d(
        int64_t n_seq_tokens,
        int64_t n_seqs,
        int64_t n_tokens,
-        int64_t cache_head) {
+        int64_t cache_head,
+        uint32_t mem_size,
+        uint32_t n_rs_seq) {
    const int64_t d_inner = head_dim * n_head;
    const int64_t conv_state_size = (d_conv - 1) * d_inner;
    const int64_t total_state_size = 3 * conv_state_size;
@@ -193,13 +197,18 @@ static ggml_tensor * bailingmoe3_causal_conv1d(
    x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
    ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0);

-    ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
-            conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]);
-    ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv_x,
-            ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
-                (d_conv - 1) * ggml_element_size(conv_states_all),
-                total_state_size * ggml_element_size(conv_states_all),
-                (cache_head * total_state_size + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
+    const int64_t K = (int64_t) n_rs_seq + 1;
+    const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
+
+    for (int64_t slot = 0; slot < n_written; ++slot) {
+        ggml_tensor * conv_snap = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
+                conv_x->nb[1], conv_x->nb[2], (conv_x->ne[0] - (d_conv - 1) - slot) * conv_x->nb[0]);
+        ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_snap,
+                ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
+                    (d_conv - 1) * ggml_element_size(conv_states_all),
+                    total_state_size * ggml_element_size(conv_states_all),
+                    ((slot * mem_size + cache_head) * total_state_size + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
+    }

    ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
    ggml_tensor * out = ggml_ssm_conv(ctx0, conv_x, conv_weight);
@@ -237,6 +246,8 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
    GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);

    for (int il = 0; il < n_layer; ++il) {
+        res->t_layer_inp[il] = inpL;
+
        const auto & layer = model.layers[il];
        ggml_tensor * inpSA = inpL;
        ggml_tensor * cur = build_norm(inpL, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
@@ -245,18 +256,19 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
        if (hparams.is_recr(il)) {
            const auto * mctx_cur = inp_rs->mctx;
            const auto cache_head = mctx_cur->get_head();
+            const auto mem_size = mctx_cur->get_size();
            ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
            ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);

            ggml_tensor * q = bailingmoe3_causal_conv1d(
                    gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv,
-                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head);
+                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
            ggml_tensor * k = bailingmoe3_causal_conv1d(
                    gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv,
-                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head);
+                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
            ggml_tensor * v = bailingmoe3_causal_conv1d(
                    gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv,
-                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head);
+                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);

            ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_f_a, cur);
            gate = ggml_add(ctx0, gate, layer.ssm_dt_b);
@@ -276,11 +288,8 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
            ggml_tensor * state = build_rs(inp_rs, states_all, hparams.n_embd_s(), n_seqs);
            state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head, n_seqs);

-            auto result = build_delta_net(q, k, v, gate, beta, state, il);
-            ggml_tensor * out = ggml_cont(ctx0, result.first);
-            ggml_build_forward_expand(gf, ggml_cpy(ctx0, result.second,
-                    ggml_view_1d(ctx0, states_all, hparams.n_embd_s() * n_seqs,
-                        cache_head * hparams.n_embd_s() * ggml_element_size(states_all))));
+            ggml_tensor * out = ggml_cont(ctx0, build_recurrent_attn(
+                    inp_rs, states_all, q, k, v, gate, beta, state, il));

            ggml_tensor * out_gate = ggml_mul_mat(ctx0, layer.ssm_g_a, cur);
            out_gate = ggml_reshape_3d(ctx0, out_gate, head_dim, n_head, n_tokens);
@@ -8,7 +8,7 @@ void test_json_serialization(testing &t) {
    auto json_serialized = original.to_json().dump();

    t.test("compare before/after", [&](testing &t) {
-        auto deserialized = common_peg_arena::from_json(nlohmann::json::parse(json_serialized));
+        auto deserialized = common_peg_arena::from_json(common_json::parse(json_serialized));

        // Test complex JSON
        std::string input = R"({"name": "test", "values": [1, 2, 3], "nested": {"a": true}})";
@@ -23,6 +23,6 @@ void test_json_serialization(testing &t) {
    });

    t.bench("deserialize", [&]() {
-        auto deserialized = common_peg_arena::from_json(nlohmann::json::parse(json_serialized));
+        auto deserialized = common_peg_arena::from_json(common_json::parse(json_serialized));
    }, 100);
 }
@@ -1,7 +1,7 @@
 #pragma once

 // Common includes for all test files
-#include <nlohmann/json.hpp>
+#include "json.h"
 #include <string>
 #include <vector>

@@ -11,9 +11,9 @@
 #include "simple-tokenize.h"

 struct bench_tool_call {
-    std::string            id;
-    std::string            name;
-    nlohmann::ordered_json args;
+    std::string id;
+    std::string name;
+    common_json args;
 };

 // Test function declarations
@@ -11,9 +11,9 @@
 #include <regex>
 #include <string>

-#include "nlohmann/json.hpp"
+#include "json.h"

-using json = nlohmann::ordered_json;
+using json = common_json;

 static json create_tools();
 static void test_example_native(testing & t);
@@ -63,10 +63,10 @@ static json create_tools() {
                          { { "type", "string" }, { "description", "The city and state, e.g. San Francisco, CA" } } },
                        { "unit",
                          { { "type", "string" },
-                            { "enum", { "celsius", "fahrenheit" } },
+                            { "enum", json::array({ "celsius", "fahrenheit" }) },
                            { "description",
                              "The temperature unit to use. Infer this from the users location." } } } } },
-                    { "required", { "location", "unit" } },
+                    { "required", json::array({ "location", "unit" }) },
                } },
          }                      }
    };
@@ -86,14 +86,14 @@ static json create_tools() {
                          { { "type", "string" }, { "description", "The city and state, e.g. San Francisco, CA" } } },
                        { "unit",
                          { { "type", "string" },
-                            { "enum", { "celsius", "fahrenheit" } },
+                            { "enum", json::array({ "celsius", "fahrenheit" }) },
                            { "description", "The temperature unit to use. Infer this from the users location." } } },
                        { "days",
                          { { "type", "integer" },
                            { "description", "Number of days to forecast (1-10)" },
                            { "minimum", 1 },
                            { "maximum", 10 } } } } },
-                    { "required", { "location", "unit" } },
+                    { "required", json::array({ "location", "unit" }) },
                } },
          }                      }
    };
@@ -114,9 +114,9 @@ static json create_tools() {
                        { "default", 5 } } },
                    { "category",
                      { { "type", "string" },
-                        { "enum", { "api", "troubleshooting", "billing", "general" } },
+                        { "enum", json::array({ "api", "troubleshooting", "billing", "general" }) },
                        { "description", "Filter search by specific category." } } } } },
-                { "required", { "query", "category" } },
+                { "required", json::array({ "query", "category" }) },
                { "additionalProperties", false } } },
            { "strict", true } } }
    };
@@ -341,7 +341,7 @@ static void test_example_native(testing & t) {
                { { "invoice_number", { { "type", "string" } } },
                  { "amount", { { "type", "number" } } },
                  { "due_date", { { "type", "string" } } } } },
-              { "required", { "invoice_number", "amount", "due_date" } } },
+              { "required", json::array({ "invoice_number", "amount", "due_date" }) } },
         /* .parallel_tool_calls =  */ false,
         /* .generation_prompt =    */ "<think>",
         /* .input =                */
@@ -406,7 +406,7 @@ static void test_example_qwen3_coder(testing & t) {

            std::set<std::string> required_properties;
            if (function.contains("required")) {
-                function.at("required").get_to(required_properties);
+                required_properties = function.at("required").get<std::set<std::string>>();
            }

            std::vector<common_peg_parser> arg_parsers;
@@ -661,8 +661,8 @@ void test_command7_parser_compare(testing & t) {
        "5. Provide a detailed cost breakdown that includes accommodation, transportation, meals, and entry fees "
        "to attractions.";

-    std::vector<std::tuple<std::string, std::string, nlohmann::json>> tool_calls = {
-        { "call_0", "plan_trip", nlohmann::json::parse(R"({
+    std::vector<std::tuple<std::string, std::string, common_json>> tool_calls = {
+        { "call_0", "plan_trip", common_json::parse(R"({
            "destination": "Japan",
            "duration": 14,
            "budget": 4000,
@@ -686,16 +686,16 @@ void test_command7_parser_compare(testing & t) {
    if (!tool_calls.empty()) {
        tokens.emplace_back("<|START_ACTION|>");

-        auto json = nlohmann::json::array();
+        auto json = common_json::array();
        for (const auto & tc : tool_calls) {
-            auto tc_json            = nlohmann::json::object();
+            auto tc_json            = common_json::object();
            tc_json["tool_call_id"] = std::get<0>(tc);
            tc_json["tool_name"]    = std::get<1>(tc);
            tc_json["parameters"]   = std::get<2>(tc);
            json.push_back(tc_json);
        }

-        auto tokenized = simple_tokenize(json.dump(-1, ' ', true));
+        auto tokenized = simple_tokenize(json.dump(-1));
        tokens.insert(tokens.end(), tokenized.begin(), tokenized.end());

        tokens.emplace_back("<|END_ACTION|>");
@@ -737,7 +737,7 @@ static void test_prefix_tool_names(testing & t) {
                      {
                          { "arg1", { { "type", "integer" } } },
                      } },
-                    { "required", { "arg1" } },
+                    { "required", json::array({ "arg1" }) },
                } },
          } }
    };
@@ -757,7 +757,7 @@ static void test_prefix_tool_names(testing & t) {
                          { "arg1", { { "type", "integer" } } },
                          { "arg2", { { "type", "integer" } } },
                      } },
-                    { "required", { "arg1" } },
+                    { "required", json::array({ "arg1" }) },
                } },
          } }
    };
@@ -7,7 +7,7 @@
 #include <fstream>
 #include <filesystem>

-#include <nlohmann/json.hpp>
+#include "json.h"

 #undef NDEBUG
 #include <cassert>
@@ -20,7 +20,7 @@
 #include "jinja/lexer.h"
 #include "jinja/caps.h"

-using json = nlohmann::ordered_json;
+using json = common_json;

 static int main_automated_tests(void);

@@ -304,8 +304,8 @@ void run_single(const std::string& contents, json input, bool use_common, bool d
        if (input.contains("eos_token")) {
            eos_token = input["eos_token"].get<std::string>();
        }
-        nlohmann::ordered_json msgs_json = input["messages"];
-        nlohmann::ordered_json tools_json = input["tools"];
+        common_json msgs_json = input["messages"];
+        common_json tools_json = input["tools"];
        auto messages = common_chat_msgs_parse_oaicompat(msgs_json);
        auto tools = common_chat_tools_parse_oaicompat(tools_json);
        auto output = format_using_common(contents, bos_token, eos_token, messages, tools);
@@ -19,12 +19,12 @@
 #include <fstream>
 #include <functional>
 #include <iostream>
-#include <nlohmann/json.hpp>
+#include "json.h"
 #include <set>
 #include <stdexcept>
 #include <string>

-using json = nlohmann::ordered_json;
+using json = common_json;

 static std::ostream & operator<<(std::ostream & os, const common_chat_msg_diff & diff) {
    os << "{ content_delta: " << diff.content_delta << "; ";
@@ -7,13 +7,13 @@
 #include "../src/unicode.h"
 #include "../src/llama-grammar.h"

-#include <nlohmann/json.hpp>
+#include "json.h"

 #include <cassert>
 #include <string>
 #include <vector>

-using json = nlohmann::ordered_json;
+using json = common_json;

 static llama_grammar * build_grammar_with_root(const std::string & grammar_str, const char * grammar_root) {
    return llama_grammar_init_impl(nullptr, grammar_str.c_str(), grammar_root, false, nullptr, 0, nullptr, 0);
@@ -3,7 +3,7 @@
 #include <random>
 #include <cstdlib>

-#include <nlohmann/json.hpp>
+#include "json.h"
 #include "subproc.h"

 #include "jinja/runtime.h"
@@ -14,7 +14,7 @@

 #include "testing.h"

-using json = nlohmann::ordered_json;
+using json = common_json;

 static void test_template(testing & t, const std::string & name, const std::string & tmpl, const json & vars, const std::string & expect);

@@ -240,7 +240,7 @@ static void test_conditionals(testing & t) {

    test_template(t, "is undefined key falsy",
        "{{ 'yes' if not y['x'] else 'no' }}",
-        {{"y", {{}}}},
+        {{"y", json::array({nullptr})}},
        "yes"
    );

@@ -282,7 +282,7 @@ static void test_conditionals(testing & t) {

    test_template(t, "is non-empty object truthy",
        "{{ 'yes' if y else 'no' }}",
-        {{"y", {"x", false}}},
+        {{"y", json::array({"x", false})}},
        "yes"
    );

@@ -6,7 +6,7 @@

 #include "../src/llama-grammar.h"

-#include <nlohmann/json.hpp>
+#include "json.h"

 #include <cassert>
 #include <fstream>
@@ -1442,7 +1442,7 @@ static void test_resolves_to_string() {
    auto test = [](const std::string & name, const std::string & schema_str, bool expected) {
        fprintf(stderr, "- %s\n", name.c_str());
        common_schema_info info;
-        auto schema = nlohmann::ordered_json::parse(schema_str);
+        auto schema = common_json::parse(schema_str);
        info.resolve_refs(schema);
        bool result = info.resolves_to_string(schema);
        if (result != expected) {
@@ -1517,7 +1517,7 @@ int main() {

    test_all("C++", [](const TestCase & tc) {
        try {
-            tc.verify(json_schema_to_grammar(nlohmann::ordered_json::parse(tc.schema), true));
+            tc.verify(json_schema_to_grammar(common_json::parse(tc.schema), true));
            tc.verify_status(SUCCESS);
        } catch (const std::invalid_argument & ex) {
            fprintf(stderr, "Error: %s\n", ex.what());
@@ -1531,7 +1531,7 @@ int main() {
        auto run = [](const TestCase & tc) {
            fprintf(stderr, "- %s\n", tc.name.c_str());
            try {
-                tc.verify(json_schema_to_grammar(nlohmann::ordered_json::parse(tc.schema), true));
+                tc.verify(json_schema_to_grammar(common_json::parse(tc.schema), true));
                tc.verify_status(SUCCESS);
            } catch (const std::invalid_argument & ex) {
                fprintf(stderr, "Error: %s\n", ex.what());
@@ -9,7 +9,7 @@
 #include "http.h"
 #include "log.h"

-#include <nlohmann/json.hpp>
+#include "json.h"

 #include <algorithm>
 #include <cstdio>
@@ -55,7 +55,7 @@ static const char * COMMIT = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
 static void serve_repos(httplib::Server & server) {
    server.Get(R"(/api/models/(.+)/refs)", [](const httplib::Request & req, httplib::Response & res) {
        if (g_repos.count(req.matches[1])) {
-            res.set_content(nlohmann::json{{"branches", {{{"name", "main"}, {"targetCommit", COMMIT}}}}}.dump(),
+            res.set_content(common_json{{"branches", common_json::array({ common_json{{"name", "main"}, {"targetCommit", COMMIT}} })}}.dump(),
                            "application/json");
        } else {
            res.status = 404;
@@ -66,7 +66,7 @@ static void serve_repos(httplib::Server & server) {
            res.status = 404;
            return;
        }
-        auto files = nlohmann::json::array();
+        auto files = common_json::array();
        size_t i = 0;
        for (const auto & p : g_repos[req.matches[1]]) {
            char oid[41];
@@ -6,8 +6,7 @@
 #include "log.h"
 #include "console.h"

-#define JSON_ASSERT GGML_ASSERT
-#include <nlohmann/json.hpp>
+#include "json.h"

 #include <algorithm>
 #include <cctype>
@@ -16,7 +15,7 @@
 #include <map>
 #include <set>

-using json = nlohmann::ordered_json;
+using json = common_json;

 struct cli_context_impl {
    json messages      = json::array();
@@ -73,7 +72,7 @@ static std::string format_error_message(const json & err) {

 // err is the raw response body of a failed request; it may or may not be JSON
 static std::string format_error_message(const std::string & err) {
-    json parsed = json::parse(err, nullptr, false);
+    json parsed = json::parse_no_throw(err);
    if (!parsed.is_discarded()) {
        return format_error_message(parsed);
    }
@@ -157,7 +156,7 @@ bool cli_context::init() {
            if (!list_and_ask_models()) {
                return false;
            }
-        } catch (const json::parse_error & e) {
+        } catch (const common_json_error & e) {
            ui::show_error(e.what());
            ui::show_message("This might be caused by an incorrect server-base endpoint URL");
            return false;
@@ -364,7 +363,7 @@ bool cli_context::generate_completion(generated_content & content_out, cli_timin
    ui::assistant_turn a;

    std::string err = client.post_sse("/v1/chat/completions", body.dump(), should_stop, [&](const std::string & payload) {
-        json chunk = json::parse(payload, nullptr, false);
+        json chunk = json::parse_no_throw(payload);
        if (chunk.is_discarded()) {
            return;
        }
@@ -33,6 +33,7 @@ int llama_fit_params(int argc, char ** argv) {
    if (!params.fit_params_print) {
        const common_params_fit_status status = common_fit_params(params.model.path.c_str(), &mparams, &cparams,
                params.tensor_split, params.tensor_buft_overrides.data(), params.fit_params_target.data(), params.fit_params_min_ctx,
+                nullptr,
                params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR);
        if (status != COMMON_PARAMS_FIT_STATUS_SUCCESS) {
            LOG_ERR("%s: failed to fit CLI arguments to free memory, exiting...\n", __func__);
@@ -2294,6 +2294,7 @@ int llama_bench(int argc, char ** argv) {
                fit_overrides.data(),
                margins.data(),
                inst.fit_min_ctx,
+                nullptr,
                params.verbose ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR);
       }

@@ -30,6 +30,7 @@ add_library(mtmd
            models/models.h
            models/cogvlm.cpp
            models/conformer.cpp
+            models/dots3note.cpp
            models/dotsocr.cpp
            models/exaone4_5.cpp
            models/gemma4a.cpp
@@ -120,6 +120,12 @@ struct clip_graph {
            ffn_op_type type_op,
            int il) const;

+    ggml_tensor * build_moe_ffn(
+            ggml_tensor * cur,
+            const clip_layer & layer,
+            ffn_op_type type_op,
+            int il) const;
+
    ggml_tensor * build_attn(
            ggml_tensor * wo,
            ggml_tensor * wo_b,
@@ -131,9 +137,15 @@ struct clip_graph {
            int il,
            ggml_tensor * sinks = nullptr) const;

-    // implementation of the 2D RoPE without adding a new op in ggml
-    // this is not efficient (use double the memory), but works on all backends
-    // TODO: there was a more efficient which relies on ggml_view and ggml_rope_ext_inplace, but the rope inplace does not work well with non-contiguous tensors ; we should fix that and revert back to the original implementation in https://github.com/ggml-org/llama.cpp/pull/13065
+    // implementation of the 2D RoPE using two ggml_rope_ext calls
+    //
+    // unlike GGML_ROPE_TYPE_VISION which forces NEOX ordering, this rotates adjacent pairs (normal ordering)
+    //
+    // example:
+    //  given a single head with size = 8 --> [00000000]
+    //  dims [0, 4) rotate with pos_a, dims [4, 8) rotate with pos_b --> [aaaabbbb]
+    //  interleave_freq = false --> both halves use the same inv_freq set (like GGML_ROPE_TYPE_VISION)
+    //  interleave_freq = true  --> first half uses even inv_freq, second half uses odd inv_freq (used by pixtral)
    ggml_tensor * build_rope_2d(
        ggml_context * ctx0,
        ggml_tensor * cur,
@@ -75,6 +75,7 @@
 #define KEY_SAM_N_HEAD             "clip.vision.sam.head_count"
 #define KEY_SAM_N_BLOCK            "clip.vision.sam.block_count"
 #define KEY_SAM_N_EMBD             "clip.vision.sam.embedding_length"
+#define KEY_VISION_N_EXPERT_USED   "clip.vision.expert_used_count"
 // audio-specific
 #define KEY_AUDIO_PROJ_TYPE        "clip.audio.projector_type" // for models with mixed modalities
 #define KEY_A_NUM_MEL_BINS         "clip.audio.num_mel_bins"
@@ -119,7 +120,11 @@
 #define TN_FFN_DOWN        "%s.blk.%d.ffn_down.%s"
 #define TN_FFN_GATE        "%s.blk.%d.ffn_gate.%s"
 #define TN_FFN_UP          "%s.blk.%d.ffn_up.%s"
-#define TN_FFN_GATE        "%s.blk.%d.ffn_gate.%s"
+#define TN_FFN_GATE_INP    "%s.blk.%d.ffn_gate_inp.%s"    // MoE router (dots3note)
+#define TN_FFN_GATE_EXPS   "%s.blk.%d.ffn_gate_exps.%s"
+#define TN_FFN_UP_EXPS     "%s.blk.%d.ffn_up_exps.%s"
+#define TN_FFN_DOWN_EXPS   "%s.blk.%d.ffn_down_exps.%s"
+#define TN_FFN_EXP_PROBS_B "%s.blk.%d.exp_probs_b.%s"
 #define TN_LN_1            "%s.blk.%d.ln1.%s" // layer norm
 #define TN_LN_2            "%s.blk.%d.ln2.%s" // layer norm
 #define TN_LS_1            "%s.blk.%d.ls1.%s"         // layer scale
@@ -471,6 +476,8 @@ enum projector_type {
    PROJECTOR_TYPE_COGVLM,
    PROJECTOR_TYPE_JANUS_PRO,
    PROJECTOR_TYPE_DOTS_OCR,
+    PROJECTOR_TYPE_DOTS3NOTE_V,
+    PROJECTOR_TYPE_DOTS3NOTE_A,
    PROJECTOR_TYPE_DEEPSEEKOCR,
    PROJECTOR_TYPE_DEEPSEEKOCR2,
    PROJECTOR_TYPE_LFM2A,
@@ -533,6 +540,8 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
    { PROJECTOR_TYPE_COGVLM,            "cogvlm"},
    { PROJECTOR_TYPE_JANUS_PRO,         "janus_pro"},
    { PROJECTOR_TYPE_DOTS_OCR,          "dots_ocr"},
+    { PROJECTOR_TYPE_DOTS3NOTE_V,       "dots3note_v"},
+    { PROJECTOR_TYPE_DOTS3NOTE_A,       "dots3note_a"},
    { PROJECTOR_TYPE_DEEPSEEKOCR,       "deepseekocr"},
    { PROJECTOR_TYPE_DEEPSEEKOCR2,      "deepseekocr2"},
    { PROJECTOR_TYPE_LFM2A,             "lfm2a"},
@@ -93,6 +93,7 @@ struct clip_hparams {

    float eps = 1e-6;
    float rope_theta = 0.0;
+    int32_t n_expert_used = 0;
    std::vector<int32_t> feature_layers;
    int32_t attn_window_size = 0;
    int32_t n_wa_pattern = 0;
@@ -259,6 +260,13 @@ struct clip_layer {
    ggml_tensor * ff_down_w = nullptr;
    ggml_tensor * ff_down_b = nullptr;

+    // MoE FFN (dots3note vision pyramid blocks)
+    ggml_tensor * ff_gate_inp_w  = nullptr;
+    ggml_tensor * ff_gate_exps_w = nullptr;
+    ggml_tensor * ff_up_exps_w   = nullptr;
+    ggml_tensor * ff_down_exps_w = nullptr;
+    ggml_tensor * ff_exp_probs_b = nullptr;
+
    // layernorm 2 (or pre-FFN norm)
    ggml_tensor * ln_2_w = nullptr;
    ggml_tensor * ln_2_b = nullptr;
@@ -514,11 +514,13 @@ ggml_tensor * clip_graph::build_vit(
        cb(cur, "ffn_inp_normed", il);

        // ffn
-        cur = build_ffn(cur,
-            layer.ff_up_w, layer.ff_up_b,
-            layer.ff_gate_w, layer.ff_gate_b,
-            layer.ff_down_w, layer.ff_down_b,
-            ffn_t, il);
+        cur = layer.ff_gate_exps_w
+            ? build_moe_ffn(cur, layer, ffn_t, il)
+            : build_ffn(cur,
+                layer.ff_up_w, layer.ff_up_b,
+                layer.ff_gate_w, layer.ff_gate_b,
+                layer.ff_down_w, layer.ff_down_b,
+                ffn_t, il);

        cb(cur, "ffn_out", il);

@@ -699,6 +701,50 @@ ggml_tensor * clip_graph::build_ffn(
    return cur;
 }

+// MoE FFN with sigmoid router and normalized top-k weights (dots3note vision)
+// the router runs in fp32; exp_probs_b only affects expert selection, not the weights
+ggml_tensor * clip_graph::build_moe_ffn(ggml_tensor * cur, const clip_layer & layer, ffn_op_type type_op, int il) const {
+    const int64_t n_tokens      = cur->ne[1];
+    const int64_t n_expert      = layer.ff_gate_exps_w->ne[2];
+    const int64_t n_expert_used = std::min((int64_t) hparams.n_expert_used, n_expert);
+    GGML_ASSERT(n_expert_used > 0);
+    GGML_ASSERT(type_op == FFN_SILU);
+
+    ggml_tensor * probs = ggml_sigmoid(ctx0, build_mm(layer.ff_gate_inp_w, cur)); // [n_expert, n_tokens]
+    cb(probs, "ffn_moe_probs", il);
+
+    ggml_tensor * sel = layer.ff_exp_probs_b
+        ? ggml_add(ctx0, probs, layer.ff_exp_probs_b)
+        : probs;
+    ggml_tensor * selected = ggml_top_k(ctx0, sel, n_expert_used); // [n_expert_used, n_tokens]
+
+    ggml_tensor * weights = ggml_get_rows(ctx0,
+        ggml_reshape_3d(ctx0, probs, 1, n_expert, n_tokens), selected);
+    weights = ggml_reshape_2d(ctx0, weights, n_expert_used, n_tokens);
+    weights = ggml_div(ctx0, weights, ggml_sum_rows(ctx0, weights));
+    weights = ggml_reshape_3d(ctx0, weights, 1, n_expert_used, n_tokens);
+    cb(weights, "ffn_moe_weights", il);
+
+    cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], 1, n_tokens);
+    ggml_tensor * gate = ggml_mul_mat_id(ctx0, layer.ff_gate_exps_w, cur, selected); // [n_ff, n_expert_used, n_tokens]
+    ggml_tensor * up   = ggml_mul_mat_id(ctx0, layer.ff_up_exps_w,   cur, selected);
+    cur = ggml_mul(ctx0, ggml_silu(ctx0, gate), up);
+    cur = ggml_mul_mat_id(ctx0, layer.ff_down_exps_w, cur, selected); // [n_embd, n_expert_used, n_tokens]
+    cur = ggml_mul(ctx0, cur, weights);
+
+    // sum over the selected experts
+    ggml_tensor * out = nullptr;
+    for (int64_t i = 0; i < n_expert_used; i++) {
+        ggml_tensor * v = ggml_view_2d(ctx0, cur, cur->ne[0], n_tokens, cur->nb[2], i * cur->nb[1]);
+        out = out ? ggml_add(ctx0, out, v) : v;
+    }
+    if (n_expert_used == 1) {
+        out = ggml_cont(ctx0, out);
+    }
+    cb(out, "ffn_moe_out", il);
+    return out;
+}
+
 ggml_tensor * clip_graph::build_attn(
        ggml_tensor * wo,
        ggml_tensor * wo_b,
@@ -773,8 +819,6 @@ ggml_tensor * clip_graph::build_attn(
 }

 // implementation of the 2D RoPE without adding a new op in ggml
-// this is not efficient (use double the memory), but works on all backends
-// TODO: there was a more efficient which relies on ggml_view and ggml_rope_ext_inplace, but the rope inplace does not work well with non-contiguous tensors ; we should fix that and revert back to the original implementation in https://github.com/ggml-org/llama.cpp/pull/13065
 ggml_tensor * clip_graph::build_rope_2d(
    ggml_context * ctx0,
    ggml_tensor * cur,
@@ -783,9 +827,7 @@ ggml_tensor * clip_graph::build_rope_2d(
    const float freq_base,
    const bool interleave_freq
 ) {
-    const int64_t n_dim  = cur->ne[0];
-    const int64_t n_head = cur->ne[1];
-    const int64_t n_pos  = cur->ne[2];
+    const int64_t n_dim = cur->ne[0];

    // for example, if we have cur tensor of shape (n_dim=8, n_head, n_pos)
    // we will have a list of 4 inv_freq: 1e-0, 1e-1, 1e-2, 1e-3
@@ -799,46 +841,30 @@ ggml_tensor * clip_graph::build_rope_2d(
                                ? std::pow(freq_base, (float)-2/n_dim)
                                : 1.0;

-    // first half
-    ggml_tensor * first;
-    {
-        first = ggml_view_3d(ctx0, cur,
-            n_dim/2, n_head, n_pos,
-            cur->nb[1],
-            cur->nb[2],
-            0);
-        first = ggml_rope_ext(
-            ctx0,
-            first,
-            pos_a,      // positions
-            nullptr,    // freq factors
-            n_dim/2,    // n_dims
-            0, 0, freq_base,
-            1.0f, 0.0f, 1.0f, 0.0f, 0.0f
-        );
-    }
+    // first half, dims [0, n_dim/2)
+    cur = ggml_rope_ext(
+        ctx0,
+        cur,
+        pos_a,      // positions
+        nullptr,    // freq factors
+        n_dim/2,    // n_dims
+        0, 0, freq_base,
+        1.0f, 0.0f, 1.0f, 0.0f, 0.0f
+    );

-    // second half
-    ggml_tensor * second;
-    {
-        second = ggml_view_3d(ctx0, cur,
-            n_dim/2, n_head, n_pos,
-            cur->nb[1],
-            cur->nb[2],
-            n_dim/2 * ggml_element_size(cur));
-        second = ggml_rope_ext(
-            ctx0,
-            second,
-            pos_b,      // positions
-            nullptr,    // freq factors
-            n_dim/2,    // n_dims
-            0, 0, freq_base,
-            freq_scale_odd,
-            0.0f, 1.0f, 0.0f, 0.0f
-        );
-    }
+    // second half, dims [n_dim/2, n_dim)
+    cur = ggml_rope_ext(
+        ctx0,
+        cur,
+        pos_b,      // positions
+        nullptr,    // freq factors
+        n_dim/2,    // n_dims
+        0, 0, freq_base,
+        freq_scale_odd,
+        0.0f, 1.0f, 0.0f, 0.0f
+    );
+    cur = ggml_rope_set_offset(cur, n_dim/2);

-    cur = ggml_concat(ctx0, first, second, 0);
    return cur;
 }

@@ -933,9 +959,14 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
                builder = std::make_unique<clip_graph_pixtral>(ctx, img);
            } break;
        case PROJECTOR_TYPE_DOTS_OCR:
+        case PROJECTOR_TYPE_DOTS3NOTE_V: // same ViT + merger; pyramid MoE is handled by build_vit
            {
                builder = std::make_unique<clip_graph_dotsocr>(ctx, img);
            } break;
+        case PROJECTOR_TYPE_DOTS3NOTE_A:
+            {
+                builder = std::make_unique<clip_graph_dots3note_a>(ctx, img);
+            } break;
        case PROJECTOR_TYPE_QWEN2VL:
        case PROJECTOR_TYPE_QWEN25VL:
            {
@@ -1510,6 +1541,25 @@ struct clip_model_loader {
                        get_u32(KEY_IMAGE_MAX_PIXELS, hparams.image_max_pixels);
                        hparams.set_warmup_n_tokens(46*46); // avoid OOM on warmup
                    } break;
+                case PROJECTOR_TYPE_DOTS3NOTE_V:
+                    {
+                        hparams.rope_theta = 10000.0f;
+                        hparams.image_resize_algo = RESIZE_ALGO_BICUBIC_PILLOW;
+                        get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge);
+                        get_u32(KEY_IMAGE_MIN_PIXELS, hparams.image_min_pixels);
+                        get_u32(KEY_IMAGE_MAX_PIXELS, hparams.image_max_pixels);
+                        get_u32(KEY_VISION_N_EXPERT_USED, hparams.n_expert_used);
+                        hparams.set_warmup_n_tokens(46*46); // avoid OOM on warmup
+                    } break;
+                case PROJECTOR_TYPE_DOTS3NOTE_A:
+                    {
+                        hparams.rope_theta = 10000.0f;
+                        hparams.audio_chunk_len   = 60; // in seconds
+                        hparams.audio_sample_rate = 16000;
+                        hparams.audio_n_fft       = 400;
+                        hparams.audio_window_len  = 400;
+                        hparams.audio_hop_len     = 160;
+                    } break;
                case PROJECTOR_TYPE_KIMIVL:
                    {
                        hparams.image_resize_algo = RESIZE_ALGO_BILINEAR;
@@ -2190,12 +2240,20 @@ struct clip_model_loader {
            layer.ln_1_b = get_tensor(string_format(TN_LN_1,        prefix, il, "bias"), false);
            layer.ln_2_b = get_tensor(string_format(TN_LN_2,        prefix, il, "bias"), false);

+            // MoE ffn (dots3note vision pyramid blocks); replaces the dense ffn when present
+            layer.ff_gate_inp_w  = get_tensor(string_format(TN_FFN_GATE_INP,  prefix, il, "weight"), false);
+            layer.ff_gate_exps_w = get_tensor(string_format(TN_FFN_GATE_EXPS, prefix, il, "weight"), false);
+            layer.ff_up_exps_w   = get_tensor(string_format(TN_FFN_UP_EXPS,   prefix, il, "weight"), false);
+            layer.ff_down_exps_w = get_tensor(string_format(TN_FFN_DOWN_EXPS, prefix, il, "weight"), false);
+            layer.ff_exp_probs_b = get_tensor(string_format(TN_FFN_EXP_PROBS_B, prefix, il, "weight"), false);
+            const bool is_moe = layer.ff_gate_exps_w != nullptr;
+
            // ffn
-            layer.ff_up_w   = get_tensor(string_format(TN_FFN_UP,   prefix, il, "weight"));
+            layer.ff_up_w   = get_tensor(string_format(TN_FFN_UP,   prefix, il, "weight"), !is_moe);
            layer.ff_up_b   = get_tensor(string_format(TN_FFN_UP,   prefix, il, "bias"),   false);
            layer.ff_gate_w = get_tensor(string_format(TN_FFN_GATE, prefix, il, "weight"), false);
            layer.ff_gate_b = get_tensor(string_format(TN_FFN_GATE, prefix, il, "bias"),   false);
-            layer.ff_down_w = get_tensor(string_format(TN_FFN_DOWN, prefix, il, "weight"));
+            layer.ff_down_w = get_tensor(string_format(TN_FFN_DOWN, prefix, il, "weight"), !is_moe);
            layer.ff_down_b = get_tensor(string_format(TN_FFN_DOWN, prefix, il, "bias"),   false);

            // mimovl per-head attention sink bias
@@ -2677,6 +2735,7 @@ struct clip_model_loader {
                    model.mm_patch_merger_w = get_tensor(string_format(TN_MM_PATCH_MERGER, "weight"), false);
                } break;
            case PROJECTOR_TYPE_DOTS_OCR:
+            case PROJECTOR_TYPE_DOTS3NOTE_V:
                {
                    model.mm_0_w = get_tensor(string_format(TN_LLAVA_PROJ, 0, "weight"));
                    model.mm_0_b = get_tensor(string_format(TN_LLAVA_PROJ, 0, "bias"));
@@ -2687,6 +2746,23 @@ struct clip_model_loader {
                    // post_trunk_norm: applied after all ViT blocks, before the merger
                    model.post_ln_w = get_tensor(string_format(TN_MM_POST_NORM, "weight"));
                } break;
+            case PROJECTOR_TYPE_DOTS3NOTE_A:
+                {
+                    model.conv2d_1_w = get_tensor(string_format(TN_CONV2D, 1, "weight"));
+                    model.conv2d_1_b = get_tensor(string_format(TN_CONV2D, 1, "bias"));
+                    model.conv2d_2_w = get_tensor(string_format(TN_CONV2D, 2, "weight"));
+                    model.conv2d_2_b = get_tensor(string_format(TN_CONV2D, 2, "bias"));
+                    model.conv2d_3_w = get_tensor(string_format(TN_CONV2D, 3, "weight"));
+                    model.conv2d_3_b = get_tensor(string_format(TN_CONV2D, 3, "bias"));
+                    model.conv_out_w = get_tensor(string_format(TN_CONV_OUT, "weight")); // no bias
+                    // adapter: LayerNorm -> Linear -> GELU -> Linear
+                    model.mm_norm_pre_w = get_tensor(string_format(TN_MM_NORM_PRE, "weight"));
+                    model.mm_norm_pre_b = get_tensor(string_format(TN_MM_NORM_PRE, "bias"));
+                    model.mm_1_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 1, "weight"));
+                    model.mm_1_b = get_tensor(string_format(TN_MM_AUDIO_MLP, 1, "bias"));
+                    model.mm_2_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 3, "weight"));
+                    model.mm_2_b = get_tensor(string_format(TN_MM_AUDIO_MLP, 3, "bias"));
+                } break;
            case PROJECTOR_TYPE_ULTRAVOX:
                {
                    model.conv1d_1_w = get_tensor(string_format(TN_CONV1D, 1, "weight"));
@@ -4075,12 +4151,18 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
            } break;
        case PROJECTOR_TYPE_PADDLEOCR:
        case PROJECTOR_TYPE_DOTS_OCR:
+        case PROJECTOR_TYPE_DOTS3NOTE_V:
            {
                // dynamic size
                int n_merge = ctx->model.hparams.n_merge;
                int stride = n_merge * n_merge;
                n_patches = CLIP_ALIGN(n_patches, stride) / stride;
            } break;
+        case PROJECTOR_TYPE_DOTS3NOTE_A:
+            {
+                // 3x stride-2 conv2d over mel frames
+                n_patches = (img->nx() + 7) / 8;
+            } break;
        case PROJECTOR_TYPE_PIXTRAL:
        case PROJECTOR_TYPE_LIGHTONOCR:
            {
@@ -4727,6 +4809,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
                set_input_i32("minimax_pos_w", pos_w);
            } break;
        case PROJECTOR_TYPE_DOTS_OCR:
+        case PROJECTOR_TYPE_DOTS3NOTE_V:
            {
                const int pw = image_size_width / patch_size;
                const int ph = image_size_height / patch_size;
@@ -5217,6 +5300,16 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
                }
                set_input_i32("pos_w", pos_data);
            } break;
+        case PROJECTOR_TYPE_DOTS3NOTE_A:
+            {
+                GGML_ASSERT(imgs.entries.size() == 1);
+                const int n_pos = (imgs.entries.front().nx() + 7) / 8; // 3x stride-2 conv2d
+                std::vector<int32_t> positions(n_pos);
+                for (int i = 0; i < n_pos; i++) {
+                    positions[i] = i;
+                }
+                set_input_i32("positions", positions);
+            } break;
        case PROJECTOR_TYPE_GEMMA4A:
            {
                GGML_ASSERT(imgs.entries.size() == 1);
@@ -5713,6 +5806,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
        case PROJECTOR_TYPE_PIXTRAL:
        case PROJECTOR_TYPE_LIGHTONOCR:
        case PROJECTOR_TYPE_DOTS_OCR:
+        case PROJECTOR_TYPE_DOTS3NOTE_V:
+        case PROJECTOR_TYPE_DOTS3NOTE_A:
            return ctx->model.mm_2_w->ne[1];
        case PROJECTOR_TYPE_MLP_NORM:
            return ctx->model.mm_3_b->ne[0];
@@ -0,0 +1,61 @@
+#include "models.h"
+
+ggml_cgraph * clip_graph_dots3note_a::build() {
+    // inp_raw: [n_frames, n_mel, 1], one 60s chunk, mel frames not padded
+    // the reference impl zero-masks conv inputs beyond the valid length at each stage;
+    // running on exactly the valid frames with the convs' zero padding is equivalent
+    ggml_tensor * inp = build_inp_raw(1);
+    GGML_ASSERT(inp->type == GGML_TYPE_F32);
+
+    // 3x conv2d (k=3, s=2, p=1) + gelu
+    {
+        auto conv_block = [&](ggml_tensor * x, ggml_tensor * w, ggml_tensor * b) {
+            x = ggml_conv_2d(ctx0, w, x, 2, 2, 1, 1, 1, 1);
+            x = ggml_add(ctx0, x, ggml_reshape_4d(ctx0, b, 1, 1, x->ne[2], 1));
+            return ggml_gelu_erf(ctx0, x);
+        };
+
+        inp = conv_block(inp, model.conv2d_1_w, model.conv2d_1_b);
+        inp = conv_block(inp, model.conv2d_2_w, model.conv2d_2_b);
+        inp = conv_block(inp, model.conv2d_3_w, model.conv2d_3_b);
+        // inp: [OW=n_frames/8, OH=n_mel/8, OC=480, 1]
+        cb(inp, "after_conv_stem", -1);
+    }
+
+    // [OW, OH, OC, 1] -> [OH*OC, OW], feature index f + OH*c (matches the reference permute+reshape)
+    inp = ggml_cont(ctx0, ggml_permute(ctx0, inp, 2, 0, 1, 3));
+    inp = ggml_reshape_2d(ctx0, inp, inp->ne[0] * inp->ne[1], inp->ne[2]);
+
+    // project to d_model (no bias)
+    inp = ggml_mul_mat(ctx0, model.conv_out_w, inp);
+    cb(inp, "after_conv_out", -1);
+
+    const int64_t n_pos = inp->ne[1];
+
+    ggml_tensor * positions = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos);
+    ggml_set_name(positions, "positions");
+    ggml_set_input(positions);
+
+    // partial rotary: first half of each head, NEOX style
+    auto add_pos = [&](ggml_tensor * cur, const clip_layer &) {
+        return ggml_rope_ext(ctx0, cur, positions, nullptr, d_head/2,
+                             GGML_ROPE_TYPE_NEOX, 0, hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+    };
+
+    ggml_tensor * cur = build_vit(inp, n_pos,
+        NORM_TYPE_RMS, hparams.ffn_op,
+        nullptr, add_pos);
+    cb(cur, "after_transformer", -1);
+
+    // adapter: LayerNorm -> Linear -> GELU -> Linear
+    cur = build_norm(cur, model.mm_norm_pre_w, model.mm_norm_pre_b, NORM_TYPE_NORMAL, 1e-5, -1);
+    cur = build_ffn(cur,
+        model.mm_1_w, model.mm_1_b,
+        nullptr, nullptr,
+        model.mm_2_w, model.mm_2_b,
+        FFN_GELU_ERF, -1);
+    cb(cur, "projected", -1);
+
+    ggml_build_forward_expand(gf, cur);
+    return gf;
+}
@@ -44,51 +44,31 @@ ggml_cgraph * clip_graph_gemma4v::build() {

    // similar to build_rope_2d, but use neox ordering
    auto add_pos = [&](ggml_tensor * cur, const clip_layer &) {
-        const int64_t n_dim  = cur->ne[0];
-        const int64_t n_head = cur->ne[1];
-        const int64_t n_pos  = cur->ne[2];
+        const int64_t n_dim = cur->ne[0];

-        // first half
-        ggml_tensor * first;
-        {
-            first = ggml_view_4d(ctx0, cur,
-                n_dim/2, n_head, n_pos, n_batch,
-                cur->nb[1],
-                cur->nb[2],
-                cur->nb[3],
-                0);
-            first = ggml_rope_ext(
-                ctx0,
-                first,
-                pos_x,      // positions
-                nullptr,    // freq factors
-                n_dim/2,    // n_dims
-                GGML_ROPE_TYPE_NEOX, 0, hparams.rope_theta,
-                1.0f, 0.0f, 1.0f, 0.0f, 0.0f
-            );
-        }
+        // first half, dims [0, n_dim/2)
+        cur = ggml_rope_ext(
+            ctx0,
+            cur,
+            pos_x,      // positions
+            nullptr,    // freq factors
+            n_dim/2,    // n_dims
+            GGML_ROPE_TYPE_NEOX, 0, hparams.rope_theta,
+            1.0f, 0.0f, 1.0f, 0.0f, 0.0f
+        );

-        // second half
-        ggml_tensor * second;
-        {
-            second = ggml_view_4d(ctx0, cur,
-                n_dim/2, n_head, n_pos, n_batch,
-                cur->nb[1],
-                cur->nb[2],
-                cur->nb[3],
-                n_dim/2 * ggml_element_size(cur));
-            second = ggml_rope_ext(
-                ctx0,
-                second,
-                pos_y,      // positions
-                nullptr,    // freq factors
-                n_dim/2,    // n_dims
-                GGML_ROPE_TYPE_NEOX, 0, hparams.rope_theta,
-                1.0f, 0.0f, 1.0f, 0.0f, 0.0f
-            );
-        }
+        // second half, dims [n_dim/2, n_dim)
+        cur = ggml_rope_ext(
+            ctx0,
+            cur,
+            pos_y,      // positions
+            nullptr,    // freq factors
+            n_dim/2,    // n_dims
+            GGML_ROPE_TYPE_NEOX, 0, hparams.rope_theta,
+            1.0f, 0.0f, 1.0f, 0.0f, 0.0f
+        );
+        cur = ggml_rope_set_offset(cur, n_dim/2);

-        cur = ggml_concat(ctx0, first, second, 0);
        return cur;
    };

@@ -2,30 +2,22 @@

 ggml_tensor * clip_graph_minimax_m3::apply_rope(
        ggml_tensor * x, ggml_tensor * pos_h, ggml_tensor * pos_w) {
-    const int64_t Hn  = x->ne[1];
-    const int64_t P   = x->ne[2];
-    const size_t  es  = ggml_element_size(x);
-    const int     dh  = (int) x->ne[0];
-    const int     axd = 2 * ((2 * (dh / 2) / 3) / 2);
+    const int dh  = (int) x->ne[0];
+    const int axd = 2 * ((2 * (dh / 2) / 3) / 2);

-    GGML_ASSERT(x->nb[0] == es);
    GGML_ASSERT(3 * axd <= dh);

    const float th  = hparams.rope_theta;

    // layout of x is [t, h, w, pad]
    // t is unrotated, h and w are rotated, pad is unrotated
-    // note: everything from n_dims onward untouched, so w and pad are rotated in one call.
-    auto sl = [&](int off, int n) {
-        return ggml_cont(ctx0, ggml_view_3d(ctx0, x, n, Hn, P, x->nb[1], x->nb[2], (size_t) off * es));
-    };
-    ggml_tensor * t = sl(0,       axd);
-    ggml_tensor * h = sl(axd,     axd);
-    ggml_tensor * w = sl(2 * axd, dh - 2 * axd); // w + pad
+    x = ggml_rope_ext(ctx0, x, pos_h, nullptr, axd, GGML_ROPE_TYPE_NEOX, 0, th, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+    x = ggml_rope_set_offset(x, axd);

-    h = ggml_rope_ext(ctx0, h, pos_h, nullptr, axd, GGML_ROPE_TYPE_NEOX, 0, th, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
-    w = ggml_rope_ext(ctx0, w, pos_w, nullptr, axd, GGML_ROPE_TYPE_NEOX, 0, th, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
-    return ggml_concat(ctx0, ggml_concat(ctx0, t, h, 0), w, 0);
+    x = ggml_rope_ext(ctx0, x, pos_w, nullptr, axd, GGML_ROPE_TYPE_NEOX, 0, th, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+    x = ggml_rope_set_offset(x, 2 * axd);
+
+    return x;
 }

 ggml_cgraph * clip_graph_minimax_m3::build() {
@@ -119,6 +119,11 @@ struct clip_graph_dotsocr : clip_graph {
    ggml_cgraph * build() override;
 };

+struct clip_graph_dots3note_a : clip_graph {
+    clip_graph_dots3note_a(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
+    ggml_cgraph * build() override;
+};
+
 struct clip_graph_cogvlm : clip_graph {
    clip_graph_cogvlm(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
    ggml_cgraph * build() override;
@@ -723,6 +723,100 @@ bool mtmd_audio_preprocessor_qwen3a::preprocess(const float *                 sa
    return true;
 }

+//
+// mtmd_audio_preprocessor_dots3note
+//
+// Matches Dots3NoteFeatureExtractor: the waveform is split into 60s chunks and each chunk gets
+// its own whisper-style log-mel (center=True, log10 + (max-8)/4). Only sample_length//hop frames
+// per chunk are valid; the reference masks everything beyond them, so we emit exactly that many.
+//
+
+void mtmd_audio_preprocessor_dots3note::initialize() {
+    cache.fill_sin_cos_table(hparams.audio_n_fft);
+    cache.fill_hann_window(hparams.audio_window_len, true);
+    cache.fill_mel_filterbank_matrix(hparams.n_mel_bins, hparams.audio_n_fft, hparams.audio_sample_rate);
+}
+
+bool mtmd_audio_preprocessor_dots3note::preprocess(const float *                 samples,
+                                                   size_t                        n_samples,
+                                                   std::vector<mtmd_audio_mel> & output) {
+    if (n_samples == 0) {
+        return false;
+    }
+
+    GGML_ASSERT(!cache.sin_vals.empty());
+    GGML_ASSERT(!cache.cos_vals.empty());
+    GGML_ASSERT(!cache.filters.data.empty());
+
+    const int    pad           = hparams.audio_n_fft / 2; // center=True padding
+    const int    hop           = hparams.audio_hop_len;
+    const size_t chunk_samples = (size_t) hparams.audio_chunk_len * hparams.audio_sample_rate;
+
+    for (size_t start = 0; start < n_samples; start += chunk_samples) {
+        const size_t n_chunk = std::min(chunk_samples, n_samples - start);
+        const float * chunk  = samples + start;
+
+        const int64_t n_valid = n_chunk / hop;
+        if (n_valid == 0) {
+            continue; // sub-hop tail, contributes no frames
+        }
+
+        // reflect-pad the start; the reference zero-pads partial chunks to 60s before the STFT,
+        // so a partial chunk sees zeros past its end while a full chunk reflects its own tail
+        std::vector<float> padded(n_chunk + 2 * pad, 0.0f);
+        for (int i = 0; i < pad; i++) {
+            int src = pad - i;
+            padded[i] = (src < (int) n_chunk) ? chunk[src] : 0.0f;
+        }
+        std::copy(chunk, chunk + n_chunk, padded.begin() + pad);
+        if (n_chunk == chunk_samples) {
+            for (int i = 0; i < pad; i++) {
+                int src = (int) n_chunk - 2 - i;
+                padded[n_chunk + pad + i] = (src >= 0) ? chunk[src] : 0.0f;
+            }
+        }
+
+        filter_params params;
+        params.n_mel            = hparams.n_mel_bins;
+        params.n_fft_bins       = 1 + (hparams.audio_n_fft / 2);
+        params.hann_window_size = hparams.audio_window_len;
+        params.hop_length       = hop;
+        params.sample_rate      = hparams.audio_sample_rate;
+        params.no_padding       = true; // padding already applied above
+        params.use_natural_log  = false;
+
+        mtmd_audio_mel mel_full;
+        if (!log_mel_spectrogram(padded.data(), (int) padded.size(), 4, params, cache, mel_full)) {
+            return false;
+        }
+        GGML_ASSERT(mel_full.n_len >= n_valid);
+
+        // per-chunk whisper-style normalization, then keep only the valid frames
+        mtmd_audio_mel out;
+        out.n_mel     = mel_full.n_mel;
+        out.n_len     = n_valid;
+        out.n_len_org = n_valid;
+        out.data.resize((size_t) out.n_mel * (size_t) out.n_len);
+
+        double mmax = -1e20;
+        for (int64_t m = 0; m < out.n_mel; m++) {
+            for (int64_t t = 0; t < n_valid; t++) {
+                mmax = std::max(mmax, (double) mel_full.data[(size_t) m * mel_full.n_len + t]);
+            }
+        }
+        mmax -= 8.0;
+        for (int64_t m = 0; m < out.n_mel; m++) {
+            for (int64_t t = 0; t < n_valid; t++) {
+                const double v = std::max((double) mel_full.data[(size_t) m * mel_full.n_len + t], mmax);
+                out.data[(size_t) m * n_valid + t] = (float) ((v + 4.0) / 4.0);
+            }
+        }
+
+        output.push_back(std::move(out));
+    }
+    return !output.empty();
+}
+
 //
 // mtmd_audio_preprocessor_mimo_audio
 //
@@ -111,6 +111,15 @@ struct mtmd_audio_preprocessor_qwen3a : mtmd_audio_preprocessor {
    mtmd_audio_cache cache;
 };

+struct mtmd_audio_preprocessor_dots3note : mtmd_audio_preprocessor {
+    mtmd_audio_preprocessor_dots3note(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {}
+    void initialize() override;
+    bool preprocess(const float * samples, size_t n_samples, std::vector<mtmd_audio_mel> & output) override;
+
+  private:
+    mtmd_audio_cache cache;
+};
+
 struct mtmd_audio_preprocessor_mimo_audio : mtmd_audio_preprocessor {
    mtmd_audio_preprocessor_mimo_audio(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {}
    void initialize() override;
@@ -825,6 +825,7 @@ struct mtmd_context {
                    image_preproc = std::make_unique<mtmd_image_preprocessor_longest_edge>(ctx_v);
                } break;
            case PROJECTOR_TYPE_DOTS_OCR:
+            case PROJECTOR_TYPE_DOTS3NOTE_V:
                {
                    // <|img|> ... (image embeddings) ... <|endofimg|>
                    img_beg = "<|img|>";
@@ -976,6 +977,13 @@ struct mtmd_context {
                    aud_end = "<audio|>";
                    audio_preproc = std::make_unique<mtmd_audio_preprocessor_gemma4ua>(ctx_a);
                } break;
+            case PROJECTOR_TYPE_DOTS3NOTE_A:
+                {
+                    // <|audio_comp_start|> ... (embeddings) ... <|audio_comp_end|>
+                    aud_beg = "<|audio_comp_start|>";
+                    aud_end = "<|audio_comp_end|>";
+                    audio_preproc = std::make_unique<mtmd_audio_preprocessor_dots3note>(ctx_a);
+                } break;
            case PROJECTOR_TYPE_MIMO_AUDIO:
                {
                    aud_beg = "<|mimo_audio_start|>";
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Xuan-Son Nguyen	b21e4de745	mtmd: use ggml_rope_set_offset (#27521 ) * mtmd: use ggml_rope_set_offset * add comment	2026-08-22 16:33:47 +02:00
Xuan-Son Nguyen	d9f918d2d0	common: add json.h abstraction (#27511 ) * add common/json * migrate common * adapt jinja * migrate server * big wip * migrate tests * wip * revert some excessive changes * wip * wip 2 * revert redundant changes * fix server crash * various fixes * fix ci * harden a bit * clean up * rm json-shim * add some comments * rm redundant decl	2026-08-22 16:28:28 +02:00
Xuan-Son Nguyen	2fb989b9e7	fit: also take into account n_streams (#27496 ) * fit: also take into account n_streams * server: make the draft context follow the target context With a non-unified KV cache the target context now holds n_ctx_train tokens per sequence, while the draft context was still created with n_ctx = 0 and fell back to n_ctx_train / n_streams per sequence. A slot filled beyond that point makes the draft batch fail to decode, and the server answers 500 on the request. The draft context now takes its size from the target context, so both hold the same number of tokens per sequence. Contexts that share their cells with the target no longer need the kv_size override. The memory reserved for the draft model before fitting is measured at the largest context the target can take, since the draft context grows with the target and a fixed byte margin cannot express that. * fit: take an optional second model into account Illustrates the alternative discussed on the draft context fix. The memory of a draft or MTP context is currently handed to the fit as a fixed byte margin, which cannot express a memory that grows with the context the fit is still deciding on. common_fit_params now takes an optional second model that shares the devices of the main one. Its context follows the main context and its memory is measured again whenever that context changes, so the reduce path stays exact instead of conservative. A model that cannot be measured on its own, such as a shared cell MTP context, is skipped with a warning and the main model is fitted alone. This drops the reservation block in the server, which no longer has to probe the trained context size of the target to guess an upper bound. --------- Co-authored-by: Pascal <admin@serveurperso.com>	2026-08-22 16:16:06 +02:00
Xuan-Son Nguyen	9fee29e943	arg: remove -no-cnv from cli [no ci] (#27542 ) * arg: remove -no-cnv from cli * clarify about not adding exccesive test cases	2026-08-22 15:53:56 +02:00
Mario Limonciello	e85caa81ea	ci : Restore ROCm job for Ubuntu (#27399 ) * Revert "ci : disable ubuntu-rocm (#26969)" This reverts commit `9558fa44c9`. * ci: set ccache compiler_check=content for ROCm build The ROCm toolchain is pip-installed fresh on every run, so the clang binary's mtime changes each time. With ccache's default compiler_check=mtime that invalidates the whole cache and warm builds only reached ~70% hits. Hash the compiler contents instead so the cache survives toolchain reinstalls. * Update ccache size to 1GB We're waivering with so many architectures built, we need a bigger ccache limit. * merge fix --------- Co-authored-by: Jim Wu <ywu@xilinx.com> Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>	2026-08-22 13:28:30 +03:00
Tiwei Bie	2115b73d8e	model : support DSpark for bailingmoe3 (#27508 )	2026-08-22 12:19:48 +03:00
Xuan-Son Nguyen	54ee5ee643	mtmd: support dots3-note vision+audio (#27524 ) * text: conversion * init impl * mtmd: conversion * impl mtmd cpp * Update gguf-py/gguf/tensor_mapping.py Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> --------- Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>	2026-08-22 10:35:50 +02:00
Georgi Gerganov	3a653fea93	ci : add older, min and dry-run options to ccache-clear (#27504 ) * ci : add older, min and dry-run options to ccache-clear Assisted-by: pi:llama.cpp/Qwen3.8-27B * pi : add note about not wrapping lines in PR descriptions [no ci] Assisted-by: pi:llama.cpp/Qwen3.8-27B	2026-08-22 11:31:30 +03:00
Kartik Sirohi	369e1cd614	ggml: optimize concat op by replacing per-element memcpy with row-level memcpy (#24575 ) * ggml: optimize concat op by replacing per-element memcpy with row-level memcpy * ggml: fix concat offsets for row-level copies * ggml: add concat row contiguity asserts * ggml: move concat block size asserts * ggml: remove redundant concat asserts * Update ggml/src/ggml-cpu/ops.cpp Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> --------- Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>	2026-08-22 11:30:31 +03:00
Shahir BIn Zulfiker	2c6b141efb	common : fix draft-mtp with embeddings (#26352 , #27299 ) (#27400 ) * common: fix draft-mtp with embeddings (#26352) * --whitespace --------- Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>	2026-08-22 09:44:22 +02:00
Sigbjørn Skjæret	8672290039	sycl : add Q2_K reordered MMVQ and ESIMD kernels (again) (#27490 ) * Revert "Revert "sycl : add Q2_K reordered MMVQ and ESIMD kernels (#26336)" (#…" This reverts commit `7a0e42fd01`. * add gate params	2026-08-22 10:09:26 +03:00
Sigbjørn Skjæret	3aeb924628	readme : fix server badge alt (#27533 )	2026-08-22 10:08:07 +03:00
Georgi Gerganov	2100e59260	readme : update badges (#27531 )	2026-08-22 08:25:00 +03:00