llama: add token ID tracking to KV cell (#27762 )

* kv: track token id * rm get_prev_tokens, move it to the main pr * nits * add get_prev_tokens
ui: Move Settings and MCP Servers routes to dialog-based views (#27744 )
2026-08-27 02:27:34 +02:00 · 2026-08-26 23:34:28 +02:00 · 2026-08-26 21:07:24 +02:00 · 2026-08-26 20:19:19 +02:00 · 2026-08-26 19:02:06 +03:00 · 2026-08-26 16:49:32 +02:00
89 changed files with 2744 additions and 1233 deletions
@@ -90,6 +90,9 @@ RUN bash -c "source ${OpenVINO_DIR}/setupvars.sh && \
    cmake -B build/ReleaseOV -G Ninja \
        -DCMAKE_BUILD_TYPE=Release \
        -DLLAMA_BUILD_TESTS=OFF \
+        -DGGML_NATIVE=OFF \
+        -DGGML_BACKEND_DL=ON \
+        -DGGML_CPU_ALL_VARIANTS=ON \
        -DGGML_OPENVINO=ON && \
    cmake --build build/ReleaseOV --parallel "

@@ -58,34 +58,38 @@ runs:
      if: ${{ inputs.save == 'true' }}
      shell: bash
      run: |
-        set +e -uo pipefail
-        source .venv-hf/bin/activate
-        CCACHE_DIR=$(ccache -k cache_dir)
-        if [[ -d "$CCACHE_DIR" ]]; then
-            ccache -s
-            if [[ -n "${{ inputs.evict-old-files }}" ]]; then
-                ccache --evict-older-than "${{ inputs.evict-old-files }}"
+        if [[ -n "$HF_TOKEN" ]]; then
+            set +e -uo pipefail
+            source .venv-hf/bin/activate
+            CCACHE_DIR=$(ccache -k cache_dir)
+            if [[ -d "$CCACHE_DIR" ]]; then
+                ccache -s
+                if [[ -n "${{ inputs.evict-old-files }}" ]]; then
+                    ccache --evict-older-than "${{ inputs.evict-old-files }}"
+                fi
+                DATESTAMP=$(date -u +'%Y-%m-%dT%H:%M:%SZ')
+                CACHEFILE="${{ inputs.key }}-$DATESTAMP.tar.gz"
+                if tar -czf ccache_bucket.tar.gz -C "$CCACHE_DIR" .; then
+                    hf buckets cp ccache_bucket.tar.gz "hf://buckets/${{ inputs.hf_bucket }}/${{ inputs.folder }}/$CACHEFILE"
+                fi
+                rm ccache_bucket.tar.gz
+            else
+                echo "'$CCACHE_DIR' not found."
            fi
-            DATESTAMP=$(date -u +'%Y-%m-%dT%H:%M:%SZ')
-            CACHEFILE="${{ inputs.key }}-$DATESTAMP.tar.gz"
-            if tar -czf ccache_bucket.tar.gz -C "$CCACHE_DIR" .; then
-                hf buckets cp ccache_bucket.tar.gz "hf://buckets/${{ inputs.hf_bucket }}/${{ inputs.folder }}/$CACHEFILE"
-            fi
-            rm ccache_bucket.tar.gz
-        else
-            echo "'$CCACHE_DIR' not found."
        fi

    - name: Remove old ccache files from buckets
      if: ${{ inputs.save == 'true' }}
      shell: bash
      run: |
-        set +e -uo pipefail
-        source .venv-hf/bin/activate
-        CACHE_FILES=$(hf buckets list "hf://buckets/${{ inputs.hf_bucket }}/${{ inputs.folder }}" --json | jq -r '[.[] | select(.type == "file") | select((.uploaded_at | .[:19]+"Z" | fromdateiso8601) < (now - 5 * 60)) | select(.path | startswith("${{ inputs.folder }}/${{ inputs.key }}") and endswith(".tar.gz"))] | sort_by(.path)[:-1] | .[] | [.path // ""] | @tsv')
-        if [[ -n "$CACHE_FILES" ]]; then
-            echo "Removing old ccache files..."
-            while IFS=$'\t' read -r CACHE_PATH; do
-                hf buckets rm "hf://buckets/${{ inputs.hf_bucket }}/$CACHE_PATH" -y
-            done <<< "$CACHE_FILES"
+        if [[ -n "$HF_TOKEN" ]]; then
+            set +e -uo pipefail
+            source .venv-hf/bin/activate
+            CACHE_FILES=$(hf buckets list "hf://buckets/${{ inputs.hf_bucket }}/${{ inputs.folder }}" --json | jq -r '[.[] | select(.type == "file") | select((.uploaded_at | .[:19]+"Z" | fromdateiso8601) < (now - 5 * 60)) | select(.path | startswith("${{ inputs.folder }}/${{ inputs.key }}") and endswith(".tar.gz"))] | sort_by(.path)[:-1] | .[] | [.path // ""] | @tsv')
+            if [[ -n "$CACHE_FILES" ]]; then
+                echo "Removing old ccache files..."
+                while IFS=$'\t' read -r CACHE_PATH; do
+                    hf buckets rm "hf://buckets/${{ inputs.hf_bucket }}/$CACHE_PATH" -y
+                done <<< "$CACHE_FILES"
+            fi
        fi
@@ -22,7 +22,8 @@ on:
    types: [opened, synchronize, reopened]
    paths: [
      '.github/workflows/build-apple.yml',
-      'ggml/src/ggml-metal/**'
+      'ggml/src/ggml-metal/**',
+      'ggml/src/ggml-rpc/**'
    ]

 concurrency:
@@ -65,7 +65,7 @@ jobs:
        with:
          key: cuda-ubuntu-24.04-cuda
          folder: llama.cpp
-          hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }}
+          hf_bucket: ggml-org/cache

      - name: Build with CMake
        # TODO: Remove GGML_CUDA_CUB_3DOT2 flag once CCCL 3.2 is bundled within CTK and that CTK version is used in this project
@@ -89,7 +89,7 @@ jobs:
          key: cuda-ubuntu-24.04-cuda
          folder: llama.cpp
          evict-old-files: 1d
-          hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }}
+          hf_bucket: ggml-org/cache
          save: true

  hip:
@@ -120,7 +120,7 @@ jobs:
        with:
          key: cuda-ubuntu-22.04-hip
          folder: llama.cpp
-          hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }}
+          hf_bucket: ggml-org/cache

      - name: Build with native CMake HIP support
        id: cmake_build
@@ -140,7 +140,7 @@ jobs:
          key: cuda-ubuntu-22.04-hip
          folder: llama.cpp
          evict-old-files: 1d
-          hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }}
+          hf_bucket: ggml-org/cache
          save: true

  musa:
@@ -171,7 +171,7 @@ jobs:
        with:
          key: cuda-ubuntu-22.04-musa
          folder: llama.cpp
-          hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }}
+          hf_bucket: ggml-org/cache

      - name: Build with native CMake MUSA support
        id: cmake_build
@@ -189,5 +189,5 @@ jobs:
          key: cuda-ubuntu-22.04-musa
          folder: llama.cpp
          evict-old-files: 1d
-          hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }}
+          hf_bucket: ggml-org/cache
          save: true
@@ -64,7 +64,7 @@ jobs:
    needs: create_tag
    uses: ./.github/workflows/ui-build.yml
    with:
-      hf_ui_version: ${{ needs.create_tag.outputs.source_tag }}
+      ui_version: ${{ needs.create_tag.outputs.source_tag }}

  prepare_matrices:
    name: Prepare Docker matrices
@@ -162,7 +162,7 @@ jobs:
        if: ${{ matrix.config.prebuilt_ui == true }}
        uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist

      - name: Set up QEMU
@@ -61,31 +61,8 @@ jobs:
            echo "should_release=false" >> $GITHUB_OUTPUT
          fi

-  get-version:
-    runs-on: ubuntu-slim
-    outputs:
-      ui_version: ${{ steps.version.outputs.ui_version }}
-    steps:
-      - uses: actions/checkout@v6
-        with:
-          fetch-depth: 0
-      - id: version
-        run: |
-          # Resolve UI version: BUILD_NUMBER from cmake/build-info.cmake > git hash + epoch > fallback
-          version=""
-          if grep -q "BUILD_NUMBER" cmake/build-info.cmake; then
-            build_number=$(grep "set(BUILD_NUMBER" cmake/build-info.cmake | grep -oP '\d+')
-            if [ -n "$build_number" ] && [ "$build_number" -gt 0 ]; then
-              version="b${build_number}"
-            fi
-          fi
-          if [ -z "$version" ]; then
-            version=$(git rev-parse --short HEAD)-$(date +%s)
-          fi
-          echo "ui_version=${version}" >> $GITHUB_OUTPUT
-
  macos-cpu:
-    needs: [check-release, get-version]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}
    strategy:
      matrix:
@@ -119,12 +96,11 @@ jobs:
        with:
          fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
@@ -141,7 +117,6 @@ jobs:
            -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
            -DLLAMA_FATAL_WARNINGS=ON \
            -DLLAMA_BUILD_BORINGSSL=ON \
-            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
            ${{ env.CMAKE_ARGS }}
          cmake --build build --config Release -j $(sysctl -n hw.logicalcpu)

@@ -167,7 +142,7 @@ jobs:
          key: release-${{ matrix.os }}-${{ matrix.arch }}

  ubuntu-cpu:
-    needs: [check-release, get-version]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}
    strategy:
      matrix:
@@ -191,12 +166,11 @@ jobs:
        with:
          fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: Dependencies
        id: depends
@@ -227,7 +201,6 @@ jobs:
            -DGGML_NATIVE=OFF \
            -DGGML_CPU_ALL_VARIANTS=ON \
            -DLLAMA_FATAL_WARNINGS=ON \
-            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
            ${{ env.CMAKE_ARGS }}
          cmake --build build --config Release -j $(nproc)

@@ -254,7 +227,7 @@ jobs:
          key: release-${{ matrix.os }}-cpu

  ubuntu-vulkan:
-    needs: [check-release, get-version]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    strategy:
@@ -277,12 +250,11 @@ jobs:
        with:
          fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: Dependencies
        id: depends
@@ -314,7 +286,6 @@ jobs:
            -DGGML_NATIVE=OFF \
            -DGGML_CPU_ALL_VARIANTS=ON \
            -DGGML_VULKAN=ON \
-            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
            ${{ env.CMAKE_ARGS }}
          cmake --build build --config Release -j $(nproc)

@@ -340,7 +311,7 @@ jobs:
          key: release-${{ matrix.os }}-vulkan

  android-arm64:
-    needs: [check-release, get-version]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    runs-on: ubuntu-latest
@@ -358,12 +329,11 @@ jobs:
        with:
          fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: Set up JDK
        uses: actions/setup-java@v5
@@ -407,7 +377,6 @@ jobs:
            -DLLAMA_FATAL_WARNINGS=ON \
            -DGGML_OPENMP=OFF \
            -DLLAMA_BUILD_BORINGSSL=ON \
-            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
            ${{ env.CMAKE_ARGS }}
          cmake --build build --config Release -j $(nproc)

@@ -433,7 +402,7 @@ jobs:
          name: llama-bin-android-arm64.tar.gz

  ubuntu-24-openvino:
-    needs: [check-release, get-version]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    runs-on: ubuntu-24.04
@@ -460,12 +429,11 @@ jobs:
        with:
          fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
@@ -508,7 +476,6 @@ jobs:
            -DGGML_OPENVINO=ON \
            -DCMAKE_INSTALL_RPATH='$ORIGIN' \
            -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
-            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
            ${{ env.CMAKE_ARGS }}
          cmake --build build/ReleaseOV --config Release --parallel

@@ -552,7 +519,7 @@ jobs:
          key: release-ubuntu-24.04-openvino-release-no-preset-v1

  windows-openvino:
-    needs: [check-release]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    runs-on: windows-2022
@@ -577,12 +544,11 @@ jobs:
        with:
            fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
@@ -682,7 +648,7 @@ jobs:

  windows-cpu:
    name: windows-cpu / ${{ matrix.arch }}
-    needs: [check-release]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    runs-on: windows-2025-vs2026
@@ -702,12 +668,11 @@ jobs:
        with:
          fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: Install Ninja
        run: |
@@ -749,8 +714,10 @@ jobs:
        with:
          key: release-windows-2025-vs2026-${{ matrix.arch }}-cpu

+  # TODO: build only the ggml-hip backend like the other windows backend jobs
+  #       (windows-cuda, windows-sycl), then drop the ui-build dependency
  windows-rocm:
-    needs: [check-release]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    runs-on: windows-2022
@@ -769,6 +736,12 @@ jobs:
        with:
          fetch-depth: 0

+      - name: Download UI build
+        uses: actions/download-artifact@v7
+        with:
+          name: llama-ui.zip
+          path: tools/ui/dist
+
      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
        with:
@@ -879,6 +852,8 @@ jobs:
        with:
          key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}

+  # note: builds only the backend library - llama-server (with the embedded UI)
+  #       is injected from the windows-cpu zip during the release "Merge artifacts" step
  windows:
    needs: [check-release]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}
@@ -909,13 +884,6 @@ jobs:
        id: checkout
        uses: actions/checkout@v6

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
-        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
-
      - name: Install Vulkan SDK
        id: get_vulkan
        if: ${{ matrix.backend == 'vulkan' }}
@@ -978,6 +946,8 @@ jobs:
          path: llama-bin-win-${{ matrix.backend }}-${{ matrix.arch }}.zip
          name: llama-bin-win-${{ matrix.backend }}-${{ matrix.arch }}.zip

+  # note: builds only the ggml-cuda backend - llama-server is injected from the
+  #       windows-cpu zip during the release "Merge artifacts" step
  windows-cuda:
    name: windows-cuda (${{ matrix.cuda }}, ${{ matrix.arch }})
    needs: [check-release]
@@ -1006,13 +976,6 @@ jobs:
        id: checkout
        uses: actions/checkout@v6

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
-        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
-
      - name: Install Cuda Toolkit
        uses: ./.github/actions/windows-setup-cuda
        with:
@@ -1084,6 +1047,8 @@ jobs:
        with:
          key: release-windows-2022-${{ matrix.arch }}-cuda-${{ matrix.cuda }}

+  # note: builds only the ggml-sycl backend - llama-server is injected from the
+  #       windows-cpu zip during the release "Merge artifacts" step
  windows-sycl:
    needs: [check-release]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}
@@ -1118,13 +1083,6 @@ jobs:
          Expand-Archive -Path "level-zero-win-sdk.zip" -DestinationPath "C:/level-zero-sdk" -Force
          "LEVEL_ZERO_V1_SDK_PATH=C:/level-zero-sdk" | Out-File -FilePath $env:GITHUB_ENV -Append

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
-        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
-
      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
        with:
@@ -1195,7 +1153,7 @@ jobs:
          key: release-windows-2022-x64-sycl

  ubuntu-24-sycl:
-    needs: [check-release]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    strategy:
@@ -1237,12 +1195,11 @@ jobs:
          wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/level-zero-devel_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O level-zero-devel.deb
          sudo apt-get install -y ./level-zero.deb ./level-zero-devel.deb

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
@@ -1287,11 +1244,11 @@ jobs:
        with:
          key: release-ubuntu-24.04-sycl-${{ matrix.build }}

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

-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04

    permissions:
      actions: write
@@ -1310,12 +1267,11 @@ jobs:
        with:
          fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: Free up disk space
        uses: ggml-org/free-disk-space@v1.3.1
@@ -1325,7 +1281,7 @@ jobs:
      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
        with:
-          key: release-ubuntu-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
+          key: release-ubuntu-24.04-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
          evict-old-files: 1d
          max-size: "1G"

@@ -1388,7 +1344,6 @@ jobs:
            -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)

@@ -1414,10 +1369,10 @@ jobs:
      - name: ccache-clear
        uses: ./.github/actions/ccache-clear
        with:
-          key: release-ubuntu-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
+          key: release-ubuntu-24.04-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}

  ios-xcode:
-    needs: [check-release, get-version]
+    needs: [check-release]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}
    runs-on: macos-26

@@ -1445,8 +1400,7 @@ jobs:
            -DLLAMA_BUILD_SERVER=OFF \
            -DCMAKE_SYSTEM_NAME=iOS \
            -DCMAKE_OSX_DEPLOYMENT_TARGET=16.0 \
-            -DCMAKE_XCODE_ATTRIBUTE_DEVELOPMENT_TEAM=ggml \
-            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }}
+            -DCMAKE_XCODE_ATTRIBUTE_DEVELOPMENT_TEAM=ggml
          cmake --build build --config Release -j $(sysctl -n hw.logicalcpu) -- CODE_SIGNING_ALLOWED=NO

      - name: xcodebuild for swift package
@@ -1569,11 +1523,9 @@ jobs:
 #          name: llama-bin-${{ matrix.chip_type }}-openEuler-${{ matrix.arch }}${{ matrix.use_acl_graph == 'on' && '-aclgraph' || '' }}.tar.gz

  ui-build:
-    needs: [check-release, get-version]
+    needs: [check-release]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}
    uses: ./.github/workflows/ui-build.yml
-    with:
-      hf_ui_version: ${{ needs.get-version.outputs.ui_version }}

  release:
    if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
@@ -1588,14 +1540,13 @@ jobs:
    runs-on: ubuntu-slim

    needs:
-      - get-version
      - windows
      - windows-cpu
      - windows-cuda
      - windows-sycl
      - windows-rocm
      - windows-openvino
-      - ubuntu-22-rocm
+      - ubuntu-24-rocm
      - ubuntu-cpu
      - ubuntu-vulkan
      - ubuntu-24-openvino
@@ -1628,24 +1579,27 @@ jobs:
          path: ./artifact
          merge-multiple: true

-      - name: Move artifacts
+      - name: Merge artifacts
        id: move_artifacts
        run: |
          mkdir -p release

-          echo "Adding CPU backend files to existing zips..."
+          # the windows-cpu zip contains the full toolset (llama-server with the embedded
+          # UI, ggml-cpu) - inject it into the other windows zips so that every archive
+          # ships the same binaries, only with a different backend library on top
+          echo "Injecting windows-cpu binaries (llama-server + CPU backend) into the backend zips..."
          for arch in x64 arm64; do
            cpu_zip="artifact/llama-bin-win-cpu-${arch}.zip"
            temp_dir=$(mktemp -d)
-            echo "Extracting CPU backend for $arch..."
+            echo "Extracting windows-cpu-${arch} package..."
            unzip "$cpu_zip" -d "$temp_dir"

-            echo "Adding CPU files to $arch zips..."
+            echo "Merging into $arch zips..."
            for target_zip in artifact/llama-bin-win-*-${arch}.zip; do
              if [[ "$target_zip" == "$cpu_zip" ]]; then
                continue
              fi
-              echo "Adding CPU backend to $(basename "$target_zip")"
+              echo "Injecting into $(basename "$target_zip")"
              realpath_target_zip=$(realpath "$target_zip")
              (cd "$temp_dir" && zip -r "$realpath_target_zip" .)
            done
@@ -1669,7 +1623,7 @@ jobs:
        id: download_ui
        uses: actions/download-artifact@v7
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: ./ui-dist

      - name: Package UI
@@ -73,13 +73,6 @@ jobs:
          fetch-depth: 0
          ref: ${{ github.event.inputs.sha || github.event.pull_request.head.sha || github.sha || github.head_ref || github.ref_name }}

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
-        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
-
      - name: Build
        id: cmake_build
        run: |
@@ -31,6 +31,6 @@ jobs:
      - name: Upload built UI
        uses: actions/upload-artifact@v6
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist/
          retention-days: 1
@@ -3,8 +3,8 @@ name: UI Build
 on:
  workflow_call:
    inputs:
-      hf_ui_version:
-        description: 'Version string for version.json (e.g. 12345)'
+      ui_version:
+        description: 'Version string embedded in build.json (e.g. b1234); defaults to b<commit-count>'
        required: false
        type: string

@@ -17,6 +17,17 @@ jobs:
    steps:
      - name: Checkout code
        uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
+
+      - name: Resolve UI version
+        id: version
+        run: |
+          version="${{ inputs.ui_version }}"
+          if [ -z "$version" ]; then
+            version="b$(git rev-list --count HEAD)"
+          fi
+          echo "ui_version=${version}" >> $GITHUB_OUTPUT

      - name: Setup Node.js
        uses: actions/setup-node@v6
@@ -31,8 +42,7 @@ jobs:

      - name: Build application
        env:
-          HF_UI_VERSION: ${{ inputs.hf_ui_version || '' }}
-          LLAMA_BUILD_NUMBER: ${{ inputs.hf_ui_version || 'b0000' }}
+          LLAMA_BUILD_NUMBER: ${{ steps.version.outputs.ui_version }}
        run: npm run build
        working-directory: tools/ui

@@ -43,6 +53,6 @@ jobs:
      - name: Upload built UI
        uses: actions/upload-artifact@v6
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist/
          retention-days: 1
@@ -37,7 +37,7 @@ jobs:
      - name: Download UI build artifact
        uses: actions/download-artifact@v7
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist/

      - name: Create distribution archive
@@ -64,7 +64,7 @@ jobs:
      - name: Download built UI artifacts
        uses: actions/download-artifact@v6
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist/

      - name: Run type checking
@@ -106,7 +106,7 @@ jobs:
      - name: Download built UI artifacts
        uses: actions/download-artifact@v6
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist/

      - name: Build Storybook
@@ -63,7 +63,7 @@ jobs:
      - name: Download built UI artifacts
        uses: actions/download-artifact@v6
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist/

      - name: Install dependencies
@@ -126,7 +126,7 @@ jobs:
      - name: Download built UI artifacts (reuses ui-build)
        uses: actions/download-artifact@v6
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist/

      - name: Install Playwright browsers
@@ -134,8 +134,8 @@ option(LLAMA_BUILD_TOOLS     "llama: build tools"
 option(LLAMA_BUILD_EXAMPLES  "llama: build examples"                                                             ${LLAMA_STANDALONE})
 option(LLAMA_BUILD_SERVER    "llama: build server example"                                                       ${LLAMA_STANDALONE})
 option(LLAMA_BUILD_APP       "llama: build the unified binary"                                                   ${LLAMA_STANDALONE})
-option(LLAMA_BUILD_UI        "llama: build the embedded Web UI for server"                                       ON)
-option(LLAMA_USE_PREBUILT_UI "llama: use prebuilt UI from HF Bucket when available (requires LLAMA_BUILD_UI=ON)" ON)
+option(LLAMA_BUILD_UI        "llama: build the embedded Web UI for server"                                       OFF)
+option(LLAMA_USE_PREBUILT_UI "llama: use prebuilt UI from HF Bucket when available" ON)

 option(LLAMA_TOOLS_INSTALL "llama: install tools" ${LLAMA_TOOLS_INSTALL_DEFAULT})
 option(LLAMA_TESTS_INSTALL "llama: install tests" ON)
@@ -202,6 +202,10 @@ class NemotronHModel(GraniteHybridModel):
    is_moe: bool = False
    supports_mtp_export = True

+    _SSM_LAYER_TYPES = {"mamba", "linear_attention"}
+    _ATTN_LAYER_TYPES = {"attention", "full_attention"}
+    _MLP_LAYER_TYPES = {"moe"}
+
    def __init__(self, *args, **kwargs):
        # We have to determine the correct model architecture (MoE vs non-MoE) before
        # calling the parent __init__. This is because the parent constructor
@@ -242,8 +246,8 @@ class NemotronHModel(GraniteHybridModel):
            self._ssm_layers = [i for i, val in enumerate(pattern) if val == "M"]
            self._mlp_layers = [i for i, val in enumerate(pattern) if val == ("E" if self.is_moe else "-")]
        else:
-            self._ssm_layers = [i for i, val in enumerate(pattern) if val == "mamba"]
-            self._mlp_layers = [i for i, val in enumerate(pattern) if val == "moe"]
+            self._ssm_layers = [i for i, val in enumerate(pattern) if val in self._SSM_LAYER_TYPES]
+            self._mlp_layers = [i for i, val in enumerate(pattern) if val in self._MLP_LAYER_TYPES]

        # `--no-mtp` drops it entirely; `--mtp` exports only the MTP head
        self._mtp_bid: int | None = None
@@ -272,7 +276,7 @@ class NemotronHModel(GraniteHybridModel):
        if isinstance(pattern, str):
            return [i for i, val in enumerate(pattern) if val == "*"]

-        return [i for i, val in enumerate(pattern) if val == "attention"]
+        return [i for i, val in enumerate(pattern) if val in self._ATTN_LAYER_TYPES]

    @classmethod
    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
@@ -35,8 +35,8 @@ Legend:
 |                              COS | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                      COUNT_EQUAL | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              CPY | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
-|               CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|          CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|               CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|          CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           CUMSUM | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             DIAG | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    DIAG_MASK_INF | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -19292,10 +19292,10 @@
 "Vulkan0","FLASH_ATTN_EXT","hsk=128,hsv=128,nh=8,nr23=[4,1],kv=4096,nb=512,mask=1,sinks=0,max_bias=0.000000,logit_softcap=0.000000,prec=f32,type_K=f16,type_V=f16,permute=[0,1,2,3]","support","1","yes","Vulkan"
 "Vulkan0","FLASH_ATTN_EXT","hsk=256,hsv=256,nh=4,nr23=[6,1],kv=16384,nb=512,mask=1,sinks=0,max_bias=0.000000,logit_softcap=0.000000,prec=f32,type_K=f16,type_V=f16,permute=[0,1,2,3]","support","1","yes","Vulkan"
 "Vulkan0","FLASH_ATTN_EXT","hsk=128,hsv=128,nh=8,nr23=[4,1],kv=16384,nb=512,mask=1,sinks=0,max_bias=0.000000,logit_softcap=0.000000,prec=f32,type_K=f16,type_V=f16,permute=[0,1,2,3]","support","1","yes","Vulkan"
-"Vulkan0","CROSS_ENTROPY_LOSS","type=f32,ne=[10,5,4,3]","support","0","no","Vulkan"
-"Vulkan0","CROSS_ENTROPY_LOSS","type=f32,ne=[30000,1,1,1]","support","0","no","Vulkan"
-"Vulkan0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[10,5,4,3]","support","0","no","Vulkan"
-"Vulkan0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[30000,1,1,1]","support","0","no","Vulkan"
+"Vulkan0","CROSS_ENTROPY_LOSS","type=f32,ne=[10,5,4,3]","support","1","yes","Vulkan"
+"Vulkan0","CROSS_ENTROPY_LOSS","type=f32,ne=[30000,1,1,1]","support","1","yes","Vulkan"
+"Vulkan0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[10,5,4,3]","support","1","yes","Vulkan"
+"Vulkan0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[30000,1,1,1]","support","1","yes","Vulkan"
 "Vulkan0","OPT_STEP_ADAMW","type=f32,ne=[10,5,4,3]","support","1","yes","Vulkan"
 "Vulkan0","OPT_STEP_SGD","type=f32,ne=[10,5,4,3]","support","1","yes","Vulkan"
 "Vulkan0","GATED_DELTA_NET","type=f32,head_count=32,head_size=128,n_seq_tokens=1,n_seqs=1,v_repeat=1,permuted=0,kda=0,K=1","support","1","yes","Vulkan"
@@ -110,6 +110,16 @@ set_and_check(GGML_INCLUDE_DIR "@PACKAGE_GGML_INCLUDE_INSTALL_DIR@")
 set_and_check(GGML_LIB_DIR "@PACKAGE_GGML_LIB_INSTALL_DIR@")
 #set_and_check(GGML_BIN_DIR "@PACKAGE_GGML_BIN_INSTALL_DIR@")

+if (NOT GGML_SHARED_LIB AND GGML_CPU_KLEIDIAI)
+    unset(KLEIDIAI_LIBRARY CACHE)
+    unset(KLEIDIAI_LIBRARY)
+    find_library(KLEIDIAI_LIBRARY kleidiai
+        REQUIRED
+        HINTS ${GGML_LIB_DIR}
+        NO_CMAKE_FIND_ROOT_PATH)
+    list(APPEND GGML_CPU_INTERFACE_LINK_LIBRARIES ${KLEIDIAI_LIBRARY})
+endif()
+
 if(NOT TARGET ggml::ggml)
    find_package(Threads REQUIRED)

@@ -6,8 +6,8 @@
 extern "C" {
 #endif

-#define RPC_PROTO_MAJOR_VERSION    5
-#define RPC_PROTO_MINOR_VERSION    1
+#define RPC_PROTO_MAJOR_VERSION    6
+#define RPC_PROTO_MINOR_VERSION    0
 #define RPC_PROTO_PATCH_VERSION    0

 #ifdef  __cplusplus
@@ -83,6 +83,7 @@ extern "C" {
    GGML_API ggml_backend_buffer_t ggml_backend_multi_buffer_alloc_buffer(ggml_backend_buffer_t * buffers, size_t n_buffers);
    GGML_API bool                  ggml_backend_buffer_is_multi_buffer(ggml_backend_buffer_t buffer);
    GGML_API void                  ggml_backend_multi_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage);
+    GGML_API void                  ggml_backend_meta_buffer_set_usage (ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage);

    //
    // Backend (meta)
@@ -1168,7 +1168,6 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
 }

 static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync) {
-    GGML_ASSERT(ggml_backend_buffer_is_meta(tensor->buffer));
    ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context;
    return ggml_backend_meta_get_split_state(buf_ctx->get_simple_tensor_container(tensor), tensor, assume_sync);
 }
@@ -1259,7 +1258,14 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m
            t_ij->data = (char *) ggml_backend_buffer_get_base(simple_buf)
                + size_t(tensor->data) - size_t(ggml_backend_buffer_get_base(tensor->buffer));
        }
-        t_ij->extra = tensor->extra;
+
+        if (simple_buf) {
+            // the backend that owns the buffer will set .extra
+            ggml_backend_buffer_init_tensor(simple_buf, t_ij);
+        } else {
+            t_ij->extra = tensor->extra;
+        }
+
        for (int i = 0; i < GGML_MAX_SRC; i++) {
            t_ij->src[i] = tensor->src[i];
            if (tensor->src[i] == tensor) {
@@ -1668,6 +1674,16 @@ bool ggml_backend_buffer_is_meta(ggml_backend_buffer_t buf) {
    return buf != nullptr && buf->iface.free_buffer == ggml_backend_meta_buffer_iface.free_buffer;
 }

+void ggml_backend_meta_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage) {
+    GGML_ASSERT(ggml_backend_buffer_is_meta(buffer));
+    ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context;
+    for (size_t i = 0; i < buf_ctx->bufs.size(); i++) {
+        if (buf_ctx->bufs[i]) {
+            ggml_backend_buffer_set_usage(buf_ctx->bufs[i].get(), usage);
+        }
+    }
+}
+
 static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
    const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft);

@@ -182,6 +182,8 @@ void ggml_backend_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backe
    // FIXME: add a generic callback to the buffer interface
    if (ggml_backend_buffer_is_multi_buffer(buffer)) {
        ggml_backend_multi_buffer_set_usage(buffer, usage);
+    } else if (ggml_backend_buffer_is_meta(buffer)) {
+        ggml_backend_meta_buffer_set_usage(buffer, usage);
    }
 }

@@ -576,10 +576,25 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
    endif()

    if (GGML_CPU_KLEIDIAI)
-        message(STATUS "Using KleidiAI optimized kernels if applicable")
+        # upstream repo requires at least cmake 3.16
+        if (CMAKE_VERSION VERSION_LESS 3.16)
+            message(FATAL_ERROR "GGML_CPU_KLEIDIAI requires CMake >= 3.16")
+        endif()

-        # Disable the KleidiAI tests
-        set(KLEIDIAI_BUILD_TESTS  OFF)
+        set(GGML_CPU_KLEIDIAI_AARCH64 OFF)
+        if (GGML_SYSTEM_ARCH STREQUAL "ARM" AND
+            (APPLE OR WIN32 OR CMAKE_SYSTEM_NAME MATCHES "^(Linux|Android)$") AND
+            (CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64|ARM64|arm64-v8a)$" OR
+             CMAKE_OSX_ARCHITECTURES MATCHES "arm64" OR
+             CMAKE_GENERATOR_PLATFORM_LWR STREQUAL "arm64" OR
+             CMAKE_ANDROID_ARCH_ABI STREQUAL "arm64-v8a"))
+            set(GGML_CPU_KLEIDIAI_AARCH64 ON)
+        endif()
+        if (NOT GGML_CPU_KLEIDIAI_AARCH64)
+            message(FATAL_ERROR "GGML_CPU_KLEIDIAI requires a Linux, Android, Apple, or Windows AArch64/arm64 target")
+        endif()
+
+        message(STATUS "Using KleidiAI optimized kernels if applicable")

        # Fetch KleidiAI sources:
        include(FetchContent)
@@ -595,31 +610,49 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
            list(APPEND KLEIDIAI_FETCH_ARGS DOWNLOAD_EXTRACT_TIMESTAMP NEW)
        endif()

-        if (CMAKE_VERSION VERSION_GREATER_EQUAL "3.28")
-            FetchContent_Declare(KleidiAI_Download
-                ${KLEIDIAI_FETCH_ARGS}
+        FetchContent_Declare(kleidiai
+            ${KLEIDIAI_FETCH_ARGS}
+        )
+
+        # Disable tests and benchmark building
+        set(KLEIDIAI_BUILD_TESTS OFF CACHE BOOL "" FORCE)
+        set(KLEIDIAI_BUILD_BENCHMARK OFF CACHE BOOL "" FORCE)
+
+        # Use the Populate/add_subdirectory flow for compatibility with CMake 3.16.
+        FetchContent_GetProperties(kleidiai
+            SOURCE_DIR KLEIDIAI_SRC
+            BINARY_DIR KLEIDIAI_BIN
+            POPULATED KLEIDIAI_POPULATED
+        )
+        if (NOT KLEIDIAI_POPULATED)
+            FetchContent_Populate(kleidiai)
+            FetchContent_GetProperties(kleidiai
+                SOURCE_DIR KLEIDIAI_SRC
+                BINARY_DIR KLEIDIAI_BIN
+            )
+        endif()
+
+        if (NOT TARGET kleidiai)
+            add_subdirectory(
+                "${CMAKE_CURRENT_SOURCE_DIR}/ggml-cpu/kleidiai"
+                "${CMAKE_CURRENT_BINARY_DIR}/kleidiai-wrapper"
                EXCLUDE_FROM_ALL
            )
-
-            FetchContent_MakeAvailable(KleidiAI_Download)
-            FetchContent_GetProperties(KleidiAI_Download SOURCE_DIR KLEIDIAI_SRC)
-        else()
-            FetchContent_Declare(KleidiAI_Download
-                ${KLEIDIAI_FETCH_ARGS}
-            )
-
-            FetchContent_GetProperties(KleidiAI_Download
-                SOURCE_DIR KLEIDIAI_SRC
-                POPULATED KLEIDIAI_POPULATED
-            )
-
-            if (NOT KLEIDIAI_POPULATED)
-                FetchContent_Populate(KleidiAI_Download)
-                FetchContent_GetProperties(KleidiAI_Download SOURCE_DIR KLEIDIAI_SRC)
+            if (NOT CMAKE_SKIP_INSTALL_RULES AND
+                (NOT DEFINED BUILD_SHARED_LIBS OR NOT BUILD_SHARED_LIBS))
+                install(TARGETS kleidiai ARCHIVE)
            endif()
        endif()

-        add_compile_definitions(GGML_USE_CPU_KLEIDIAI)
+        if (NOT TARGET kleidiai)
+            message(FATAL_ERROR "KleidiAI target was not created")
+        endif()
+
+        set_target_properties(kleidiai PROPERTIES POSITION_INDEPENDENT_CODE ON)
+
+        target_link_libraries(${GGML_CPU_NAME} PRIVATE kleidiai)
+
+        target_compile_definitions(${GGML_CPU_NAME} PRIVATE GGML_USE_CPU_KLEIDIAI)

        list(APPEND GGML_CPU_SOURCES
            ggml-cpu/kleidiai/kleidiai.cpp
@@ -627,108 +660,6 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
            ggml-cpu/kleidiai/kleidiai.h
            ggml-cpu/kleidiai/kernels.h
            )
-
-        # KleidiAI
-        include_directories(
-            ${KLEIDIAI_SRC}/
-            ${KLEIDIAI_SRC}/kai/
-            ${KLEIDIAI_SRC}/kai/ukernels/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_fp32_bf16p_bf16p/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/)
-
-        set(ARCH_FLAGS_TEMP "${ARCH_FLAGS}")
-        if (NOT ARCH_FLAGS_TEMP)
-            string(REGEX MATCH "-march=[^ ]+" ARCH_FLAGS_TEMP "${CMAKE_C_FLAGS}")
-        endif()
-        string(FIND "${ARCH_FLAGS_TEMP}" "+dotprod" DOTPROD_ENABLED)
-        string(FIND "${ARCH_FLAGS_TEMP}" "+i8mm" I8MM_ENABLED)
-        string(FIND "${ARCH_FLAGS_TEMP}" "+sme" SME_ENABLED)
-        string(FIND "${ARCH_FLAGS_TEMP}" "+sve" SVE_ENABLED)
-
-        set(PRIVATE_ARCH_FLAGS ${ARCH_FLAGS_TEMP})
-
-        list(APPEND GGML_KLEIDIAI_SOURCES
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32p_f32.c
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32p4x8sb_f32_neon.c
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon.c
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32p_f32_neon.c
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0.c
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_quant_pack_qai8dxp_f32.c
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi8cxp_qsi8cx_neon.c)
-
-        if (NOT DOTPROD_ENABLED MATCHES -1)
-            list(APPEND GGML_KLEIDIAI_SOURCES
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p4x8_1x4x32_neon_dotprod.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4x4_1x4_neon_dotprod.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x4_qsi4c32p4x4_16x4_neon_dotprod.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp4x4_qsi8cxp4x4_16x4_neon_dotprod.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4x4_1x4_neon_dotprod.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x8_qsi8cxp4x8_1x4_neon_dotprod.c)
-        endif()
-
-        if (NOT I8MM_ENABLED MATCHES -1)
-            list(APPEND GGML_KLEIDIAI_SOURCES
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p4x8_16x4_neon_i8mm.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp4x8_qsi8cxp4x8_16x4_neon_i8mm.c)
-        endif()
-
-        if (NOT SME_ENABLED MATCHES -1)
-            list(APPEND GGML_KLEIDIAI_SME_SOURCES
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa_asm.S)
-            set_source_files_properties(${GGML_KLEIDIAI_SME_SOURCES}
-                PROPERTIES COMPILE_OPTIONS "-fno-tree-vectorize;${ARCH_FLAGS_TEMP}+sve+sve2+sme")
-            list(APPEND GGML_CPU_SOURCES ${GGML_KLEIDIAI_SME_SOURCES})
-
-            list(APPEND GGML_KLEIDIAI_SME2_SOURCES
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4vlx4_1x4vl_sme2_sdot.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_fp32_bf16p_bf16p/kai_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_fp32_bf16p_bf16p/kai_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_bf16p2vlx2_f32_sme.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_f16pmrx2_f32_neon.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_f32p2vlx1_f32_sme.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_f32p2vlx1_f32_sme_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme_asm.S
-                ${KLEIDIAI_SRC}/kai/kai_common_sme_asm.S)
-            set_source_files_properties(${GGML_KLEIDIAI_SME2_SOURCES}
-                PROPERTIES COMPILE_OPTIONS "-fno-tree-vectorize;${ARCH_FLAGS_TEMP}+sve+sve2+sme2+fp16")
-            list(APPEND GGML_CPU_SOURCES ${GGML_KLEIDIAI_SME2_SOURCES})
-            set(PRIVATE_ARCH_FLAGS "-fno-tree-vectorize;${PRIVATE_ARCH_FLAGS}")
-        endif()
-
-        if (NOT SVE_ENABLED MATCHES -1)
-            list(APPEND GGML_KLEIDIAI_SOURCES
-                ${KLEIDIAI_SRC}/kai/kai_common_sve_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p8x8_1x8_sve_dotprod_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p8x8_1x8_sve_dotprod.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p8x8_16x8_sve_i8mm_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p8x8_16x8_sve_i8mm.c)
-        endif()
-
-        set_source_files_properties(${GGML_KLEIDIAI_SOURCES} PROPERTIES COMPILE_OPTIONS "${PRIVATE_ARCH_FLAGS}")
-        list(APPEND GGML_CPU_SOURCES ${GGML_KLEIDIAI_SOURCES})
    endif()

    message(STATUS "Adding CPU backend variant ${GGML_CPU_NAME}: ${ARCH_FLAGS} ${ARCH_DEFINITIONS}")
@@ -0,0 +1,14 @@
+set(BUILD_SHARED_LIBS OFF)
+set(CMAKE_SKIP_INSTALL_RULES TRUE)
+
+add_subdirectory("${KLEIDIAI_SRC}" "${KLEIDIAI_BIN}" EXCLUDE_FROM_ALL)
+
+if (NOT TARGET kleidiai)
+    message(FATAL_ERROR "KleidiAI target was not created")
+endif()
+
+if (MSVC)
+    target_compile_options(kleidiai PRIVATE $<$<COMPILE_LANGUAGE:C,CXX>:/WX->)
+else()
+    target_compile_options(kleidiai PRIVATE $<$<COMPILE_LANGUAGE:C,CXX>:-Wno-error>)
+endif()
@@ -3,44 +3,44 @@
 //

 // KleidiAI micro-kernels
-#include "kai_matmul_clamp_f32_qsi8d32p_qsi4c32p_interface.h"
-#include "kai_matmul_clamp_f32_qai8dxp_qsi8cxp_interface.h"
-#include "kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p4x8_1x4x32_neon_dotprod.h"
-#include "kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4x4_1x4_neon_dotprod.h"
-#include "kai_matmul_clamp_f32_qsi8d32p4x4_qsi4c32p4x4_16x4_neon_dotprod.h"
-#include "kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p4x8_16x4_neon_i8mm.h"
-#include "kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4vlx4_1x4vl_sme2_sdot.h"
-#include "kai_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa.h"
-#include "kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa.h"
-#include "kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot.h"
-#include "kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa.h"
-#include "kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot.h"
-#include "kai_matmul_clamp_f32_qai8dxp1x8_qsi8cxp4x8_1x4_neon_dotprod.h"
-#include "kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4x4_1x4_neon_dotprod.h"
-#include "kai_matmul_clamp_f32_qai8dxp4x4_qsi8cxp4x4_16x4_neon_dotprod.h"
-#include "kai_matmul_clamp_f32_qai8dxp4x8_qsi8cxp4x8_16x4_neon_i8mm.h"
-#include "kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p8x8_16x8_sve_i8mm.h"
-#include "kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p8x8_1x8_sve_dotprod.h"
-#include "kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa.h"
-#include "kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.h"
-#include "kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla.h"
-#include "kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p_qsi4c32p_interface.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp_qsi8cxp_interface.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p4x8_1x4x32_neon_dotprod.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4x4_1x4_neon_dotprod.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x4_qsi4c32p4x4_16x4_neon_dotprod.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p4x8_16x4_neon_i8mm.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4vlx4_1x4vl_sme2_sdot.h"
+#include "kai/ukernels/matmul/matmul_clamp_fp32_bf16p_bf16p/kai_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x8_qsi8cxp4x8_1x4_neon_dotprod.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4x4_1x4_neon_dotprod.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp4x4_qsi8cxp4x4_16x4_neon_dotprod.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp4x8_qsi8cxp4x8_16x4_neon_i8mm.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p8x8_16x8_sve_i8mm.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p8x8_1x8_sve_dotprod.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa.h"

-#include "kai_lhs_pack_bf16p2vlx2_f32_sme.h"
-#include "kai_lhs_pack_f32p2vlx1_f32_sme.h"
-#include "kai_lhs_quant_pack_qsi8d32p_f32.h"
-#include "kai_lhs_quant_pack_qsi8d32p4x8sb_f32_neon.h"
-#include "kai_lhs_quant_pack_qsi8d32p_f32_neon.h"
-#include "kai_lhs_quant_pack_qai8dxp_f32.h"
+#include "kai/ukernels/matmul/pack/kai_lhs_pack_bf16p2vlx2_f32_sme.h"
+#include "kai/ukernels/matmul/pack/kai_lhs_pack_f32p2vlx1_f32_sme.h"
+#include "kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32p_f32.h"
+#include "kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32p4x8sb_f32_neon.h"
+#include "kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32p_f32_neon.h"
+#include "kai/ukernels/matmul/pack/kai_lhs_quant_pack_qai8dxp_f32.h"

-#include "kai_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme.h"
-#include "kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.h"
-#include "kai_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0.h"
-#include "kai_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon.h"
-#include "kai_rhs_pack_nxk_qsi8cxp_qsi8cx_neon.h"
-#include "kai_lhs_pack_f16pmrx2_f32_neon.h"
+#include "kai/ukernels/matmul/pack/kai_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme.h"
+#include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.h"
+#include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0.h"
+#include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon.h"
+#include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi8cxp_qsi8cx_neon.h"
+#include "kai/ukernels/matmul/pack/kai_lhs_pack_f16pmrx2_f32_neon.h"

-#include "kai_common.h"
+#include "kai/kai_common.h"

 #include "simd-mappings.h"

@@ -328,9 +328,8 @@ static void dequantize_row_qsi8cxp(
 }

 static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
-#if defined(__ARM_FEATURE_SME)
    {
-        /* SME GEMM */
+        /* SME2 GEMM */
        /* .kern_info = */ {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa,
@@ -351,7 +350,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_size_ex        = */ &lhs_ps_fn6<kai_get_lhs_packed_size_lhs_pack_f16pmrx2_f32_neon>,
            /* .pack_func_ex          = */ &lhs_pack_void_fn10<kai_run_lhs_pack_f16pmrx2_f32_neon>,
        },
-        /* SME GEMV */
+        /* SME2 GEMV */
        /* .kern_info = */ {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4vlx4_1x4vl_sme2_sdot,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4vlx4_1x4vl_sme2_sdot,
@@ -378,13 +377,13 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_stride_ex      = */ &rhs_stride_fn4<kai_get_rhs_packed_stride_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon>,
            /* .pack_func_ex          = */ &rhs_pack_fn12<kai_run_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon>,
        },
-        /* .required_cpu       = */ CPU_FEATURE_SME2,
+        /* .required_cpu       = */ CPU_FEATURE_SME2 | CPU_FEATURE_FP16,
        /* .lhs_type           = */ GGML_TYPE_F32,
        /* .rhs_type           = */ GGML_TYPE_Q4_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
    {
-        /* SME GEMM */
+        /* SME2 GEMM */
        /* .kern_info = */ {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa,
@@ -404,7 +403,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_pack_bf16p2vlx2_f32_sme>,
            /* .pack_func_ex          = */ &lhs_pack_void_fn9<kai_run_lhs_pack_bf16p2vlx2_f32_sme>,
        },
-        /* SME GEMV */
+        /* SME2 GEMV */
        /* .kern_info = */ {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa,
@@ -436,9 +435,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
        /* .rhs_type           = */ GGML_TYPE_F16,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
 #if defined(__APPLE__)
-#if defined(__ARM_FEATURE_DOTPROD)
    {
        /* DOTPROD GEMM */
        /* .kern_info = */ {
@@ -492,8 +489,6 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
        /* .rhs_type           = */ GGML_TYPE_Q4_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
-#if defined(__ARM_FEATURE_MATMUL_INT8)
    {
        /* i8mm GEMM */
        /* .kern_info = */ {
@@ -515,7 +510,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_size_ex        = */ &lhs_ps_fn6<kai_get_lhs_packed_size_lhs_quant_pack_qsi8d32p4x8sb_f32_neon>,
            /* .pack_func_ex          = */ &lhs_pack_float_fn10<kai_run_lhs_quant_pack_qsi8d32p4x8sb_f32_neon>,
        },
-        /* i8mm GEMV */
+        /* DOTPROD GEMV */
        /* .kern_info = */ {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p4x8_1x4x32_neon_dotprod,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p4x8_1x4x32_neon_dotprod,
@@ -542,14 +537,12 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_stride_ex      = */ &rhs_stride_fn4<kai_get_rhs_packed_stride_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
            /* .pack_func_ex          = */ &rhs_pack_fn12<kai_run_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
        },
-        /* .required_cpu       = */ CPU_FEATURE_I8MM,
+        /* .required_cpu       = */ CPU_FEATURE_I8MM | CPU_FEATURE_DOTPROD,
        /* .lhs_type           = */ GGML_TYPE_F32,
        /* .rhs_type           = */ GGML_TYPE_Q4_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
 #else
-#if defined(__ARM_FEATURE_SVE)
    {
        /* SVE i8mm GEMM */
        /* .kern_info = */ {
@@ -603,8 +596,6 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
        /* .rhs_type           = */ GGML_TYPE_Q4_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
-#if defined(__ARM_FEATURE_MATMUL_INT8)
    {
        /* i8mm GEMM */
        /* .kern_info = */ {
@@ -626,7 +617,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_size_ex        = */ &lhs_ps_fn6<kai_get_lhs_packed_size_lhs_quant_pack_qsi8d32p4x8sb_f32_neon>,
            /* .pack_func_ex          = */ &lhs_pack_float_fn10<kai_run_lhs_quant_pack_qsi8d32p4x8sb_f32_neon>,
        },
-        /* i8mm GEMV */
+        /* DOTPROD GEMV */
        /* .kern_info = */ {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p4x8_1x4x32_neon_dotprod,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p4x8_1x4x32_neon_dotprod,
@@ -653,13 +644,11 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_stride_ex      = */ &rhs_stride_fn4<kai_get_rhs_packed_stride_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
            /* .pack_func_ex          = */ &rhs_pack_fn12<kai_run_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
        },
-        /* .required_cpu       = */ CPU_FEATURE_I8MM,
+        /* .required_cpu       = */ CPU_FEATURE_I8MM | CPU_FEATURE_DOTPROD,
        /* .lhs_type           = */ GGML_TYPE_F32,
        /* .rhs_type           = */ GGML_TYPE_Q4_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif // __ARM_FEATURE_MATMUL_INT8
-#if defined(__ARM_FEATURE_DOTPROD)
    {
        /* DOTPROD GEMM */
        /* .kern_info = */ {
@@ -713,15 +702,13 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
        /* .rhs_type           = */ GGML_TYPE_Q4_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
 #endif
    { /* Sentinel */ }
 };

 static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
-#if defined(__ARM_FEATURE_SME)
    {
-        /* SME GEMM */
+        /* SME2 GEMM */
        {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa,
@@ -741,7 +728,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32>,
            /* .pack_func_ex          = */ &lhs_pack_float_fn9_no_bl<kai_run_lhs_quant_pack_qai8dxp_f32>,
        },
-        /* SME GEMV */
+        /* SME2 GEMV */
        {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot,
@@ -826,8 +813,6 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
        /* .rhs_type           = */ GGML_TYPE_Q8_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
-#if defined(__ARM_FEATURE_MATMUL_INT8)
    {
        /* I8MM GEMM */
        {
@@ -876,13 +861,11 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
            /* .packed_stride_ex      = */ &rhs_stride_fn4<kai_get_rhs_packed_stride_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
            /* .pack_func_ex          = */ &rhs_pack_scale_fn12<kai_run_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
        },
-        /* .required_cpu       = */ CPU_FEATURE_I8MM,
+        /* .required_cpu       = */ CPU_FEATURE_I8MM | CPU_FEATURE_DOTPROD,
        /* .lhs_type           = */ GGML_TYPE_F32,
        /* .rhs_type           = */ GGML_TYPE_Q8_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
-#if defined(__ARM_FEATURE_DOTPROD)
    {
        /* DOTPROD GEMM */
        {
@@ -936,12 +919,10 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
        /* .rhs_type           = */ GGML_TYPE_Q8_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
    { /* Sentinel */ }
 };

 static ggml_kleidiai_kernels ggml_kleidiai_kernels_f32[] = {
-#if defined(__ARM_FEATURE_SME)
    {
        /* SME2 GEMM */
        {
@@ -1048,7 +1029,6 @@ static ggml_kleidiai_kernels ggml_kleidiai_kernels_f32[] = {
        /* .rhs_type           = */ GGML_TYPE_F32,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
    { /* Sentinel */ }
 };

@@ -1056,10 +1036,6 @@ ggml_kleidiai_kernels * ggml_kleidiai_select_kernels(cpu_feature cpu_features, c
    ggml_kleidiai_kernels * kernel = nullptr;

    if (tensor->op == GGML_OP_MUL_MAT && tensor->src[0] != nullptr && tensor->src[1] != nullptr) {
-#if defined(__ARM_FEATURE_SME)          ||  \
-    defined(__ARM_FEATURE_DOTPROD)      ||  \
-    defined(__ARM_FEATURE_MATMUL_INT8)  ||  \
-    defined(__ARM_FEATURE_SVE)
        auto try_table = [&](auto & table) {
            for (size_t i = 0; i < NELEMS(table) - 1; ++i) {
                if ((cpu_features & table[i].required_cpu) == table[i].required_cpu &&
@@ -1080,12 +1056,6 @@ ggml_kleidiai_kernels * ggml_kleidiai_select_kernels(cpu_feature cpu_features, c
        } else {
            try_table(gemm_gemv_kernels);
        }
-#else
-    GGML_UNUSED(gemm_gemv_kernels);
-    GGML_UNUSED(gemm_gemv_kernels_q8);
-    GGML_UNUSED(ggml_kleidiai_kernels_f32);
-    GGML_UNUSED(cpu_features);
-#endif
    }

    return kernel;
@@ -1094,19 +1064,13 @@ ggml_kleidiai_kernels * ggml_kleidiai_select_kernels(cpu_feature cpu_features, c
 ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q4_0(cpu_feature features) {
    ggml_kleidiai_kernels * kernels = nullptr;

-#if defined(__ARM_FEATURE_SME)          ||  \
-    defined(__ARM_FEATURE_DOTPROD)      ||  \
-    defined(__ARM_FEATURE_MATMUL_INT8)  ||  \
-    defined(__ARM_FEATURE_SVE)
    for (size_t i = 0; i < NELEMS(gemm_gemv_kernels) - 1; ++i) {
-        if ((features & gemm_gemv_kernels[i].required_cpu) == gemm_gemv_kernels[i].required_cpu) {
+        if ((features & gemm_gemv_kernels[i].required_cpu) == gemm_gemv_kernels[i].required_cpu &&
+            gemm_gemv_kernels[i].rhs_type == GGML_TYPE_Q4_0) {
            kernels = &gemm_gemv_kernels[i];
            break;
        }
    }
-#else
-    GGML_UNUSED(features);
-#endif

    return kernels;
 }
@@ -1114,16 +1078,12 @@ ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q4_0(cpu_feature features)
 ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q8_0(cpu_feature features) {
    ggml_kleidiai_kernels * kernels = nullptr;

-#if defined(__ARM_FEATURE_SME) || defined(__ARM_FEATURE_DOTPROD) || defined(__ARM_FEATURE_MATMUL_INT8)
    for (size_t i = 0; i < NELEMS(gemm_gemv_kernels_q8) - 1; ++i) {
        if ((features & gemm_gemv_kernels_q8[i].required_cpu) == gemm_gemv_kernels_q8[i].required_cpu) {
            kernels = &gemm_gemv_kernels_q8[i];
            break;
        }
    }
-#else
-    GGML_UNUSED(features);
-#endif

    return kernels;
 }
@@ -1131,16 +1091,11 @@ ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q8_0(cpu_feature features)
 ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_f32(cpu_feature features) {
    ggml_kleidiai_kernels * kernels = nullptr;

-#if defined(__ARM_FEATURE_SME)
    for (size_t i = 0; i < NELEMS(ggml_kleidiai_kernels_f32) - 1; ++i) {
        if ((features & ggml_kleidiai_kernels_f32[i].required_cpu) == ggml_kleidiai_kernels_f32[i].required_cpu) {
            kernels = &ggml_kleidiai_kernels_f32[i];
            break;
        }
    }
-#else
-    GGML_UNUSED(features);
-#endif
-
    return kernels;
 }
@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: Copyright 2025 Arm Limited and/or its affiliates <open-source-office@arm.com>
+// SPDX-FileCopyrightText: Copyright 2025-2026 Arm Limited and/or its affiliates <open-source-office@arm.com>
 // SPDX-License-Identifier: MIT
 //

@@ -12,7 +12,8 @@ enum cpu_feature {
    CPU_FEATURE_I8MM    = 2,
    CPU_FEATURE_SVE     = 4,
    CPU_FEATURE_SME     = 8,
-    CPU_FEATURE_SME2    = 16
+    CPU_FEATURE_SME2    = 16,
+    CPU_FEATURE_FP16    = 32
 };

 inline cpu_feature& operator|=(cpu_feature& lhs, cpu_feature rhs) {
@@ -48,7 +48,7 @@

 #include "kernels.h"

-#include "kai_common.h"
+#include "kai/kai_common.h"

 #define GGML_COMMON_DECL_CPP
 #include "ggml-common.h"
@@ -316,6 +316,7 @@ static void init_kleidiai_context(void) {

        ctx.features  = (runtime_feat.has_dotprod  ? CPU_FEATURE_DOTPROD : CPU_FEATURE_NONE) |
                        (runtime_feat.has_i8mm     ? CPU_FEATURE_I8MM    : CPU_FEATURE_NONE) |
+                        (runtime_feat.has_fp16     ? CPU_FEATURE_FP16    : CPU_FEATURE_NONE) |
                        (runtime_feat.sve_cnt == QK8_0 ? CPU_FEATURE_SVE : CPU_FEATURE_NONE);

        if (env_threads) {
@@ -1,4 +1,4 @@
-static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal(ggml_type type, int J, bool fallback) {
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal_dp4a(ggml_type type, int J, bool fallback) {
    CASE(GGML_TYPE_Q1_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
@@ -0,0 +1,273 @@
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal_older(ggml_type type, int J, bool fallback) {
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
@@ -314,7 +314,9 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t
    }

    if (ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_DP4A) {
-        return false;
+        // for MoE, mmq is faster even without native dp4a
+        // TODO: check if cards older than pascal might benefit from this as well
+        return cc >= GGML_CUDA_CC_PASCAL && n_experts > 0;
    }

 #ifdef GGML_CUDA_FORCE_MMQ
@@ -213,7 +213,8 @@ struct ggml_cuda_mmq_config {
        return ggml_cuda_mmq_config((type_), (nthreads_), (occupancy_), (I_), (J_), (sram_layout_), (K_vram_), (stream_k_), (fallback_)); \
    }                                                                                                                                     \

-#include "mmq-config-pascal.cuh"
+#include "mmq-config-pascal-older.cuh"
+#include "mmq-config-pascal-dp4a.cuh"
 #include "mmq-config-ampere.cuh"
 #include "mmq-config-blackwell.cuh"

@@ -247,7 +248,10 @@ static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type ty
    if (ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) {
        return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
    }
-    return ggml_cuda_mmq_get_config_pascal(type, J, fallback);
+    if (ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_DP4A) {
+        return ggml_cuda_mmq_get_config_pascal_dp4a(type, J, fallback);
+    }
+    return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback);
 }

 static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_type type, int J, bool fallback) {
@@ -268,8 +272,10 @@ static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_t
    return ggml_cuda_mmq_get_config_blackwell(type, J, fallback);
 #elif __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
    return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
+#elif __CUDA_ARCH__ >= GGML_CUDA_CC_DP4A
+    return ggml_cuda_mmq_get_config_pascal_dp4a(type, J, fallback);
 #else
-    return ggml_cuda_mmq_get_config_pascal(type, J, fallback);
+    return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback);
 #endif // BLACKWELL_MMA_AVAILABLE
 #endif // GGML_USE_HIP
    GGML_UNUSED_VARS(type, J, fallback);
@@ -572,7 +572,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched
    return res;
 }

-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_metal_library_t lib, const ggml_tensor * op)  {
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_metal_library_t lib, const ggml_tensor * op, bool tail)  {
    GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);

    char base[256];
@@ -580,7 +580,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_me

    const int nsg = (ne00 + 31)/32;

-    snprintf(base, 256, "kernel_ssm_scan_%s", ggml_type_name(op->src[0]->type));
+    snprintf(base, 256, "kernel_ssm_scan_%s%s", ggml_type_name(op->src[0]->type), tail ? "_tail" : "");
    snprintf(name, 256, "%s_nsg=%d", base, nsg);

    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
@@ -598,6 +598,27 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_me
    return res;
 }

+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan_ssd_mma(ggml_metal_library_t lib, const ggml_tensor * op)  {
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_ssm_scan_ssd_mma_%s", ggml_type_name(op->src[0]->type));
+    snprintf(name, 256, "%s", base);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+    }
+
+    // acs/exp(acs)/state-decay vectors + dtX + SAM rows + two 8x8 tiles per simdgroup
+    res.smem = (3*OP_SSM_SCAN_SSD_CS +
+                OP_SSM_SCAN_SSD_CS*OP_SSM_SCAN_SSD_HD +
+                OP_SSM_SCAN_SSD_NSG*8*OP_SSM_SCAN_SSD_CS +
+                OP_SSM_SCAN_SSD_NSG*2*8*8)*sizeof(float);
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv(ggml_metal_library_t lib, const ggml_tensor * op) {
    char base[256];
    char name[256];
@@ -129,7 +129,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc           (ggml_metal_library_t lib, enum ggml_op op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv          (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched  (ggml_metal_library_t lib, const struct ggml_tensor * op, int ssm_conv_bs);
-struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan          (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan          (ggml_metal_library_t lib, const struct ggml_tensor * op, bool tail);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan_ssd_mma  (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv              (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net   (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri         (ggml_metal_library_t lib, const struct ggml_tensor * op);
@@ -158,6 +158,10 @@
 #define OP_SUM_ROWS_NUM_SUM_ROWS 10
 #define OP_SUM_ROWS_NUM_MEAN     11

+#define OP_SSM_SCAN_SSD_CS  64 // Metal-specific; Chunk Size; 64 is largest multiple of 8 (simdgroup tile) fitting into 32 KiB Metal threadgroup mem limit (~26.75 KiB shared mem; see smem layout comment in kernel_ssm_scan_ssd_mma_f32)
+#define OP_SSM_SCAN_SSD_HD  64 // Metal-specific; Head Dim the MMA kernel is specialized for (Mamba-2); use_mma gates on d_inner == this
+#define OP_SSM_SCAN_SSD_NSG 4  // Metal-specific; Number of SimdGroups per threadgroup; NSG*32 == threads dispatched per threadgroup
+
 // kernel argument structs
 //
 // - element counters (e.g. ne00) typically use int32_t to reduce register usage
@@ -893,6 +897,8 @@ typedef struct {
    int64_t  n_head;
    int64_t  n_group;
    int64_t  n_seq_tokens;
+    int64_t  n_seq_tokens_total;
+    int64_t  token_offset;
    int64_t  n_seqs;
    int64_t  K;
    uint64_t s_off;
@@ -1677,6 +1677,7 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) {

    ggml_metal_library_t lib = ctx->lib;
    ggml_metal_encoder_t enc = ctx->enc;
+    const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev);

    GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
    GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
@@ -1722,6 +1723,8 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) {
        /*.n_head       =*/ n_head,
        /*.n_group      =*/ n_group,
        /*.n_seq_tokens =*/ n_seq_tokens,
+        /*.n_seq_tokens_total =*/ n_seq_tokens,
+        /*.token_offset =*/ 0,
        /*.n_seqs       =*/ n_seqs,
        /*.K            =*/ K,
        /*.s_off        =*/ ggml_nelements(op->src[1]) * sizeof(float),
@@ -1751,26 +1754,53 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) {
        /*.nb0          =*/ nb0,
    };

-    auto pipeline = ggml_metal_library_get_pipeline_ssm_scan(lib, op);
+    constexpr int64_t CHUNK = OP_SSM_SCAN_SSD_CS;

-    GGML_ASSERT(d_state <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
+    const int64_t snap_reserve = K > 1 ? K : 0; // tokens reserved for sequential kernel rollback snapshots
+    const int64_t mma_tokens = ((n_seq_tokens - snap_reserve) / CHUNK) * CHUNK; // largest multiple of CHUNK that leaves snap_reserve for the tail
+    const bool use_mma =
+        mma_tokens > 0 &&
+        ne30 == 1 && // checks that A tensor is set to scalar decay per head (A shape {1, n_head})
+        props_dev->has_simdgroup_mm && // hardware check for M1 or newer
+        d_state % 8 == 0 && // d_state must be multiple of 8 to align with simdgroup_float 8x8 tiles
+        d_inner == OP_SSM_SCAN_SSD_HD; // mma kernel is specialized for the Mamba-2 head dim; this checks it

-    const size_t smem = pipeline.smem;
+    const auto dispatch = [&](ggml_metal_pipeline_with_params pipeline, int64_t nth, int64_t n_tg_x) {
+        GGML_ASSERT(nth <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
+        GGML_ASSERT(pipeline.smem <= props_dev->max_theadgroup_memory_size);

-    ggml_metal_encoder_set_pipeline(enc, pipeline);
-    ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[1]), 2);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[2]), 3);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[3]), 4);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[4]), 5);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[5]), 6);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[6]), 7);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         8);
+        ggml_metal_encoder_set_pipeline(enc, pipeline);
+        ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[1]), 2);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[2]), 3);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[3]), 4);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[4]), 5);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[5]), 6);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[6]), 7);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         8);
+        ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline.smem, 0);
+        ggml_metal_encoder_dispatch_threadgroups(enc, n_tg_x, n_head, n_seqs, nth, 1, 1);
+    };

-    ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
+    if (!use_mma) {
+        dispatch(ggml_metal_library_get_pipeline_ssm_scan(lib, op, false), d_state, d_inner);
+        return 1;
+    }

-    ggml_metal_encoder_dispatch_threadgroups(enc, d_inner, n_head, n_seqs, d_state, 1, 1);
+    args.n_seq_tokens = mma_tokens;
+    dispatch(
+        ggml_metal_library_get_pipeline_ssm_scan_ssd_mma(lib, op),
+        OP_SSM_SCAN_SSD_NSG*32,
+        1);
+
+    if (mma_tokens < n_seq_tokens) {
+        ggml_metal_op_concurrency_reset(ctx);
+
+        args.n_seq_tokens = n_seq_tokens - mma_tokens;
+        args.token_offset = mma_tokens;
+        dispatch(ggml_metal_library_get_pipeline_ssm_scan(lib, op, true), d_state, d_inner);
+    }

    return 1;
 }
@@ -159,7 +159,9 @@ kernel void kernel_ssm_conv_f32_f32_batched_4(

 // 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(
+// TAIL == false is the whole-sequence / decode path: token_offset folds away at compile time.
+template<bool TAIL>
+kernel void kernel_ssm_scan_impl(
        constant ggml_metal_kargs_ssm_scan & args,
        device const void * src0,
        device const void * src1,
@@ -200,13 +202,17 @@ kernel void kernel_ssm_scan_f32(
    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 n_t_total = TAIL ? args.n_seq_tokens_total : n_t;
+    const int32_t t_off     = TAIL ? args.token_offset       : 0;

    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);
+    device const float * s0_buff = t_off != 0 ?
+        s_buff :
+        (device const float *) ((device const char *) src0 + ir*args.nb02 + ids[i3]*args.nb03);

    const int32_t i = i0 + i1*nc;
    const int32_t g = ir / (nh / ng); // repeat_interleave
@@ -218,12 +224,12 @@ kernel void kernel_ssm_scan_f32(

    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 const float * x  = (device const float *)((device const char *) src1 + i1*args.nb10 + ir*args.nb11 + t_off*args.nb12 + i3*args.nb13); // {dim, nh, nt, ns}
+    device const float * dt = (device const float *)((device const char *) src2 + ir*args.nb20 + t_off*args.nb21 + i3*args.nb22);                 // {nh, nt, ns}
+    device const float * B  = (device const float *)((device const char *) src4 + g*args.nb41 + t_off*args.nb42 + i3*args.nb43);                  // {d_state, ng, nt, ns}
+    device const float * C  = (device const float *)((device const char *) src5 + g*args.nb51 + t_off*args.nb52 + i3*args.nb53);                  // {d_state, ng, nt, ns}

-    device float * y = dst + (i1 + ir*(nr) + i3*(n_t*nh*nr)); // {dim, nh, nt, ns}
+    device float * y = dst + (i1 + ir*nr + t_off*nh*nr + i3*(n_t_total*nh*nr)); // {dim, nh, nt, ns}

    for (int i2 = 0; i2 < n_t; i2 += sgptg) {
        threadgroup_barrier(mem_flags::mem_threadgroup);
@@ -285,3 +291,183 @@ kernel void kernel_ssm_scan_f32(

    s_buff[i] = s;
 }
+
+typedef decltype(kernel_ssm_scan_impl<false>) kernel_ssm_scan_t;
+
+template [[host_name("kernel_ssm_scan_f32")]]      kernel kernel_ssm_scan_t kernel_ssm_scan_impl<false>;
+template [[host_name("kernel_ssm_scan_f32_tail")]] kernel kernel_ssm_scan_t kernel_ssm_scan_impl<true>;
+
+// Chunked SSD SSM scan via Metal simdgroup MMatrix Multiply-Accumulate (simdgroup_float8x8) fast path.
+// One threadgroup per (head, sequence) and tokens are processed in chunks.
+// C*B^T computed in each chunk one time and reused across the head_dim channel tiles.
+kernel void kernel_ssm_scan_ssd_mma_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]],
+        ushort  tiitg[[thread_index_in_threadgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]]) {
+    constexpr short CS  = OP_SSM_SCAN_SSD_CS;
+    constexpr short TC  = 8; // Tile Count of each edge in a simdgroup 8x8 tile
+    constexpr short HD  = OP_SSM_SCAN_SSD_HD;
+    constexpr short NSG = OP_SSM_SCAN_SSD_NSG;
+
+    // acs/exp(acs)/state-decay vectors, dtX[CS][HD], four private SAM row tiles [8][CS],
+    // and two 8x8 scratch tiles per simdgroup. Total: 26.75 KiB.
+    threadgroup float * shared_acs         = shared;
+    threadgroup float * shared_exp_acs     = shared + CS;
+    threadgroup float * shared_state_decay = shared + 2*CS;
+    threadgroup float * shared_dtx         = shared + 3*CS;
+    threadgroup float * shared_sam         = shared + 3*CS + CS*HD;
+    threadgroup float * sam_rows    = shared_sam + sgitg*TC*CS;
+    threadgroup float * shared_tile = shared_sam + NSG*TC*CS;
+    threadgroup float * tile0       = shared_tile + sgitg*2*TC*TC;
+    threadgroup float * tile1       = tile0 + TC*TC;
+
+    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_t_total = args.n_seq_tokens_total;
+    const int32_t g   = ir / (nh / ng);
+
+    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 + args.s_off);
+
+    device const float * A  = (device const float *) ((device const char *) src3 + ir*args.nb31);
+    device const float * x  = (device const float *) ((device const char *) src1 + ir*args.nb11 + i3*args.nb13);
+    device const float * dt = (device const float *) ((device const char *) src2 + ir*args.nb20 + i3*args.nb22);
+    device const float * B  = (device const float *) ((device const char *) src4 +  g*args.nb41 + i3*args.nb43);
+    device const float * C  = (device const float *) ((device const char *) src5 +  g*args.nb51 + i3*args.nb53);
+
+    device float * y = dst + (ir*nr + i3*(n_t_total*nh*nr));
+
+    for (int32_t t0 = 0; t0 < n_t; t0 += CS) {
+        for (int32_t idx = tiitg; idx < CS*HD; idx += NSG*N_SIMDWIDTH) {
+            const int32_t t = idx / HD;
+            const int32_t c = idx % HD;
+            const float dt0  = dt[(t0 + t) * (int32_t) args.ns21];
+            const float dtsp = dt0 <= 20.0f ? log(1.0f + exp(dt0)) : dt0;
+            shared_dtx[idx] = x[(t0 + t) * (int32_t) args.ns12 + c] * dtsp;
+        }
+        if (tiitg < CS) {
+            const float dt0  = dt[(t0 + tiitg) * (int32_t) args.ns21];
+            const float dtsp = dt0 <= 20.0f ? log(1.0f + exp(dt0)) : dt0;
+            shared_acs[tiitg] = dtsp * A[0];
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiitg == 0) {
+            float acc = 0.0f;
+            for (short t = 0; t < CS; ++t) {
+                acc += shared_acs[t];
+                shared_acs[t] = acc;
+            }
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (tiitg < CS) {
+            shared_exp_acs[tiitg] = exp(shared_acs[tiitg]);
+            shared_state_decay[tiitg] = exp(shared_acs[CS - 1] - shared_acs[tiitg]);
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        device const float * state = t0 == 0 ? s0_buff : s_buff;
+
+        // Build one 8x64 row tile of SAM per simdgroup, then reuse it across every channel tile.
+        for (short ib = sgitg; ib < CS/TC; ib += NSG) {
+            for (short jb = 0; jb <= ib; ++jb) {
+                simdgroup_float8x8 cb = make_filled_simdgroup_matrix<float, 8>(0.0f);
+
+                for (int32_t k0 = 0; k0 < nc; k0 += TC) {
+                    simdgroup_float8x8 mc;
+                    simdgroup_float8x8 mb;
+                    simdgroup_load(mc, C + (t0 + ib*TC)*(int32_t) args.ns52 + k0, args.ns52);
+                    simdgroup_load(mb, B + (t0 + jb*TC)*(int32_t) args.ns42 + k0, args.ns42, 0, true);
+                    simdgroup_multiply_accumulate(cb, mc, mb, cb);
+                }
+
+                threadgroup float * sam = sam_rows + jb*TC;
+                simdgroup_store(cb, sam, CS);
+                simdgroup_barrier(mem_flags::mem_threadgroup);
+                for (short e = tiisg; e < TC*TC; e += N_SIMDWIDTH) {
+                    const short ri = e / TC;
+                    const short rj = e % TC;
+                    const short i  = ib*TC + ri;
+                    const short j  = jb*TC + rj;
+                    sam[ri*CS + rj] = j <= i ?
+                        sam[ri*CS + rj] * exp(shared_acs[i] - shared_acs[j]) : 0.0f;
+                }
+                simdgroup_barrier(mem_flags::mem_threadgroup);
+            }
+
+            for (short ch = 0; ch < HD/TC; ++ch) {
+                simdgroup_float8x8 y_diag  = make_filled_simdgroup_matrix<float, 8>(0.0f);
+                simdgroup_float8x8 y_inter = make_filled_simdgroup_matrix<float, 8>(0.0f);
+
+                for (short jb = 0; jb <= ib; ++jb) {
+                    simdgroup_float8x8 sam;
+                    simdgroup_float8x8 mdtx;
+                    simdgroup_load(sam,  sam_rows + jb*TC,                   CS);
+                    simdgroup_load(mdtx, shared_dtx + jb*TC*HD + ch*TC,     HD);
+                    simdgroup_multiply_accumulate(y_diag, sam, mdtx, y_diag);
+                }
+
+                for (int32_t k0 = 0; k0 < nc; k0 += TC) {
+                    simdgroup_float8x8 mc;
+                    simdgroup_float8x8 ms;
+                    simdgroup_load(mc, C + (t0 + ib*TC)*(int32_t) args.ns52 + k0, args.ns52);
+                    simdgroup_load(ms, state + ch*TC*nc + k0, nc, 0, true);
+                    simdgroup_multiply_accumulate(y_inter, mc, ms, y_inter);
+                }
+
+                simdgroup_store(y_diag,  tile0, TC);
+                simdgroup_store(y_inter, tile1, TC);
+                simdgroup_barrier(mem_flags::mem_threadgroup);
+                for (short e = tiisg; e < TC*TC; e += N_SIMDWIDTH) {
+                    const short ri = e / TC;
+                    const short ci = e % TC;
+                    const int32_t token = t0 + ib*TC + ri;
+                    y[token*nh*nr + ch*TC + ci] =
+                        tile0[e] + shared_exp_acs[ib*TC + ri] * tile1[e];
+                }
+                simdgroup_barrier(mem_flags::mem_threadgroup);
+            }
+        }
+
+        // All simdgroups must finish reading s_buff before any thread overwrites it.
+        threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup);
+
+        // Keep the carried-state reduction in token order. Reassociating this particular product
+        // with MMA compounds rounding differences at every chunk boundary; CB, y_diag, and C*S
+        // remain on the matrix unit.
+        const float chunk_decay = exp(shared_acs[CS - 1]);
+        for (int32_t idx = tiitg; idx < nc*HD; idx += NSG*N_SIMDWIDTH) {
+            const int32_t ci = idx / nc;
+            const int32_t si = idx % nc;
+            float state_c = 0.0f;
+            for (short t = 0; t < CS; ++t) {
+                state_c += shared_state_decay[t] *
+                    B[(t0 + t)*(int32_t) args.ns42 + si] *
+                    shared_dtx[t*HD + ci];
+            }
+            s_buff[idx] = chunk_decay * state[idx] + state_c;
+        }
+
+        // All state tiles must be visible before the next chunk consumes s_buff as S_prev.
+        threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup);
+    }
+}
@@ -9,10 +9,18 @@ if (WIN32)
    target_link_libraries(ggml-rpc PRIVATE ws2_32)
 endif()

-# RDMA auto-detection (Linux only, requires libibverbs)
-if (NOT WIN32 AND NOT APPLE)
-    find_library(IBVERBS_LIB ibverbs)
-    if (IBVERBS_LIB)
+# RDMA auto-detection: Linux RoCE/IB via libibverbs, Apple RDMA-over-Thunderbolt via librdma
+if (APPLE)
+    set(RDMA_LIB_NAME rdma)
+    set(RDMA_DESC "Apple RDMA-over-Thunderbolt, UC")
+elseif (NOT WIN32)
+    set(RDMA_LIB_NAME ibverbs)
+    set(RDMA_DESC "auto-detected")
+endif()
+
+if (RDMA_LIB_NAME)
+    find_library(RDMA_LIB ${RDMA_LIB_NAME})
+    if (RDMA_LIB)
        option(GGML_RPC_RDMA "ggml: enable RDMA transport for RPC" ON)
    else()
        option(GGML_RPC_RDMA "ggml: enable RDMA transport for RPC" OFF)
@@ -22,12 +30,16 @@ else()
 endif()

 if (GGML_RPC_RDMA)
-    if (NOT IBVERBS_LIB)
-        find_library(IBVERBS_LIB ibverbs REQUIRED)
+    if (NOT RDMA_LIB)
+        find_library(RDMA_LIB ${RDMA_LIB_NAME} REQUIRED)
    endif()
    target_compile_definitions(ggml-rpc PRIVATE GGML_RPC_RDMA)
-    target_link_libraries(ggml-rpc PRIVATE ${IBVERBS_LIB})
-    message(STATUS "  RDMA transport enabled (auto-detected)")
+    target_link_libraries(ggml-rpc PRIVATE ${RDMA_LIB})
+    if (APPLE)
+        target_compile_definitions(ggml-rpc PRIVATE GGML_RPC_RDMA_APPLE)
+        target_sources(ggml-rpc PRIVATE transport-apple.cpp)
+    endif()
+    message(STATUS "  RDMA transport enabled (${RDMA_DESC})")
 else()
    message(STATUS "  RDMA transport disabled")
 endif()
@@ -9,6 +9,9 @@
 #include <optional>
 #include <string>
 #include <vector>
+#include <queue>
+#include <condition_variable>
+#include <future>
 #include <memory>
 #include <mutex>
 #include <unordered_map>
@@ -17,6 +20,8 @@
 #include <fstream>
 #include <filesystem>
 #include <algorithm>
+#include <atomic>
+#include <thread>

 static const char * RPC_DEBUG = std::getenv("GGML_RPC_DEBUG");

@@ -72,6 +77,7 @@ enum rpc_cmd {
    RPC_CMD_DEVICE_COUNT,
    RPC_CMD_GRAPH_RECOMPUTE,
    RPC_CMD_MEMSET_TENSOR,
+    RPC_CMD_NONE,
    RPC_CMD_COUNT,
 };

@@ -223,24 +229,24 @@ struct ggml_backend_rpc_buffer_type_context {
    size_t      max_size;
 };

+class rpc_dispatcher;
 struct ggml_backend_rpc_context {
-    std::string endpoint;
-    uint32_t    device;
-    std::string name;
+    std::shared_ptr<rpc_dispatcher> dispatcher;
+    uint32_t                        device;
+    std::string                     name;
 };

 struct ggml_backend_rpc_buffer_context {
-    std::shared_ptr<socket_t> sock;
-    void * base_ptr;
-    uint64_t remote_ptr;
+    std::shared_ptr<rpc_dispatcher>   dispatcher;
+    void                            * base_ptr;
+    uint64_t                          remote_ptr;
 };

 // RPC helper functions

 // Computes FNV-1a hash of the data
-static uint64_t fnv_hash(const uint8_t * data, size_t len) {
+static uint64_t fnv_hash(const uint8_t * data, size_t len, uint64_t hash = 0xcbf29ce484222325ULL) {
    const uint64_t fnv_prime = 0x100000001b3ULL;
-    uint64_t hash = 0xcbf29ce484222325ULL;

    for (size_t i = 0; i < len; ++i) {
        hash ^= data[i];
@@ -253,7 +259,10 @@ static bool send_msg(socket_ptr sock, const void * msg, size_t msg_size) {
    if (!sock->send_data(&msg_size, sizeof(msg_size))) {
        return false;
    }
-    return sock->send_data(msg, msg_size);
+    if (!sock->send_data(msg, msg_size)) {
+        return false;
+    }
+    return sock->flush();
 }

 static bool recv_msg(socket_ptr sock, void * msg, size_t msg_size) {
@@ -308,7 +317,7 @@ static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input,
    if (!sock->send_data(input, input_size)) {
        return false;
    }
-    return true;
+    return sock->flush();
 }

 // RPC request : | rpc_cmd (1 byte) | request_size (8 bytes) | request_data (request_size bytes) |
@@ -354,44 +363,248 @@ static bool negotiate_hello(const std::shared_ptr<socket_t> & sock) {
    return true;
 }

-static std::shared_ptr<socket_t> get_socket(const std::string & endpoint) {
-    static std::mutex mutex;
-    std::lock_guard<std::mutex> lock(mutex);
-    static std::unordered_map<std::string, std::weak_ptr<socket_t>> sockets;
+template <typename T>
+class message_queue {
+public:
+    message_queue() {}

-    auto it = sockets.find(endpoint);
-    if (it != sockets.end()) {
-        if (auto sock = it->second.lock()) {
-            return sock;
+    bool push(const T &value) {
+        std::unique_lock<std::mutex> lock(mutex);
+        if (interrupted) {
+            return false;
        }
+        queue.push(value);
+        cvar.notify_all();
+        return true;
    }
+
+    bool pop(T* out) {
+        std::unique_lock<std::mutex> lock(mutex);
+        cvar.wait(lock, [this] { return !queue.empty() || interrupted; });
+        if (interrupted) {
+            return false;
+        }
+        *out = queue.front();
+        queue.pop();
+        return true;
+    }
+
+    void interrupt() {
+        std::unique_lock<std::mutex> lock(mutex);
+        interrupted = true;
+        lock.unlock();
+        cvar.notify_all();
+    }
+
+private:
+    bool interrupted = false;
+    std::queue<T> queue;
+    std::mutex mutex;
+    std::condition_variable cvar;
+};
+
+class rpc_dispatcher {
+public:
+    rpc_dispatcher() {
+    }
+
+    void send(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size);
+    void send(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size, void * output, size_t output_size);
+    void send_async(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size);
+    void send_async(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size, void * output, size_t output_size);
+
+    ggml_backend_event_t event_new(ggml_backend_dev_t dev);
+    void event_free(ggml_backend_event_t event);
+    void event_synchronize(ggml_backend_event_t event);
+    void event_record(ggml_backend_event_t event);
+    void synchronize();
+
+    void start(const std::string & endpoint);
+    void work();
+
+    ~rpc_dispatcher();
+
+private:
+    struct rpc_msg {
+        rpc_cmd                       cmd;
+        std::shared_ptr<const void>   input;
+        size_t                        input_size;
+        void                        * output;
+        size_t                        output_size;
+        std::promise<void>            completion;
+    };
+    using rpc_msg_ptr   = std::shared_ptr<rpc_msg>;
+    using rpc_msg_queue = message_queue<rpc_msg_ptr>;
+    struct rpc_event {
+        rpc_msg_ptr              msg;
+        std::shared_future<void> sf;
+    };
+    rpc_msg_queue    queue;
+    socket_ptr       sock;
+    std::atomic_bool running;
+    std::thread      thread;
+};
+
+static void rpc_dispatcher_trampoline(rpc_dispatcher * dispatcher)
+{
+    dispatcher->work();
+}
+
+void rpc_dispatcher::send(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size) {
+    auto msg = std::make_shared<rpc_msg>();
+    msg->cmd = cmd;
+    msg->input = input;
+    msg->input_size = input_size;
+    msg->output = nullptr;
+    msg->output_size = 0;
+    GGML_ASSERT(queue.push(msg));
+    auto future = msg->completion.get_future();
+    future.wait();
+}
+
+void rpc_dispatcher::send_async(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size) {
+    auto msg = std::make_shared<rpc_msg>();
+    msg->cmd = cmd;
+    msg->input = input;
+    msg->input_size = input_size;
+    msg->output = nullptr;
+    msg->output_size = 0;
+    GGML_ASSERT(queue.push(msg));
+}
+
+void rpc_dispatcher::send(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size, void * output, size_t output_size) {
+    auto msg = std::make_shared<rpc_msg>();
+    msg->cmd = cmd;
+    msg->input = input;
+    msg->input_size = input_size;
+    msg->output = output;
+    msg->output_size = output_size;
+    GGML_ASSERT(queue.push(msg));
+    auto future = msg->completion.get_future();
+    future.wait();
+}
+
+void rpc_dispatcher::send_async(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size, void * output, size_t output_size) {
+    auto msg = std::make_shared<rpc_msg>();
+    msg->cmd = cmd;
+    msg->input = input;
+    msg->input_size = input_size;
+    msg->output = output;
+    msg->output_size = output_size;
+    GGML_ASSERT(queue.push(msg));
+}
+
+ggml_backend_event_t rpc_dispatcher::event_new(ggml_backend_dev_t dev) {
+    rpc_event * ev = new rpc_event;
+    ev->msg = std::make_shared<rpc_msg>();
+    ev->msg->cmd = RPC_CMD_NONE;
+    ev->sf = ev->msg->completion.get_future().share();
+    GGML_ASSERT(queue.push(ev->msg));
+    return new ggml_backend_event {
+        /* .device  = */ dev,
+        /* .context = */ ev,
+    };
+}
+
+void rpc_dispatcher::event_free(ggml_backend_event_t event) {
+    rpc_event * ev = (rpc_event *)event->context;
+    delete ev;
+}
+
+void rpc_dispatcher::event_synchronize(ggml_backend_event_t event) {
+    rpc_event * ev = (rpc_event *)event->context;
+    ev->sf.wait();
+}
+
+void rpc_dispatcher::event_record(ggml_backend_event_t event) {
+    rpc_event * ev = (rpc_event *)event->context;
+    ev->msg = std::make_shared<rpc_msg>();
+    ev->msg->cmd = RPC_CMD_NONE;
+    ev->sf = ev->msg->completion.get_future().share();
+    GGML_ASSERT(queue.push(ev->msg));
+}
+
+void rpc_dispatcher::synchronize() {
+    // to ensure all messages are processed, submit dummy message and wait for it to complete
+    auto msg = std::make_shared<rpc_msg>();
+    msg->cmd = RPC_CMD_NONE;
+    GGML_ASSERT(queue.push(msg));
+    msg->completion.get_future().wait();
+}
+
+void rpc_dispatcher::start(const std::string & endpoint) {
    std::string host;
    int port;
    if (!parse_endpoint(endpoint, host, port)) {
-        GGML_LOG_ERROR("Failed to parse endpoint: %s\n", endpoint.c_str());
-        return nullptr;
+        GGML_ABORT("Failed to parse endpoint: %s\n", endpoint.c_str());
+    }
+    if (!rpc_transport_init()) {
+        GGML_ABORT("RPC transport initialization failed\n");
    }

-    if (!rpc_transport_init()) {
-        return nullptr;
-    }
-    auto sock = socket_t::connect(host.c_str(), port);
+    sock = socket_t::connect(host.c_str(), port);
    if (sock == nullptr) {
-        return nullptr;
+        GGML_ABORT("Failed to connect to %s\n", endpoint.c_str());
    }
    if (!negotiate_hello(sock)) {
-        return nullptr;
+        GGML_ABORT("RPC handshake failed for %s\n", endpoint.c_str());
    }
    LOG_DBG("[%s] connected to %s\n", __func__, endpoint.c_str());
-    sockets[endpoint] = sock;
-    return sock;
+    running = true;
+    thread = std::thread(rpc_dispatcher_trampoline, this);
+}
+
+void rpc_dispatcher::work() {
+    while (running) {
+        rpc_msg_ptr msg_ptr;
+        if (!queue.pop(&msg_ptr)) {
+            break;
+        }
+        if (msg_ptr->cmd != RPC_CMD_NONE) {
+            if (msg_ptr->output) {
+                bool status = send_rpc_cmd(sock, msg_ptr->cmd, msg_ptr->input.get(), msg_ptr->input_size, msg_ptr->output, msg_ptr->output_size);
+                RPC_STATUS_ASSERT(status);
+            } else {
+                bool status = send_rpc_cmd(sock, msg_ptr->cmd, msg_ptr->input.get(), msg_ptr->input_size);
+                RPC_STATUS_ASSERT(status);
+            }
+        }
+        msg_ptr->completion.set_value();
+    }
+}
+
+rpc_dispatcher::~rpc_dispatcher() {
+    running = false;
+    queue.interrupt();
+    sock = nullptr;
+    if (thread.joinable()) {
+        thread.join();
+    }
+}
+
+static std::shared_ptr<rpc_dispatcher> get_dispatcher(const std::string & endpoint) {
+    static std::mutex mutex;
+    std::lock_guard<std::mutex> lock(mutex);
+    static std::unordered_map<std::string, std::weak_ptr<rpc_dispatcher>> dispatchers;
+
+    auto it = dispatchers.find(endpoint);
+    if (it != dispatchers.end()) {
+        if (auto dispatcher = it->second.lock()) {
+            return dispatcher;
+        }
+    }
+
+    auto dispatcher = std::make_shared<rpc_dispatcher>();
+    dispatcher->start(endpoint);
+    dispatchers[endpoint] = dispatcher;
+    return dispatcher;
 }

 static void ggml_backend_rpc_buffer_free_buffer(ggml_backend_buffer_t buffer) {
    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
-    rpc_msg_free_buffer_req request = {ctx->remote_ptr};
-    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_FREE_BUFFER, &request, sizeof(request), nullptr, 0);
-    RPC_STATUS_ASSERT(status);
+    auto request = std::make_shared<rpc_msg_free_buffer_req>();
+    request->remote_ptr = ctx->remote_ptr;
+    ctx->dispatcher->send(RPC_CMD_FREE_BUFFER, request, sizeof(*request));
    delete ctx;
 }

@@ -400,10 +613,10 @@ static void * ggml_backend_rpc_buffer_get_base(ggml_backend_buffer_t buffer) {
    if (ctx->base_ptr != nullptr) {
        return ctx->base_ptr;
    }
-    rpc_msg_buffer_get_base_req request = {ctx->remote_ptr};
+    auto request = std::make_shared<rpc_msg_buffer_get_base_req>();
+    request->remote_ptr = ctx->remote_ptr;
    rpc_msg_buffer_get_base_rsp response;
-    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_BUFFER_GET_BASE, &request, sizeof(request), &response, sizeof(response));
-    RPC_STATUS_ASSERT(status);
+    ctx->dispatcher->send(RPC_CMD_BUFFER_GET_BASE, request, sizeof(*request), &response, sizeof(response));
    ctx->base_ptr = reinterpret_cast<void *>(response.base_ptr);
    return ctx->base_ptr;
 }
@@ -460,12 +673,9 @@ static enum ggml_status ggml_backend_rpc_buffer_init_tensor(ggml_backend_buffer_
    // Due to bandwidth constraints, we only call the server init tensor functions if necessary.
    // In particular, only quantized tensors need padding
    if (ggml_is_quantized(tensor->type) && (tensor->ne[0] % 512 != 0) && (tensor->view_src == nullptr)) {
-        rpc_msg_init_tensor_req request;
-
-        request.tensor = serialize_tensor(tensor);
-
-        bool status = send_rpc_cmd(ctx->sock, RPC_CMD_INIT_TENSOR, &request, sizeof(request), nullptr, 0);
-        RPC_STATUS_ASSERT(status);
+        auto request = std::make_shared<rpc_msg_init_tensor_req>();
+        request->tensor = serialize_tensor(tensor);
+        ctx->dispatcher->send(RPC_CMD_INIT_TENSOR, request, sizeof(*request));
    }
    return GGML_STATUS_SUCCESS;
 }
@@ -473,27 +683,24 @@ static enum ggml_status ggml_backend_rpc_buffer_init_tensor(ggml_backend_buffer_
 static void ggml_backend_rpc_buffer_memset_tensor(
        ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) {
    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
-    rpc_msg_memset_tensor_req request = {
-        /* .tensor = */ serialize_tensor(tensor),
-        /* .offset = */ offset,
-        /* .size   = */ size,
-        /* .value  = */ value,
-    };
-    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_MEMSET_TENSOR, &request, sizeof(request), nullptr, 0);
-    RPC_STATUS_ASSERT(status);
+    auto request = std::make_shared<rpc_msg_memset_tensor_req>();
+    request->tensor = serialize_tensor(tensor);
+    request->offset = offset;
+    request->size   = size;
+    request->value  = value;
+    ctx->dispatcher->send(RPC_CMD_MEMSET_TENSOR, request, sizeof(*request));
 }

 static void ggml_backend_rpc_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
    rpc_tensor rpc_tensor = serialize_tensor(tensor);
    if (size > HASH_THRESHOLD) {
-        rpc_msg_set_tensor_hash_req request;
-        request.tensor = rpc_tensor;
-        request.offset = offset;
-        request.hash = fnv_hash((const uint8_t*)data, size);
+        auto request = std::make_shared<rpc_msg_set_tensor_hash_req>();
+        request->tensor = rpc_tensor;
+        request->offset = offset;
+        request->hash = fnv_hash((const uint8_t*)data, size);
        rpc_msg_set_tensor_hash_rsp response;
-        bool status = send_rpc_cmd(ctx->sock, RPC_CMD_SET_TENSOR_HASH, &request, sizeof(request), &response, sizeof(response));
-        RPC_STATUS_ASSERT(status);
+        ctx->dispatcher->send(RPC_CMD_SET_TENSOR_HASH, request, sizeof(*request), &response, sizeof(response));
        if (response.result) {
            // the server has the same data, no need to send it
            return;
@@ -501,22 +708,21 @@ static void ggml_backend_rpc_buffer_set_tensor(ggml_backend_buffer_t buffer, ggm
    }
    // input serialization format: | rpc_tensor | offset (8 bytes) | data (size bytes)
    size_t input_size = sizeof(rpc_tensor) + sizeof(uint64_t) + size;
-    std::vector<uint8_t> input(input_size, 0);
-    memcpy(input.data(), &rpc_tensor, sizeof(rpc_tensor));
-    memcpy(input.data() + sizeof(rpc_tensor), &offset, sizeof(offset));
-    memcpy(input.data() + sizeof(rpc_tensor) + sizeof(offset), data, size);
-    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_SET_TENSOR, input.data(), input.size());
-    RPC_STATUS_ASSERT(status);
+    uint8_t * input = new uint8_t[input_size]();
+    memcpy(input, &rpc_tensor, sizeof(rpc_tensor));
+    memcpy(input + sizeof(rpc_tensor), &offset, sizeof(offset));
+    memcpy(input + sizeof(rpc_tensor) + sizeof(offset), data, size);
+    std::shared_ptr<uint8_t> input_ptr(input, std::default_delete<uint8_t[]>());
+    ctx->dispatcher->send(RPC_CMD_SET_TENSOR, input_ptr, input_size);
 }

 static void ggml_backend_rpc_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) {
    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
-    rpc_msg_get_tensor_req request;
-    request.tensor = serialize_tensor(tensor);
-    request.offset = offset;
-    request.size = size;
-    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_GET_TENSOR, &request, sizeof(request), data, size);
-    RPC_STATUS_ASSERT(status);
+    auto request = std::make_shared<rpc_msg_get_tensor_req>();
+    request->tensor = serialize_tensor(tensor);
+    request->offset = offset;
+    request->size = size;
+    ctx->dispatcher->send(RPC_CMD_GET_TENSOR, request, sizeof(*request), data, size);
 }

 static bool ggml_backend_rpc_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) {
@@ -526,16 +732,15 @@ static bool ggml_backend_rpc_buffer_cpy_tensor(ggml_backend_buffer_t buffer, con
        ggml_backend_rpc_buffer_context * src_ctx = (ggml_backend_rpc_buffer_context *)src_buffer->context;
        ggml_backend_buffer_t dst_buffer = dst->buffer;
        ggml_backend_rpc_buffer_context * dst_ctx = (ggml_backend_rpc_buffer_context *)dst_buffer->context;
-        if (src_ctx->sock != dst_ctx->sock) {
+        if (src_ctx->dispatcher != dst_ctx->dispatcher) {
            return false;
        }
        ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
-        rpc_msg_copy_tensor_req request;
-        request.src = serialize_tensor(src);
-        request.dst = serialize_tensor(dst);
+        auto request = std::make_shared<rpc_msg_copy_tensor_req>();
+        request->src = serialize_tensor(src);
+        request->dst = serialize_tensor(dst);
        rpc_msg_copy_tensor_rsp response;
-        bool status = send_rpc_cmd(ctx->sock, RPC_CMD_COPY_TENSOR, &request, sizeof(request), &response, sizeof(response));
-        RPC_STATUS_ASSERT(status);
+        ctx->dispatcher->send(RPC_CMD_COPY_TENSOR, request, sizeof(*request), &response, sizeof(response));
        return response.result;
    }
    return false;
@@ -543,9 +748,10 @@ static bool ggml_backend_rpc_buffer_cpy_tensor(ggml_backend_buffer_t buffer, con

 static void ggml_backend_rpc_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) {
    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
-    rpc_msg_buffer_clear_req request = {ctx->remote_ptr, value};
-    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_BUFFER_CLEAR, &request, sizeof(request), nullptr, 0);
-    RPC_STATUS_ASSERT(status);
+    auto request = std::make_shared<rpc_msg_buffer_clear_req>();
+    request->remote_ptr = ctx->remote_ptr;
+    request->value = value;
+    ctx->dispatcher->send(RPC_CMD_BUFFER_CLEAR, request, sizeof(*request));
 }

 static ggml_backend_buffer_i ggml_backend_rpc_buffer_interface = {
@@ -569,15 +775,17 @@ static const char * ggml_backend_rpc_buffer_type_name(ggml_backend_buffer_type_t

 static ggml_backend_buffer_t ggml_backend_rpc_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
    ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context;
-    rpc_msg_alloc_buffer_req request = {buft_ctx->device, size};
+    auto request = std::make_shared<rpc_msg_alloc_buffer_req>();
+    request->device = buft_ctx->device;
+    request->size = size;
    rpc_msg_alloc_buffer_rsp response;
-    auto sock = get_socket(buft_ctx->endpoint);
-    bool status = send_rpc_cmd(sock, RPC_CMD_ALLOC_BUFFER, &request, sizeof(request), &response, sizeof(response));
-    RPC_STATUS_ASSERT(status);
+
+    auto dispatcher = get_dispatcher(buft_ctx->endpoint);
+    dispatcher->send(RPC_CMD_ALLOC_BUFFER, request, sizeof(*request), &response, sizeof(response));
    if (response.remote_ptr != 0) {
        ggml_backend_buffer_t buffer = ggml_backend_buffer_init(buft,
            ggml_backend_rpc_buffer_interface,
-            new ggml_backend_rpc_buffer_context{sock, nullptr, response.remote_ptr},
+            new ggml_backend_rpc_buffer_context{dispatcher, nullptr, response.remote_ptr},
            response.remote_size);
        return buffer;
    } else {
@@ -585,11 +793,11 @@ static ggml_backend_buffer_t ggml_backend_rpc_buffer_type_alloc_buffer(ggml_back
    }
 }

-static size_t get_alignment(const std::shared_ptr<socket_t> & sock, uint32_t device) {
-    rpc_msg_get_alignment_req request = {device};
+static size_t get_alignment(const std::shared_ptr<rpc_dispatcher> & dispatcher, uint32_t device) {
+    auto request = std::make_shared<rpc_msg_get_alignment_req>();
+    request->device = device;
    rpc_msg_get_alignment_rsp response;
-    bool status = send_rpc_cmd(sock, RPC_CMD_GET_ALIGNMENT, &request, sizeof(request), &response, sizeof(response));
-    RPC_STATUS_ASSERT(status);
+    dispatcher->send(RPC_CMD_GET_ALIGNMENT, request, sizeof(*request), &response, sizeof(response));
    return response.alignment;
 }

@@ -598,11 +806,11 @@ static size_t ggml_backend_rpc_buffer_type_get_alignment(ggml_backend_buffer_typ
    return buft_ctx->alignment;
 }

-static size_t get_max_size(const std::shared_ptr<socket_t> & sock, uint32_t device) {
-    rpc_msg_get_max_size_req request = {device};
+static size_t get_max_size(const std::shared_ptr<rpc_dispatcher> & dispatcher, uint32_t device) {
+    auto request = std::make_shared<rpc_msg_get_max_size_req>();
+    request->device = device;
    rpc_msg_get_max_size_rsp response;
-    bool status = send_rpc_cmd(sock, RPC_CMD_GET_MAX_SIZE, &request, sizeof(request), &response, sizeof(response));
-    RPC_STATUS_ASSERT(status);
+    dispatcher->send(RPC_CMD_GET_MAX_SIZE, request, sizeof(*request), &response, sizeof(response));
    return response.max_size;
 }

@@ -625,23 +833,63 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty

    if (rpc_get) {
        ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context;
-        auto sock = get_socket(buft_ctx->endpoint);

-        rpc_msg_get_alloc_size_req request = {
-            /*.device =*/ buft_ctx->device,
-            /*.tensor =*/ serialize_tensor(tensor),
-            /*.srcs   =*/ {},
+        // Cache key for calls to read the alloc_size.
+        // We deliberately exclude src tensor dimensions from the key because:
+        // 1. For CPU backends, alloc_size = ggml_nbytes(output) regardless of src shapes
+        // 2. For GPU backends, the reservation graph uses max dimensions, so the
+        //    cached value from reservation is always >= any subsequent request
+        // 3. Including src dims causes cache misses per-ubatch (e.g. growing KV cache)
+        //    which blocks the main thread behind in-flight GRAPH_COMPUTE commands
+        struct alloc_size_cache_key {
+            uint32_t device;
+            uint32_t type;
+            uint32_t op;
+            int32_t  op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)];
+            uint32_t ne[GGML_MAX_DIMS];
        };

+        alloc_size_cache_key key = {};
+        key.device = buft_ctx->device;
+        key.type = tensor->type;
+        key.op = tensor->op;
+        memcpy(key.op_params, tensor->op_params, sizeof(key.op_params));
+        for (int i = 0; i < GGML_MAX_DIMS; i++) {
+            key.ne[i] = (uint32_t)tensor->ne[i];
+        }
+
+        uint64_t cache_hash = fnv_hash((const uint8_t *)&key, sizeof(key));
+        cache_hash = fnv_hash((const uint8_t *)buft_ctx->endpoint.data(), buft_ctx->endpoint.size(), cache_hash);
+
+        // alloc sizes are immutable for a given tensor configuration
+        static std::mutex cache_mutex;
+        static std::unordered_map<uint64_t, size_t> cache;
+
+        {
+            std::lock_guard<std::mutex> lock(cache_mutex);
+            auto it = cache.find(cache_hash);
+            if (it != cache.end()) {
+                return it->second;
+            }
+        }
+
+        auto request = std::make_shared<rpc_msg_get_alloc_size_req>();
+        request->device = buft_ctx->device;
+        request->tensor = serialize_tensor(tensor);
+
        // .get_alloc_size could be a function of the tensor's srcs, so we must serialize them as well
        for (int i = 0; i < GGML_MAX_SRC; i++) {
-            request.srcs[i] = serialize_tensor(tensor->src[i]);
+            request->srcs[i] = serialize_tensor(tensor->src[i]);
        }

-        // TODO: cache the alloc responses to avoid extra RPC calls?
        rpc_msg_get_alloc_size_rsp response;
-        bool status = send_rpc_cmd(sock, RPC_CMD_GET_ALLOC_SIZE, &request, sizeof(request), &response, sizeof(response));
-        RPC_STATUS_ASSERT(status);
+        auto dispatcher = get_dispatcher(buft_ctx->endpoint);
+        dispatcher->send(RPC_CMD_GET_ALLOC_SIZE, request, sizeof(*request), &response, sizeof(response));
+
+        {
+            std::lock_guard<std::mutex> lock(cache_mutex);
+            cache[cache_hash] = response.alloc_size;
+        }

        return response.alloc_size;
    }
@@ -670,9 +918,44 @@ static void ggml_backend_rpc_free(ggml_backend_t backend) {
    delete backend;
 }

+static void ggml_backend_rpc_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
+    ggml_backend_rpc_context * ctx = (ggml_backend_rpc_context *)backend->context;
+    rpc_tensor rpc_tensor = serialize_tensor(tensor);
+    if (size > HASH_THRESHOLD) {
+        auto request = std::make_shared<rpc_msg_set_tensor_hash_req>();
+        request->tensor = rpc_tensor;
+        request->offset = offset;
+        request->hash = fnv_hash((const uint8_t*)data, size);
+        rpc_msg_set_tensor_hash_rsp response;
+        // TODO: make this async
+        ctx->dispatcher->send(RPC_CMD_SET_TENSOR_HASH, request, sizeof(*request), &response, sizeof(response));
+        if (response.result) {
+            // the server has the same data, no need to send it
+            return;
+        }
+    }
+    // input serialization format: | rpc_tensor | offset (8 bytes) | data (size bytes)
+    size_t input_size = sizeof(rpc_tensor) + sizeof(uint64_t) + size;
+    uint8_t * input = new uint8_t[input_size]();
+    memcpy(input, &rpc_tensor, sizeof(rpc_tensor));
+    memcpy(input + sizeof(rpc_tensor), &offset, sizeof(offset));
+    memcpy(input + sizeof(rpc_tensor) + sizeof(offset), data, size);
+    std::shared_ptr<uint8_t> input_ptr(input, std::default_delete<uint8_t[]>());
+    ctx->dispatcher->send_async(RPC_CMD_SET_TENSOR, input_ptr, input_size);
+}
+
+static void ggml_backend_rpc_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) {
+    ggml_backend_rpc_context * ctx = (ggml_backend_rpc_context *)backend->context;
+    auto request = std::make_shared<rpc_msg_get_tensor_req>();
+    request->tensor = serialize_tensor(tensor);
+    request->offset = offset;
+    request->size = size;
+    ctx->dispatcher->send_async(RPC_CMD_GET_TENSOR, request, sizeof(*request), data, size);
+}
+
 static void ggml_backend_rpc_synchronize(ggml_backend_t backend) {
-    GGML_UNUSED(backend);
-    // this is no-op because we don't have any async operations
+    ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context;
+    rpc_ctx->dispatcher->synchronize();
 }

 static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::vector<rpc_tensor> & tensors, std::unordered_set<ggml_tensor*> & visited) {
@@ -695,7 +978,7 @@ static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::ve
    tensors.push_back(result);
 }

-static void serialize_graph(uint32_t device, const ggml_cgraph * cgraph, std::vector<uint8_t> & output) {
+static uint8_t * serialize_graph(uint32_t device, const ggml_cgraph * cgraph, size_t * output_size) {
    uint32_t n_nodes = cgraph->n_nodes;
    std::vector<rpc_tensor> tensors;
    std::unordered_set<ggml_tensor*> visited;
@@ -705,9 +988,9 @@ static void serialize_graph(uint32_t device, const ggml_cgraph * cgraph, std::ve
    // serialization format:
    // | device (4 bytes) | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(rpc_tensor)) |
    uint32_t n_tensors = tensors.size();
-    int output_size = 2*sizeof(uint32_t) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t) + n_tensors * sizeof(rpc_tensor);
-    output.resize(output_size, 0);
-    uint8_t * dest = output.data();
+    *output_size = 2*sizeof(uint32_t) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t) + n_tensors * sizeof(rpc_tensor);
+    uint8_t * output = new uint8_t[*output_size]();
+    uint8_t * dest = output;
    memcpy(dest, &device, sizeof(device));
    dest += sizeof(device);
    memcpy(dest, &n_nodes, sizeof(n_nodes));
@@ -720,6 +1003,7 @@ static void serialize_graph(uint32_t device, const ggml_cgraph * cgraph, std::ve
    dest += sizeof(n_tensors);
    rpc_tensor * out_tensors = (rpc_tensor *)dest;
    memcpy(out_tensors, tensors.data(), n_tensors * sizeof(rpc_tensor));
+    return output;
 }

 static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) {
@@ -730,27 +1014,35 @@ static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, g
    GGML_ASSERT(cgraph->n_nodes > 0);
    bool reuse = cgraph->uid != 0 && rpc_dev_ctx->last_graph_uid == cgraph->uid;
    if (reuse) {
-        rpc_msg_graph_recompute_req request;
-        request.device = rpc_ctx->device;
-        auto sock = get_socket(rpc_ctx->endpoint);
-        bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_RECOMPUTE, &request, sizeof(request));
-        RPC_STATUS_ASSERT(status);
+        auto request = std::make_shared<rpc_msg_graph_recompute_req>();
+        request->device = rpc_ctx->device;
+        rpc_ctx->dispatcher->send_async(RPC_CMD_GRAPH_RECOMPUTE, request, sizeof(*request));
    } else {
        rpc_dev_ctx->last_graph_uid = cgraph->uid;
-        std::vector<uint8_t> input;
-        serialize_graph(rpc_ctx->device, cgraph, input);
-        auto sock = get_socket(rpc_ctx->endpoint);
-        bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_COMPUTE, input.data(), input.size());
-        RPC_STATUS_ASSERT(status);
+        size_t input_size = 0;
+        uint8_t * input = serialize_graph(rpc_ctx->device, cgraph, &input_size);
+        std::shared_ptr<uint8_t> input_ptr(input, std::default_delete<uint8_t[]>());
+        rpc_ctx->dispatcher->send_async(RPC_CMD_GRAPH_COMPUTE, input_ptr, input_size);
    }
    return GGML_STATUS_SUCCESS;
 }

+static void ggml_backend_rpc_event_record(ggml_backend_t backend, ggml_backend_event_t event) {
+    ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context;
+    rpc_ctx->dispatcher->event_record(event);
+}
+
+static void ggml_backend_rpc_event_wait(ggml_backend_t backend, ggml_backend_event_t event) {
+    // this is noop for RPC as we have a single stream
+    GGML_UNUSED(backend);
+    GGML_UNUSED(event);
+}
+
 static ggml_backend_i ggml_backend_rpc_interface = {
    /* .get_name                = */ ggml_backend_rpc_name,
    /* .free                    = */ ggml_backend_rpc_free,
-    /* .set_tensor_async        = */ NULL,
-    /* .get_tensor_async        = */ NULL,
+    /* .set_tensor_async        = */ ggml_backend_rpc_set_tensor_async,
+    /* .get_tensor_async        = */ ggml_backend_rpc_get_tensor_async,
    /* .set_tensor_2d_async     = */ NULL,
    /* .get_tensor_2d_async     = */ NULL,
    /* .cpy_tensor_async        = */ NULL,
@@ -760,8 +1052,8 @@ static ggml_backend_i ggml_backend_rpc_interface = {
    /* .graph_plan_update       = */ NULL,
    /* .graph_plan_compute      = */ NULL,
    /* .graph_compute           = */ ggml_backend_rpc_graph_compute,
-    /* .event_record            = */ NULL,
-    /* .event_wait              = */ NULL,
+    /* .event_record            = */ ggml_backend_rpc_event_record,
+    /* .event_wait              = */ ggml_backend_rpc_event_wait,
    /* .graph_optimize          = */ NULL,
 };

@@ -775,13 +1067,9 @@ ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, u
    if (it != buft_map.end()) {
        return it->second;
    }
-    auto sock = get_socket(endpoint);
-    if (sock == nullptr) {
-        GGML_LOG_ERROR("Failed to connect to %s\n", endpoint);
-        return nullptr;
-    }
-    size_t alignment = get_alignment(sock, device);
-    size_t max_size = get_max_size(sock, device);
+    auto dispatcher = get_dispatcher(endpoint);
+    size_t alignment = get_alignment(dispatcher, device);
+    size_t max_size = get_max_size(dispatcher, device);
    ggml_backend_rpc_buffer_type_context * buft_ctx = new ggml_backend_rpc_buffer_type_context {
        /* .endpoint  = */ endpoint,
        /* .device    = */ device,
@@ -801,10 +1089,11 @@ ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, u

 ggml_backend_t ggml_backend_rpc_init(const char * endpoint, uint32_t device) {
    std::string dev_name = "RPC" + std::to_string(device) + "[" + std::string(endpoint) + "]";
+    auto dispatcher = get_dispatcher(endpoint);
    ggml_backend_rpc_context * ctx = new ggml_backend_rpc_context {
-        /* .endpoint       = */ endpoint,
-        /* .device         = */ device,
-        /* .name           = */ dev_name,
+        /* .dispatcher = */ dispatcher,
+        /* .device     = */ device,
+        /* .name       = */ dev_name,
    };
    auto reg = ggml_backend_rpc_add_server(endpoint);
    ggml_backend_t backend = new ggml_backend {
@@ -820,26 +1109,16 @@ bool ggml_backend_is_rpc(ggml_backend_t backend) {
    return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_rpc_guid());
 }

-static void get_device_memory(const std::shared_ptr<socket_t> & sock, uint32_t device, size_t * free, size_t * total) {
-    rpc_msg_get_device_memory_req request;
-    request.device = device;
+void ggml_backend_rpc_get_device_memory(const char * endpoint, uint32_t device, size_t * free, size_t * total) {
+    auto dispatcher = get_dispatcher(endpoint);
+    auto request = std::make_shared<rpc_msg_get_device_memory_req>();
+    request->device = device;
    rpc_msg_get_device_memory_rsp response;
-    bool status = send_rpc_cmd(sock, RPC_CMD_GET_DEVICE_MEMORY, &request, sizeof(request), &response, sizeof(response));
-    RPC_STATUS_ASSERT(status);
+    dispatcher->send(RPC_CMD_GET_DEVICE_MEMORY, request, sizeof(*request), &response, sizeof(response));
    *free = response.free_mem;
    *total = response.total_mem;
 }

-void ggml_backend_rpc_get_device_memory(const char * endpoint, uint32_t device, size_t * free, size_t * total) {
-    auto sock = get_socket(endpoint);
-    if (sock == nullptr) {
-        *free = 0;
-        *total = 0;
-        return;
-    }
-    get_device_memory(sock, device, free, total);
-}
-
 // RPC server-side implementation

 class rpc_server {
@@ -1644,9 +1923,6 @@ static void rpc_serve_client(const std::vector<ggml_backend_t> & backends, const
                if (!server.free_buffer(request)) {
                    return;
                }
-                if (!send_msg(sock, nullptr, 0)) {
-                    return;
-                }
                break;
            }
            case RPC_CMD_BUFFER_CLEAR: {
@@ -1657,9 +1933,6 @@ static void rpc_serve_client(const std::vector<ggml_backend_t> & backends, const
                if (!server.buffer_clear(request)) {
                    return;
                }
-                if (!send_msg(sock, nullptr, 0)) {
-                    return;
-                }
                break;
            }
            case RPC_CMD_MEMSET_TENSOR: {
@@ -1670,9 +1943,6 @@ static void rpc_serve_client(const std::vector<ggml_backend_t> & backends, const
                if (!server.memset_tensor(request)) {
                    return;
                }
-                if (!send_msg(sock, nullptr, 0)) {
-                    return;
-                }
                break;
            }
            case RPC_CMD_SET_TENSOR: {
@@ -1707,9 +1977,6 @@ static void rpc_serve_client(const std::vector<ggml_backend_t> & backends, const
                if (!server.init_tensor(request)) {
                    return;
                }
-                if (!send_msg(sock, nullptr, 0)) {
-                    return;
-                }
                break;
            }
            case RPC_CMD_GET_TENSOR: {
@@ -1886,10 +2153,10 @@ static void ggml_backend_rpc_device_get_props(ggml_backend_dev_t dev, struct ggm
    props->type        = ggml_backend_rpc_device_get_type(dev);
    ggml_backend_rpc_device_get_memory(dev, &props->memory_free, &props->memory_total);
    props->caps = {
-        /* .async                 = */ false,
+        /* .async                 = */ true,
        /* .host_buffer           = */ false,
        /* .buffer_from_host_ptr  = */ false,
-        /* .events                = */ false,
+        /* .events                = */ true,
        /* .mmap_support          = */ true,
    };
 }
@@ -1926,6 +2193,24 @@ static bool ggml_backend_rpc_device_supports_buft(ggml_backend_dev_t dev, ggml_b
    return buft_ctx->endpoint == dev_ctx->endpoint && buft_ctx->device == dev_ctx->device;
 }

+static ggml_backend_event_t ggml_backend_rpc_device_event_new(ggml_backend_dev_t dev) {
+    ggml_backend_rpc_device_context * ctx = (ggml_backend_rpc_device_context *)dev->context;
+    auto dispatcher = get_dispatcher(ctx->endpoint);
+    return dispatcher->event_new(dev);
+}
+
+static void ggml_backend_rpc_device_event_free(ggml_backend_dev_t dev, ggml_backend_event_t event) {
+    ggml_backend_rpc_device_context * ctx = (ggml_backend_rpc_device_context *)dev->context;
+    auto dispatcher = get_dispatcher(ctx->endpoint);
+    dispatcher->event_free(event);
+}
+
+static void ggml_backend_rpc_device_event_synchronize(ggml_backend_dev_t dev, ggml_backend_event_t event) {
+    ggml_backend_rpc_device_context * ctx = (ggml_backend_rpc_device_context *)dev->context;
+    auto dispatcher = get_dispatcher(ctx->endpoint);
+    dispatcher->event_synchronize(event);
+}
+
 static const struct ggml_backend_device_i ggml_backend_rpc_device_i = {
    /* .get_name             = */ ggml_backend_rpc_device_get_name,
    /* .get_description      = */ ggml_backend_rpc_device_get_description,
@@ -1939,9 +2224,9 @@ static const struct ggml_backend_device_i ggml_backend_rpc_device_i = {
    /* .supports_op          = */ ggml_backend_rpc_device_supports_op,
    /* .supports_buft        = */ ggml_backend_rpc_device_supports_buft,
    /* .offload_op           = */ NULL,
-    /* .event_new            = */ NULL,
-    /* .event_free           = */ NULL,
-    /* .event_synchronize    = */ NULL,
+    /* .event_new            = */ ggml_backend_rpc_device_event_new,
+    /* .event_free           = */ ggml_backend_rpc_device_event_free,
+    /* .event_synchronize    = */ ggml_backend_rpc_device_event_synchronize,
 };

 // backend reg interface
@@ -2001,14 +2286,9 @@ ggml_backend_reg_t ggml_backend_rpc_reg(void) {
 }

 static uint32_t ggml_backend_rpc_get_device_count(const char * endpoint) {
-    auto sock = get_socket(endpoint);
-    if (sock == nullptr) {
-        GGML_LOG_ERROR("Failed to connect to %s\n", endpoint);
-        return 0;
-    }
+    auto dispatcher = get_dispatcher(endpoint);
    rpc_msg_device_count_rsp response;
-    bool status = send_rpc_cmd(sock, RPC_CMD_DEVICE_COUNT, nullptr, 0, &response, sizeof(response));
-    RPC_STATUS_ASSERT(status);
+    dispatcher->send(RPC_CMD_DEVICE_COUNT, nullptr, 0, &response, sizeof(response));
    return response.device_count;
 }

@@ -0,0 +1,470 @@
+#include "transport-apple.h"
+#include "transport.h"
+#include "ggml-impl.h"
+
+#include <infiniband/verbs.h>
+
+#include <cerrno>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+#include <poll.h>
+#include <sys/socket.h>
+#include <unistd.h>
+
+// Apple RDMA-over-Thunderbolt (see Apple TN3205).
+//
+// Apple's RDMA is quite different from what's supported in Linux - deserving of its own transport implementation.
+// see https://developer.apple.com/documentation/technotes/tn3205-low-latency-communication-with-rdma-over-thunderbolt for details
+// at a high level the main differences are:
+// UC(unreliable connection) on Apple vs RC(reliable connection) QP transport types on Linux (though in practice UC on Apple is still lossless)
+// fixed 128KiB stride on Apple vs variable chunk size on Linux
+// relying on Apple's hardware credit based flow control vs RNR NAKs + retries on Linux
+//
+// on Apple a SEND and its corresponding RECV must cover the same number of 4 KiB Thunderbolt frames,
+// so every SEND posts a whole 128KiB stride over the wire, even when partially filled.
+// (In testing 128KiB was the best performing among 32, 64, 128, 256)
+
+static constexpr uint32_t RDMA_SEG_MAGIC   = 0x52534547u; // "RSEG"
+static constexpr int      RDMA_NBUF        = 16;          // ring depth (frames per direction)
+static constexpr size_t   RDMA_FRAME       = 4096;        // Thunderbolt frame (fixed on Apple)
+static constexpr size_t   RDMA_STRIDE      = 128 * 1024;  // 32 Thunderbolt frames; NBUF x this = 2 MiB pinned per direction
+static constexpr uint32_t RDMA_PSN         = 0;           // any value works if both sides match: UC has no retransmit
+static constexpr size_t   RDMA_GID_SIZE    = 16;
+
+static_assert(RDMA_STRIDE % RDMA_FRAME == 0, "RDMA_STRIDE must be a whole number of frames");
+// TN3205 counts queue depth in Thunderbolt frames, not work requests.
+static constexpr uint32_t RDMA_QP_WR       = (uint32_t)RDMA_NBUF * (RDMA_STRIDE / RDMA_FRAME);
+static constexpr uint64_t RDMA_RECV_WR     = 1ull << 20;  // wr_id bit tagging recv completions
+static constexpr uint64_t RDMA_WR_IDX_MASK = 0xffff;      // buffer index in the low bits of wr_id
+static constexpr uint8_t  RDMA_SYNC_READY  = 0x2A;        // readiness-handshake byte (peer activated)
+
+struct rdma_seg_hdr {
+    uint32_t magic; // RDMA_SEG_MAGIC; a mismatch means the stream desynced
+    uint32_t len;   // payload bytes in this frame; the rest of the stride is padding
+};
+static constexpr size_t RDMA_PAYLOAD = RDMA_STRIDE - sizeof(rdma_seg_hdr);
+
+struct apple_rdma_caps {
+    uint32_t qpn;
+    uint16_t lid;
+    uint16_t reserved;
+    uint8_t  gid[RDMA_GID_SIZE];
+};
+
+static_assert(sizeof(apple_rdma_caps) == RPC_CONN_CAPS_SIZE, "apple_rdma_caps must match conn_caps size");
+
+struct apple_rdma::impl {
+    int fd = -1;                      // bootstrap TCP socket, kept as the liveness anchor
+
+    struct ibv_context * ctx = nullptr;
+    struct ibv_pd * pd = nullptr;
+    struct ibv_cq * cq = nullptr;     // one CQ for both directions; RDMA_RECV_WR tags recv completions
+    struct ibv_qp * qp = nullptr;
+
+    uint8_t       * send_mem = nullptr;
+    struct ibv_mr * send_mr  = nullptr;
+    uint8_t       * recv_mem = nullptr;
+    struct ibv_mr * recv_mr  = nullptr;
+
+    int      send_busy[RDMA_NBUF] = {};  // 1 while this buffer has a send in flight
+    // completed recv frames, oldest first: ring index, bytes already handed to
+    // the reader, and total payload length
+    struct { int buf; uint32_t off; uint32_t len; } inq[RDMA_NBUF] = {};
+    int      inq_head = 0;
+    int      inq_count = 0;
+    int      pend_buf = -1;
+    uint32_t pend_len = 0;
+    bool     broken = false;
+
+    uint32_t     qpn = 0;
+    uint8_t      port = 0;
+    int          gid_idx = 0;
+    enum ibv_mtu path_mtu = IBV_MTU_1024;
+
+    int  progress();
+    bool acquire_pending();
+    bool post_pending();
+
+    bool post_recv(int i) {
+        struct ibv_sge sge = {};
+        sge.addr   = (uintptr_t)(recv_mem + (size_t)i * RDMA_STRIDE);
+        sge.length = (uint32_t)RDMA_STRIDE;
+        sge.lkey   = recv_mr->lkey;
+        struct ibv_recv_wr wr = {}, * bad = nullptr;
+        wr.wr_id   = RDMA_RECV_WR | (uint64_t)i;
+        wr.sg_list = &sge;
+        wr.num_sge = 1;
+        return ibv_post_recv(qp, &wr, &bad) == 0;
+    }
+
+    bool post_send(int i, size_t len) {
+        struct ibv_sge sge = {};
+        sge.addr   = (uintptr_t)(send_mem + (size_t)i * RDMA_STRIDE);
+        sge.length = (uint32_t)len;
+        sge.lkey   = send_mr->lkey;
+        struct ibv_send_wr wr = {}, * bad = nullptr;
+        wr.wr_id   = (uint64_t)i;
+        wr.sg_list = &sge;
+        wr.num_sge = 1;
+        wr.opcode  = IBV_WR_SEND;
+        wr.send_flags = IBV_SEND_SIGNALED;
+        return ibv_post_send(qp, &wr, &bad) == 0;
+    }
+
+    ~impl() {
+        broken = true;
+        // the QP must be destroyed before the memory it can still write to is
+        // deregistered and freed: ERR only starts flushing the posted WQEs
+        if (qp) {
+            struct ibv_qp_attr a = {};
+            a.qp_state = IBV_QPS_ERR;
+            ibv_modify_qp(qp, &a, IBV_QP_STATE);
+            struct ibv_wc wc[RDMA_NBUF * 2];
+            while (ibv_poll_cq(cq, RDMA_NBUF * 2, wc) > 0) {}
+            ibv_destroy_qp(qp);
+        }
+        if (send_mr) ibv_dereg_mr(send_mr);
+        if (recv_mr) ibv_dereg_mr(recv_mr);
+        free(send_mem);
+        free(recv_mem);
+        if (cq)  ibv_destroy_cq(cq);
+        if (pd)  ibv_dealloc_pd(pd);
+        if (ctx) ibv_close_device(ctx);
+    }
+};
+
+apple_rdma::apple_rdma(std::unique_ptr<impl> p) : pimpl(std::move(p)) {}
+
+apple_rdma::~apple_rdma() = default;
+
+bool apple_rdma::broken() const {
+    return pimpl->broken;
+}
+
+// The readiness handshake below still runs over the bootstrap socket, one byte
+// each way, before the transport is declared live.
+static bool tcp_send_byte(int fd, uint8_t b) {
+    ssize_t n;
+    do { n = ::send(fd, &b, sizeof(b), 0); } while (n < 0 && errno == EINTR);
+    return n == sizeof(b);
+}
+
+static bool tcp_recv_byte(int fd, uint8_t * b) {
+    ssize_t n;
+    do { n = ::recv(fd, b, sizeof(*b), 0); } while (n < 0 && errno == EINTR);
+    return n == (ssize_t)sizeof(*b);
+}
+
+// Index of the GID on this port equal to the target, or -1. Thunderbolt GIDs are
+// RoCEv2 IPv4-mapped (::ffff:a.b.c.d), so this matches the local TCP address.
+static int rdma_match_gid(struct ibv_context * ctx, uint8_t port, int gid_tbl_len,
+                          const uint8_t * target, union ibv_gid * out) {
+    for (int i = 0; i < gid_tbl_len; i++) {
+        union ibv_gid g;
+        if (ibv_query_gid(ctx, port, i, &g) != 0) continue;
+        if (memcmp(g.raw, target, RDMA_GID_SIZE) != 0) continue;
+        if (out) *out = g;
+        return i;
+    }
+    return -1;
+}
+
+// First ACTIVE port on the device. Only a cabled, up Thunderbolt link reports
+// ACTIVE, and it is not always port 1, so the port cannot be hardcoded the way
+// the Linux path does. Returns 0 if none.
+static uint8_t rdma_first_active_port(struct ibv_context * ctx, struct ibv_port_attr * out) {
+    struct ibv_device_attr da;
+    if (ibv_query_device(ctx, &da) != 0) return 0;
+    for (uint8_t p = 1; p <= da.phys_port_cnt; p++) {
+        struct ibv_port_attr pa;
+        if (ibv_query_port(ctx, p, &pa) != 0) continue;
+        if (pa.state == IBV_PORT_ACTIVE) { if (out) *out = pa; return p; }
+    }
+    return 0;
+}
+
+// Called before the endpoints are exchanged: pick the local device facing this
+// peer, create a UC QP and register the frame rings. RDMA is point-to-point, so
+// the device is the one whose GID equals the bootstrap connection's local
+// address, i.e. the one cabled to the peer.
+std::unique_ptr<apple_rdma> apple_rdma::probe(int fd, const uint8_t * target_gid, uint8_t * caps) {
+    int ndev = 0;
+    ibv_device ** devs = ibv_get_device_list(&ndev);
+    if (!devs) return nullptr;
+
+    ibv_context * ctx = nullptr;
+    uint8_t port = 0;
+    struct ibv_port_attr pa = {};
+    union ibv_gid gid = {};
+    int gid_idx = -1;
+    std::string matched;
+    for (int d = 0; d < ndev; d++) {
+        ibv_context * c = ibv_open_device(devs[d]);
+        if (!c) continue;
+        struct ibv_port_attr p = {};
+        uint8_t pt = rdma_first_active_port(c, &p);
+        int gi = pt ? rdma_match_gid(c, pt, p.gid_tbl_len, target_gid, &gid) : -1;
+        if (gi < 0) { ibv_close_device(c); continue; }
+        ctx = c; port = pt; pa = p; gid_idx = gi;
+        const char * name = ibv_get_device_name(devs[d]);
+        matched = name ? name : "";
+        break;
+    }
+    ibv_free_device_list(devs);
+    if (!ctx) return nullptr;
+
+    std::unique_ptr<impl> c(new impl());
+    c->fd  = fd;
+    c->ctx = ctx;
+    c->port = port;
+    c->gid_idx = gid_idx;
+    c->path_mtu = pa.active_mtu;
+
+    c->pd = ibv_alloc_pd(ctx);
+    if (!c->pd) return nullptr;
+
+    c->cq = ibv_create_cq(ctx, 2 * RDMA_QP_WR + 1, nullptr, nullptr, 0);
+    if (!c->cq) return nullptr;
+
+    ibv_qp_init_attr qia = {};
+    qia.send_cq = c->cq;
+    qia.recv_cq = c->cq;
+    qia.qp_type = IBV_QPT_UC;
+    qia.cap.max_send_wr  = RDMA_QP_WR;
+    qia.cap.max_recv_wr  = RDMA_QP_WR;
+    qia.cap.max_send_sge = 1;
+    qia.cap.max_recv_sge = 1;
+    c->qp = ibv_create_qp(c->pd, &qia);
+    if (!c->qp) return nullptr;
+
+    {
+        ibv_qp_attr a = {};
+        a.qp_state = IBV_QPS_INIT;
+        a.pkey_index = 0;
+        a.port_num = port;
+        a.qp_access_flags = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_WRITE;
+        if (ibv_modify_qp(c->qp, &a,
+                IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS) != 0) {
+            return nullptr;
+        }
+    }
+
+    long page = sysconf(_SC_PAGESIZE);
+    if (page <= 0) page = 4096;
+    const size_t ring_bytes = (size_t)RDMA_NBUF * RDMA_STRIDE;
+    if (posix_memalign((void **)&c->send_mem, (size_t)page, ring_bytes) != 0) c->send_mem = nullptr;
+    if (posix_memalign((void **)&c->recv_mem, (size_t)page, ring_bytes) != 0) c->recv_mem = nullptr;
+    if (!c->send_mem || !c->recv_mem) return nullptr;
+
+    // Apple's provider rejects LOCAL_WRITE-only MRs even for two-sided SEND/RECV.
+    const int mr_flags = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_WRITE;
+    c->send_mr = ibv_reg_mr(c->pd, c->send_mem, ring_bytes, mr_flags);
+    c->recv_mr = ibv_reg_mr(c->pd, c->recv_mem, ring_bytes, mr_flags);
+    if (!c->send_mr || !c->recv_mr) return nullptr;
+
+    // Recvs are posted in activate() after the RTS transition, not here: Apple's
+    // provider rejects ibv_post_recv on a QP that has not reached RTS.
+
+    c->qpn = c->qp->qp_num;
+
+    apple_rdma_caps rc = {};
+    rc.qpn = c->qpn;
+    rc.lid = pa.lid;
+    memcpy(rc.gid, gid.raw, RDMA_GID_SIZE);
+    memcpy(caps, &rc, sizeof(rc));
+
+    GGML_LOG_INFO("RDMA(Apple/UC) probed: dev=%s port=%u gid=%d qpn=%u lid=%u mtu=%d ring=%d x %zu KiB\n",
+                  matched.c_str(), port, gid_idx, c->qpn, (unsigned)pa.lid, 128 << c->path_mtu,
+                  RDMA_NBUF, RDMA_STRIDE / 1024);
+    return std::unique_ptr<apple_rdma>(new apple_rdma(std::move(c)));
+}
+
+// Called once the peer's endpoint has arrived: INIT -> RTR -> RTS (UC: GID/GRH
+// addressing, no timeout/retry/rnr/rd_atomic), then the readiness handshake.
+bool apple_rdma::activate(const uint8_t * caps) {
+    impl * c = pimpl.get();
+
+    apple_rdma_caps rc = {};
+    memcpy(&rc, caps, sizeof(rc));
+
+    bool ok = true;
+    {
+        ibv_qp_attr a = {};
+        a.qp_state   = IBV_QPS_RTR;
+        a.path_mtu   = c->path_mtu;
+        a.rq_psn     = RDMA_PSN;
+        a.dest_qp_num = rc.qpn;
+        a.ah_attr.is_global     = 1;
+        a.ah_attr.port_num      = c->port;
+        a.ah_attr.sl            = 0;
+        a.ah_attr.src_path_bits = 0;
+        a.ah_attr.dlid          = rc.lid;
+        a.ah_attr.grh.hop_limit  = 1;
+        a.ah_attr.grh.sgid_index = (uint8_t)c->gid_idx;
+        memcpy(&a.ah_attr.grh.dgid, rc.gid, RDMA_GID_SIZE);
+        if (ibv_modify_qp(c->qp, &a,
+                IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN | IBV_QP_RQ_PSN) != 0) {
+            GGML_LOG_ERROR("RDMA(Apple/UC) RTR failed: %s\n", strerror(errno));
+            ok = false;
+        }
+    }
+    if (ok) {
+        ibv_qp_attr a = {};
+        a.qp_state = IBV_QPS_RTS;
+        a.sq_psn   = RDMA_PSN;
+        if (ibv_modify_qp(c->qp, &a, IBV_QP_STATE | IBV_QP_SQ_PSN) != 0) {
+            GGML_LOG_ERROR("RDMA(Apple/UC) RTS failed: %s\n", strerror(errno));
+            ok = false;
+        }
+    }
+
+    // Recvs are posted only now: the controller starts processing them at RTR.
+    for (int i = 0; ok && i < RDMA_NBUF; i++) {
+        if (!c->post_recv(i)) {
+            GGML_LOG_ERROR("RDMA(Apple/UC) post_recv %d/%d failed\n", i, RDMA_NBUF);
+            ok = false;
+        }
+    }
+
+    // A queue pair processes receives only after RTR and the transitions above can
+    // fail on one side alone, so neither peer sends a frame until both report their
+    // recvs posted.
+    uint8_t peer_ready = 0;
+    if (!tcp_send_byte(c->fd, ok ? RDMA_SYNC_READY : 0) || !tcp_recv_byte(c->fd, &peer_ready)) {
+        return false;
+    }
+    if (!ok || peer_ready != RDMA_SYNC_READY) {
+        return false;
+    }
+
+    GGML_LOG_INFO("RDMA(Apple/UC) activated: qpn=%u->%u mtu=%d rx_depth=%d\n",
+                  c->qpn, rc.qpn, 128 << c->path_mtu, RDMA_NBUF);
+    return true;
+}
+
+// Drain the CQ: release completed send buffers, queue completed recv frames for
+// the reader. Returns the number of completions reaped, or -1 on error.
+int apple_rdma::impl::progress() {
+    struct ibv_wc wc[RDMA_NBUF * 2];
+    int n = ibv_poll_cq(cq, RDMA_NBUF * 2, wc);
+    if (n < 0) { GGML_LOG_ERROR("RDMA(Apple/UC) poll_cq failed\n"); broken = true; return -1; }
+    for (int j = 0; j < n; j++) {
+        uint64_t id = wc[j].wr_id;
+        bool is_recv = (id & RDMA_RECV_WR) != 0;
+        if (wc[j].status != IBV_WC_SUCCESS) {
+            GGML_LOG_ERROR("RDMA(Apple/UC) %s wc error: status=%d\n", is_recv ? "recv" : "send", wc[j].status);
+            broken = true;
+            return -1;
+        }
+        if (is_recv) {
+            int b = (int)(id & RDMA_WR_IDX_MASK);
+            const rdma_seg_hdr * h = (const rdma_seg_hdr *)(recv_mem + (size_t)b * RDMA_STRIDE);
+            if (h->magic != RDMA_SEG_MAGIC) { GGML_LOG_ERROR("RDMA(Apple/UC) bad frame magic\n"); broken = true; return -1; }
+            if (h->len > RDMA_PAYLOAD) { GGML_LOG_ERROR("RDMA(Apple/UC) frame len %u exceeds payload\n", h->len); broken = true; return -1; }
+            int slot = (inq_head + inq_count) % RDMA_NBUF;
+            inq[slot].buf  = b;
+            inq[slot].off  = 0;
+            inq[slot].len  = h->len;
+            inq_count++;
+        } else {
+            send_busy[(int)(id & RDMA_WR_IDX_MASK)] = 0;
+        }
+    }
+    return n;
+}
+
+// Reserve a free send buffer to coalesce into, waiting on progress if none free.
+bool apple_rdma::impl::acquire_pending() {
+    if (pend_buf >= 0) return true;
+    for (;;) {
+        if (broken) return false;
+        for (int k = 0; k < RDMA_NBUF; k++) if (!send_busy[k]) { pend_buf = k; pend_len = 0; return true; }
+        if (progress() < 0) return false;
+    }
+}
+
+// Post the pending frame. The whole STRIDE goes out even when only partly filled:
+// TN3205 requires a SEND and its matching RECV to cover the same number of
+// Thunderbolt frames, so a short send would fail the peer's receive.
+bool apple_rdma::impl::post_pending() {
+    if (pend_buf < 0) return true;
+    int i = pend_buf;
+    rdma_seg_hdr * h = (rdma_seg_hdr *)(send_mem + (size_t)i * RDMA_STRIDE);
+    h->magic = RDMA_SEG_MAGIC;
+    h->len   = pend_len;
+    if (!post_send(i, RDMA_STRIDE)) { broken = true; return false; }
+    send_busy[i] = 1;
+    pend_buf = -1;
+    pend_len = 0;
+    return true;
+}
+
+// Coalescing write: append into the pending frame, posting a full frame when it
+// fills. The trailing partial is posted by flush() at each message boundary.
+bool apple_rdma::send(const void * data, size_t size) {
+    impl * c = pimpl.get();
+    const uint8_t * p = (const uint8_t *)data;
+    while (size > 0) {
+        if (c->broken) return false;
+        if (!c->acquire_pending()) return false;
+        uint8_t * sb = c->send_mem + (size_t)c->pend_buf * RDMA_STRIDE;
+        size_t space = RDMA_PAYLOAD - c->pend_len;
+        size_t chunk = size < space ? size : space;
+        memcpy(sb + sizeof(rdma_seg_hdr) + c->pend_len, p, chunk);
+        c->pend_len += (uint32_t)chunk;
+        p += chunk;
+        size -= chunk;
+        if (c->pend_len == RDMA_PAYLOAD) { if (!c->post_pending()) return false; }
+    }
+    return true;
+}
+
+bool apple_rdma::recv(void * data, size_t size) {
+    impl * c = pimpl.get();
+    uint8_t * p = (uint8_t *)data;
+    if (!c->post_pending()) return false;   // turnaround: flush the coalesced request
+    unsigned idle = 0;
+    while (size > 0) {
+        if (c->inq_count == 0) {
+            if (c->broken) return false;
+            int n = c->progress();
+            if (n < 0) return false;
+            if (n == 0) {
+                // UC gives no disconnect notification, so the bootstrap TCP fd is
+                // the liveness anchor: nothing crosses it once RDMA is up, so any
+                // readability means the peer's FIN (macOS has no POLLRDHUP).
+                // Same idle interval as the Linux path.
+                if ((++idle & 0xFFFFF) == 0) {
+                    struct pollfd pfd = { c->fd, POLLIN, 0 };
+                    if (poll(&pfd, 1, 0) > 0 &&
+                        (pfd.revents & (POLLIN | POLLHUP | POLLERR | POLLNVAL))) {
+                        return false;
+                    }
+                }
+            } else {
+                idle = 0;
+            }
+            continue;
+        }
+        idle = 0;
+        int slot = c->inq_head;
+        int b = c->inq[slot].buf;
+        uint32_t avail = c->inq[slot].len - c->inq[slot].off;
+        uint32_t take = (size < (size_t)avail) ? (uint32_t)size : avail;
+        memcpy(p, c->recv_mem + (size_t)b * RDMA_STRIDE + sizeof(rdma_seg_hdr) + c->inq[slot].off, take);
+        p += take;
+        size -= take;
+        c->inq[slot].off += take;
+        if (c->inq[slot].off == c->inq[slot].len) {
+            if (!c->post_recv(b)) { c->broken = true; return false; }
+            c->inq_head = (c->inq_head + 1) % RDMA_NBUF;
+            c->inq_count--;
+        }
+    }
+    return true;
+}
+
+bool apple_rdma::flush() {
+    return pimpl->post_pending();
+}
@@ -0,0 +1,27 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+struct apple_rdma {
+    // target_gid is 16 bytes in, caps is RPC_CONN_CAPS_SIZE bytes out.
+    static std::unique_ptr<apple_rdma> probe(int fd, const uint8_t * target_gid, uint8_t * caps);
+    ~apple_rdma();
+
+    // Peer endpoint from its caps, which must be non-zero: this blocks on a
+    // readiness handshake over fd that the peer only joins if it also has RDMA.
+    bool activate(const uint8_t * caps);
+
+    bool send(const void * data, size_t size);
+    bool recv(void * data, size_t size);
+    // Post the trailing partial frame; must be called at every message boundary.
+    bool flush();
+    // True once the connection has failed; the caller should drop the socket.
+    bool broken() const;
+
+private:
+    struct impl;
+    explicit apple_rdma(std::unique_ptr<impl> p);
+    std::unique_ptr<impl> pimpl;
+};
@@ -18,15 +18,20 @@
 #  include <unistd.h>
 #endif
 #include <cstdlib>
+#include <cstring>
 #include <mutex>
 #include <optional>

 #ifdef GGML_RPC_RDMA
 #  include <infiniband/verbs.h>
+#  include <array>
 #  include <time.h>
 #  ifndef _WIN32
 #    include <poll.h>
 #  endif
+#  ifdef GGML_RPC_RDMA_APPLE
+#    include "transport-apple.h"
+#  endif
 #endif // GGML_RPC_RDMA

 #ifdef _WIN32
@@ -42,10 +47,13 @@ static const char * RPC_DEBUG = std::getenv("GGML_RPC_DEBUG");
    do { if (RPC_DEBUG) GGML_LOG_DEBUG(__VA_ARGS__); } while (0)

 #ifdef GGML_RPC_RDMA
-static constexpr size_t RDMA_CHUNK    = 256 * 1024;   // 256 KiB per send/recv (fits default 8 MiB memlock)
-static constexpr int    RDMA_RX_DEPTH = 24;            // pre-posted recv ring: 24 × 256 KiB = 6 MiB
 static constexpr size_t RDMA_GID_SIZE = 16;            // RoCE GID / IB GID is always 16 bytes
 using rdma_gid_t = std::array<uint8_t, RDMA_GID_SIZE>;
+#endif // GGML_RPC_RDMA
+
+#if defined(GGML_RPC_RDMA) && !defined(GGML_RPC_RDMA_APPLE)
+static constexpr size_t RDMA_CHUNK    = 256 * 1024;   // 256 KiB per send/recv (fits default 8 MiB memlock)
+static constexpr int    RDMA_RX_DEPTH = 24;            // pre-posted recv ring: 24 × 256 KiB = 6 MiB

 struct rdma_conn {
    struct ibv_context * ctx = nullptr;
@@ -111,27 +119,33 @@ struct rdma_caps {

 static_assert(sizeof(rdma_caps) == RPC_CONN_CAPS_SIZE, "rdma_caps must match conn_caps size");

-#endif // GGML_RPC_RDMA
+#endif // GGML_RPC_RDMA && !GGML_RPC_RDMA_APPLE

 struct socket_t::impl {
    impl(sockfd_t fd) : use_rdma(false), fd(fd) {}
    ~impl();
    bool send_data(const void * data, size_t size);
    bool recv_data(void * data, size_t size);
+    bool flush();
    void get_caps(uint8_t * local_caps);
    void update_caps(const uint8_t * remote_caps);

 #ifdef GGML_RPC_RDMA
-    bool tcp_peer_closed();
    std::optional<rdma_gid_t> rdma_build_target_gid();
+
+#  ifdef GGML_RPC_RDMA_APPLE
+    std::unique_ptr<apple_rdma> rdma;
+#  else
    bool rdma_probe();
-    bool rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, const uint8_t * remote_gid);
-    bool rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc);
    bool rdma_send(const void * data, size_t size);
    bool rdma_recv(void * data, size_t size);
+    bool tcp_peer_closed();
+    bool rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, const uint8_t * remote_gid);
+    bool rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc);

    std::unique_ptr<rdma_conn> rdma;
    rdma_local_info            rdma_local = {};
+#  endif
 #endif // GGML_RPC_RDMA
    bool     use_rdma;
    sockfd_t fd;
@@ -151,17 +165,6 @@ socket_t::impl::~impl() {

 #ifdef GGML_RPC_RDMA

-bool socket_t::impl::tcp_peer_closed() {
-    if (fd < 0) return false;
-#ifndef _WIN32
-    struct pollfd pfd = { fd, POLLIN | POLLRDHUP, 0 };
-    int r = poll(&pfd, 1, 0);
-    return r > 0 && (pfd.revents & (POLLHUP | POLLERR | POLLRDHUP));
-#else
-    return false;
-#endif
-}
-
 // Build a RoCE GID-shaped 16-byte target from a TCP socket's local address.
 // Used to match the socket's local IP against the kernel's GID table so that
 // a single memcmp handles IPv4, IPv4-mapped IPv6, and native IPv6 uniformly:
@@ -191,6 +194,19 @@ std::optional<rdma_gid_t> socket_t::impl::rdma_build_target_gid() {
    return std::nullopt;
 }

+#ifndef GGML_RPC_RDMA_APPLE
+
+bool socket_t::impl::tcp_peer_closed() {
+    if (fd < 0) return false;
+#ifndef _WIN32
+    struct pollfd pfd = { fd, POLLIN | POLLRDHUP, 0 };
+    int r = poll(&pfd, 1, 0);
+    return r > 0 && (pfd.revents & (POLLHUP | POLLERR | POLLRDHUP));
+#else
+    return false;
+#endif
+}
+
 bool socket_t::impl::rdma_probe() {
    const char * dev_env = std::getenv("GGML_RDMA_DEV");
    const char * gid_env = std::getenv("GGML_RDMA_GID");
@@ -457,10 +473,16 @@ bool socket_t::impl::rdma_recv(void * data, size_t size) {
    return true;
 }

+#endif // !GGML_RPC_RDMA_APPLE (Linux RC transport)
+
 #endif // GGML_RPC_RDMA

 bool socket_t::impl::send_data(const void * data, size_t size) {
-#ifdef GGML_RPC_RDMA
+#ifdef GGML_RPC_RDMA_APPLE
+    if (use_rdma) {
+        return rdma->send(data, size);
+    }
+#elif defined(GGML_RPC_RDMA)
    if (use_rdma) {
        return rdma_send(data, size);
    }
@@ -480,7 +502,11 @@ bool socket_t::impl::send_data(const void * data, size_t size) {
 }

 bool socket_t::impl::recv_data(void * data, size_t size) {
-#ifdef GGML_RPC_RDMA
+#ifdef GGML_RPC_RDMA_APPLE
+    if (use_rdma) {
+        return rdma->recv(data, size);
+    }
+#elif defined(GGML_RPC_RDMA)
    if (use_rdma) {
        return rdma_recv(data, size);
    }
@@ -506,6 +532,15 @@ bool socket_t::impl::recv_data(void * data, size_t size) {
 void socket_t::impl::get_caps(uint8_t * local_caps) {
    memset(local_caps, 0, RPC_CONN_CAPS_SIZE);
 #ifdef GGML_RPC_RDMA
+    if (std::getenv("GGML_RPC_NO_RDMA")) {
+        return;
+    }
+#  ifdef GGML_RPC_RDMA_APPLE
+    auto target_gid = rdma_build_target_gid();
+    if (target_gid) {
+        rdma = apple_rdma::probe(fd, target_gid->data(), local_caps);
+    }
+#  else
    rdma_local = {};
    if (rdma_probe()) {
        rdma_caps rc = {};
@@ -516,21 +551,30 @@ void socket_t::impl::get_caps(uint8_t * local_caps) {
    } else {
        rdma.reset();
    }
+#  endif
 #endif // GGML_RPC_RDMA
 }

 void socket_t::impl::update_caps(const uint8_t * remote_caps) {
 #ifdef GGML_RPC_RDMA
-    if (!rdma) {
-        return;
+    // a peer that has no RDMA advertises all-zero caps and takes no further part
+    // in the negotiation, so drop to TCP without reporting a failure
+    bool remote_rdma = false;
+    for (size_t i = 0; i < RPC_CONN_CAPS_SIZE; i++) {
+        remote_rdma |= remote_caps[i] != 0;
    }
-    rdma_caps rc = {};
-    memcpy(&rc, remote_caps, sizeof(rc));
-    if (rc.qpn == 0) {
+    if (!rdma || !remote_rdma) {
        rdma.reset();
        return;
    }
-    if (rdma_activate(rc.qpn, rc.psn, rc.gid)) {
+#  ifdef GGML_RPC_RDMA_APPLE
+    bool activated = rdma->activate(remote_caps);
+#  else
+    rdma_caps rc = {};
+    memcpy(&rc, remote_caps, sizeof(rc));
+    bool activated = rdma_activate(rc.qpn, rc.psn, rc.gid);
+#  endif
+    if (activated) {
        use_rdma = true;
    } else {
        GGML_LOG_ERROR("RDMA activate failed, staying on TCP\n");
@@ -541,6 +585,14 @@ void socket_t::impl::update_caps(const uint8_t * remote_caps) {
 #endif // GGML_RPC_RDMA
 }

+bool socket_t::impl::flush() {
+#ifdef GGML_RPC_RDMA_APPLE
+    if (use_rdma) {
+        return rdma->flush();
+    }
+#endif
+    return true;
+}

 /////////////////////////////////////////////////////////////////////////////

@@ -556,6 +608,10 @@ bool socket_t::recv_data(void * data, size_t size) {
    return pimpl->recv_data(data, size);
 }

+bool socket_t::flush() {
+    return pimpl->flush();
+}
+
 void socket_t::get_caps(uint8_t * local_caps) {
    return pimpl->get_caps(local_caps);
 }
@@ -15,6 +15,10 @@ struct socket_t {

    bool send_data(const void * data, size_t size);
    bool recv_data(void * data, size_t size);
+    // Must be called at every message boundary: the RDMA transport coalesces
+    // writes into fixed-size frames and posts the trailing partial frame only
+    // here. No-op on TCP.
+    bool flush();

    socket_ptr accept();

@@ -1042,6 +1042,8 @@ struct vk_device_struct {
    vk_pipeline pipeline_argsort_large_f32[num_argsort_pipelines];
    vk_pipeline pipeline_topk_f32[num_topk_pipelines];
    vk_pipeline pipeline_sum_rows_f32;
+    vk_pipeline pipeline_cross_entropy_loss_f32, pipeline_cross_entropy_loss_f32_wg512;
+    vk_pipeline pipeline_cross_entropy_loss_back_f32, pipeline_cross_entropy_loss_back_f32_wg512;
    vk_pipeline pipeline_fwht_f32[4];
    vk_pipeline pipeline_cumsum_f32;
    vk_pipeline pipeline_cumsum_small_f32;
@@ -4169,10 +4171,16 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
    const uint32_t subgroup_size_16 = std::max(device->subgroup_size, 16u);
    const uint32_t subgroup_size_32 = std::max(device->subgroup_size, 32u);

+    // clamp WARP for l_/m_ warptiles so WM <= BM (breaks on subgroupSize > 64)
+    const uint32_t mm_warp_8  = std::min(subgroup_size_8,  64u);
+    const uint32_t mm_warp_16 = std::min(subgroup_size_16, 64u);
+
    const uint32_t mul_mat_subgroup_size = (device->vendor_id == VK_VENDOR_ID_INTEL && device->subgroup_size_control) ? device->subgroup_min_size : device->subgroup_size;
    const uint32_t mul_mat_subgroup_size_8 = std::max(mul_mat_subgroup_size, 8u);
    const uint32_t mul_mat_subgroup_size_16 = std::max(mul_mat_subgroup_size, 16u);
    const uint32_t mul_mat_subgroup_size_32 = std::max(mul_mat_subgroup_size, 32u);
+    const uint32_t mul_mat_mm_warp_8  = std::min(mul_mat_subgroup_size_8,  64u);
+    const uint32_t mul_mat_mm_warp_16 = std::min(mul_mat_subgroup_size_16, 64u);

    const bool subgroup_min_size_16 = (!device->subgroup_size_control && device->subgroup_size >= 16) ||
                                      (device->subgroup_size_control && device->subgroup_max_size >= 16);
@@ -4253,39 +4261,39 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {

        const uint32_t s_warptile_wm = device->subgroup_size == 8 ? 8 : 32;

-        l_warptile = { 128,             128, 128, 16, subgroup_size_8 * 2, 64, 2, tm_l, tn_l, tk_l, subgroup_size_8 };
-        m_warptile = { 128,              64,  64, 16, subgroup_size_8,     32, 2, tm_m, tn_m, tk_m, subgroup_size_8 };
-        s_warptile = { subgroup_size_32, 32,  32, 16, s_warptile_wm,       32, 2, tm_s, tn_s, tk_s, subgroup_size_8 };
+        l_warptile = { 128,             128, 128, 16, mm_warp_8 * 2, 64, 2, tm_l, tn_l, tk_l, mm_warp_8 };
+        m_warptile = { 128,              64,  64, 16, mm_warp_8,     32, 2, tm_m, tn_m, tk_m, mm_warp_8 };
+        s_warptile = { subgroup_size_32, 32,  32, 16, s_warptile_wm, 32, 2, tm_s, tn_s, tk_s, subgroup_size_8 };

-        l_warptile_mmq = { 128,             128, 128, 32, subgroup_size_8 * 2, 64, 2, tm_l, tn_l, tk_l, subgroup_size_8 };
-        m_warptile_mmq = { 128,              64,  64, 32, subgroup_size_8,     32, 2, tm_m, tn_m, tk_m, subgroup_size_8 };
-        s_warptile_mmq = { subgroup_size_32, 32,  32, 32, s_warptile_wm,       32, 2, tm_s, tn_s, tk_s, subgroup_size_8 };
+        l_warptile_mmq = { 128,             128, 128, 32, mm_warp_8 * 2, 64, 2, tm_l, tn_l, tk_l, mm_warp_8 };
+        m_warptile_mmq = { 128,              64,  64, 32, mm_warp_8,     32, 2, tm_m, tn_m, tk_m, mm_warp_8 };
+        s_warptile_mmq = { subgroup_size_32, 32,  32, 32, s_warptile_wm, 32, 2, tm_s, tn_s, tk_s, subgroup_size_8 };

        // Integer MMQ has a smaller shared memory profile, but heavier register use
-        l_warptile_mmq_int = { 128,             128, 128, 32, subgroup_size_8 * 2, 64, 2, 4, 4, 1, subgroup_size_8 };
-        m_warptile_mmq_int = { 128,              64,  64, 32, subgroup_size_8,     32, 2, 2, 2, 1, subgroup_size_8 };
-        s_warptile_mmq_int = { subgroup_size_32, 32,  32, 32, s_warptile_wm,       32, 2, 2, 1, 1, subgroup_size_8 };
+        l_warptile_mmq_int = { 128,             128, 128, 32, mm_warp_8 * 2, 64, 2, 4, 4, 1, mm_warp_8 };
+        m_warptile_mmq_int = { 128,              64,  64, 32, mm_warp_8,     32, 2, 2, 2, 1, mm_warp_8 };
+        s_warptile_mmq_int = { subgroup_size_32, 32,  32, 32, s_warptile_wm, 32, 2, 2, 1, 1, subgroup_size_8 };

        // K-quants use even more registers, mitigate by setting WMITER to 1
-        l_warptile_mmq_int_k = { 128,               128, 128, 32, subgroup_size_8 * 2, 64, 1, 4, 4, 1, subgroup_size_8 };
-        m_warptile_mmq_int_k = { 128,                64,  64, 32, subgroup_size_8,     32, 1, 2, 2, 1, subgroup_size_8 };
-        s_warptile_mmq_int_k = { subgroup_size_32,   32,  32, 32, s_warptile_wm,       32, 1, 2, 1, 1, subgroup_size_8 };
+        l_warptile_mmq_int_k = { 128,               128, 128, 32, mm_warp_8 * 2, 64, 1, 4, 4, 1, mm_warp_8 };
+        m_warptile_mmq_int_k = { 128,                64,  64, 32, mm_warp_8,     32, 1, 2, 2, 1, mm_warp_8 };
+        s_warptile_mmq_int_k = { subgroup_size_32,   32,  32, 32, s_warptile_wm, 32, 1, 2, 1, 1, subgroup_size_8 };

-        l_warptile_id = { 128,                      128, 128, 16, mul_mat_subgroup_size_16 * 2, 64, 2, tm_l, tn_l, tk_l, mul_mat_subgroup_size_16 };
-        m_warptile_id = { 128,                       64,  64, 16, mul_mat_subgroup_size_16,     32, 2, tm_m, tn_m, tk_m, mul_mat_subgroup_size_16 };
-        s_warptile_id = { mul_mat_subgroup_size_16,  32,  32, 16, s_warptile_wm,                32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_16 };
+        l_warptile_id = { 128,                      128, 128, 16, mul_mat_mm_warp_16 * 2, 64, 2, tm_l, tn_l, tk_l, mul_mat_mm_warp_16 };
+        m_warptile_id = { 128,                       64,  64, 16, mul_mat_mm_warp_16,     32, 2, tm_m, tn_m, tk_m, mul_mat_mm_warp_16 };
+        s_warptile_id = { mul_mat_subgroup_size_16,  32,  32, 16, s_warptile_wm,          32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_16 };

-        l_warptile_mmqid = { 128,                       128, 128, 32, mul_mat_subgroup_size_8 * 2, 64, 2, tm_l, tn_l, tk_l, mul_mat_subgroup_size_8 };
-        m_warptile_mmqid = { 128,                        64,  64, 32, mul_mat_subgroup_size_8,     32, 2, tm_m, tn_m, tk_m, mul_mat_subgroup_size_8 };
-        s_warptile_mmqid = { mul_mat_subgroup_size_32,   32,  32, 32, s_warptile_wm,               32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_8 };
+        l_warptile_mmqid = { 128,                       128, 128, 32, mul_mat_mm_warp_8 * 2, 64, 2, tm_l, tn_l, tk_l, mul_mat_mm_warp_8 };
+        m_warptile_mmqid = { 128,                        64,  64, 32, mul_mat_mm_warp_8,     32, 2, tm_m, tn_m, tk_m, mul_mat_mm_warp_8 };
+        s_warptile_mmqid = { mul_mat_subgroup_size_32,   32,  32, 32, s_warptile_wm,         32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_8 };

-        l_warptile_mmqid_int = { 128,                       128, 128, 32, mul_mat_subgroup_size_8 * 2, 64, 2, 4, 4, 1, mul_mat_subgroup_size_8 };
-        m_warptile_mmqid_int = { 128,                        64,  64, 32, mul_mat_subgroup_size_8,     32, 2, 2, 2, 1, mul_mat_subgroup_size_8 };
-        s_warptile_mmqid_int = { mul_mat_subgroup_size_32,   32,  32, 32, s_warptile_wm,               32, 2, 2, 1, 1, mul_mat_subgroup_size_8 };
+        l_warptile_mmqid_int = { 128,                       128, 128, 32, mul_mat_mm_warp_8 * 2, 64, 2, 4, 4, 1, mul_mat_mm_warp_8 };
+        m_warptile_mmqid_int = { 128,                        64,  64, 32, mul_mat_mm_warp_8,     32, 2, 2, 2, 1, mul_mat_mm_warp_8 };
+        s_warptile_mmqid_int = { mul_mat_subgroup_size_32,   32,  32, 32, s_warptile_wm,         32, 2, 2, 1, 1, mul_mat_subgroup_size_8 };

-        l_warptile_mmqid_int_k = { 128,                     128, 128, 32, mul_mat_subgroup_size_16 * 2, 64, 1, 4, 4, 1, mul_mat_subgroup_size_16 };
-        m_warptile_mmqid_int_k = { 128,                      64,  64, 32, mul_mat_subgroup_size_16,     32, 1, 2, 2, 1, mul_mat_subgroup_size_16 };
-        s_warptile_mmqid_int_k = { mul_mat_subgroup_size_32, 32,  32, 32, s_warptile_wm,                32, 1, 2, 1, 1, mul_mat_subgroup_size_16 };
+        l_warptile_mmqid_int_k = { 128,                     128, 128, 32, mul_mat_mm_warp_16 * 2, 64, 1, 4, 4, 1, mul_mat_mm_warp_16 };
+        m_warptile_mmqid_int_k = { 128,                      64,  64, 32, mul_mat_mm_warp_16,     32, 1, 2, 2, 1, mul_mat_mm_warp_16 };
+        s_warptile_mmqid_int_k = { mul_mat_subgroup_size_32, 32,  32, 32, s_warptile_wm,          32, 1, 2, 1, 1, mul_mat_subgroup_size_16 };

        // chip specific tuning
        if ((device->architecture == AMD_GCN) && (device->driver_id != vk::DriverId::eAmdProprietary)) {
@@ -4293,13 +4301,13 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
            m_warptile_mmqid = m_warptile_mmqid_int = { 256, 64, 64, 32, 16, 16, 2, 2, 2, 1, 16 };
        } else if (device->vendor_id == VK_VENDOR_ID_AMD && device->coopmat_support && device->driver_id != vk::DriverId::eAmdProprietary) {
            // This is intentionally using tx_m values, slight performance increase
-            l_warptile = { 256, 128, 128, 16, subgroup_size_8, 64, 2, tm_m, tn_m, tk_m, subgroup_size_8 };
-            l_warptile_mmq = l_warptile_mmq_int = { 256, 128, 128, 32, subgroup_size_8, 64, 2, tm_m, tn_m, tk_m, subgroup_size_8 };
-            l_warptile_mmq_int_k = { 256, 128, 128, 32, subgroup_size_16, 64, 1, 4, 2, 1, subgroup_size_16 };
+            l_warptile = { 256, 128, 128, 16, mm_warp_8, 64, 2, tm_m, tn_m, tk_m, mm_warp_8 };
+            l_warptile_mmq = l_warptile_mmq_int = { 256, 128, 128, 32, mm_warp_8, 64, 2, tm_m, tn_m, tk_m, mm_warp_8 };
+            l_warptile_mmq_int_k = { 256, 128, 128, 32, mm_warp_16, 64, 1, 4, 2, 1, mm_warp_16 };
        } else if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support) {
            // Xe2/Xe3 with coopmat enabled - warptile performance tuning
-            l_warptile = { 512, 128, 128, 16, subgroup_size_8, 32, 2, tm_m, tn_m, tk_m, subgroup_size_8 };
-            l_warptile_mmq = { 512, 128, 128, 32, subgroup_size_8, 32, 2, tm_m, tn_m, tk_m, subgroup_size_8 };
+            l_warptile = { 512, 128, 128, 16, mm_warp_8, 32, 2, tm_m, tn_m, tk_m, mm_warp_8 };
+            l_warptile_mmq = { 512, 128, 128, 32, mm_warp_8, 32, 2, tm_m, tn_m, tk_m, mm_warp_8 };
        }

        l_mmq_wg_denoms = l_wg_denoms = {128, 128, 1 };
@@ -5172,8 +5180,8 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
        const uint32_t s_warptile_wm = device->subgroup_size == 8 ? 8 : 32;

        // use scalar tile sizes
-        l_warptile = { 128, 128, 128, 16, subgroup_size_8 * 2, 64, 2, 4, 4, 1, subgroup_size_8 };
-        m_warptile = { 128,  64,  64, 16, subgroup_size_8, 32, 2, 4, 2, 1, subgroup_size_8 };
+        l_warptile = { 128, 128, 128, 16, mm_warp_8 * 2, 64, 2, 4, 4, 1, mm_warp_8 };
+        m_warptile = { 128,  64,  64, 16, mm_warp_8, 32, 2, 4, 2, 1, mm_warp_8 };
        s_warptile = { subgroup_size_32, 32, 32, 16, s_warptile_wm, 32, 2, 2, 2, 1, subgroup_size_8 };

        l_wg_denoms = {128, 128, 1 };
@@ -5758,6 +5766,10 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
    ggml_vk_create_pipeline(device, device->pipeline_argmax_f32, "argmax_f32", argmax_f32_len, argmax_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);

    ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_sum_rows_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_f32, "cross_entropy_loss_f32", cross_entropy_loss_f32_len, cross_entropy_loss_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_f32_wg512, "cross_entropy_loss_f32_wg512", cross_entropy_loss_f32_len, cross_entropy_loss_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, { 512 }, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_back_f32, "cross_entropy_loss_back_f32", cross_entropy_loss_back_f32_len, cross_entropy_loss_back_f32_data, "main", 4, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_back_f32_wg512, "cross_entropy_loss_back_f32_wg512", cross_entropy_loss_back_f32_len, cross_entropy_loss_back_f32_data, "main", 4, sizeof(vk_op_push_constants), {1, 1, 1}, { 512 }, 1);
    // Intel Windows driver in range [32.0.101.8509, 32.0.101.8860) will crash when using fwht kernels so we gate that here
    const bool can_use_fwht = device->driver_id != vk::DriverId::eIntelProprietaryWindows ||
        !ggml_vk_intel_windows_driver_in_range(device->properties.driverVersion, 101, 8509, 101, 8860);
@@ -11577,6 +11589,17 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
            return ctx->device->pipeline_sum_rows_f32;
        }
        return nullptr;
+    case GGML_OP_CROSS_ENTROPY_LOSS:
+        if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+            return src0->ne[0] > 1024 ? ctx->device->pipeline_cross_entropy_loss_f32_wg512 : ctx->device->pipeline_cross_entropy_loss_f32;
+        }
+        return nullptr;
+    case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
+        // src0 is the scalar grad; src1 is logits
+        if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && src2 && src2->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+            return src1->ne[0] > 1024 ? ctx->device->pipeline_cross_entropy_loss_back_f32_wg512 : ctx->device->pipeline_cross_entropy_loss_back_f32;
+        }
+        return nullptr;
    case GGML_OP_CUMSUM:
        if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
            if (src0->ne[0] <= 512) {
@@ -13942,6 +13965,103 @@ static void ggml_vk_cumsum(ggml_backend_vk_context * ctx, vk_context& subctx, co
    ctx->prealloc_split_k_need_sync = true;
 }

+static std::array<uint32_t, 3> ggml_vk_nrows_elements(uint32_t nr) {
+    if (nr > 262144) {
+        return { 512, 512, CEIL_DIV(nr, 262144) };
+    }
+    if (nr > 512) {
+        return { 512, CEIL_DIV(nr, 512), 1 };
+    }
+    return { nr, 1, 1 };
+}
+
+static void ggml_vk_cross_entropy_loss(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
+
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_is_contiguous(src0));
+    GGML_ASSERT(ggml_is_contiguous(src1));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+    GGML_ASSERT(ggml_are_same_shape(src0, src1));
+    GGML_ASSERT(ggml_is_scalar(dst));
+
+    const uint32_t nclasses = (uint32_t)src0->ne[0];
+    const uint32_t nrows    = (uint32_t)ggml_nrows(src0);
+
+    vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, nullptr, dst, GGML_OP_CROSS_ENTROPY_LOSS);
+    GGML_ASSERT(pipeline != nullptr);
+
+    ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+    ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_sum_rows_f32, 1);
+
+    vk_subbuffer src0_buf = ggml_vk_tensor_subbuffer(ctx, src0);
+    vk_subbuffer src1_buf = ggml_vk_tensor_subbuffer(ctx, src1);
+    vk_subbuffer dst_buf  = ggml_vk_tensor_subbuffer(ctx, dst, true);
+
+    const vk_op_push_constants pc = { nclasses, nrows, 0.0f, 0.0f, 0.0f, 0.0f };
+
+    const size_t tmp_size = (size_t)nrows * sizeof(float);
+    if (ctx->prealloc_size_x < tmp_size) {
+        ctx->prealloc_size_x = tmp_size;
+        ggml_vk_preallocate_buffers(ctx, subctx);
+    }
+    if (ctx->prealloc_x_need_sync) {
+        ggml_vk_sync_buffers(ctx, subctx);
+    }
+
+    vk_subbuffer tmp_buf = { ctx->prealloc_x, 0, tmp_size };
+    ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { src0_buf, src1_buf, tmp_buf }, pc, ggml_vk_nrows_elements(nrows));
+    ggml_vk_sync_buffers(ctx, subctx);
+
+    vk_op_sum_rows_push_constants sp = {};
+    sp.n_cols = nrows;
+    sp.ne01 = 1;
+    sp.ne02 = 1;
+    sp.weight = 1.0f;
+    init_pushconst_fastdiv(sp);
+    sp.misalign_offsets = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type);
+
+    ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_sum_rows_f32, { tmp_buf, dst_buf }, sp, { 1, 1, 1 });
+    ctx->prealloc_x_need_sync = true;
+}
+
+static void ggml_vk_cross_entropy_loss_back(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
+    const ggml_tensor * grad   = dst->src[0];
+    const ggml_tensor * logits = dst->src[1];
+    const ggml_tensor * labels = dst->src[2];
+
+    GGML_ASSERT(grad->type   == GGML_TYPE_F32);
+    GGML_ASSERT(logits->type == GGML_TYPE_F32);
+    GGML_ASSERT(labels->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type    == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_is_scalar(grad));
+    GGML_ASSERT(ggml_is_contiguous(grad));
+    GGML_ASSERT(ggml_is_contiguous(logits));
+    GGML_ASSERT(ggml_is_contiguous(labels));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+    GGML_ASSERT(ggml_are_same_shape(logits, labels));
+    GGML_ASSERT(ggml_are_same_shape(logits, dst));
+
+    const uint32_t nclasses = (uint32_t)logits->ne[0];
+    const uint32_t nrows    = (uint32_t)ggml_nrows(logits);
+
+    vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, grad, logits, labels, dst, GGML_OP_CROSS_ENTROPY_LOSS_BACK);
+    GGML_ASSERT(pipeline != nullptr);
+
+    ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+
+    vk_subbuffer grad_buf   = ggml_vk_tensor_subbuffer(ctx, grad);
+    vk_subbuffer logits_buf = ggml_vk_tensor_subbuffer(ctx, logits);
+    vk_subbuffer labels_buf = ggml_vk_tensor_subbuffer(ctx, labels);
+    vk_subbuffer dst_buf    = ggml_vk_tensor_subbuffer(ctx, dst);
+
+    const vk_op_push_constants pc = { nclasses, nrows, 0.0f, 0.0f, 0.0f, 0.0f };
+    ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { grad_buf, logits_buf, labels_buf, dst_buf }, pc, ggml_vk_nrows_elements(nrows));
+}
+
 static void ggml_vk_argmax(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
    ggml_vk_op_f32<vk_op_push_constants>(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_ARGMAX, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], 0.0f, 0.0f, 0.0f, 0.0f });
 }
@@ -15687,6 +15807,14 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
    case GGML_OP_ARGMAX:
        ggml_vk_argmax(ctx, compute_ctx, src0, node);

+        break;
+    case GGML_OP_CROSS_ENTROPY_LOSS:
+        ggml_vk_cross_entropy_loss(ctx, compute_ctx, node);
+
+        break;
+    case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
+        ggml_vk_cross_entropy_loss_back(ctx, compute_ctx, node);
+
        break;
    case GGML_OP_COUNT_EQUAL:
        ggml_vk_count_equal(ctx, compute_ctx, src0, src1, node);
@@ -18511,6 +18639,18 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
            }
        case GGML_OP_ARGMAX:
            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
+        case GGML_OP_CROSS_ENTROPY_LOSS:
+            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32
+                && ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32
+                && ggml_are_same_shape(op->src[0], op->src[1])
+                && ggml_is_contiguous(op) && ggml_is_scalar(op) && op->type == GGML_TYPE_F32;
+        case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
+            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32 && ggml_is_scalar(op->src[0])
+                && ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32
+                && ggml_is_contiguous(op->src[2]) && op->src[2]->type == GGML_TYPE_F32
+                && ggml_are_same_shape(op->src[1], op->src[2])
+                && ggml_are_same_shape(op->src[1], op)
+                && ggml_is_contiguous(op) && op->type == GGML_TYPE_F32;
        case GGML_OP_COUNT_EQUAL:
            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_I32
                && ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_I32;
@@ -19437,6 +19577,10 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
            tensor_clone = ggml_mean(ggml_ctx, src_clone[0]);
        } else if (tensor->op == GGML_OP_ARGMAX) {
            tensor_clone = ggml_argmax(ggml_ctx, src_clone[0]);
+        } else if (tensor->op == GGML_OP_CROSS_ENTROPY_LOSS) {
+            tensor_clone = ggml_cross_entropy_loss(ggml_ctx, src_clone[0], src_clone[1]);
+        } else if (tensor->op == GGML_OP_CROSS_ENTROPY_LOSS_BACK) {
+            tensor_clone = ggml_cross_entropy_loss_back(ggml_ctx, src_clone[0], src_clone[1], src_clone[2]);
        } else if (tensor->op == GGML_OP_COUNT_EQUAL) {
            tensor_clone = ggml_count_equal(ggml_ctx, src_clone[0], src_clone[1]);
        } else if (tensor->op == GGML_OP_SOLVE_TRI) {
@@ -0,0 +1,78 @@
+#version 450
+
+#include "generic_head.glsl"
+#include "types.glsl"
+
+#extension GL_EXT_control_flow_attributes : enable
+
+layout(constant_id = 0) const uint BLOCK_SIZE = 32;
+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
+layout (binding = 1) readonly buffer B {B_TYPE data_b[];};
+layout (binding = 2) writeonly buffer D {D_TYPE data_d[];};
+
+shared FLOAT_TYPE tmp[BLOCK_SIZE];
+
+FLOAT_TYPE wg_reduce_max(FLOAT_TYPE v) {
+    const uint tid = gl_LocalInvocationID.x;
+    tmp[tid] = v;
+    barrier();
+    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            tmp[tid] = max(tmp[tid], tmp[tid + s]);
+        }
+        barrier();
+    }
+    v = tmp[0];
+    barrier();
+    return v;
+}
+
+FLOAT_TYPE wg_reduce_sum(FLOAT_TYPE v) {
+    const uint tid = gl_LocalInvocationID.x;
+    tmp[tid] = v;
+    barrier();
+    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            tmp[tid] += tmp[tid + s];
+        }
+        barrier();
+    }
+    v = tmp[0];
+    barrier();
+    return v;
+}
+
+void main() {
+    const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
+    const uint tid = gl_LocalInvocationID.x;
+
+    if (row >= p.KY) {
+        return;
+    }
+
+    const uint off = row * p.KX;
+
+    FLOAT_TYPE max_logit = FLOAT_TYPE(uintBitsToFloat(0xFF800000));
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        max_logit = max(max_logit, FLOAT_TYPE(data_a[off + i]));
+    }
+    max_logit = wg_reduce_max(max_logit);
+
+    FLOAT_TYPE sum_exp = FLOAT_TYPE(0.0f);
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        sum_exp += exp(FLOAT_TYPE(data_a[off + i]) - max_logit);
+    }
+    const FLOAT_TYPE log_sum = log(wg_reduce_sum(sum_exp));
+
+    FLOAT_TYPE loss = FLOAT_TYPE(0.0f);
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        loss += (FLOAT_TYPE(data_a[off + i]) - max_logit - log_sum) * FLOAT_TYPE(data_b[off + i]);
+    }
+    loss = -wg_reduce_sum(loss) / FLOAT_TYPE(p.KY);
+
+    if (tid == 0) {
+        data_d[row] = D_TYPE(loss);
+    }
+}
@@ -0,0 +1,75 @@
+#version 450
+
+#include "generic_head.glsl"
+#include "types.glsl"
+
+#extension GL_EXT_control_flow_attributes : enable
+
+layout(constant_id = 0) const uint BLOCK_SIZE = 32;
+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer G {A_TYPE data_g[];};
+layout (binding = 1) readonly buffer X {B_TYPE data_x[];};
+layout (binding = 2) readonly buffer Y {B_TYPE data_y[];};
+layout (binding = 3) writeonly buffer D {D_TYPE data_d[];};
+
+shared FLOAT_TYPE tmp[BLOCK_SIZE];
+
+FLOAT_TYPE wg_reduce_max(FLOAT_TYPE v) {
+    const uint tid = gl_LocalInvocationID.x;
+    tmp[tid] = v;
+    barrier();
+    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            tmp[tid] = max(tmp[tid], tmp[tid + s]);
+        }
+        barrier();
+    }
+    v = tmp[0];
+    barrier();
+    return v;
+}
+
+FLOAT_TYPE wg_reduce_sum(FLOAT_TYPE v) {
+    const uint tid = gl_LocalInvocationID.x;
+    tmp[tid] = v;
+    barrier();
+    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            tmp[tid] += tmp[tid + s];
+        }
+        barrier();
+    }
+    v = tmp[0];
+    barrier();
+    return v;
+}
+
+void main() {
+    const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
+    const uint tid = gl_LocalInvocationID.x;
+
+    if (row >= p.KY) {
+        return;
+    }
+
+    const uint off = row * p.KX;
+    const FLOAT_TYPE d_by_nrows = FLOAT_TYPE(data_g[0]) / FLOAT_TYPE(p.KY);
+
+    FLOAT_TYPE max_logit = FLOAT_TYPE(uintBitsToFloat(0xFF800000));
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        max_logit = max(max_logit, FLOAT_TYPE(data_x[off + i]));
+    }
+    max_logit = wg_reduce_max(max_logit);
+
+    FLOAT_TYPE sum_exp = FLOAT_TYPE(0.0f);
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        sum_exp += exp(FLOAT_TYPE(data_x[off + i]) - max_logit);
+    }
+    const FLOAT_TYPE inv_sum = FLOAT_TYPE(1.0f) / wg_reduce_sum(sum_exp);
+
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        const FLOAT_TYPE sm = exp(FLOAT_TYPE(data_x[off + i]) - max_logit) * inv_sum;
+        data_d[off + i] = D_TYPE((sm - FLOAT_TYPE(data_y[off + i])) * d_by_nrows);
+    }
+}
@@ -1029,6 +1029,8 @@ void process_shaders() {

    string_to_spv("argmax_f32", "argmax.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "int"}}));
    string_to_spv("sum_rows_f32", "sum_rows.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
+    string_to_spv("cross_entropy_loss_f32", "cross_entropy_loss.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
+    string_to_spv("cross_entropy_loss_back_f32", "cross_entropy_loss_back.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
    string_to_spv("fwht_f32", "fwht.comp", {});
    string_to_spv("fwht_shmem_f32", "fwht.comp", {{"FWHT_SHMEM", "1"}});
    string_to_spv("count_equal_i32", "count_equal.comp", merge_maps(base_dict, {{"A_TYPE", "int"}, {"B_TYPE", "int"}, {"D_TYPE", "int"}}));
@@ -43,10 +43,10 @@
 #define LLAMA_FILE_MAGIC_GGSQ 0x67677371u // 'ggsq'

 #define LLAMA_SESSION_MAGIC   LLAMA_FILE_MAGIC_GGSN
-#define LLAMA_SESSION_VERSION 9
+#define LLAMA_SESSION_VERSION 10

 #define LLAMA_STATE_SEQ_MAGIC   LLAMA_FILE_MAGIC_GGSQ
-#define LLAMA_STATE_SEQ_VERSION 2
+#define LLAMA_STATE_SEQ_VERSION 3

 #ifdef __cplusplus
 extern "C" {
@@ -12,6 +12,7 @@
 #include <limits>
 #include <map>
 #include <stdexcept>
+#include <unordered_map>

 static bool ggml_is_power_of_2(int n) {
    return (n & (n - 1)) == 0;
@@ -1128,11 +1129,18 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &

            cells.pos_set(idx, ubatch.pos[i]);

-            if (ubatch.is_pos_2d()) {
-                llama_kv_cell_ext ext {
-                    /*.x =*/ ubatch.pos[i + ubatch.n_tokens*2],
-                    /*.y =*/ ubatch.pos[i + ubatch.n_tokens],
-                };
+            if (ubatch.is_pos_2d() || ubatch.token) {
+                llama_kv_cell_ext ext;
+
+                if (ubatch.is_pos_2d()) {
+                    ext.x = ubatch.pos[i + ubatch.n_tokens*2];
+                    ext.y = ubatch.pos[i + ubatch.n_tokens];
+                }
+
+                if (ubatch.token) {
+                    ext.tok = ubatch.token[i];
+                }
+
                cells.ext_set(idx, ext);
            }

@@ -1805,6 +1813,69 @@ void llama_kv_cache::set_input_v_rot(ggml_tensor * dst) const {
    memcpy(dst->data, attn_rot_hadamard.at(n_rot).data(), ggml_nbytes(dst));
 }

+bool llama_kv_cache::has_cell_ext() const {
+    return hparams.n_pos_per_embd() > 1;
+}
+
+void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const {
+    const uint32_t n_tokens = ubatch.n_tokens;
+
+    res.clear();
+    res.resize(n_tokens*n, LLAMA_TOKEN_NULL);
+
+    if (n == 0) {
+        return;
+    }
+
+    // note: apply_ubatch() has already stored the current ubatch
+    //       the window below thus covers tokens of this very ubatch as well, which is what we want
+    llama_pos p_min = std::numeric_limits<llama_pos>::max();
+    llama_pos p_max = std::numeric_limits<llama_pos>::min();
+
+    std::bitset<LLAMA_MAX_SEQ> seqs;
+
+    for (uint32_t i = 0; i < n_tokens; ++i) {
+        p_min = std::min(p_min, ubatch.pos[i]);
+        p_max = std::max(p_max, ubatch.pos[i]);
+    }
+
+    for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
+        seqs.set(ubatch.seq_id_unq[s]);
+    }
+
+    // (seq_id, pos) -> token, for every cell that could be a predecessor of a ubatch token
+    std::unordered_map<uint64_t, llama_token> hist;
+
+    const auto key = [](llama_seq_id seq_id, llama_pos pos) {
+        return ((uint64_t) seq_id << 32) | (uint32_t) pos;
+    };
+
+    for (uint32_t s = 0; s < n_stream; ++s) {
+        v_cells[s].for_each_token_in(seqs, p_min - (llama_pos) n, p_max,
+            [&](llama_seq_id seq_id, llama_pos pos, llama_token tok) {
+                hist[key(seq_id, pos)] = tok;
+            });
+    }
+
+    for (uint32_t i = 0; i < n_tokens; ++i) {
+        // TODO: a token that belongs to more than one sequence has an ambiguous history.
+        //       the n-gram architectures have to reject such batches
+        const llama_seq_id seq_id = ubatch.seq_id[i][0];
+
+        for (uint32_t j = 0; j < n; ++j) {
+            const llama_pos p = ubatch.pos[i] - (llama_pos) (n - j);
+            if (p < 0) {
+                continue;
+            }
+
+            const auto it = hist.find(key(seq_id, p));
+            if (it != hist.end()) {
+                res[i*n + j] = it->second;
+            }
+        }
+    }
+}
+
 size_t llama_kv_cache::total_size() const {
    size_t size = 0;

@@ -2106,7 +2177,7 @@ void llama_kv_cache::state_write_meta(llama_io_write_i & io, const cell_ranges_t
            io.write(&pos,      sizeof(pos));
            io.write(&n_seq_id, sizeof(n_seq_id));

-            if (hparams.n_pos_per_embd() > 1) {
+            if (has_cell_ext()) {
                const llama_kv_cell_ext ext = cells.ext_get(i);
                io.write(&ext, sizeof(ext));
            }
@@ -2243,12 +2314,17 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
                return false;
            }

-            if (hparams.n_pos_per_embd() > 1) {
+            if (has_cell_ext()) {
                llama_kv_cell_ext ext;
                io.read(&ext, sizeof(ext));

-                ubatch.pos[i + ubatch.n_tokens]   = ext.y;
-                ubatch.pos[i + ubatch.n_tokens*2] = ext.x;
+                if (hparams.n_pos_per_embd() > 1) {
+                    ubatch.pos[i + ubatch.n_tokens]   = ext.y;
+                    ubatch.pos[i + ubatch.n_tokens*2] = ext.x;
+                }
+
+                // apply_ubatch() below restores ext.tok from the ubatch tokens
+                ubatch.token[i] = ext.tok;
            }

            // read the sequence id, but directly discard it - we will use dest_seq_id instead
@@ -2268,7 +2344,8 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
            return false;
        }

-        // TODO: we cannot yet restore llama_kv_cell_ext as the apply_ubatch() does not support it yet
+        // note: apply_ubatch() rebuilds llama_kv_cell_ext from the ubatch
+        //       only ext.tok and the M-RoPE 2D position round-trip through it
        //       see: https://github.com/ggml-org/llama.cpp/pull/16825#issuecomment-3460868350
        apply_ubatch(sinfo, ubatch);

@@ -2301,7 +2378,7 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32

            cells.pos_set(i, pos);

-            if (hparams.n_pos_per_embd() > 1) {
+            if (has_cell_ext()) {
                llama_kv_cell_ext ext;
                io.read(&ext, sizeof(ext));
                cells.ext_set(i, ext);
@@ -2652,3 +2729,7 @@ void llama_kv_cache_context::set_input_k_rot(ggml_tensor * dst) const {
 void llama_kv_cache_context::set_input_v_rot(ggml_tensor * dst) const {
    kv->set_input_v_rot(dst);
 }
+
+void llama_kv_cache_context::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const {
+    kv->get_prev_tokens(ubatch, n, res);
+}
@@ -219,6 +219,14 @@ public:
    void set_input_k_rot(ggml_tensor * dst) const;
    void set_input_v_rot(ggml_tensor * dst) const;

+    // true if llama_kv_cell_ext holds information that has to survive a state save/restore
+    bool has_cell_ext() const;
+
+    // for every token of the ubatch, the ids of the n tokens that precede it in its sequence
+    // entries with no matching cell are set to LLAMA_TOKEN_NULL
+    // note: used by n-gram input embeddings
+    void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const;
+
 private:
    const llama_model & model;
    const llama_hparams & hparams;
@@ -401,6 +409,9 @@ public:
    void set_input_k_rot(ggml_tensor * dst) const;
    void set_input_v_rot(ggml_tensor * dst) const;

+    // see llama_kv_cache::get_prev_tokens()
+    void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const;
+
 private:
    llama_memory_status status;

@@ -15,6 +15,10 @@ struct llama_kv_cell_ext {
    llama_pos x = 0;
    llama_pos y = 0;

+    // when tok = LLAMA_TOKEN_NULL when the cell is produced by embedding input (i.e. multimodal)
+    // use case: n-gram embeddings hash
+    llama_token tok = LLAMA_TOKEN_NULL;
+
    // return true if the current 2D spatial position is greater than other
    bool is_2d_gt(llama_pos ox, llama_pos oy) const {
        return (y > oy) || (y == oy && x > ox);
@@ -23,7 +27,7 @@ struct llama_kv_cell_ext {
    void reset() {
        static_assert(std::is_trivially_copyable_v<llama_kv_cell_ext>);

-        memset(this, 0, sizeof(*this));
+        *this = llama_kv_cell_ext{};
    }
 };

@@ -305,6 +309,29 @@ public:
        return seq[i].test(seq_id);
    }

+    // gather the token ids of the cells in `seqs` with position in [p0, p1)
+    // the callback receives (seq_id, pos, token) for every such (cell, seq) pair
+    // note: used by n-gram input embeddings to recover the tokens preceding a ubatch
+    template<typename F>
+    void for_each_token_in(const std::bitset<LLAMA_MAX_SEQ> & seqs, llama_pos p0, llama_pos p1, F && f) const {
+        for (const auto & i : used) {
+            if (pos[i] < p0 || pos[i] >= p1) {
+                continue;
+            }
+
+            const auto m = seq[i] & seqs;
+            if (m.none()) {
+                continue;
+            }
+
+            for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) {
+                if (m.test(s)) {
+                    f(s, pos[i], ext[i].tok);
+                }
+            }
+        }
+    }
+
    // note: call only if the cell is not empty and the seq_id is not in the cell
    void seq_add(uint32_t i, llama_seq_id seq_id) {
        assert(i < pos.size());
@@ -4120,9 +4120,10 @@ struct test_ssm_scan : public test_case {
    const int64_t n_seqs;
    const bool    xbc_overlap;
    const int64_t K;
+    const bool    weak_decay;

    std::string vars() override {
-        return VARS_TO_STR9(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, xbc_overlap, K);
+        return VARS_TO_STR10(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, xbc_overlap, K, weak_decay);
    }

    test_ssm_scan(ggml_type type = GGML_TYPE_F32,
@@ -4133,8 +4134,9 @@ struct test_ssm_scan : public test_case {
            int64_t n_seq_tokens = 32,
            int64_t n_seqs = 32,
            bool xbc_overlap = false,
-            int64_t K = 1)
-        : type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), xbc_overlap(xbc_overlap), K(K) {}
+            int64_t K = 1,
+            bool weak_decay = false)
+        : type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), xbc_overlap(xbc_overlap), K(K), weak_decay(weak_decay) {}

    double max_nmse_err() override {
        // SSD path (head_dim > 1) uses FP16 intermediates (M matrix, X_dt); Mamba-1 is pure FP32.
@@ -4187,7 +4189,7 @@ struct test_ssm_scan : public test_case {
                continue;
            } else if (t->ne[1] == n_head && t->ne[2] == 1) {
                // A {1 or d_state, n_head}: negative decay (2-D tensor, ne[2]==1 distinguishes from 3-D/4-D tensors)
-                init_tensor_uniform(t, -1.0f, -0.5f);
+                init_tensor_uniform(t, weak_decay ? -0.02f : -1.0f, weak_decay ? -0.005f : -0.5f);
            } else {
                init_tensor_uniform(t);
            }
@@ -9111,6 +9113,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 4, 2, false, /*K=*/4)); // Mamba-2 rollback snapshots
    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 8, 2, false, /*K=*/3)); // Mamba-2 rollback overflow
    test_cases.emplace_back(new test_ssm_scan_rollback(GGML_TYPE_F32, 128, 64, 16, 2, 8, 2, /*K=*/3)); // rollback snapshots match prefix states
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 64, 4)); // Metal SSD one chunk MMA only, no seq tail
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 65, 2)); // SSD one chunk + 1-token sequential tail
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 128, 2)); // SSD multi-chunk, no tail (exercises the chunk-to-chunk state handoff)
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 128, 2, false, /*K=*/1, /*weak_decay=*/true)); // SSD multi-chunk, carried state not numerically negligible

    test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 1, 1));
    test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 32, 1));
@@ -97,9 +97,19 @@ By default, the cache is stored in the `$HOME/.cache/llama.cpp/rpc` directory an

 ### RDMA transport

-On Linux systems with RoCEv2-capable NICs (e.g. Mellanox ConnectX), the RPC backend can use RDMA instead of TCP for lower latency and higher throughput. The transport is negotiated automatically -- no changes to command-line usage are required.
+The RPC backend can use RDMA instead of TCP for lower latency and higher throughput. The transport is negotiated during the initial handshake -- no changes to command-line usage are required, and the connection falls back to TCP unless both peers can use RDMA.

-RDMA is enabled by default when `libibverbs` is found at build time.
+Two providers are supported, each enabled by default when its library is found at build time:
+
+- **Linux**: RoCEv2-capable NICs (e.g. Mellanox ConnectX), via `libibverbs`.
+- **macOS**: RDMA over Thunderbolt on Apple silicon Macs with Thunderbolt 5, via `librdma`. Requires macOS 26.2 or later, with RDMA enabled once from macOS Recovery via `rdma_ctl enable`. See [TN3205](https://developer.apple.com/documentation/technotes/tn3205-low-latency-communication-with-rdma-over-thunderbolt).
+
+RDMA is point-to-point, so each side uses the local device whose GID matches the address the connection was made on. Connect over the RDMA-capable link -- with Thunderbolt, use the peer's Thunderbolt address in `--rpc`; a connection made over another interface stays on TCP.
+
+To force plain TCP without rebuilding, set `GGML_RPC_NO_RDMA` on either peer:
+```bash
+$ GGML_RPC_NO_RDMA=1 bin/ggml-rpc-server
+```

 ### Troubleshooting

@@ -1280,6 +1280,12 @@ json oaicompat_chat_params_parse(
    if (inputs.continue_final_message != COMMON_CHAT_CONTINUATION_NONE && inputs.add_generation_prompt) {
        throw std::invalid_argument("Cannot set both add_generation_prompt and continue_final_message to true.");
    }
+    if (inputs.continue_final_message != COMMON_CHAT_CONTINUATION_NONE
+        && !inputs.messages.empty()
+        && inputs.messages.back().role == "assistant"
+        && !inputs.messages.back().tool_calls.empty()) {
+        throw std::invalid_argument("Cannot continue an assistant message that contains tool calls.");
+    }
    inputs.reasoning_format = opt.reasoning_format;
    if (body.contains("reasoning_format")) {
        inputs.reasoning_format = common_reasoning_format_from_name(body.at("reasoning_format").get<std::string>());
@@ -8,7 +8,8 @@
 		ChatFormInputFileInputInvisible,
 		ChatFormMcpResourcesList,
 		ChatFormPickers,
-		DialogMcpResourcesBrowser
+		DialogMcpResourcesBrowser,
+		DialogMcpServers
 	} from '$lib/components/app';
 	import {
 		CLIPBOARD_CONTENT_QUOTE_PREFIX,
@@ -183,6 +184,9 @@
 	let isResourceDialogOpen = $state(false);
 	let preSelectedResourceUri = $state<string | undefined>(undefined);

+	// MCP Servers Dialog State
+	let isMcpServersDialogOpen = $state(false);
+
 	let currentConfig = $derived(settingsStore.config);

 	let pasteLongTextToFileLength = $derived.by(() => {
@@ -618,6 +622,7 @@
 				onFileUpload={handleFileUpload}
 				onMcpPromptClick={showMcpPromptButton ? () => pickers.openPromptPicker() : undefined}
 				onMcpResourcesClick={() => (isResourceDialogOpen = true)}
+				onMcpSettingsClick={() => (isMcpServersDialogOpen = true)}
 				onMicClick={handleMicClick}
 				{onStop}
 				onSystemPromptClick={() => onSystemPromptClick?.({ files: uploadedFiles, message: value })}
@@ -656,3 +661,5 @@
 	}}
 	preSelectedUri={preSelectedResourceUri}
 />
+
+<DialogMcpServers bind:open={isMcpServersDialogOpen} />
@@ -1,10 +1,6 @@
 <script lang="ts">
-	import { File, FolderOpen, MessageSquare, Plus, Zap } from '@lucide/svelte';
-	import {
-		ChatFormActionAddMcpServersSubmenu,
-		ChatFormActionAddReasoningSubmenu,
-		ChatFormActionAddToolsSubmenu
-	} from '$lib/components/app';
+	import { File, MessageSquare, Plus } from '@lucide/svelte';
+	import { ChatFormActionAddToolsSubmenu, McpLogo } from '$lib/components/app';
 	import { buttonVariants } from '$lib/components/ui/button';
 	import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
 	import * as Tooltip from '$lib/components/ui/tooltip';
@@ -31,11 +27,6 @@
 	// must not restore focus to the trigger on close
 	let suppressCloseAutoFocus = false;

-	function handleMcpSettingsClick() {
-		dropdownOpen = false;
-		chatFormActions.onMcpSettingsClick?.();
-	}
-
 	const attachmentMenu = useAttachmentMenu(
 		() => ({
 			hasAudioModality: chatFormActions.hasAudioModality,
@@ -93,10 +84,6 @@
 				}
 			}}
 		>
-			<ChatFormActionAddReasoningSubmenu />
-
-			<DropdownMenu.Separator />
-
 			<DropdownMenu.Sub>
 				<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
 					<File class={ICON_CLASS_DEFAULT} />
@@ -156,31 +143,14 @@

 			<ChatFormActionAddToolsSubmenu />

-			<ChatFormActionAddMcpServersSubmenu onMcpSettingsClick={handleMcpSettingsClick} />
+			<DropdownMenu.Item
+				class="flex cursor-pointer items-center gap-2"
+				onclick={chatFormActions.onMcpSettingsClick}
+			>
+				<McpLogo class={ICON_CLASS_DEFAULT} />

-			{#if chatFormActions.hasMcpPromptsSupport}
-				<DropdownMenu.Separator />
-
-				<DropdownMenu.Item
-					class="flex cursor-pointer items-center gap-2"
-					onclick={chatFormActions.onMcpPromptClick}
-				>
-					<Zap class={ICON_CLASS_DEFAULT} />
-
-					<span>MCP Prompt</span>
-				</DropdownMenu.Item>
-			{/if}
-
-			{#if chatFormActions.hasMcpResourcesSupport}
-				<DropdownMenu.Item
-					class="flex cursor-pointer items-center gap-2"
-					onclick={chatFormActions.onMcpResourcesClick}
-				>
-					<FolderOpen class={ICON_CLASS_DEFAULT} />
-
-					<span>MCP Resources</span>
-				</DropdownMenu.Item>
-			{/if}
+				<span>MCP Servers</span>
+			</DropdownMenu.Item>
 		</DropdownMenu.Content>
 	</DropdownMenu.Root>
 </div>
@@ -1,152 +0,0 @@
-<script lang="ts">
-	import { Plus, Settings } from '@lucide/svelte';
-	import { goto } from '$app/navigation';
-	import { DropdownMenuSearchable, McpLogo, McpServerIdentity } from '$lib/components/app';
-	import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
-	import { Switch } from '$lib/components/ui/switch';
-	import { ICON_CLASS_DEFAULT, ROUTES } from '$lib/constants';
-	import { HealthCheckStatus } from '$lib/enums';
-	import { conversationsStore, mcpStore } from '$lib/stores';
-	import type { MCPServerSettingsEntry } from '$lib/types';
-
-	interface Props {
-		onMcpSettingsClick?: () => void;
-	}
-
-	let { onMcpSettingsClick }: Props = $props();
-
-	let mcpSearchQuery = $state('');
-	// Every configured server is listed; `enabled` is an on/off state,
-	// not a visibility filter, so a disabled server stays toggleable.
-	let mcpServers = $derived(mcpStore.getServers());
-	let hasMcpServers = $derived(mcpServers.length > 0);
-	let filteredMcpServers = $derived.by(() => {
-		const query = mcpSearchQuery.toLowerCase().trim();
-
-		if (!query) return mcpServers;
-
-		return mcpServers.filter((s) => {
-			const name = getServerLabel(s).toLowerCase();
-			const url = s.url.toLowerCase();
-
-			return name.includes(query) || url.includes(query);
-		});
-	});
-
-	function getServerLabel(server: MCPServerSettingsEntry): string {
-		return mcpStore.getServerLabel(server);
-	}
-
-	function isServerEnabledForChat(serverId: string): boolean {
-		return conversationsStore.preferences.isMcpServerEnabledForChat(serverId);
-	}
-
-	async function toggleServerForChat(serverId: string) {
-		await conversationsStore.preferences.toggleMcpServerForChat(serverId);
-	}
-
-	function handleMcpSubMenuOpen(open: boolean) {
-		if (open) {
-			mcpSearchQuery = '';
-			mcpStore.runHealthChecksForServers(mcpServers);
-		}
-	}
-
-	function handleMcpSettingsClick() {
-		onMcpSettingsClick?.();
-
-		goto(`${hasMcpServers ? '' : '?add'}${ROUTES.MCP_SERVERS}`);
-	}
-</script>
-
-<DropdownMenu.Root>
-	<DropdownMenu.Sub onOpenChange={handleMcpSubMenuOpen}>
-		<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
-			<McpLogo class={ICON_CLASS_DEFAULT} />
-
-			<span>MCP Servers</span>
-		</DropdownMenu.SubTrigger>
-
-		<DropdownMenu.SubContent class="w-72 pt-0">
-			{#if hasMcpServers}
-				<DropdownMenuSearchable
-					bind:searchValue={mcpSearchQuery}
-					emptyMessage="No servers found"
-					isEmpty={filteredMcpServers.length === 0}
-					placeholder="Search servers..."
-				>
-					<div class="max-h-64 overflow-y-auto">
-						{#each filteredMcpServers as server (server.id)}
-							{@const healthState = mcpStore.getHealthCheckState(server.id)}
-							{@const hasError = healthState.status === HealthCheckStatus.ERROR}
-							{@const isEnabledForChat = isServerEnabledForChat(server.id)}
-							{@const displayName = getServerLabel(server)}
-							{@const faviconUrl = mcpStore.getServerFavicon(server.id)}
-
-							<button
-								class="flex w-full items-center justify-between gap-2 rounded-sm px-2 py-2 text-left transition-colors hover:bg-accent disabled:cursor-not-allowed disabled:opacity-50"
-								disabled={hasError}
-								onclick={() => !hasError && toggleServerForChat(server.id)}
-								type="button"
-							>
-								<div class="flex min-w-0 flex-1 items-center gap-2">
-									<div class="min-w-0 flex-1">
-										<McpServerIdentity
-											{displayName}
-											{faviconUrl}
-											iconClass={ICON_CLASS_DEFAULT}
-											iconRounded="rounded-sm"
-											nameClass="text-sm"
-											showVersion={false}
-										/>
-									</div>
-
-									{#if hasError}
-										<span
-											class="shrink-0 rounded bg-destructive/15 px-1.5 py-0.5 text-xs text-destructive"
-										>
-											Error
-										</span>
-									{/if}
-								</div>
-
-								<Switch
-									checked={isEnabledForChat}
-									disabled={hasError}
-									onCheckedChange={() => toggleServerForChat(server.id)}
-									onclick={(e) => e.stopPropagation()}
-								/>
-							</button>
-						{/each}
-					</div>
-
-					{#snippet footer()}
-						<DropdownMenu.Item
-							class="flex cursor-pointer items-center gap-2"
-							onclick={handleMcpSettingsClick}
-						>
-							<Settings class={ICON_CLASS_DEFAULT} />
-
-							<span>Manage MCP Servers</span>
-						</DropdownMenu.Item>
-					{/snippet}
-				</DropdownMenuSearchable>
-			{:else}
-				<div class="px-2 py-3 text-center text-sm text-muted-foreground">
-					No MCP servers configured
-				</div>
-
-				<DropdownMenu.Separator />
-
-				<DropdownMenu.Item
-					class="flex cursor-pointer items-center gap-2"
-					onclick={handleMcpSettingsClick}
-				>
-					<Plus class={ICON_CLASS_DEFAULT} />
-
-					<span>Add MCP Servers</span>
-				</DropdownMenu.Item>
-			{/if}
-		</DropdownMenu.SubContent>
-	</DropdownMenu.Sub>
-</DropdownMenu.Root>
@@ -0,0 +1,51 @@
+<script lang="ts">
+	import { FolderOpen, Server, Zap } from '@lucide/svelte';
+	import { McpLogo } from '$lib/components/app';
+	import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
+	import { ICON_CLASS_DEFAULT } from '$lib/constants';
+	import { getChatFormActionsContext } from '$lib/contexts';
+
+	const chatFormActions = getChatFormActionsContext();
+
+	function handleServersClick() {
+		chatFormActions.onMcpSettingsClick?.();
+	}
+</script>
+
+<DropdownMenu.Sub>
+	<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
+		<McpLogo class={ICON_CLASS_DEFAULT} />
+
+		<span>MCP</span>
+	</DropdownMenu.SubTrigger>
+
+	<DropdownMenu.SubContent class="w-48">
+		<DropdownMenu.Item class="flex cursor-pointer items-center gap-2" onclick={handleServersClick}>
+			<Server class={ICON_CLASS_DEFAULT} />
+
+			<span>Servers</span>
+		</DropdownMenu.Item>
+
+		{#if chatFormActions.hasMcpPromptsSupport}
+			<DropdownMenu.Item
+				class="flex cursor-pointer items-center gap-2"
+				onclick={chatFormActions.onMcpPromptClick}
+			>
+				<Zap class={ICON_CLASS_DEFAULT} />
+
+				<span>Prompts</span>
+			</DropdownMenu.Item>
+		{/if}
+
+		{#if chatFormActions.hasMcpResourcesSupport}
+			<DropdownMenu.Item
+				class="flex cursor-pointer items-center gap-2"
+				onclick={chatFormActions.onMcpResourcesClick}
+			>
+				<FolderOpen class={ICON_CLASS_DEFAULT} />
+
+				<span>Resources</span>
+			</DropdownMenu.Item>
+		{/if}
+	</DropdownMenu.SubContent>
+</DropdownMenu.Sub>
@@ -1,6 +1,5 @@
 <script lang="ts">
 	import { SkipForward, Square } from '@lucide/svelte';
-	import { goto } from '$app/navigation';
 	import { page } from '$app/state';
 	import {
 		ChatFormActionModels,
@@ -10,7 +9,7 @@
 		ChatFormContextGauge
 	} from '$lib/components/app';
 	import { Button } from '$lib/components/ui/button';
-	import { ICON_CLASS_DEFAULT, ROUTES } from '$lib/constants';
+	import { ICON_CLASS_DEFAULT } from '$lib/constants';
 	import { setChatFormActionsContext } from '$lib/contexts';
 	import { FileTypeCategory, MessageRole } from '$lib/enums';
 	import { ChatService } from '$lib/services';
@@ -34,6 +33,7 @@
 		onSystemPromptClick?: () => void;
 		onMcpPromptClick?: () => void;
 		onMcpResourcesClick?: () => void;
+		onMcpSettingsClick?: () => void;
 	}

 	let {
@@ -47,6 +47,7 @@
 		onFileUpload,
 		onMcpPromptClick,
 		onMcpResourcesClick,
+		onMcpSettingsClick,
 		onMicClick,
 		onStop,
 		onSystemPromptClick,
@@ -163,7 +164,7 @@
 			return onMcpResourcesClick;
 		},
 		get onMcpSettingsClick() {
-			return () => goto(ROUTES.MCP_SERVERS);
+			return onMcpSettingsClick;
 		},
 		get onSystemPromptClick() {
 			return onSystemPromptClick;
@@ -221,25 +221,17 @@ export { default as ChatFormActionModels } from './ChatForm/ChatFormActions/Chat
 export { default as ChatFormActionAddToolsSubmenu } from './ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte';

 /**
- * Dropdown submenu for managing MCP servers in the chat form.
+ * Dropdown submenu for MCP prompts and resources in the chat form.
 *
- * Displays a searchable list of enabled MCP servers with toggle switches
- * to enable/disable each server for chat. Shows server favicon, health status,
- * and a "Manage MCP Servers" settings link.
- *
- * Features:
- * - Search/filter servers by name or URL
- * - Per-server toggle to enable/disable for chat
- * - Health check indicator (shows "Error" badge for failed servers)
- * - Server favicon display
- * - Settings link to manage MCP server configuration
+ * Shows an "MCP" sub-menu item with entries for MCP Prompts and MCP
+ * Resources. Only visible when the server supports them.
 *
 * @example
 * ```svelte
- * <ChatFormActionAddMcpServersSubmenu onMcpSettingsClick={handleMcpSettingsClick} />
+ * <ChatFormActionAddMcpSubmenu />
 * ```
 */
-export { default as ChatFormActionAddMcpServersSubmenu } from './ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpServersSubmenu.svelte';
+export { default as ChatFormActionAddMcpSubmenu } from './ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpSubmenu.svelte';

 /**
 * Dropdown submenu for selecting reasoning effort level.
@@ -253,7 +253,7 @@
 </script>

 <Dialog.Root onOpenChange={handleOpenChange} {open}>
-	<Dialog.Content class="max-h-[80vh] !max-w-4xl overflow-hidden p-0">
+	<Dialog.Content class="max-h-[80vh] md:max-w-4xl! w-full! overflow-hidden p-0">
 		<Dialog.Header class="border-b border-border/30 px-6 py-4">
 			<Dialog.Title class="flex items-center gap-2">
 				<FolderOpen class="h-5 w-5" />
@@ -246,7 +246,7 @@
 </script>

 <Dialog.Root onOpenChange={handleOpenChange} {open}>
-	<Dialog.Content class="sm:max-w-2xl">
+	<Dialog.Content class="max-w-2xl!">
 		<Dialog.Header>
 			<Dialog.Title class="select-none">Add New MCP Server</Dialog.Title>
 		</Dialog.Header>
@@ -0,0 +1,33 @@
+<script lang="ts">
+	import { McpLogo } from '$lib/components/app';
+	import { SettingsMcpServers } from '$lib/components/app/settings';
+	import * as Dialog from '$lib/components/ui/dialog';
+
+	interface Props {
+		open?: boolean;
+		onOpenChange?: (open: boolean) => void;
+	}
+
+	let { onOpenChange, open = $bindable(false) }: Props = $props();
+
+	function handleOpenChange(value: boolean) {
+		open = value;
+		onOpenChange?.(value);
+	}
+</script>
+
+<Dialog.Root onOpenChange={handleOpenChange} {open}>
+	<Dialog.Content
+		class="md:h-[calc(100vh-4rem)]! md:max-h-240! md:w-[calc(100vw-4rem)]! md:max-w-360! flex flex-col"
+	>
+		<Dialog.Header>
+			<Dialog.Title class="flex items-center gap-2">
+				<McpLogo class="h-5 w-5" />
+
+				<span>MCP Servers</span>
+			</Dialog.Title>
+		</Dialog.Header>
+
+		<SettingsMcpServers class="mt-4" />
+	</Dialog.Content>
+</Dialog.Root>
@@ -14,6 +14,7 @@
 <Dialog.Root bind:open {onOpenChange}>
 	<Dialog.Content
 		class="z-999999 grid max-h-full max-w-full! grid-rows-[1fr_auto] overflow-hidden p-0 md:h-[90vh] md:max-w-[90vw]!"
+		showCloseButton
 	>
 		<MermaidPreview {svgHtml} />
 	</Dialog.Content>
@@ -0,0 +1,34 @@
+<script lang="ts">
+	import { Settings } from '@lucide/svelte';
+	import { SettingsChat } from '$lib/components/app/settings';
+	import * as Dialog from '$lib/components/ui/dialog';
+
+	interface Props {
+		open?: boolean;
+		onOpenChange?: (open: boolean) => void;
+		initialSection?: string;
+	}
+
+	let { initialSection, onOpenChange, open = $bindable(false) }: Props = $props();
+
+	function handleOpenChange(value: boolean) {
+		open = value;
+		onOpenChange?.(value);
+	}
+</script>
+
+<Dialog.Root onOpenChange={handleOpenChange} {open}>
+	<Dialog.Content
+		class="md:h-[calc(100vh-4rem)]! md:max-h-240! md:w-[calc(100vw-4rem)]! md:max-w-6xl! flex flex-col p-0 md:p-6 gap-0"
+	>
+		<Dialog.Header class="md:p-0 p-4">
+			<Dialog.Title class="flex items-center gap-2">
+				<Settings class="h-5 w-5" />
+
+				<span>Settings</span>
+			</Dialog.Title>
+		</Dialog.Header>
+
+		<SettingsChat {initialSection} onClose={() => (open = false)} onSectionChange={() => {}} />
+	</Dialog.Content>
+</Dialog.Root>
@@ -18,6 +18,23 @@
 */
 export { default as DialogMcpServerAddNew } from './DialogMcpServerAddNew.svelte';

+/**
+ * **DialogMcpServers** - MCP servers dialog shown from the chat form
+ *
+ * Shows the same MCP server list as the `/mcp-servers` route inside a modal
+ * dialog.
+ */
+export { default as DialogMcpServers } from './DialogMcpServers.svelte';
+
+/**
+ * **DialogSettingsChat** - Chat settings shown in a modal dialog
+ *
+ * Wraps the full SettingsChat layout (sidebar, mobile header, fields, footer)
+ * inside a ShadCN Dialog instead of a dedicated route. Section switching is
+ * handled in-app via `onSectionChange` rather than URL navigation.
+ */
+export { default as DialogSettingsChat } from './DialogSettingsChat.svelte';
+
 /**
 * **DialogExportSettings** - Settings export dialog with sensitive data warning
 *
@@ -1,13 +1,22 @@
 <script lang="ts">
-	import { Database, FileText, ListChecks, MessageSquare, Sparkles, Wrench } from '@lucide/svelte';
+	import {
+		Database,
+		ExternalLink,
+		FileText,
+		ListChecks,
+		MessageSquare,
+		Sparkles,
+		Wrench
+	} from '@lucide/svelte';
 	import { Badge } from '$lib/components/ui/badge';
 	import type { MCPCapabilitiesInfo } from '$lib/types';

 	interface Props {
 		capabilities?: MCPCapabilitiesInfo;
+		onBrowseResources?: () => void;
 	}

-	let { capabilities }: Props = $props();
+	let { capabilities, onBrowseResources }: Props = $props();
 </script>

 {#if capabilities}
@@ -20,10 +29,24 @@
 	{/if}

 	{#if capabilities.server.resources}
-		<Badge class="h-5 gap-1 bg-blue-50 px-1.5 text-[10px] dark:bg-blue-950" variant="outline">
+		<Badge
+			class="h-5 cursor-pointer gap-1 bg-blue-50 px-1.5 text-[10px] transition-colors hover:bg-blue-100 dark:bg-blue-950 dark:hover:bg-blue-900"
+			onclick={onBrowseResources}
+			onkeydown={(e) => {
+				if (e.key === 'Enter' || e.key === ' ') {
+					e.preventDefault();
+					onBrowseResources?.();
+				}
+			}}
+			role="button"
+			tabindex={0}
+			variant="outline"
+		>
 			<Database class="h-3 w-3 text-blue-600 dark:text-blue-400" />

 			Resources
+
+			<ExternalLink class="h-3 w-3 text-blue-600 dark:text-blue-400" />
 		</Badge>
 	{/if}

@@ -22,9 +22,10 @@
 		onToggle: (enabled: boolean) => void;
 		onUpdate: (updates: Partial<MCPServerSettingsEntry>) => void;
 		onDelete: () => void;
+		onBrowseResources?: () => void;
 	}

-	let { enabled, onDelete, onToggle, onUpdate, server }: Props = $props();
+	let { enabled, onBrowseResources, onDelete, onToggle, onUpdate, server }: Props = $props();

 	let healthState = $derived<HealthCheckState>(mcpStore.getHealthCheckState(server.id));
 	let displayName = $derived(mcpStore.getServerLabel(server));
@@ -125,6 +126,7 @@
 			{displayName}
 			enabled={enabled ?? server.enabled}
 			{faviconUrl}
+			{onBrowseResources}
 			{onToggle}
 			{serverInfo}
 			{transportType}
@@ -191,12 +193,14 @@
 				</div>
 			{/if}

-			<McpServerCardActions
-				{isHealthChecking}
-				onDelete={handleDeleteClick}
-				onEdit={startEditing}
-				onRefresh={handleHealthCheck}
-			/>
+			<div class="flex items-center gap-2">
+				<McpServerCardActions
+					{isHealthChecking}
+					onDelete={handleDeleteClick}
+					onEdit={startEditing}
+					onRefresh={handleHealthCheck}
+				/>
+			</div>
 		</div>
 	{/if}
 </Card.Root>
@@ -12,6 +12,7 @@
 		enabled: boolean;
 		disabled?: boolean;
 		onToggle: (enabled: boolean) => void;
+		onBrowseResources?: () => void;
 		serverInfo?: MCPServerInfo;
 		capabilities?: MCPCapabilitiesInfo;
 		transportType?: MCPTransportType;
@@ -23,6 +24,7 @@
 		displayName,
 		enabled,
 		faviconUrl,
+		onBrowseResources,
 		onToggle,
 		serverInfo,
 		transportType
@@ -57,7 +59,7 @@
 					{/if}

 					{#if capabilities}
-						<McpCapabilitiesBadges {capabilities} />
+						<McpCapabilitiesBadges {capabilities} {onBrowseResources} />
 					{/if}
 				</div>
 			{/if}
@@ -5,6 +5,7 @@
 	import {
 		ActionIcon,
 		DialogConversationRename,
+		DialogSettingsChat,
 		Logo,
 		SidebarNavigationActions,
 		SidebarNavigationConversationList
@@ -91,6 +92,7 @@
 	let selectedIds = new SvelteSet<string>();

 	let renameDialogOpen = $state(false);
+	let settingsDialogOpen = $state(false);
 	let renameTargetConversationId = $state<string | null>(null);
 	let renameDraft = $state('');
 	let renameOriginalTitle = $state('');
@@ -308,7 +310,7 @@

 <svelte:window bind:innerWidth onkeydown={handleKeydown} />

-{#if innerWidth > 768 || (!page.url.hash.includes(ROUTES.SETTINGS) && !page.url.hash.includes(ROUTES.MCP_SERVERS) && !page.url.hash.includes(ROUTES.SEARCH))}
+{#if innerWidth > 768 || !page.url.hash.includes(ROUTES.SEARCH)}
 	<aside
 		class={[
 			'fixed md:sticky top-2 left-2 md:left-0 md:ml-2 md:mt-2 pt-2 z-10 w-[calc(100dvw-1rem)]',
@@ -400,6 +402,7 @@
 					isSearchModeActive = false;
 					searchQuery = '';
 				}}
+				onSettingsClick={() => (settingsDialogOpen = true)}
 			/>

 			{#if uiStore.isSidebarExpanded || isOnMobile}
@@ -447,6 +450,8 @@
 	onConfirm={handleRenameConfirm}
 />

+<DialogSettingsChat bind:open={settingsDialogOpen} />
+
 <style>
 	aside {
 		@media (max-width: 768px) {
@@ -26,6 +26,7 @@
 		onSearchDeactivated?: () => void;
 		onSearchClick?: () => void;
 		onNewChat?: () => void;
+		onSettingsClick?: () => void;
 	}

 	let {
@@ -35,6 +36,7 @@
 		onNewChat,
 		onSearchClick,
 		onSearchDeactivated,
+		onSettingsClick,
 		searchQuery = $bindable('')
 	}: Props = $props();

@@ -115,14 +117,16 @@
 							onNewChat?.();
 							void conversationsStore.openNewChat();
 						}
-					: item.route
-						? () => {
-								onNewChat?.();
-								goto(item.route!);
-							}
-						: isSearchOnMobile
-							? undefined
-							: onSearchClick}
+					: item.action === SidebarAction.SETTINGS
+						? () => onSettingsClick?.()
+						: item.route
+							? () => {
+									onNewChat?.();
+									goto(item.route!);
+								}
+							: isSearchOnMobile
+								? undefined
+								: onSearchClick}
 			{@const itemTransition = {
 				delay: !initialized ? i * ICON_STRIP_TRANSITION_DELAY_MULTIPLIER : 0,
 				duration: ICON_STRIP_TRANSITION_DURATION,
@@ -169,14 +173,16 @@
 							onNewChat?.();
 							void conversationsStore.openNewChat();
 						}
-					: item.route
-						? () => {
-								onNewChat?.();
-								goto(item.route!);
-							}
-						: isSearchOnMobile
-							? undefined
-							: onSearchClick}
+					: item.action === SidebarAction.SETTINGS
+						? () => onSettingsClick?.()
+						: item.route
+							? () => {
+									onNewChat?.();
+									goto(item.route!);
+								}
+							: isSearchOnMobile
+								? undefined
+								: onSearchClick}
 			{@const itemTransition = {
 				delay: !initialized ? i * ICON_STRIP_TRANSITION_DELAY_MULTIPLIER : 0,
 				duration: ICON_STRIP_TRANSITION_DURATION,
@@ -1,7 +1,5 @@
 <script lang="ts">
 	import { RefreshCw } from '@lucide/svelte';
-	import { goto } from '$app/navigation';
-	import { page } from '$app/state';
 	import {
 		SettingsChatDesktopSidebar,
 		SettingsChatFields,
@@ -18,21 +16,29 @@
 		SETTINGS_SECTION_SLUGS
 	} from '$lib/constants';
 	import { ColorMode } from '$lib/enums/ui.enums';
-	import { RouterService } from '$lib/services/router.service';
-	import { modelsStore, serverStore, settingsReferrer, settingsStore } from '$lib/stores';
-	import type { SettingsSection } from '$lib/types';
+	import { modelsStore, serverStore, settingsStore } from '$lib/stores';
+	import type { SettingsSection, SettingsSectionTitle } from '$lib/types';
 	import { setMode } from 'mode-watcher';
 	import { fade } from 'svelte/transition';
 	interface Props {
 		initialSection?: string;
-		getSectionHref?: (section: SettingsSection) => string;
+		onSectionChange?: (section: SettingsSectionTitle) => void;
+		onClose?: () => void;
 	}

-	let { getSectionHref, initialSection }: Props = $props();
+	let { initialSection, onClose, onSectionChange }: Props = $props();

-	let activeSlug = $derived(
-		initialSection ?? (page.params as Record<string, string | undefined>).section ?? 'general'
-	);
+	let activeSlug = $derived(initialSection ?? 'general');
+
+	function handleSectionChange(section: SettingsSectionTitle) {
+		const found = SETTINGS_CHAT_SECTIONS.find((s) => s.title === section);
+
+		if (found) {
+			activeSlug = found.slug;
+		}
+
+		onSectionChange?.(section);
+	}

 	let currentSection = $derived(
 		SETTINGS_CHAT_SECTIONS.find((section) => section.slug === activeSlug) ||
@@ -115,7 +121,7 @@
 		}

 		settingsStore.updateMultipleConfig(processedConfig);
-		goto(settingsReferrer.url);
+		onClose?.();
 	}

 	export function reset() {
@@ -123,32 +129,24 @@
 	}
 </script>

-<div in:fade={{ duration: 150 }} class="mx-auto flex h-full w-full flex-col md:pl-8">
-	<div class="flex flex-1 flex-col gap-4 md:flex-row">
+<div in:fade={{ duration: 150 }} class="mx-auto flex h-full w-full flex-col">
+	<div class="flex flex-1 flex-col md:flex-row md:gap-4">
 		<SettingsChatDesktopSidebar
-			getHref={getSectionHref ??
-				((section: SettingsSection) => RouterService.settings(section.slug))}
 			isActive={(section: SettingsSection) => section.slug === activeSlug}
+			onSectionChange={handleSectionChange}
 			sections={SETTINGS_CHAT_SECTIONS}
 		/>

 		<SettingsChatMobileHeader
 			bind:this={mobileHeader}
-			getHref={getSectionHref ??
-				((section: SettingsSection) => RouterService.settings(section.slug))}
 			isActive={(section: SettingsSection) => section.slug === activeSlug}
+			onSectionChange={handleSectionChange}
 			sections={SETTINGS_CHAT_SECTIONS}
 		/>

-		<div class="mx-auto max-w-3xl flex-1">
-			<div class="space-y-6 p-4 md:p-6 md:pt-28">
+		<div class="mx-auto max-w-2xl px-4 flex-1 md:mt-4">
+			<div class="space-y-6 pt-3">
 				<div class="grid">
-					<div class="mb-6 flex items-center gap-2 border-b border-border/30 pb-6 md:flex">
-						<currentSection.icon class="h-5 w-5" />
-
-						<h3 class="text-lg font-semibold">{currentSection.title}</h3>
-					</div>
-
 					{#if currentSection.slug === SETTINGS_SECTION_SLUGS.TOOLS}
 						<SettingsChatToolsTab />
 					{:else if currentSection.slug === SETTINGS_SECTION_SLUGS.IMPORT_EXPORT}
@@ -1,54 +1,31 @@
 <script lang="ts">
-	import { Settings } from '@lucide/svelte';
 	import { ICON_CLASS_DEFAULT } from '$lib/constants';
 	import type { SettingsSection, SettingsSectionTitle } from '$lib/types';

 	interface Props {
 		sections: SettingsSection[];
 		isActive: (section: SettingsSection) => boolean;
-		getHref?: (section: SettingsSection) => string;
 		onSectionChange?: (section: SettingsSectionTitle) => void;
 	}

-	let { getHref, isActive, onSectionChange, sections }: Props = $props();
+	let { isActive, onSectionChange, sections }: Props = $props();
 </script>

-<div class="sticky top-2 hidden w-64 flex-col self-start bg-background py-4 md:flex gap-6">
-	<div class="flex items-center gap-2 py-2">
-		<Settings class="h-5 w-5 md:h-6 md:w-6" />
-
-		<h1 class="text-xl font-semibold md:text-2xl">Settings</h1>
-	</div>
-
+<div class="sticky top-12 hidden w-64 flex-col self-start bg-background md:flex gap-6">
 	<nav class="space-y-1">
 		{#each sections as section (section.title)}
-			{#if getHref}
-				<a
-					class="flex w-full cursor-pointer items-center gap-3 rounded-lg px-3 py-2 text-left text-sm no-underline transition-colors hover:bg-accent {isActive(
-						section
-					)
-						? 'bg-accent text-accent-foreground'
-						: 'text-muted-foreground'}"
-					href={getHref(section)}
-				>
-					<section.icon class={ICON_CLASS_DEFAULT} />
+			<button
+				class="flex w-full cursor-pointer items-center gap-3 rounded-lg px-3 py-2 text-left text-sm transition-colors hover:bg-accent {isActive(
+					section
+				)
+					? 'bg-accent text-accent-foreground'
+					: 'text-muted-foreground'}"
+				onclick={() => onSectionChange?.(section.title)}
+			>
+				<section.icon class={ICON_CLASS_DEFAULT} />

-					<span class="ml-2">{section.title}</span>
-				</a>
-			{:else}
-				<button
-					class="flex w-full cursor-pointer items-center gap-3 rounded-lg px-3 py-2 text-left text-sm transition-colors hover:bg-accent {isActive(
-						section
-					)
-						? 'bg-accent text-accent-foreground'
-						: 'text-muted-foreground'}"
-					onclick={() => onSectionChange?.(section.title)}
-				>
-					<section.icon class={ICON_CLASS_DEFAULT} />
-
-					<span class="ml-2">{section.title}</span>
-				</button>
-			{/if}
+				<span class="ml-2">{section.title}</span>
+			</button>
 		{/each}
 	</nav>
 </div>
@@ -1,5 +1,4 @@
 <script lang="ts">
-	import { Settings } from '@lucide/svelte';
 	import { ScrollCarousel } from '$lib/components/app';
 	import { ICON_CLASS_DEFAULT, UI_DATA_ATTRS } from '$lib/constants';
 	import { BooleanString } from '$lib/enums';
@@ -10,11 +9,10 @@
 	interface Props {
 		sections: SettingsSection[];
 		isActive: (section: SettingsSection) => boolean;
-		getHref?: (section: SettingsSection) => string;
 		onSectionChange?: (section: SettingsSectionTitle) => void;
 	}

-	let { getHref, isActive, onSectionChange, sections }: Props = $props();
+	let { isActive, onSectionChange, sections }: Props = $props();

 	const carousel = useScrollCarousel();

@@ -37,51 +35,26 @@
 	}
 </script>

-<div class="sticky top-0 z-10 flex flex-col bg-background md:hidden">
-	<div class="flex items-center gap-2 px-4 pt-4 pb-2 md:pt-6">
-		<Settings class="h-5 w-5 md:h-6 md:w-6" />
-
-		<h1 class="text-xl font-semibold md:text-2xl">Settings</h1>
-	</div>
-
-	<div class="border-b border-border/30 py-2">
+<div class="flex flex-col bg-background md:hidden sticky top-13 z-50">
+	<div class="border-b border-border/30">
 		<ScrollCarousel alwaysShowArrows {carousel} containerClass="py-2" innerClass="gap-2">
 			{#each sections as section (section.title)}
-				{#if getHref}
-					<a
-						class="flex cursor-pointer items-center gap-2 rounded-lg px-3 py-2 text-sm whitespace-nowrap no-underline transition-colors first:ml-4 last:mr-4 hover:bg-accent {isActive(
-							section
-						)
-							? 'bg-accent text-accent-foreground'
-							: 'text-muted-foreground'}"
-						{...{ [UI_DATA_ATTRS.ACTIVE]: isActive(section) }}
-						href={getHref(section)}
-						onclick={(e: MouseEvent) => {
-							carousel.scrollToCenter(e.currentTarget as HTMLElement);
-						}}
-					>
-						<section.icon class="{ICON_CLASS_DEFAULT} flex-shrink-0" />
+				<button
+					class="flex cursor-pointer items-center gap-2 rounded-lg px-3 py-2 text-sm whitespace-nowrap transition-colors first:ml-4 last:mr-4 hover:bg-accent {isActive(
+						section
+					)
+						? 'bg-accent text-accent-foreground'
+						: 'text-muted-foreground'}"
+					{...{ [UI_DATA_ATTRS.ACTIVE]: isActive(section) }}
+					onclick={(e: MouseEvent) => {
+						onSectionChange?.(section.title);
+						carousel.scrollToCenter(e.currentTarget as HTMLElement);
+					}}
+				>
+					<section.icon class="{ICON_CLASS_DEFAULT} flex-shrink-0" />

-						<span>{section.title}</span>
-					</a>
-				{:else}
-					<button
-						class="flex cursor-pointer items-center gap-2 rounded-lg px-3 py-2 text-sm whitespace-nowrap transition-colors first:ml-4 last:mr-4 hover:bg-accent {isActive(
-							section
-						)
-							? 'bg-accent text-accent-foreground'
-							: 'text-muted-foreground'}"
-						{...{ [UI_DATA_ATTRS.ACTIVE]: isActive(section) }}
-						onclick={(e: MouseEvent) => {
-							onSectionChange?.(section.title);
-							carousel.scrollToCenter(e.currentTarget as HTMLElement);
-						}}
-					>
-						<section.icon class="{ICON_CLASS_DEFAULT} flex-shrink-0" />
-
-						<span>{section.title}</span>
-					</button>
-				{/if}
+					<span>{section.title}</span>
+				</button>
 			{/each}
 		</ScrollCarousel>
 	</div>
@@ -29,7 +29,7 @@
 	}
 </script>

-<div class="sticky bottom-0 mx-auto mt-4 flex w-full justify-between p-6">
+<div class="sticky bottom-0 mx-auto mt-4 flex w-full justify-between pb-4 md:pb-0">
 	<div class="flex gap-2">
 		<Button onclick={handleResetClick} variant="outline">
 			<RotateCcw class="h-3 w-3" />
@@ -1,14 +1,11 @@
 <script lang="ts">
-	import McpLogo from '../mcp/McpLogo.svelte';
-	import { Plus, X } from '@lucide/svelte';
-	import { browser } from '$app/environment';
-	import { goto, replaceState } from '$app/navigation';
+	import { Plus } from '@lucide/svelte';
+	import { replaceState } from '$app/navigation';
 	import { page } from '$app/state';
-	import { ActionIcon, McpServerCard, McpServerCardSkeleton } from '$lib/components/app';
-	import { DialogMcpServerAddNew } from '$lib/components/app/dialogs';
+	import { McpServerCard, McpServerCardSkeleton } from '$lib/components/app';
+	import { DialogMcpResourcesBrowser, DialogMcpServerAddNew } from '$lib/components/app/dialogs';
 	import { Button } from '$lib/components/ui/button';
 	import * as Empty from '$lib/components/ui/empty';
-	import { ROUTES } from '$lib/constants';
 	import { HealthCheckStatus } from '$lib/enums';
 	import { conversationsStore, mcpStore, toolsStore } from '$lib/stores';
 	import { onMount } from 'svelte';
@@ -23,26 +20,7 @@
 	let servers = $derived(mcpStore.getServers());

 	let isAddingServer = $state(false);
-
-	let previousRouteId = $state<string | null>(null);
-
-	$effect(() => {
-		const currentId = page.route.id;
-
-		return () => {
-			previousRouteId = currentId;
-		};
-	});
-
-	function handleClose() {
-		const prevIsMcpServers = previousRouteId === '/mcp-servers';
-
-		if (browser && window.history.length > 1 && !prevIsMcpServers) {
-			history.back();
-		} else {
-			goto(ROUTES.START);
-		}
-	}
+	let isResourcesDialogOpen = $state(false);

 	onMount(() => {
 		if (page.url.searchParams.has('add')) {
@@ -71,25 +49,13 @@
 	}
 </script>

-<div in:fade={{ duration: 150 }} class="flex min-h-[calc(100dvh-4rem)] flex-col">
-	<div class="fixed top-4.5 right-4 z-50 md:hidden">
-		<ActionIcon icon={X} onclick={handleClose} tooltip="Close" />
-	</div>
-
-	<div
-		class="sticky top-0 z-10 mt-4 mb-2 flex items-start gap-4 md:p-4 p-0 px-4 md:justify-between md:px-8"
-	>
-		<div class="flex items-center gap-2">
-			<McpLogo class="h-5 w-5 md:h-6 md:w-6" />
-
-			<h1 class="text-lg font-semibold md:text-2xl">MCP Servers</h1>
-		</div>
-	</div>
-
+<div in:fade={{ duration: 150 }} class="flex flex-col h-full">
 	<DialogMcpServerAddNew bind:open={isAddingServer} />

+	<DialogMcpResourcesBrowser bind:open={isResourcesDialogOpen} />
+
 	{#if servers.length === 0}
-		<div class="flex flex-1 items-center justify-center py-16">
+		<div class="flex flex-1 items-center justify-center pb-20 pt-10 my-auto">
 			<Empty.Root class="max-w-md">
 				<Empty.Header>
 					<Empty.Media variant="icon">
@@ -112,8 +78,8 @@
 		</div>
 	{:else}
 		<div
-			class="grid gap-3 {className}"
-			style="grid-template-columns: repeat(auto-fill, minmax(min(32rem, calc(100dvw - 2rem)), 1fr));"
+			class="grid gap-4 {className}"
+			style="grid-template-columns: repeat(auto-fill, minmax(min(25rem, calc(100dvw - 4rem)), 1fr));"
 		>
 			{#each servers as server (server.id)}
 				{#if isServerPending(server.id, server.enabled)}
@@ -121,6 +87,7 @@
 				{:else}
 					<McpServerCard
 						enabled={conversationsStore.preferences.isMcpServerEnabledForChat(server.id)}
+						onBrowseResources={() => (isResourcesDialogOpen = true)}
 						onDelete={() => mcpStore.removeServer(server.id)}
 						onToggle={async () => {
 							const wasEnabled = conversationsStore.preferences.isMcpServerEnabledForChat(
@@ -1,21 +1,21 @@
 /**
 * Full chat settings page layout with sidebar, mobile header, and content area.
 * Manages local configuration state, section navigation, and context setup.
- * Accepts an optional `initialSection` prop to override the URL-based section resolution.
+ * Accepts an optional `initialSection` prop to set the initial active section.
 */
 export { default as SettingsChat } from './SettingsChat/SettingsChat.svelte';

 /**
 * Desktop sidebar navigation for chat settings.
 * Displays a list of settings sections with icons and titles.
- * Supports both hash-link navigation (via `getHref`) and in-app section switching (via `onSectionChange`).
+ * Switches sections in-app via `onSectionChange`.
 */
 export { default as SettingsChatDesktopSidebar } from './SettingsChatDesktopSidebar.svelte';

 /**
 * Mobile header with a horizontally scrollable section picker for chat settings.
 * Shows chevron buttons for scroll navigation and highlights the active section.
- * Supports both hash-link navigation (via `getHref`) and in-app section switching (via `onSectionChange`).
+ * Switches sections in-app via `onSectionChange`.
 */
 export { default as SettingsChatMobileHeader } from './SettingsChatMobileHeader.svelte';

@@ -5,3 +5,9 @@
 </script>

 <DialogPrimitive.Close bind:ref data-slot="dialog-close" {...restProps} />
+
+<style>
+	:global([data-dialog-close]) {
+		z-index: 999;
+	}
+</style>
@@ -10,7 +10,7 @@
 		class: className,
 		portalProps,
 		ref = $bindable(null),
-		showCloseButton = true,
+		showCloseButton = false,
 		...restProps
 	}: WithoutChildrenOrChild<DialogPrimitive.ContentProps> & {
 		portalProps?: DialogPrimitive.PortalProps;
@@ -25,7 +25,8 @@
 	<DialogPrimitive.Content
 		bind:ref
 		class={cn(
-			`fixed top-[50%] left-[50%] z-50 grid max-h-[100dvh] w-full max-w-[calc(100%-2rem)] translate-x-[-50%] translate-y-[-50%] gap-4 overflow-y-auto rounded-lg border border-border/30 bg-background p-6 shadow-lg duration-200 data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=closed]:fill-mode-forwards data-[state=closed]:zoom-out-95 data-[state=open]:animate-in data-[state=open]:fade-in-0 data-[state=open]:zoom-in-95 sm:max-w-lg md:max-h-[100vh]`,
+			`fixed top-[50%] left-[50%] z-50 grid translate-x-[-50%] translate-y-[-50%] gap-4 overflow-y-auto rounded-lg border border-border/30 bg-background p-6 shadow-lg duration-200 data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=closed]:fill-mode-forwards data-[state=closed]:zoom-out-95 data-[state=open]:animate-in data-[state=open]:fade-in-0 data-[state=open]:zoom-in-95`,
+			'max-h-[100dvh] max-w-full size-auto sm:max-w-lg md:w-auto md:size-auto md:max-w-[calc(100%-2rem)]',
 			className
 		)}
 		data-slot="dialog-content"
@@ -1,20 +1,43 @@
 <script lang="ts">
+	import XIcon from '@lucide/svelte/icons/x';
 	import { cn, type WithElementRef } from '$lib/components/ui/utils';
+	import { Dialog as DialogPrimitive } from 'bits-ui';
 	import type { HTMLAttributes } from 'svelte/elements';

 	let {
 		children,
 		class: className,
 		ref = $bindable(null),
+		showCloseButton = true,
 		...restProps
-	}: WithElementRef<HTMLAttributes<HTMLDivElement>> = $props();
+	}: WithElementRef<HTMLAttributes<HTMLDivElement>> & {
+		showCloseButton?: boolean;
+	} = $props();
 </script>

+<!--
+	Header is `sticky`, so it stays at the top while the dialog body scrolls. The close
+	button lives here (not in the body) so it sticks together with the title. `sticky`
+	makes it the containing block, so the close can be absolutely placed at its corner.
+-->
 <div
 	bind:this={ref}
-	class={cn('flex flex-col gap-2 text-center sm:text-left', className)}
+	class={cn(
+		'flex flex-col gap-2 text-center sm:text-left sticky top-0 z-50 bg-background md:bg-transparent',
+		className
+	)}
 	data-slot="dialog-header"
 	{...restProps}
 >
 	{@render children?.()}
+
+	{#if showCloseButton}
+		<DialogPrimitive.Close
+			class="absolute top-0 right-0 rounded-xs opacity-70 ring-offset-background transition-opacity hover:opacity-100 focus:ring-2 focus:ring-ring focus:ring-offset-2 focus:outline-hidden disabled:pointer-events-none [&_svg]:pointer-events-none [&_svg]:shrink-0 [&_svg:not([class*='size-'])]:size-4"
+		>
+			<XIcon />
+
+			<span class="sr-only">Close</span>
+		</DialogPrimitive.Close>
+	{/if}
 </div>
@@ -15,10 +15,6 @@ export const ROUTES = {
 	MCP_SERVERS: '#/mcp-servers',
 	/** Search — mobile-only full-page conversation search. */
 	SEARCH: '#/search',
-	/** Settings base — for dynamic settings URLs use RouterService. */
-	SETTINGS: '#/settings',
-	/** Exit destination for the settings view (fallback when no referrer). */
-	SETTINGS_EXIT: '#/',
 	/** Root — start of the app. */
 	START: '#/'
 } as const;
@@ -1,6 +1,4 @@
-import { ROUTES } from './routes.constants';
 import { Package, Search, Settings, SquarePen } from '@lucide/svelte';
-import McpLogo from '$lib/components/app/mcp/McpLogo.svelte';
 import { SidebarAction, ToolSource } from '$lib/enums';
 import type { DesktopIconStripItem } from '$lib/types';

@@ -64,15 +62,8 @@ export const SIDEBAR_ACTIONS_ITEMS: DesktopIconStripItem[] = [
 	},
 	{ icon: Search, keys: ['cmd', 'k'], tooltip: 'Search' },
 	{
-		activeRouteId: '/mcp-servers',
-		icon: McpLogo,
-		route: ROUTES.MCP_SERVERS,
-		tooltip: 'MCP Servers'
-	},
-	{
-		activeUrlIncludes: '#/settings',
+		action: SidebarAction.SETTINGS,
 		icon: Settings,
-		route: `${ROUTES.SETTINGS}/general`,
 		tooltip: 'Settings'
 	}
 ];
@@ -23,7 +23,8 @@ export enum ScrollCarouselVariant {
 * Sidebar icon strip actions handled directly by the sidebar.
 */
 export enum SidebarAction {
-	NEW_CHAT = 'new-chat'
+	NEW_CHAT = 'new-chat',
+	SETTINGS = 'settings'
 }

 /**
@@ -1,45 +0,0 @@
-import { beforeNavigate } from '$app/navigation';
-import { page } from '$app/state';
-import { ROUTES } from '$lib/constants';
-import { settingsReferrer } from '$lib/stores';
-
-export interface ChatSettings {
-	reset: () => void;
-}
-
-export function useSettingsNavigation() {
-	const subroute = $state({
-		activePanel: 'chat' as 'chat' | 'settings' | 'mcp',
-		chatSettingsRef: undefined as ChatSettings | undefined
-	});
-	const isSettingsRoute = $derived(!!page.route.id?.startsWith('/settings'));
-
-	beforeNavigate(({ from, to }) => {
-		if (to?.route?.id?.startsWith('/settings') && !from?.route?.id?.startsWith('/settings')) {
-			settingsReferrer.url = window.location.hash || ROUTES.START;
-		}
-	});
-
-	$effect(() => {
-		if (subroute.activePanel === 'settings' && subroute.chatSettingsRef) {
-			subroute.chatSettingsRef.reset();
-		}
-	});
-
-	// Return to chat when navigating to a new route
-	$effect(() => {
-		void page.url;
-
-		subroute.activePanel = 'chat';
-	});
-
-	return {
-		get isSettingsRoute() {
-			return isSettingsRoute;
-		},
-
-		get panel() {
-			return subroute;
-		}
-	};
-}
@@ -307,7 +307,6 @@ export { SandboxService } from './sandbox.service';
 *
 * **Key Responsibilities:**
 * - Build chat URLs for specific conversations: `RouterService.chat(id)` → `#/chat/:id`
- * - Build settings URLs for sections: `RouterService.settings(section)` → `#/settings/:section`
 *
 * @see ROUTES in constants/routes.ts — static route base paths
 */
@@ -1,8 +1,7 @@
 /**
 * RouterService - Builds app route paths
 *
- * Returns chat and settings route strings from a single source of truth
- * (ROUTES). No state.
+ * Returns chat route strings from a single source of truth (ROUTES). No state.
 */

 import { ROUTES } from '$lib/constants';
@@ -11,8 +10,4 @@ export class RouterService {
 	static chat(id: string): string {
 		return `${ROUTES.CHAT}/${id}`;
 	}
-
-	static settings(section: string): string {
-		return `${ROUTES.SETTINGS}/${section}`;
-	}
 }
@@ -49,8 +49,6 @@ export { uiStore } from './ui.svelte';
 // SETTINGS / UI PREFERENCES
 export { settingsStore } from './settings/index.svelte';

-export { settingsReferrer } from './settings/referrer.svelte';
-
 export { permissionsStore } from './permissions.svelte';

 // TOOLS
@@ -1,19 +0,0 @@
-/**
- * settingsReferrer - Remembers the settings route to return to after exit
- *
- * Tracks the last settings section the user was on so the app can return
- * there after a fallback exit. Standalone reactive value, no host.
- */
-
-import { ROUTES } from '$lib/constants';
-
-let _url = $state<string>(ROUTES.SETTINGS_EXIT);
-
-export const settingsReferrer = {
-	get url() {
-		return _url;
-	},
-	set url(value: string) {
-		_url = value;
-	}
-};
@@ -1,5 +0,0 @@
-<script lang="ts">
-	import { SettingsMcpServers } from '$lib/components/app/settings';
-</script>
-
-<SettingsMcpServers class="mx-auto w-full p-4 md:p-8 md:py-8" />
@@ -1,38 +0,0 @@
-<script lang="ts">
-	import { X } from '@lucide/svelte';
-	import { browser } from '$app/environment';
-	import { goto } from '$app/navigation';
-	import { page } from '$app/state';
-	import { ActionIcon } from '$lib/components/app';
-	import { ROUTES } from '$lib/constants';
-
-	let { children } = $props();
-
-	let previousRouteId = $state<string | null>(null);
-
-	$effect(() => {
-		const currentId = page.route.id;
-
-		return () => {
-			previousRouteId = currentId;
-		};
-	});
-
-	function handleClose() {
-		const prevIsSettings = previousRouteId?.startsWith('/settings');
-
-		if (browser && window.history.length > 1 && !prevIsSettings) {
-			history.back();
-		} else {
-			goto(ROUTES.SETTINGS_EXIT);
-		}
-	}
-</script>
-
-<div class="fixed top-4.5 right-4 z-50 md:hidden">
-	<ActionIcon icon={X} onclick={handleClose} tooltip="Close" />
-</div>
-
-<div class="min-h-full">
-	{@render children?.()}
-</div>
@@ -1,15 +0,0 @@
-<script lang="ts">
-	import { afterNavigate, replaceState } from '$app/navigation';
-	import { page } from '$app/state';
-	import { SettingsChat } from '$lib/components/app/settings';
-	import { SETTINGS_SECTION_SLUGS } from '$lib/constants';
-	import { RouterService } from '$lib/services';
-
-	afterNavigate(() => {
-		if (!page.params.section) {
-			replaceState(RouterService.settings(SETTINGS_SECTION_SLUGS.GENERAL), {});
-		}
-	});
-</script>
-
-<SettingsChat initialSection={(page.params as Record<string, string | undefined>).section} />
Author	SHA1	Message	Date
Xuan-Son NguyenandGitHub	925e117994	llama: add token ID tracking to KV cell (#27762 ) * kv: track token id * rm get_prev_tokens, move it to the main pr * nits * add get_prev_tokens	2026-08-26 23:34:28 +02:00
Aleksander GrygierandGitHub	539f24529b	ui: Move Settings and MCP Servers routes to dialog-based views (#27744 ) * ui : open MCP servers in a dialog from the chat form Replace the MCP servers submenu with a single "MCP Servers" item that opens a new DialogMcpServers dialog instead of navigating to the /mcp-servers route. Assisted-by: pi * ui : browse MCP resources from the server card Make the Resources capability badge clickable so it opens the MCP resources browser dialog, and drop the page-only chrome from SettingsMcpServers. Assisted-by: pi * ui : remove mcp-servers route and sidebar entry MCP servers are now managed in a dialog, so drop the dedicated route and the sidebar icon that navigated to it. Assisted-by: pi * ui : remove unused MCP servers submenu component The submenu was replaced by the MCP servers dialog, so delete the component and its export. Assisted-by: pi * feat(ui): add DialogSettingsChat dialog * refactor(ui): switch SettingsChat to in-app section navigation * feat(ui): open settings as dialog from sidebar * refactor(ui): remove settings route and URL-based settings navigation * fix(ui): adjust MCP dialogs for new base sizing * chore: Formatting & linting	2026-08-26 21:07:24 +02:00
Aleksander GrygierandGitHub	0379a19f09	ui: Update Dialog component styling (#27743 ) * feat(ui): make base dialog responsive and support sticky headers * ui: move dialog close button to the sticky header Assisted-by: pi * chore: Formatting & linting	2026-08-26 20:19:19 +02:00
Ruben OrtlamandGitHub	5e6a37cb11	vulkan: warptiles currently assume warp sizes <= 64, clamp to work around larger warps (#27726 )	2026-08-26 19:02:06 +03:00
Pranav UttarkarandGitHub	bf94216469	Implemented vulkan cross_entropy_loss and cross_entropy_loss_back (#27216 )	2026-08-26 16:49:32 +02:00
Radoslav GerganovandGitHub	d0132a680a	rpc : implement event and async backend APIs (#18626 ) * rpc : implement event and async backend APIs * cache responses from RPC_CMD_GET_ALLOC_SIZE	2026-08-26 17:34:46 +03:00
Aleksander GrygierandGitHub	4d19b28769	ci: Clean up UI builds from releases (#27706 ) * ci : inline UI version resolution into ui-build.yml * ci : build UI once and reuse the artifact in release jobs Server jobs now extract the ui-build artifact into tools/ui/dist instead of npm-building the UI. Also removes the get-version job and the no-op -DHF_UI_VERSION flags. Assisted-by: pi:Kimi-K3 * ui : disable the npm UI build by default (LLAMA_BUILD_UI=OFF) The flag now only controls building the UI from source via npm. The UI is still embedded by default from local tools/ui/dist or the prebuilt download (LLAMA_USE_PREBUILT_UI=ON). CI jobs no longer npm-build the UI; server-sanitize does not need node anymore. Assisted-by: pi:Kimi-K3 * ci : rename the ui-build artifact to llama-ui.zip Consistent with the other artifact names in the Actions summary. Assisted-by: pi:Kimi-K3 * ci : clarify the windows artifact merge in release.yml The windows-cuda/vulkan/sycl jobs build only the backend library; llama-server (with the embedded UI) is injected into their zips from the windows-cpu package during the release. State this in the job comments and use accurate wording in the merge step. Assisted-by: pi:Kimi-K3	2026-08-26 14:12:09 +02:00
David FriehsandGitHub	fc35562ba4	cuda: unblock mmq for MoE on sm_60 (#26264 ) * cuda: unblock mmq for MoE on sm_60 * cuda: duplicate mmq-config-pascal for dp4a and older * cuda: reduce occupancy on non-dp4a pascal for Q2_K, Q4_K, Q5_K, Q6_K	2026-08-26 18:35:54 +08:00
Sigbjørn SkjæretandGitHub	da9b5d68c3	ci : make cache bucket public (#27728 ) * check for hf token * make bucket public	2026-08-26 12:08:23 +02:00
Daniel BeveniusandGitHub	dac869b0a0	conversion : fix Nemotron 3.5 Lightning layers (#27729 ) This commit contains a fix for the conversion of NVIDIA Nemotron 3.5 Lightning which currently incorrectly converts when using a transformers version later than 5.5.1. When converting using [convert](https://github.com/ggml-org/convert) the transformers version is 5.13.1 and this produces the following: ```console WARNING:gguf.gguf_writer:Duplicated key name 'nemotron_h_moe.attention.head_count_kv', overwriting it with new value [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] of type ARRAY ``` This does not happen with transformers 5.5.1. The reason seems to be that the configuration is different in later versions, for example when using 5.13.1 the configuration block looks like this: ```console transformers 5.13.1 raw has layers_block_type: True autoconfig has layers_block_type: True autoconfig layers_block_type: [ 'linear_attention', 'moe', 'linear_attention', 'moe', 'linear_attention', 'full_attention', 'moe', ... ] ``` And with 5.5.1 we get: ```console transformers 5.5.1 raw has layers_block_type: True autoconfig has layers_block_type: True autoconfig layers_block_type: [ 'mamba', 'moe', 'mamba', 'moe', 'mamba', 'attention', 'moe' ... ] ``` In our conversion script we only match for attention, not full attention which is causing this issue. With the changes in this commit the output with transformers 5.13.1 will be: ```console (venv) $ gguf-dump models/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16.gguf \| grep head_count_kv INFO:gguf-dump:* Loading: models/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16.gguf 29: [INT32] \| 52 \| nemotron_h_moe.attention.head_count_kv = [0, 0, 0, 0, 0, 2, ...] ``` Resolves: https://github.com/ggml-org/llama.cpp/issues/27718 Refs: https://github.com/ggml-org/convert/actions/runs/32949047680/job/98116096069#step:5:2391	2026-08-26 12:05:31 +02:00
Dominik Pantaleoni GitHub dpantaleoni forforever73	11cd988428	ggml-metal: add chunked SSD MMA for Mamba-2 prefill optimization (#26647 ) * metal: WIP chunked SSD SSM_SCAN kernels for multi-token prefill * metal: drop scalar SSD path; MMA + sequential tail * drop WIP ssm scan test noise * remove state_from_dst and rename CS and NSG constants * remove unrelated added whitespace padding * added clarity to mma_tokens calculation * added clarity to use_mma bool checks * added comments to metal ssd op constants for clarity * reserve K tokens for sequential kernel rollback snapshots * reset concurrency between mma and seq tail * remove print args no longer used * fixed comment to no longer point to specific line * add FC_SSM_SCAN so seq path skips token offlset unless it's mma tail * added changes to new ssm.metal for rebase after ggml-metal.metal refactor * specialize ssm_scan tail with a template instead of a function constant --------- Co-authored-by: dpantaleoni <dominikpantaleoni@gmail.com> Co-authored-by: forforever73 <690105611@qq.com>	2026-08-26 11:57:07 +03:00
Max KrasnyanskyandGitHub	5d5cb4c3a4	ggml-meta: propagate buffer usage and call init on the new tensors (#27586 )	2026-08-26 08:27:51 +03:00
Jonathan ClohessyandGitHub	d222767c7a	kleidiai: Rework KleidiAI Build System/Integration (#26077 ) * Rework KleidiAI Build System/Integration Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com> * Add fp16 guard, and fix cmake caching issue Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com> * Fix formatting, and rebase issue Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com> --------- Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>	2026-08-25 14:07:29 -07:00
Mario LimoncielloandGitHub	eab8ee41f8	ci : update OS used for ROCM to Ubuntu 24.04 (#27681 ) This matches what other build targets use and also what AMD advertises wheels as supporting.	2026-08-25 20:14:16 +03:00
Ryan C GitHub Ryan Churaman	b114b47397	rpc: support apple RDMA as an RPC transport (#26421 ) * rpc: support apple RDMA as an RPC transport * remove set_tensor micro optimization, rpc socket pinning per CR * remove transparent reconnect * trigger apple builds on RPC changes --------- Co-authored-by: Ryan Churaman <rschu@meta.com>	2026-08-25 20:12:15 +03:00
krzsztf GitHub Sigbjørn Skjæret	0a5ac49bce	devops: use GGML_NATIVE=OFF for OpenVINO (#27338 ) * devops: use GGML_NATIVE=OFF for OpenVINO Same as in other Dockerfiles. Should fix #23100 * enable backend dl and cpu all variants --------- Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>	2026-08-25 19:37:11 +03:00
KyozzzandGitHub	1729ed5371	server : reject prefilled assistant messages with tool calls (#27626 ) * server: fix tool calls getting silently stripped with --prefill-assistant Last assistant carries tool_calls + --prefill-assistant is on → request flips into continuation mode, add_generation_prompt forced off, tail rebuilt from reasoning_content + content only. Tool calls just vanish. - Auto-continuation now skips trailing assistant msgs that have tool calls - continue_final_message on those throws a clear error instead of silently corrupting the prompt - Regression tests included, red before / green after Fixes #27588 Developed with AI assistance, disclosed per the contribution policy. * server : address review: fail on prefill-assistant + trailing tool_calls Move validation into oaicompat_chat_params_parse (next to the existing two-or-more-assistant check) and remove it from common_chat_templates_apply, which has no precedent for validation. Drop the regression tests. Per review: --prefill-assistant with a trailing assistant message containing tool calls is not supported and should fail loudly.	2026-08-25 09:35:22 -05:00