quantize: cap working memory size to avoid loading big tensors onto RAM

* args: add --video-* CLI arguments

* gen docs

* nits

* add mtmd_helper_init_opt

.devops/openvino.Dockerfile

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

.github/actions/ccache-buckets/action.yml

@@ -0,0 +1,95 @@
+name: "ccache-buckets"
+description: "Save/restore latest GitHub Actions ccache matching a key prefix to/from HF buckets"
+inputs:
+  key:
+    description: "Cache key prefix to match and load"
+    required: true
+  folder:
+    description: "Bucket folder containing ccache files"
+    required: true
+  evict-old-files:
+    description: "Corresponds to the ccache --evict-older-than AGE option, where AGE is the number of seconds or days followed by the 's' or 'd' suffix respectively."
+    default: ''
+  save:
+    description: "Save ccache"
+    required: false
+    default: false
+    type: boolean
+  hf_bucket:
+    description: 'Hugging Face buckets path'
+    required: true
+
+runs:
+  using: "composite"
+  steps:
+    - name: Install Hugging Face Hub CLI
+      shell: bash
+      run: |
+        python3 -m venv .venv-hf
+        .venv-hf/bin/pip install -U huggingface_hub==1.28.0
+
+    - name: Restore ccache from buckets
+      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
+            CACHE_PATH=$(hf buckets list "hf://buckets/${{ inputs.hf_bucket }}/${{ inputs.folder }}" --json | jq -r '[.[] | select(.type == "file") | select(.path | startswith("${{ inputs.folder }}/${{ inputs.key }}") and endswith(".tar.gz"))] | sort_by(.path) | last | .path // ""')
+            if [[ -n "$CACHE_PATH" ]]; then
+                echo "Restoring ccache from '$CACHE_PATH'."
+                hf buckets cp "hf://buckets/${{ inputs.hf_bucket }}/$CACHE_PATH" ccache_bucket.tar.gz
+                mkdir -p ccache_bucket
+                if tar -xzf ccache_bucket.tar.gz -C ccache_bucket; then
+                    rm -rf "$CCACHE_DIR"
+                    mv ccache_bucket "$CCACHE_DIR"
+                    ccache -z
+                fi
+                rm ccache_bucket.tar.gz
+            else
+                echo "No ccache found."
+            fi
+        else
+            echo "'$CCACHE_DIR' not found."
+        fi
+
+    - name: Save ccache to buckets
+      if: ${{ inputs.save == 'true' }}
+      shell: bash
+      run: |
+        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
+        fi
+
+    - name: Remove old ccache files from buckets
+      if: ${{ inputs.save == 'true' }}
+      shell: bash
+      run: |
+        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

.github/actions/ccache-clear/action.yml

@@ -1,23 +1,50 @@
 # note: place this as the last step of the job, so the new cache is saved by "Post ccache" right after the old one is cleared
 name: "ccache-clear"
-description: "Delete all GitHub Actions caches matching a key prefix"
+description: "Delete GitHub Actions caches matching a key prefix, oldest first"
 inputs:
   key:
     description: "Cache key prefix to match and delete"
     required: true
+  older:
+    description: "Only delete caches created more than this long ago (e.g. 90m, 1h, 1d). By default all matching caches are deleted"
+    required: false
+    default: ""
+  min:
+    description: "Stop deleting if fewer than this many caches would remain (e.g. 1). By default there is no minimum"
+    required: false
+    default: "0"
+  dry-run:
+    description: "Only print the caches that would be deleted, without deleting them"
+    required: false
+    default: "false"
 
 runs:
   using: "composite"
   steps:
+    - name: Install GitHub CLI if missing
+      shell: bash
+      run: |
+        # e.g. in container jobs, where it is not preinstalled
+        if ! command -v gh >/dev/null 2>&1; then
+          echo "GitHub CLI not found, installing..."
+          if ! command -v curl >/dev/null 2>&1; then
+            apt-get update >/dev/null 2>&1 || true
+            apt-get install -y curl >/dev/null 2>&1 || true
+          fi
+          mkdir -p -m 755 /etc/apt/keyrings
+          curl -fsSL https://cli.github.com/packages/githubcli-archive-keyring.gpg | tee /etc/apt/keyrings/githubcli-archive-keyring.gpg >/dev/null
+          chmod go+r /etc/apt/keyrings/githubcli-archive-keyring.gpg
+          echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/githubcli-archive-keyring.gpg] https://cli.github.com/packages stable main" > /etc/apt/sources.list.d/github-cli.list
+          apt-get update >/dev/null 2>&1 || true
+          apt-get install -y gh || { echo "Failed to install GitHub CLI (gh)" >&2; exit 1; }
+        fi
+        command -v gh >/dev/null 2>&1 || { echo "GitHub CLI (gh) is required but could not be installed" >&2; exit 1; }
+
     - name: Clear caches
       shell: bash
       run: |
-        CACHES=$(gh cache list --key "ccache-${{ inputs.key }}" --json id,key --jq '.[] | "\(.id) \(.key)"' 2>/dev/null)
-        if [ -z "$CACHES" ]; then
-          echo "No caches found with key prefix: ${{ inputs.key }}"
-          exit 0
-        fi
-        while read -r id key; do
-          echo "Deleting cache: $id ($key)"
-          gh cache delete "$id"
-        done <<< "$CACHES"
+        bash scripts/ccache-clear.sh \
+          --key "${{ inputs.key }}" \
+          --older "${{ inputs.older }}" \
+          --min "${{ inputs.min }}" \
+          ${{ inputs.dry-run == 'true' && '--dry-run' || '' }}

.github/workflows/build-apple.yml

@@ -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:
@@ -73,6 +74,16 @@ jobs:
           cd build
           ctest -L main -E "test-llama-archs" --verbose --timeout 900
 
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: apple-arm64
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
+
   macos-latest-x64:
     runs-on: macos-15-intel
 
@@ -109,6 +120,16 @@ jobs:
           cd build
           ctest -L main --verbose --timeout 900
 
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: apple-x64
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
+
   macos-latest-ios-xcode:
     runs-on: macos-latest
 
@@ -163,14 +184,6 @@ jobs:
         id: checkout
         uses: actions/checkout@v6
 
-      # TODO: this likely does not do anything - if yes, remove it
-      - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
-        with:
-          key: apple-tvos
-          evict-old-files: 1d
-          save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
-
       - name: Build
         id: cmake_build
         run: |
@@ -196,14 +209,6 @@ jobs:
         id: checkout
         uses: actions/checkout@v6
 
-      # TODO: this likely does not do anything - if yes, remove it
-      - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
-        with:
-          key: apple-visionos
-          evict-old-files: 1d
-          save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
-
       - name: Build
         id: cmake_build
         run: |
@@ -234,14 +239,6 @@ jobs:
         id: checkout
         uses: actions/checkout@v6
 
-      # TODO: this likely does not do anything - if yes, remove it
-      - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
-        with:
-          key: apple-swift
-          evict-old-files: 1d
-          save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
-
       - name: Download xcframework artifact
         uses: actions/download-artifact@v7
         with:

.github/workflows/build-cpu.yml

@@ -117,6 +117,18 @@ jobs:
           ./bin/llama-convert-llama2c-to-ggml --copy-vocab-from-model ./tok512.bin --llama2c-model stories260K.bin --llama2c-output-model stories260K.gguf
           ./bin/llama-completion -m stories260K.gguf -p "One day, Lily met a Shoggoth" -n 500 -c 256
 
+      # note: real deletion only on push to master (same condition as the ccache save),
+      #       dry-run otherwise (the token is read-only on PRs from forks)
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: cpu-${{ matrix.os }}
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
+
   windows:
     name: windows / ${{ matrix.build }}
     runs-on: windows-2025
@@ -203,3 +215,13 @@ jobs:
       #     cd build
       #     $env:LLAMA_SKIP_TESTS_SLOW_ON_EMULATOR = 1
       #     & $sde -future -- ctest -L main -C Release --verbose --timeout 900
+
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: cpu-windows-2025-${{ matrix.build }}
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}

.github/workflows/build-cuda-ubuntu.yml

@@ -50,14 +50,22 @@ jobs:
           DEBIAN_FRONTEND: noninteractive
         run: |
           apt update
-          apt install -y cmake build-essential ninja-build libgomp1 git libssl-dev
+          apt install -y cmake build-essential ninja-build libgomp1 git libssl-dev jq python3 python3-venv python3-pip
 
       - name: ccache
         uses: ggml-org/ccache-action@v1.2.21
         with:
           key: cuda-ubuntu-24.04-cuda
-          evict-old-files: 1d
-          save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+          save: false
+
+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: cuda-ubuntu-24.04-cuda
+          folder: llama.cpp
+          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
@@ -72,6 +80,18 @@ jobs:
             -DGGML_CUDA_CUB_3DOT2=ON
           cmake --build build
 
+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: cuda-ubuntu-24.04-cuda
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true
+
   hip:
     runs-on: ubuntu-22.04
     container: rocm/dev-ubuntu-22.04:6.1.2
@@ -85,14 +105,22 @@ jobs:
         id: depends
         run: |
           sudo apt-get update
-          sudo apt-get install -y build-essential git cmake rocblas-dev hipblas-dev libssl-dev rocwmma-dev
+          sudo apt-get install -y build-essential git cmake rocblas-dev hipblas-dev libssl-dev rocwmma-dev jq python3-venv
 
       - name: ccache
         uses: ggml-org/ccache-action@v1.2.21
         with:
           key: cuda-ubuntu-22.04-hip
-          evict-old-files: 1d
-          save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+          save: false
+
+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: cuda-ubuntu-22.04-hip
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache
 
       - name: Build with native CMake HIP support
         id: cmake_build
@@ -103,6 +131,18 @@ jobs:
             -DGGML_HIP=ON
           cmake --build build --config Release -j $(nproc)
 
+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: cuda-ubuntu-22.04-hip
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true
+
   musa:
     runs-on: ubuntu-22.04
     container: mthreads/musa:rc4.3.0-devel-ubuntu22.04-amd64
@@ -116,14 +156,22 @@ jobs:
         id: depends
         run: |
           apt-get update
-          apt-get install -y build-essential git cmake libssl-dev
+          apt-get install -y build-essential git cmake libssl-dev jq
 
       - name: ccache
         uses: ggml-org/ccache-action@v1.2.21
         with:
           key: cuda-ubuntu-22.04-musa
-          evict-old-files: 1d
-          save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+          save: false
+
+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: cuda-ubuntu-22.04-musa
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache
 
       - name: Build with native CMake MUSA support
         id: cmake_build
@@ -131,3 +179,15 @@ jobs:
           cmake -B build -S . \
             -DGGML_MUSA=ON
           time cmake --build build --config Release -j $(nproc)
+
+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: cuda-ubuntu-22.04-musa
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true

.github/workflows/build-opencl.yml

@@ -80,3 +80,13 @@ jobs:
         run: |
           cmake -S . -B build -G "Ninja Multi-Config" -D CMAKE_TOOLCHAIN_FILE=cmake/arm64-windows-llvm.cmake -DCMAKE_PREFIX_PATH="$env:RUNNER_TEMP/opencl-arm64-release" -DGGML_OPENCL=ON -DGGML_OPENCL_USE_ADRENO_KERNELS=ON -DLLAMA_BUILD_BORINGSSL=ON
           cmake --build build --config Release -j ${env:NUMBER_OF_PROCESSORS}
+
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: opencl-windows-2025-x64
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}

.github/workflows/build-openvino.yml

@@ -167,3 +167,13 @@ jobs:
 
           cd build
           ctest --test-dir ReleaseOV -L main -E "test-llama-archs|test-recurrent-state-rollback-nemotron-h" -C Release --verbose --timeout 3000
+
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: openvino-windows-2022
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}

.github/workflows/build-sycl.yml

@@ -96,6 +96,16 @@ jobs:
             -DGGML_SYCL_F16=${{ matrix.fp16 }}
           time cmake --build build --config Release -j $(nproc)
 
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: sycl-ubuntu-24-${{ matrix.build }}
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
+
   windows-latest-sycl:
     runs-on: windows-2022
 
@@ -139,3 +149,13 @@ jobs:
       - name: Build
         id: cmake_build
         run:  examples/sycl/win-build-sycl.bat
+
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: sycl-windows-latest
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}

.github/workflows/build-vulkan.yml

@@ -55,7 +55,7 @@ jobs:
       - name: ccache
         uses: ggml-org/ccache-action@v1.2.21
         with:
-          key: vulkan-ubuntu-24.04-arm-new
+          key: vulkan-ubuntu-24.04-arm
           variant: ccache
           evict-old-files: 1d
           save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
@@ -73,6 +73,16 @@ jobs:
         run: |
           time cmake --build build -j $(nproc)
 
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: vulkan-ubuntu-24.04-arm
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
+
   ubuntu-llvmpipe:
     runs-on: ubuntu-24.04
 
@@ -128,6 +138,16 @@ jobs:
           # test-backend-ops is too slow on llvmpipe, skip it
           ctest -L main -E test-backend-ops --verbose --timeout 900
 
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: vulkan-ubuntu-24.04-llvmpipe
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
+
   windows:
     runs-on: windows-2025
 
@@ -180,3 +200,13 @@ jobs:
         run: |
           cd build
           ctest -L main -C Release --verbose --timeout 900
+
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: cpu-windows-2025-x64-vulkan
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}

.github/workflows/build-wasm.yml

@@ -88,3 +88,13 @@ jobs:
             -DEMDAWNWEBGPU_DIR=emdawnwebgpu_pkg
 
           time cmake --build build-wasm --config Release --target test-backend-ops -j $(nproc)
+
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: webgpu-ubuntu-24.04-arm-wasm
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}

.github/workflows/build-webgpu.yml

@@ -101,6 +101,16 @@ jobs:
           cd build
           ctest -L main --verbose --timeout 900
 
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: webgpu-macos-latest
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
+
   ubuntu:
     runs-on: ubuntu-24.04
 
@@ -153,3 +163,13 @@ jobs:
           # This is using llvmpipe and runs slower than other backends
           # test-backend-ops is too slow on llvmpipe, skip it
           ctest -L main -E test-backend-ops --verbose --timeout 900
+
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: webgpu-ubuntu-24.04
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}

.github/workflows/docker.yml

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

.github/workflows/hip-quality-check.yml

@@ -84,3 +84,13 @@ jobs:
           cd build
           make -j $(nproc) 2>&1 | tee metrics.log | grep -v 'Rpass-analysis=kernel-resource-usage\|remark:\|^$'
           python3 ../scripts/hip/gcn-cdna-vgpr-check.py metrics.log
+
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: hip-quality-check-ubuntu-22.04
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}

.github/workflows/make-release.yml

@@ -70,6 +70,7 @@ jobs:
           cat nightly-tag.txt
 
       - name: Create release
+        id: create_release
         if: ${{ github.event.inputs.dry_run == 'false' }}
         uses: ggml-org/action-create-release@v1
         env:
@@ -83,11 +84,13 @@ jobs:
 
             New version has been released.
 
+            ## Assets
+
             ${{ steps.desc.outputs.nightly }}
 
-            **Web UI:** the `nightly-tag.txt` asset contains the tag of the corresponding nightly release
+            ## More info
 
-            **More info:** [dist : releases and versioning of ggml-org projects](https://github.com/ggml-org/ggml/discussions/1579)
+            - [Releases and versioning of `ggml-org` projects](https://github.com/ggml-org/ggml/discussions/1579)
 
             ## ${{ steps.desc.outputs.changelog_title }}
 

.github/workflows/release.yml

@@ -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,11 +736,18 @@ 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:
           key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
           evict-old-files: 1d
+          max-size: "1G"
 
       # - name: Cache ROCm Installation
       #   id: cache-rocm
@@ -878,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' }}
@@ -908,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' }}
@@ -977,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]
@@ -1005,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:
@@ -1083,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' }}
@@ -1117,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:
@@ -1194,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:
@@ -1236,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
@@ -1286,126 +1244,135 @@ jobs:
         with:
           key: release-ubuntu-24.04-sycl-${{ matrix.build }}
 
-#   ubuntu-22-rocm:
-#     needs: [check-release, get-version]
-#     if: ${{ needs.check-release.outputs.should_release == 'true' }}
+  ubuntu-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
+    permissions:
+      actions: write
 
-#     strategy:
-#       matrix:
-#         include:
-#           - ROCM_VERSION: "7.14.0"
-#             gpu_targets: "gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201"
-#             build: 'x64'
+    strategy:
+      matrix:
+        include:
+          - ROCM_VERSION: "7.14.0"
+            gpu_targets: "gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201"
+            build: 'x64'
 
-#     steps:
-#       - name: Clone
-#         id: checkout
-#         uses: actions/checkout@v6
-#         with:
-#           fetch-depth: 0
+    steps:
+      - name: Clone
+        id: checkout
+        uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
 
-#       - name: Setup Node.js
-#         uses: actions/setup-node@v6
-#         with:
-#           node-version: "24"
-#           cache: "npm"
-#           cache-dependency-path: "tools/ui/package-lock.json"
+      - name: Download UI build
+        uses: actions/download-artifact@v7
+        with:
+          name: llama-ui.zip
+          path: tools/ui/dist
 
-#       - name: Free up disk space
-#         uses: ggml-org/free-disk-space@v1.3.1
-#         with:
-#           tool-cache: true
+      - name: Free up disk space
+        uses: ggml-org/free-disk-space@v1.3.1
+        with:
+          tool-cache: true
 
-#       # - name: ccache
-#       #   uses: ggml-org/ccache-action@v1.2.21
-#       #   with:
-#       #     key: release-ubuntu-22.04-rocm-${{ matrix.ROCM_VERSION }}
+      - name: ccache
+        uses: ggml-org/ccache-action@v1.2.21
+        with:
+          key: release-ubuntu-24.04-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
+          evict-old-files: 1d
+          max-size: "1G"
 
-#       - name: Dependencies
-#         id: depends
-#         run: |
-#           sudo apt install -y build-essential git cmake wget
+      - name: Tune ccache for reinstalled ROCm toolchain
+        run: |
+          # ROCm is pip-installed fresh each run, so the clang binary's mtime
+          # changes every time. With the default compiler_check=mtime that
+          # invalidates the cache; hash compiler contents instead so warm
+          # builds hit.
+          ccache --set-config=compiler_check=content
+          ccache --set-config=sloppiness=time_macros,include_file_mtime,include_file_ctime
 
-#       - name: Setup TheRock with Wheels
-#         id: therock_env
-#         run: |
-#           # Create Python virtual environment
-#           python3 -m venv .venv
-#           source .venv/bin/activate
+      - name: Dependencies
+        id: depends
+        run: |
+          sudo apt install -y build-essential git cmake wget
 
-#           # Install ROCm wheels for build
-#           # libraries = HIP runtime and CMake configs needed for linking
-#           # devel = compilers, headers, static libs
-#           python -m pip install --upgrade pip
-#           python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ matrix.ROCM_VERSION }}"
+      - name: Setup TheRock with Wheels
+        id: therock_env
+        run: |
+          # Create Python virtual environment
+          python3 -m venv .venv
+          source .venv/bin/activate
 
-#           # Get ROCm installation paths using the rocm-sdk CLI tool
-#           ROCM_PATH=$(rocm-sdk path --root)
-#           CMAKE_PATH=$(rocm-sdk path --cmake)
-#           BIN_PATH=$(rocm-sdk path --bin)
-#           echo "ROCM_PATH=$ROCM_PATH"
-#           echo "CMAKE_PATH=$CMAKE_PATH"
-#           echo "BIN_PATH=$BIN_PATH"
+          # Install ROCm wheels for build
+          # libraries = HIP runtime and CMake configs needed for linking
+          # devel = compilers, headers, static libs
+          python -m pip install --upgrade pip
+          python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ matrix.ROCM_VERSION }}"
 
-#           # Set environment variables
-#           echo "ROCM_PATH=$ROCM_PATH" >> $GITHUB_ENV
-#           echo "CMAKE_PREFIX_PATH=$CMAKE_PATH" >> $GITHUB_ENV
-#           echo "HIP_PATH=$ROCM_PATH" >> $GITHUB_ENV
-#           echo "PATH=$BIN_PATH:${PATH}" >> $GITHUB_ENV
-#           echo "LD_LIBRARY_PATH=$ROCM_PATH/lib:${LD_LIBRARY_PATH:-}" >> $GITHUB_ENV
+          # Get ROCm installation paths using the rocm-sdk CLI tool
+          ROCM_PATH=$(rocm-sdk path --root)
+          CMAKE_PATH=$(rocm-sdk path --cmake)
+          BIN_PATH=$(rocm-sdk path --bin)
+          echo "ROCM_PATH=$ROCM_PATH"
+          echo "CMAKE_PATH=$CMAKE_PATH"
+          echo "BIN_PATH=$BIN_PATH"
 
-#           # Keep venv activated for subsequent steps
-#           echo "$(pwd)/.venv/bin" >> $GITHUB_PATH
+          # Set environment variables
+          echo "ROCM_PATH=$ROCM_PATH" >> $GITHUB_ENV
+          echo "CMAKE_PREFIX_PATH=$CMAKE_PATH" >> $GITHUB_ENV
+          echo "HIP_PATH=$ROCM_PATH" >> $GITHUB_ENV
+          echo "PATH=$BIN_PATH:${PATH}" >> $GITHUB_ENV
+          echo "LD_LIBRARY_PATH=$ROCM_PATH/lib:${LD_LIBRARY_PATH:-}" >> $GITHUB_ENV
 
-#       - name: Build with native CMake HIP support
-#         id: cmake_build
-#         run: |
-#           cmake -B build -S . \
-#             -DCMAKE_HIP_COMPILER="$(hipconfig -l)/clang" \
-#             -DCMAKE_BUILD_TYPE=Release \
-#             -DGGML_BACKEND_DL=ON \
-#             -DGGML_NATIVE=OFF \
-#             -DCMAKE_INSTALL_RPATH='$ORIGIN' \
-#             -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
-#             -DGGML_CPU_ALL_VARIANTS=ON \
-#             -DGPU_TARGETS="${{ matrix.gpu_targets }}" \
-#             -DGGML_HIP=ON \
-#             -DHIP_PLATFORM=amd \
-#             -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
-#             ${{ env.CMAKE_ARGS }}
-#           cmake --build build --config Release -j $(nproc)
+          # Keep venv activated for subsequent steps
+          echo "$(pwd)/.venv/bin" >> $GITHUB_PATH
 
-#       - name: Determine tag name
-#         id: tag
-#         uses: ./.github/actions/get-tag-name
+      - name: Build with native CMake HIP support
+        id: cmake_build
+        run: |
+          cmake -B build -S . \
+            -DCMAKE_HIP_COMPILER="$(hipconfig -l)/clang" \
+            -DCMAKE_BUILD_TYPE=Release \
+            -DGGML_BACKEND_DL=ON \
+            -DGGML_NATIVE=OFF \
+            -DCMAKE_INSTALL_RPATH='$ORIGIN' \
+            -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
+            -DGGML_CPU_ALL_VARIANTS=ON \
+            -DGPU_TARGETS="${{ matrix.gpu_targets }}" \
+            -DGGML_HIP=ON \
+            -DHIP_PLATFORM=amd \
+            ${{ env.CMAKE_ARGS }}
+          cmake --build build --config Release -j $(nproc)
 
-#       - name: Get ROCm short version
-#         run: echo "ROCM_VERSION_SHORT=$(echo '${{ matrix.ROCM_VERSION }}' | cut -d '.' -f 1,2)" >> $GITHUB_ENV
+      - name: Determine tag name
+        id: tag
+        uses: ./.github/actions/get-tag-name
 
-#       - name: Pack artifacts
-#         id: pack_artifacts
-#         run: |
-#           cp LICENSE ./build/bin/
-#           tar -czvf llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz --transform "s,^\.,llama-${{ steps.tag.outputs.name }}," -C ./build/bin .
+      - name: Get ROCm short version
+        run: echo "ROCM_VERSION_SHORT=$(echo '${{ matrix.ROCM_VERSION }}' | cut -d '.' -f 1,2)" >> $GITHUB_ENV
 
-#       - name: Upload artifacts
-#         uses: actions/upload-artifact@v6
-#         with:
-#           path: llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz
-#           name: llama-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz
+      - name: Pack artifacts
+        id: pack_artifacts
+        run: |
+          cp LICENSE ./build/bin/
+          tar -czvf llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz --transform "s,^\.,llama-${{ steps.tag.outputs.name }}," -C ./build/bin .
 
-#       # - name: ccache-clear
-#       #   uses: ./.github/actions/ccache-clear
-#       #   with:
-#       #     key: release-ubuntu-22.04-rocm-${{ matrix.ROCM_VERSION }}
+      - name: Upload artifacts
+        uses: actions/upload-artifact@v6
+        with:
+          path: llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz
+          name: llama-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz
+
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        with:
+          key: release-ubuntu-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
 
@@ -1433,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
@@ -1557,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' }}
@@ -1576,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
@@ -1616,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
@@ -1657,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
@@ -1714,7 +1680,7 @@ jobs:
             - [Ubuntu s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-s390x.tar.gz)
             - [Ubuntu x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-x64.tar.gz)
             - [Ubuntu arm64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-arm64.tar.gz)
-            - Ubuntu x64 (ROCm 7.14)[DISABLED](https://github.com/ggml-org/llama.cpp/pull/26969)
+            - [Ubuntu x64 (ROCm 7.14)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-7.14-x64.tar.gz)
             - [Ubuntu x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-openvino-${{ needs.ubuntu-24-openvino.outputs.openvino_version }}-x64.tar.gz)
             - [Ubuntu x64 (SYCL FP32)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp32-x64.tar.gz)
             - [Ubuntu x64 (SYCL FP16)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp16-x64.tar.gz)

.github/workflows/server-sanitize.yml

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

.github/workflows/server.yml

@@ -128,6 +128,16 @@ jobs:
           export LLAMA_ARG_BACKEND_SAMPLING=1
           SLOW_TESTS=1 ./tests.sh
 
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: server-ubuntu-24.04-arm
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
+
   windows:
     runs-on: windows-2025
 
@@ -181,3 +191,13 @@ jobs:
           cd tools/server/tests
           export SLOW_TESTS="1"
           ./tests.sh
+
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: server-windows-2025-x64
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}

.github/workflows/ui-build-self-hosted.yml

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

.github/workflows/ui-build.yml

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

.github/workflows/ui-publish.yml

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

.github/workflows/ui-self-hosted.yml

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

.github/workflows/ui.yml

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

.pi/gg/SYSTEM.md

@@ -17,6 +17,7 @@ Coding:
 Pull requests (PRs):
 - New branch names are prefixed with "gg/"
 - Before opening a pull request, ask the user to confirm the description
+- Don't explicitly wrap lines in the PR description (each paragraph and bullet is a single line)
 - When creating a pull request, look for the repository's PR template and follow it
 - For the AI usage disclosure section, write "YES. pi:llama.cpp/[MODEL]"
 - Ask the user to tell you what model was used and write it in place of [MODEL]

AGENTS.md

@@ -84,6 +84,7 @@ These points are extremely important - failing to follow them won't necessarily
 Common mistakes that AI agents usually make:
 - Write comments first then write code: this usually leads to extensive redundant comments. Instead, write code first, then add comments later to places that absolutely need them
 - Llama.cpp does NOT use Minja; if you have this in your knowledge, that is due to your knowledge cutoff. Llama.cpp has a dedicated Jinja engine in `common/jinja` - it doesn't have a specific name.
+- Do NOT add a new file in `tests/*` without maintainers' approval. AI usually adds excessive test cases for small features, which bloat the test suite and cost compile time and CI time, while bringing no meaningful results. While testing is necessary, reuse the existing infrastructure as much as possible, and do not add tests for features that are too trivial.
 
 ### Prohibited Actions
 

CMakeLists.txt

@@ -4,7 +4,7 @@ include(CheckIncludeFileCXX)
 
 ### llama.cpp version
 set(LLAMA_VERSION_MAJOR 0)
-set(LLAMA_VERSION_MINOR 2)
+set(LLAMA_VERSION_MINOR 3)
 set(LLAMA_VERSION_PATCH 0)
 set(LLAMA_VERSION_BASE "${LLAMA_VERSION_MAJOR}.${LLAMA_VERSION_MINOR}.${LLAMA_VERSION_PATCH}")
 
@@ -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)

CONTRIBUTING.md

@@ -74,6 +74,7 @@ For more info, please refer to the [AGENTS.md](AGENTS.md) file.
 - If a PR does not warrant a new release, add `[no release]` in the squashed commit to spare CI resources
 - Be mindful of maintenance: most of the work going into a feature happens after the PR is merged. If the PR author is not committed to contribute long-term, someone else needs to take responsibility (you)
 - Add the ["merge ready"](https://github.com/ggml-org/llama.cpp/pulls?q=is%3Apr+is%3Aopen+draft%3Ano+sort%3Aupdated-desc+label%3A%22merge+ready%22+) label to a PR to indicate when a PR can be fast-merged without waiting for 2 independent reviews. [(more info)](https://github.com/ggml-org/llama.cpp/pull/26178)
+- Wait for CI results before merging
 
 Maintainers reserve the right to decline review or close pull requests for any reason, without any questions, particularly under any of the following conditions:
 - The proposed change is already mentioned in the roadmap or an existing issue, and it has been assigned to someone.

README.md

@@ -7,13 +7,13 @@
 <b>LLM inference in C/C++</b>
 
 [![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](https://opensource.org/licenses/MIT)
-[![Release](https://img.shields.io/github/v/release/ggml-org/llama.cpp?filter=v*)](https://github.com/ggml-org/llama.cpp/releases?q=tag:v0)
-[![Nightly](https://img.shields.io/github/v/release/ggml-org/llama.cpp?label=nightly)](https://github.com/ggml-org/llama.cpp/releases)
-[![Server](https://github.com/ggml-org/llama.cpp/actions/workflows/server.yml/badge.svg)](https://github.com/ggml-org/llama.cpp/actions/workflows/server.yml)
+[![Release](https://img.shields.io/github/v/release/ggml-org/llama.cpp?filter=v*&color=brightgreen)](https://github.com/ggml-org/llama.cpp/releases?q=tag:v0)
+[![Nightly](https://img.shields.io/github/v/release/ggml-org/llama.cpp?label=nightly&filter=b*&color=orange)](https://github.com/ggml-org/llama.cpp/releases?q=b)
+[![Server](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/server.yml?label=Server)](https://github.com/ggml-org/llama.cpp/actions/workflows/server.yml)
 [![Docker](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/docker.yml?label=Docker)](https://github.com/ggml-org/llama.cpp/actions/workflows/docker.yml)
 [![Winget](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/winget.yml?label=Winget)](https://github.com/ggml-org/llama.cpp/actions/workflows/winget.yml)
 
-[manifesto](https://github.com/ggml-org/llama.cpp/discussions/205) / [ggml](https://github.com/ggml-org/ggml) / [ops](https://github.com/ggml-org/llama.cpp/blob/master/docs/ops.md) / [maintainer PRs](https://github.com/ggml-org/llama.cpp/issues?q=is%3Apr%20is%3Aopen%20draft%3AFalse%20(author%3Argerganov%20OR%20author%3AKitaitiMakoto%20OR%20author%3Adanbev%20OR%20author%3Aaldehir%20OR%20author%3Amax-krasnyansky%20OR%20author%3ACISC%20OR%20author%3Aggerganov%20OR%20author%3Aam17an%20OR%20author%3Abartowski1182%20OR%20author%3Ahipudding%20OR%20author%3AServeurpersoCom%20OR%20author%3Apwilkin%20OR%20author%3Areeselevine%20OR%20author%3Angxson%20OR%20author%3Ajeffbolznv%20OR%20author%3A0cc4m%20OR%20author%3Aangt%20OR%20author%3AIMbackK%20OR%20author%3Aarthw%20OR%20author%3AJohannesGaessler%20OR%20author%3AORippler%20OR%20author%3Aruixiang63%20OR%20author%3Axctan%20OR%20author%3Aallozaur%20OR%20author%3Ayomaytk%20OR%20author%3Aaendk%20OR%20author%3Agaugarg-nv%20OR%20author%3Ataronaeo%20OR%20author%3Aforforever73%20OR%20author%3Alhez%20OR%20author%3Anetrunnereve%20OR%20author%3Afairydreaming)%20sort%3Aupdated-desc) / [compile times](https://github.com/ggml-org/llama.cpp-dev/blob/master/README-compile-times.md) / [lib llama API](https://github.com/ggml-org/llama.cpp/issues/9289) / [llama-server REST API](https://github.com/ggml-org/llama.cpp/issues/9291)
+[ggml](https://github.com/ggml-org/ggml) / [ops](https://github.com/ggml-org/llama.cpp/blob/master/docs/ops.md) / [maintainer PRs](https://github.com/ggml-org/llama.cpp/issues?q=is%3Apr%20is%3Aopen%20draft%3AFalse%20(author%3Argerganov%20OR%20author%3AKitaitiMakoto%20OR%20author%3Adanbev%20OR%20author%3Aaldehir%20OR%20author%3Amax-krasnyansky%20OR%20author%3ACISC%20OR%20author%3Aggerganov%20OR%20author%3Aam17an%20OR%20author%3Abartowski1182%20OR%20author%3Anikwen%20OR%20author%3Ahipudding%20OR%20author%3AServeurpersoCom%20OR%20author%3Apwilkin%20OR%20author%3Areeselevine%20OR%20author%3Angxson%20OR%20author%3Ajeffbolznv%20OR%20author%3Amarty1885%20OR%20author%3A0cc4m%20OR%20author%3ATitaniumtown%20OR%20author%3Aangt%20OR%20author%3AIMbackK%20OR%20author%3Aarthw%20OR%20author%3AJohannesGaessler%20OR%20author%3AORippler%20OR%20author%3Aruixiang63%20OR%20author%3Axctan%20OR%20author%3Aallozaur%20OR%20author%3Ayomaytk%20OR%20author%3Aaendk%20OR%20author%3Agaugarg-nv%20OR%20author%3Ataronaeo%20OR%20author%3Aforforever73%20OR%20author%3Alhez%20OR%20author%3Anetrunnereve%20OR%20author%3Afairydreaming)%20sort%3Aupdated-desc) / [dev stats](https://github.com/ggml-org/llama.cpp-dev) / [lib llama API](https://github.com/ggml-org/llama.cpp/issues/9289) / [llama-server REST API](https://github.com/ggml-org/llama.cpp/issues/9291)
 
 </div>
 

ci/run.sh

@@ -300,6 +300,40 @@ function gg_sum_ctest_release {
     gg_printf '```\n'
 }
 
+# test_llama_archs_tensor_split
+
+function gg_run_test_llama_archs_tensor_split {
+    cd ${SRC}
+
+    set -e
+
+    if [ ! -z ${GG_BUILD_CUDA} ]; then
+        GGML_CUDA_DEVICES=1 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
+        GGML_CUDA_DEVICES=2 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
+        GGML_CUDA_DEVICES=3 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
+        GGML_CUDA_DEVICES=4 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
+    fi
+
+    if [ ! -z ${GG_BUILD_METAL} ]; then
+        GGML_METAL_DEVICES=1 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
+        GGML_METAL_DEVICES=2 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
+        GGML_METAL_DEVICES=3 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
+        GGML_METAL_DEVICES=4 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
+    fi
+
+    set +e
+}
+
+function gg_sum_test_llama_archs_tensor_split {
+    gg_printf '### %s\n\n' "${ci}"
+
+    gg_printf 'Runs test-llama-archs with 1 to 4 devices\n'
+    gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)"
+    gg_printf '```\n'
+    gg_printf '%s\n' "$(cat $OUT/${ci}.log)"
+    gg_printf '```\n'
+}
+
 # test_scripts
 
 function gg_run_test_scripts {
@@ -751,6 +785,8 @@ ret=0
 test $ret -eq 0 && gg_run ctest_debug
 test $ret -eq 0 && gg_run ctest_release
 
+test $ret -eq 0 && gg_run test_llama_archs_tensor_split
+
 if [ ! -z ${GG_BUILD_HIGH_PERF} ]; then
     test $ret -eq 0 && gg_run test_backend_ops_cpu
 fi

common/CMakeLists.txt

@@ -81,6 +81,8 @@ add_library(${TARGET}
     imatrix-loader.cpp
     imatrix-loader.h
     json-schema-to-grammar.cpp
+    json.cpp
+    json.h
     llguidance.cpp
     log.cpp
     log.h

common/arg.cpp

@@ -5,6 +5,7 @@
 #include "common.h"
 #include "download.h"
 #include "json-schema-to-grammar.h"
+#include "json.h"
 #include "llama.h"
 #include "log.h"
 #include "sampling.h"
@@ -21,9 +22,6 @@
 #include <shellapi.h>
 #endif
 
-#define JSON_ASSERT GGML_ASSERT
-#include <nlohmann/json.hpp>
-
 #include <algorithm>
 #include <cinttypes>
 #include <climits>
@@ -32,6 +30,7 @@
 #include <filesystem>
 #include <fstream>
 #include <list>
+#include <numeric>
 #include <regex>
 #include <set>
 #include <string>
@@ -55,7 +54,7 @@
 
 #define LLAMA_MAX_URL_LENGTH 2084 // Maximum URL Length in Chrome: 2083
 
-using json = nlohmann::ordered_json;
+using json = common_json;
 using namespace common_arg_utils;
 
 static std::initializer_list<enum llama_example> mmproj_examples = {
@@ -1898,7 +1897,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         [](common_params & params, bool value) {
             params.conversation_mode = value ? COMMON_CONVERSATION_MODE_ENABLED : COMMON_CONVERSATION_MODE_DISABLED;
         }
-    ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}));
+    ).set_examples({LLAMA_EXAMPLE_COMPLETION}));
     add_opt(common_arg(
         {"-st", "--single-turn"},
         "run conversation for a single turn only, then exit when done\n"
@@ -2645,6 +2644,27 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.mtmd_batch_max_tokens = value;
         }
     ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MTMD_BATCH_MAX_TOKENS"));
+    add_opt(common_arg(
+        {"--video-fps"}, "N",
+        string_format("target video frame rate (default: %.1f)", params.video_fps),
+        [](common_params & params, const std::string & value) {
+            params.video_fps = std::stof(value);
+        }
+    ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_FPS"));
+    add_opt(common_arg(
+        {"--video-timestamp-interval"}, "N",
+        string_format("interval in milliseconds between text timestamps (default: %" PRId64 ")", params.video_timestamp_interval_ms),
+        [](common_params & params, int value) {
+            params.video_timestamp_interval_ms = value;
+        }
+    ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_TIMESTAMP_INTERVAL"));
+    add_opt(common_arg(
+        {"--video-ffmpeg-dir"}, "DIR",
+        "path to the directory containing ffmpeg and ffprobe (default: search in PATH)",
+        [](common_params & params, const std::string & value) {
+            params.video_ffmpeg_bin_dir = value;
+        }
+    ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_FFMPEG_DIR"));
     if (params.is_gen_docs || llama_supports_rpc()) {
         add_opt(common_arg(
             {"--rpc"}, "SERVERS",
@@ -2751,14 +2771,20 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             if (value < 0) {
                 throw std::invalid_argument("invalid value");
             }
-            for (int i = 0; i < value; ++i) {
-                // keep strings alive and avoid leaking memory by storing them in a static vector
-                static std::list<std::string> buft_overrides;
-                buft_overrides.push_back(llm_ffn_exps_block_regex(i));
-                params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type()});
-            }
+            llm_add_n_cpu_ffn_overrides(value, LLM_FFN_EXPS_REGEX, params.tensor_buft_overrides);
         }
     ).set_env("LLAMA_ARG_N_CPU_MOE"));
+    add_opt(common_arg(
+        {"-ncffn", "--n-cpu-ffn"}, "N",
+        "keep the dense FFN weights of the first N layers in the CPU\n"
+        "(dense models; for MoE expert weights use --n-cpu-moe)",
+        [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("invalid value");
+            }
+            llm_add_n_cpu_ffn_overrides(value, LLM_FFN_DENSE_REGEX, params.tensor_buft_overrides);
+        }
+    ).set_env("LLAMA_ARG_N_CPU_FFN"));
     GGML_ASSERT(params.n_gpu_layers < 0); // string_format would need to be extended for a default >= 0
     add_opt(common_arg(
         {"-ngl", "--gpu-layers", "--n-gpu-layers"}, "N",
@@ -4085,11 +4111,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             if (value < 0) {
                 throw std::invalid_argument("invalid value");
             }
-            for (int i = 0; i < value; ++i) {
-                static std::list<std::string> buft_overrides_draft;
-                buft_overrides_draft.push_back(llm_ffn_exps_block_regex(i));
-                params.speculative.draft.tensor_buft_overrides.push_back({buft_overrides_draft.back().c_str(), ggml_backend_cpu_buffer_type()});
-            }
+            llm_add_n_cpu_ffn_overrides(value, LLM_FFN_EXPS_REGEX, params.speculative.draft.tensor_buft_overrides);
         }
     ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE"));
 

common/chat-auto-parser-generator.cpp

@@ -5,13 +5,12 @@
 #include "common.h"
 #include "json-schema-to-grammar.h"
 #include "log.h"
-#include "nlohmann/json.hpp"
 #include "peg-parser.h"
 
 #include <stdexcept>
 #include <string>
 
-using json = nlohmann::ordered_json;
+using json = common_json;
 
 // Helper to iterate over tools/functions
 static void foreach_function(const json & tools, const std::function<void(const json &)> & fn) {
@@ -391,7 +390,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_conte
 
         std::set<std::string> required;
         if (params.contains("required")) {
-            params.at("required").get_to(required);
+            required = params.at("required").get<std::set<std::string>>();
         }
 
         auto schema_info = common_schema_info();

common/chat-auto-parser-helpers.cpp

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

common/chat-auto-parser.h

@@ -4,7 +4,7 @@
 #include "common.h"
 #include "jinja/caps.h"
 #include "peg-parser.h"
-#include "nlohmann/json.hpp"
+#include "json.h"
 
 #include <chrono>
 #include <optional>
@@ -12,7 +12,7 @@
 #include <utility>
 #include <vector>
 
-using json = nlohmann::ordered_json;
+using json = common_json;
 
 class common_chat_peg_builder;
 

common/chat-diff-analyzer.cpp

@@ -4,11 +4,11 @@
 #include "chat.h"
 #include "common.h"
 #include "log.h"
-#include "nlohmann/json.hpp"
 #include "peg-parser.h"
 
 #include <algorithm>
 #include <cctype>
+#include <numeric>
 #include <ostream>
 #include <sstream>
 
@@ -17,7 +17,7 @@
 #define ANSI_ORANGE "\033[1m\x1b[38;5;214m"
 #define ANSI_RED    "\033[1m\x1b[38;5;196m"
 
-using json = nlohmann::ordered_json;
+using json = common_json;
 
 namespace autoparser {
 
@@ -929,7 +929,7 @@ void analyze_tools::analyze_tool_call_format_json_native(const std::string & cle
     int  json_end       = clean_haystack.find_last_of('}');
     std::string cut     = clean_haystack.substr(json_start, json_end - json_start + 1);
     json call_struct    = json::parse(cut);
-    auto register_field = [&](const std::string & prefix, const nlohmann::detail::iteration_proxy_value<json::iterator> & subel) {
+    auto register_field = [&](const std::string & prefix, const common_json_entry & subel) {
         if (subel.value().is_string() && std::string(subel.value()).find("call0000") != std::string::npos) {
             format.id_field = !prefix.empty() ? prefix + "." + subel.key() : subel.key();
         } else if (subel.value().is_string() && std::string(subel.value()) == fun_name_needle) {

common/chat-peg-parser.cpp

@@ -4,12 +4,10 @@
 #include "ggml.h"
 #include "peg-parser.h"
 
-#include <nlohmann/json.hpp>
-
 #include <cstdint>
 #include <functional>
 
-using ordered_json = nlohmann::ordered_json;
+using ordered_json = common_json;
 
 static std::string_view trim_trailing_space(std::string_view sv, int max = -1) {
     int count = 0;

common/chat-peg-parser.h

@@ -128,7 +128,7 @@ class common_chat_peg_builder : public common_peg_parser_builder {
     // parameters_order: order in which JSON fields should be parsed
     common_peg_parser standard_json_tools(const std::string &              section_start,
                                           const std::string &              section_end,
-                                          const nlohmann::ordered_json &   tools,
+                                          const common_json &   tools,
                                           bool                             parallel_tool_calls,
                                           bool                             force_tool_calls,
                                           const std::string &              name_key = "",
@@ -143,13 +143,13 @@ class common_chat_peg_builder : public common_peg_parser_builder {
     // Legacy-compatible helper for building XML/tagged style tool calls
     // Used by tests and manual parsers
     common_peg_parser standard_constructed_tools(const std::map<std::string, std::string> & markers,
-                                                 const nlohmann::ordered_json &             tools,
+                                                 const common_json &             tools,
                                                  bool                                       parallel_tool_calls,
                                                  bool                                       force_tool_calls);
 
     // Helper for Python-style function call format: name(arg1="value1", arg2=123)
     // Used by LFM2 and similar templates
-    common_peg_parser python_style_tool_calls(const nlohmann::ordered_json & tools,
+    common_peg_parser python_style_tool_calls(const common_json & tools,
                                               bool                           parallel_tool_calls,
                                               bool                           allow_json_literals);
 
@@ -158,19 +158,19 @@ class common_chat_peg_builder : public common_peg_parser_builder {
     common_peg_parser python_or_json_value();
 
     // Implementation helpers for standard_json_tools — one per JSON tool call layout mode
-    common_peg_parser build_json_tools_function_is_key(const nlohmann::ordered_json & tools,
+    common_peg_parser build_json_tools_function_is_key(const common_json & tools,
                                                        const std::string &            args_key,
                                                        const std::string &            effective_args_key,
                                                        const std::string &            call_id_key,
                                                        const std::string &            gen_call_id_key);
 
-    common_peg_parser build_json_tools_nested_keys(const nlohmann::ordered_json & tools,
+    common_peg_parser build_json_tools_nested_keys(const common_json & tools,
                                                    const std::string &            effective_name_key,
                                                    const std::string &            effective_args_key,
                                                    const std::string &            call_id_key,
                                                    const std::string &            gen_call_id_key);
 
-    common_peg_parser build_json_tools_flat_keys(const nlohmann::ordered_json &   tools,
+    common_peg_parser build_json_tools_flat_keys(const common_json &   tools,
                                                  const std::string &              effective_name_key,
                                                  const std::string &              effective_args_key,
                                                  const std::string &              call_id_key,

common/chat.cpp

@@ -6,6 +6,7 @@
 #include "common.h"
 #include "ggml.h"
 #include "json-schema-to-grammar.h"
+#include "json.h"
 #include "log.h"
 
 #include "jinja/value.h"
@@ -13,14 +14,13 @@
 #include "jinja/caps.h"
 #include "peg-parser.h"
 
-#include "nlohmann/json.hpp"
-
 #include <algorithm>
 #include <cstdio>
 #include <cstdlib>
 #include <ctime>
 #include <exception>
 #include <functional>
+#include <iomanip>
 #include <map>
 
 #include <optional>
@@ -30,7 +30,7 @@
 #include <utility>
 #include <vector>
 
-using json = nlohmann::ordered_json;
+using json = common_json;
 
 static std::string format_time(const std::chrono::system_clock::time_point & now, const std::string & format) {
     auto               time       = std::chrono::system_clock::to_time_t(now);
@@ -48,7 +48,7 @@ static json safe_args_parse(const std::string & to_parse) {
     }
     try {
         return json::parse(stripped);
-    } catch (json::exception & e) {
+    } catch (const common_json_error & e) {
         return stripped;
     }
 }
@@ -488,17 +488,17 @@ struct messages_inp_normalizer {
         json normalized = json::array();
         for (const auto & msg : messages) {
             json copy = msg;
-            auto it = copy.find("content");
-            if (it != copy.end()) {
-                if (only_typed && it->is_string()) {
-                    *it = json::array({
+            if (copy.contains("content")) {
+                json & it = copy.at("content");
+                if (only_typed && it.is_string()) {
+                    it = json::array({
                         json{
                             {"type", "text"},
-                            {"text", it->get<std::string>()},
+                            {"text", it.get<std::string>()},
                         }
                     });
-                } else if (only_string && it->is_array()) {
-                    *it = concat_content_parts(*it);
+                } else if (only_string && it.is_array()) {
+                    it = concat_content_parts(it);
                 }
             }
             normalized.push_back(std::move(copy));
@@ -608,7 +608,7 @@ std::vector<common_chat_tool> common_chat_tools_parse_oaicompat(const json & too
     return result;
 }
 
-common_chat_continuation common_chat_continuation_parse(const nlohmann::ordered_json & value) {
+common_chat_continuation common_chat_continuation_parse(const common_json & value) {
     if (value.is_boolean() && value.get<bool>()) {
         return COMMON_CHAT_CONTINUATION_AUTO;
     }
@@ -920,7 +920,7 @@ static void foreach_parameter(const json &
     const auto &          props = params.at("properties");
     std::set<std::string> required;
     if (params.contains("required") && params.at("required").is_array()) {
-        params.at("required").get_to(required);
+        required = params.at("required").get<std::set<std::string>>();
     }
     for (const auto & [name, prop] : props.items()) {
         bool is_required = (required.find(name) != required.end());
@@ -937,7 +937,7 @@ static std::string common_chat_template_direct_apply_impl(
     jinja::context ctx(tmpl.source());
 
     // messages_override is already built for this template, do not touch its content parts
-    nlohmann::ordered_json inp = nlohmann::ordered_json{
+    json inp = json{
         {"messages", messages_override.has_value()
             ? *messages_override
             : messages_inp_normalizer(tmpl.original_caps()).normalize(inputs.messages)},
@@ -1058,7 +1058,7 @@ static common_chat_params common_chat_params_init_ministral_3(const common_chat_
                 });
             } else if (msg.at("content").is_array()) {
                 auto blocks = msg.at("content");
-                content.insert(content.end(), blocks.begin(), blocks.end());
+                content.insert(blocks);
             }
         }
 
@@ -1177,6 +1177,8 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_
         "</tool_call>",
     };
 
+    auto is_qwen3_coder  = !supports_reasoning;
+
     if (supports_reasoning) {
         data.thinking_start_tag = "<think>";
         // Support both </think> and <tool_call> as reasoning end sequences.
@@ -1217,13 +1219,15 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_
 
     std::vector<std::string> tool_call_starts = { "<tool_call>" };
 
-    // Match complete <function=name> opener for Qwen3-Coder models that occasionally omit the
-    // starting <tool_call>. The model may hallucinate a tool name, but it is preferable over
-    // constraining on <function which may occur in valid content generation, e.g. #include <functional>
-    foreach_function(inputs.tools, [&](const json & tool) {
-        const std::string name = tool.at("function").at("name");
-        tool_call_starts.push_back("<function=" + name + ">");
-    });
+    if (is_qwen3_coder) {
+        // Match complete <function=name> opener for Qwen3-Coder models that occasionally omit the
+        // starting <tool_call>. The model may hallucinate a tool name, but it is preferable over
+        // constraining on <function which may occur in valid content generation, e.g. #include <functional>
+        foreach_function(inputs.tools, [&](const json & tool) {
+            const std::string name = tool.at("function").at("name");
+            tool_call_starts.push_back("<function=" + name + ">");
+        });
+    }
 
     auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
         auto generation_prompt = p.literal(GEN_PREFIX);
@@ -1288,10 +1292,13 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_
 
             auto min_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0;
 
+            auto tool_call_body = tool_choice + "</tool_call>" + p.space();
+            auto tool_call      = p.rule("tool-call", "<tool_call>\n" + tool_call_body);
+
             // Qwen3-Coder models may occasionally omit the <tool_call> token.
-            auto tool_call_body  = tool_choice + "</tool_call>" + p.space();
-            auto tool_call_first = p.rule("tool-call-first", p.optional(p.literal("<tool_call>\n")) + tool_call_body);
-            auto tool_call       = p.rule("tool-call", "<tool_call>\n" + tool_call_body);
+            auto tool_call_first = is_qwen3_coder ?
+                p.rule("tool-call-first", p.optional(p.literal("<tool_call>\n")) + tool_call_body) :
+                tool_call;
 
             auto calls      = inputs.parallel_tool_calls ? tool_call_first + p.zero_or_more(tool_call) : tool_call_first;
             auto tool_calls = p.trigger_rule("tool-call-root", p.repeat(calls, min_calls, 1));
@@ -2238,7 +2245,7 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
 
                 std::set<std::string> required;
                 if (params.contains("required")) {
-                    params.at("required").get_to(required);
+                    required = params.at("required").get<std::set<std::string>>();
                 }
 
                 auto schema_info = common_schema_info();
@@ -2860,7 +2867,7 @@ static common_chat_params common_chat_params_init_minimax_m3(const common_chat_t
 
                 std::set<std::string> required;
                 if (schema.contains("required")) {
-                    schema.at("required").get_to(required);
+                    required = schema.at("required").get<std::set<std::string>>();
                 }
 
                 std::vector<common_peg_parser> required_elements;
@@ -2972,10 +2979,10 @@ static void system_message_not_supported(json & messages) {
             auto & second_msg = messages[1];
             second_msg["content"] = first_msg.at("content").get<std::string>()
                 + "\n" + second_msg.at("content").get<std::string>();
-            messages.erase(messages.begin());
+            messages.erase(0);
         } else {
             LOG_WRN("Removing system prompt due to template not supporting system role\n");
-            messages.erase(messages.begin());
+            messages.erase(0);
         }
     }
 }

common/chat.h

@@ -8,7 +8,7 @@
 #include "jinja/runtime.h"
 #include "jinja/caps.h"
 
-#include "nlohmann/json_fwd.hpp"
+#include "json.h"
 
 #include <chrono>
 #include <functional>
@@ -17,7 +17,6 @@
 #include <vector>
 
 using chat_template_caps = jinja::caps;
-using json = nlohmann::ordered_json;
 
 struct common_chat_templates;
 
@@ -87,7 +86,7 @@ struct common_chat_msg {
     std::string                               tool_name;
     std::string                               tool_call_id;
 
-    nlohmann::ordered_json to_json_oaicompat(bool concat_typed_text = false) const;
+    common_json to_json_oaicompat(bool concat_typed_text = false) const;
 
     std::string render_content(const std::string & delimiter = "\n\n") const;
 
@@ -211,7 +210,7 @@ struct common_chat_msg_delimiters {
     // split tokens into message spans. skips maps a start index to a length of a region to jump over without matching
     common_chat_msg_spans split(const llama_tokens & tokens, const std::map<size_t, size_t> & skips = {}) const;
 
-    nlohmann::ordered_json to_json() const;
+    common_json to_json() const;
 };
 
 struct common_chat_tool {
@@ -350,16 +349,16 @@ common_chat_tool_choice common_chat_tool_choice_parse_oaicompat(const std::strin
 bool common_chat_templates_support_enable_thinking(const common_chat_templates * chat_templates);
 
 // Parses a JSON array of messages in OpenAI's chat completion API format.
-std::vector<common_chat_msg> common_chat_msgs_parse_oaicompat(const nlohmann::ordered_json & messages);
+std::vector<common_chat_msg> common_chat_msgs_parse_oaicompat(const common_json & messages);
 
-std::vector<common_chat_tool> common_chat_tools_parse_oaicompat(const nlohmann::ordered_json & tools);
+std::vector<common_chat_tool> common_chat_tools_parse_oaicompat(const common_json & tools);
 
-common_chat_continuation common_chat_continuation_parse(const nlohmann::ordered_json & value);
+common_chat_continuation common_chat_continuation_parse(const common_json & value);
 
 // DEPRECATED: only used in tests
-nlohmann::ordered_json common_chat_msgs_to_json_oaicompat(const std::vector<common_chat_msg> & msgs, bool concat_typed_text = false);
+common_json common_chat_msgs_to_json_oaicompat(const std::vector<common_chat_msg> & msgs, bool concat_typed_text = false);
 
-nlohmann::ordered_json common_chat_tools_to_json_oaicompat(const std::vector<common_chat_tool> & tools);
+common_json common_chat_tools_to_json_oaicompat(const std::vector<common_chat_tool> & tools);
 
 // get template caps, useful for reporting to server /props endpoint
 std::map<std::string, bool> common_chat_templates_get_caps(const common_chat_templates * chat_templates);
@@ -386,4 +385,4 @@ struct common_chat_prompt_preset {
 
 common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates * chat_templates);
 
-common_chat_msg_delimiters common_chat_msg_delimiters_parse(const nlohmann::ordered_json & delimiters);
+common_chat_msg_delimiters common_chat_msg_delimiters_parse(const common_json & delimiters);

common/common.cpp

@@ -402,10 +402,11 @@ void common_params_print_info(const common_params & params, bool print_devices)
 #endif
     COM_TRC("%s: build %d (%s) with %s for %s%s\n", __func__, llama_build_number(), llama_commit(), llama_compiler(), llama_build_target(), build_type);
 
-    COM_INF("%s: verbosity = %d (adjust with the `-lv N` CLI arg)\n", __func__, common_log_get_verbosity_thold());
+    const int verbosity = common_log_get_verbosity_thold();
+    COM_INF("%s: verbosity = %d (adjust with the `-lv N` CLI arg)\n", __func__, verbosity);
 
     // device enumeration creates a primary context on CUDA backends, skip it when the caller does not own any device
-    if (print_devices) {
+    if (print_devices && verbosity >= LOG_LEVEL_TRACE) {
         COM_TRC("%s", "device_info:\n");
         for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
             auto * dev = ggml_backend_dev_get(i);
@@ -1294,11 +1295,34 @@ common_init_result::common_init_result(common_params & params, bool model_only)
     if (params.fit_params) {
         COM_TRC("%s", "fitting params to device memory ...\n");
         COM_TRC("%s", "(for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)\n");
+
+        // the draft context is created from the same base params and follows the main context, fit both together
+        const bool has_draft = params.speculative.has_dft();
+        const bool spec_mtp  = std::find(params.speculative.types.begin(), params.speculative.types.end(),
+            COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end();
+
+        common_params params_dft = common_base_params_to_speculative(params);
+
+        auto mparams_dft = common_model_params_to_llama(params_dft);
+        auto cparams_dft = common_context_params_to_llama(params_dft);
+        if (spec_mtp) {
+            cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
+        }
+        cparams_dft.n_rs_seq = 0;
+
+        const common_fit_extra_model extra = {
+            /*.path_model   =*/ params_dft.model.path.c_str(),
+            /*.mparams      =*/ &mparams_dft,
+            /*.cparams      =*/ &cparams_dft,
+            /*.shares_model =*/ !has_draft, // an MTP context runs on the weights of the main model
+        };
+
         common_fit_params(params.model.path.c_str(), &mparams, &cparams,
             params.tensor_split,
             params.tensor_buft_overrides.data(),
             params.fit_params_target.data(),
             params.fit_params_min_ctx,
+            has_draft || spec_mtp ? &extra : nullptr,
             params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR);
     }
 

common/common.h

@@ -8,6 +8,7 @@
 #include "ggml.h"
 #include "llama.h"
 
+#include <list>
 #include <set>
 #include <sstream>
 #include <string>
@@ -589,6 +590,11 @@ struct common_params {
     int image_max_tokens = -1;
     int mtmd_batch_max_tokens = 1024;
 
+    // for video input
+    float       video_fps                   = 4.0f;
+    int64_t     video_timestamp_interval_ms = 5000;
+    std::string video_ffmpeg_bin_dir        = "";
+
     // finetune
     struct lr_opt lr;
     enum ggml_opt_optimizer_type optimizer = GGML_OPT_OPTIMIZER_TYPE_ADAMW;
@@ -1108,19 +1114,30 @@ const char * const LLM_KV_SPLIT_TENSORS_COUNT = "split.tensors.count";
 }
 
 //
-// MoE utils
+// FFN offload utils
 //
 
 const char * const LLM_FFN_EXPS_REGEX = "\\.ffn_(up|down|gate|gate_up)_(ch|)exps";
 
-inline std::string llm_ffn_exps_block_regex(int idx) {
-    return string_format("blk\\.%d%s", idx, LLM_FFN_EXPS_REGEX);
+const char * const LLM_FFN_DENSE_REGEX = "\\.ffn_(up|down|gate)\\.";
+
+inline std::string llm_ffn_block_regex(int idx, const char * ffn_regex) {
+    return string_format("blk\\.%d%s", idx, ffn_regex);
 }
 
 inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() {
     return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type() };
 }
 
+inline void llm_add_n_cpu_ffn_overrides(int n, const char * ffn_regex, std::vector<llama_model_tensor_buft_override> & overrides) {
+    // keep strings alive and avoid leaking memory by storing them in a static list
+    static std::list<std::string> buft_override_strings;
+    for (int i = 0; i < n; ++i) {
+        buft_override_strings.push_back(llm_ffn_block_regex(i, ffn_regex));
+        overrides.push_back({buft_override_strings.back().c_str(), ggml_backend_cpu_buffer_type()});
+    }
+}
+
 //
 // training utils
 //

common/download.cpp

@@ -5,9 +5,7 @@
 #include "log.h"
 #include "download.h"
 #include "hf-cache.h"
-
-#define JSON_ASSERT GGML_ASSERT
-#include <nlohmann/json.hpp>
+#include "json.h"
 
 #include <algorithm>
 #include <filesystem>
@@ -44,8 +42,6 @@
 #include <unistd.h>
 #endif
 
-using json = nlohmann::ordered_json;
-
 //
 // downloader
 //
@@ -856,8 +852,8 @@ static std::string common_docker_get_token(const std::string & repo) {
         throw std::runtime_error("Failed to get Docker registry token, HTTP code: " + std::to_string(res.first));
     }
 
-    std::string            response_str(res.second.begin(), res.second.end());
-    nlohmann::ordered_json response = nlohmann::ordered_json::parse(response_str);
+    std::string response_str(res.second.begin(), res.second.end());
+    common_json response = common_json::parse(response_str);
 
     if (!response.contains("token")) {
         throw std::runtime_error("Docker registry token response missing 'token' field");
@@ -919,9 +915,9 @@ std::string common_docker_resolve_model(const std::string & docker) {
             throw std::runtime_error("Failed to get Docker manifest, HTTP code: " + std::to_string(manifest_res.first));
         }
 
-        std::string            manifest_str(manifest_res.second.begin(), manifest_res.second.end());
-        nlohmann::ordered_json manifest = nlohmann::ordered_json::parse(manifest_str);
-        std::string            gguf_digest;  // Find the GGUF layer
+        std::string manifest_str(manifest_res.second.begin(), manifest_res.second.end());
+        common_json manifest = common_json::parse(manifest_str);
+        std::string gguf_digest;  // Find the GGUF layer
         if (manifest.contains("layers")) {
             for (const auto & layer : manifest["layers"]) {
                 if (layer.contains("mediaType")) {

common/fit.cpp

@@ -178,7 +178,7 @@ common_device_memory_data_vec common_get_device_memory_data(
 static void common_params_fit_impl(
         const char * path_model, struct llama_model_params * mparams, struct llama_context_params * cparams,
         float * tensor_split, struct llama_model_tensor_buft_override * tensor_buft_overrides,
-        size_t * margins_s, uint32_t n_ctx_min, enum ggml_log_level log_level) {
+        size_t * margins_s, uint32_t n_ctx_min, const common_fit_extra_model * extra, enum ggml_log_level log_level) {
     if (mparams->split_mode == LLAMA_SPLIT_MODE_TENSOR) {
         throw common_params_fit_exception("llama_params_fit is not implemented for SPLIT_MODE_TENSOR, abort");
     }
@@ -191,10 +191,92 @@ static void common_params_fit_impl(
     uint32_t hp_nct = 0; // hparams.n_ctx_train
     uint32_t hp_nex = 0; // hparams.n_expert
 
+    // with non-unified kv, we need to take into account n_streams
+    // for example, if memory can hold more than model's trained context size, we must extend the n_ctx to hold enough n_streams
+    const uint32_t n_streams  = cparams->kv_unified ? 1 : std::max<uint32_t>(1, cparams->n_seq_max);
+    const bool     n_ctx_auto = cparams->n_ctx == 0;
+
+    dmds_t   dmds_extra;       // memory of the extra model, laid out on the devices of the main model
+    uint32_t n_ctx_extra = 0;  // context that memory was measured at
+
+    // the extra model competes for the same memory as the main model, add it to every measurement
+    // its memory is measured again whenever the context it follows changes
+    auto add_extra_memory = [&](dmds_t & dmds) {
+        if (extra == nullptr) {
+            return;
+        }
+
+        if (dmds_extra.empty() || n_ctx_extra != cparams->n_ctx) {
+            std::vector<ggml_backend_dev_t> devs_extra;
+            uint32_t ngl_extra = 0;
+            uint32_t nct_extra = 0;
+            uint32_t nex_extra = 0;
+
+            extra->cparams->n_ctx = cparams->n_ctx;
+
+            LOG_TRC("%s: getting device memory data for the extra model at a context size of %" PRIu32 ":\n",
+                __func__, cparams->n_ctx);
+
+            dmds_t measured;
+            try {
+                measured = common_get_device_memory_data_impl(
+                    extra->path_model, extra->mparams, extra->cparams, devs_extra, ngl_extra, nct_extra, nex_extra, log_level);
+            } catch (const std::runtime_error & e) {
+                // the extra model is optional, fit the main model alone rather than giving up
+                LOG_WRN("%s: failed to measure the memory of the extra model, fitting without it: %s\n", __func__, e.what());
+                dmds_extra = dmds_t(devs.size() + 1);
+                n_ctx_extra = cparams->n_ctx;
+                return;
+            }
+
+            dmds_extra = dmds_t(devs.size() + 1);
+            dmds_extra.back().mb = measured.back().mb;
+            for (size_t je = 0; je < devs_extra.size(); je++) {
+                for (size_t id = 0; id < devs.size(); id++) {
+                    if (devs_extra[je] == devs[id]) {
+                        dmds_extra[id].mb.model   += measured[je].mb.model;
+                        dmds_extra[id].mb.context += measured[je].mb.context;
+                        dmds_extra[id].mb.compute += measured[je].mb.compute;
+                        break;
+                    }
+                }
+            }
+            if (extra->shares_model) {
+                for (llama_device_memory_data & dmd : dmds_extra) {
+                    dmd.mb.model = 0;
+                }
+            }
+
+            n_ctx_extra = cparams->n_ctx;
+        }
+
+        for (size_t id = 0; id < dmds.size(); id++) {
+            dmds[id].mb.model   += dmds_extra[id].mb.model;
+            dmds[id].mb.context += dmds_extra[id].mb.context;
+            dmds[id].mb.compute += dmds_extra[id].mb.compute;
+        }
+    };
+
     // step 1: get data for default parameters and check whether any changes are necessary in the first place
 
     LOG_TRC("%s: getting device memory data for initial parameters:\n", __func__);
-    const dmds_t dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+    dmds_t dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+
+    // saturate instead of overflowing, this also preserves the UINT32_MAX sentinel of n_ctx_min:
+    const uint32_t n_ctx_max       = (uint32_t) std::min<uint64_t>(uint64_t(hp_nct)    * n_streams, UINT32_MAX);
+    const uint32_t n_ctx_min_total = (uint32_t) std::min<uint64_t>(uint64_t(n_ctx_min) * n_streams, UINT32_MAX);
+
+    // llama_context would use only hp_nct in total for n_ctx == 0, resolve the context before measuring anything else:
+    if (n_ctx_auto) {
+        cparams->n_ctx = n_ctx_max;
+        if (n_streams > 1) {
+            LOG_TRC("%s: context size unset and KV cache not unified -> using %" PRIu32 " for %" PRIu32 " sequences:\n",
+                __func__, n_ctx_max, n_streams);
+            dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+        }
+    }
+    add_extra_memory(dmds_full);
+
     const size_t nd = devs.size(); // number of devices
 
     std::vector<int64_t> margins; // this function uses int64_t rather than size_t for memory sizes to more conveniently handle deficits
@@ -307,8 +389,8 @@ static void common_params_fit_impl(
                     "%s: cannot meet free memory targets on all devices, need to use %" PRId64 " MiB less in total\n",
                     __func__, -global_surplus/MiB);
             }
-            if (cparams->n_ctx == 0) {
-                if (hp_nct > n_ctx_min) {
+            if (n_ctx_auto) {
+                if (n_ctx_max > n_ctx_min_total) {
                     int64_t sum_used_target = sum_free;
                     if (nd == 0) {
                         sum_used_target -= margins[0];
@@ -328,8 +410,9 @@ static void common_params_fit_impl(
                     }
 
                     int64_t sum_projected_used_min_ctx = 0;
-                    cparams->n_ctx = n_ctx_min;
-                    const dmds_t dmds_min_ctx = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+                    cparams->n_ctx = n_ctx_min_total;
+                    dmds_t dmds_min_ctx = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+                    add_extra_memory(dmds_min_ctx);
                     if (nd == 0) {
                         sum_projected_used_min_ctx = dmds_min_ctx.back().mb.total();
                     } else {
@@ -339,14 +422,16 @@ static void common_params_fit_impl(
                     }
                     if (sum_used_target > sum_projected_used_min_ctx) {
                         // linear interpolation between minimum and maximum context size:
-                        cparams->n_ctx += (hp_nct - n_ctx_min) * (sum_used_target - sum_projected_used_min_ctx)
+                        cparams->n_ctx += (n_ctx_max - n_ctx_min_total) * (sum_used_target - sum_projected_used_min_ctx)
                             / (sum_projected_used - sum_projected_used_min_ctx);
-                        cparams->n_ctx = std::max(cparams->n_ctx - cparams->n_ctx % 256, n_ctx_min); // round down context for CUDA backend
+                        // round down context for CUDA backend, keep it divisible by the number of streams:
+                        const uint32_t align = 256 * n_streams;
+                        cparams->n_ctx = std::max(cparams->n_ctx - cparams->n_ctx % align, n_ctx_min_total);
 
-                        const int64_t bytes_per_ctx = (sum_projected_used - sum_projected_used_min_ctx) / (hp_nct - n_ctx_min);
-                        const int64_t memory_reduction = (hp_nct - cparams->n_ctx) * bytes_per_ctx;
+                        const int64_t bytes_per_ctx = (sum_projected_used - sum_projected_used_min_ctx) / (n_ctx_max - n_ctx_min_total);
+                        const int64_t memory_reduction = (n_ctx_max - cparams->n_ctx) * bytes_per_ctx;
                         LOG_TRC("%s: context size reduced from %" PRIu32 " to %" PRIu32 " -> need %" PRId64 " MiB less memory in total\n",
-                            __func__, hp_nct, cparams->n_ctx, memory_reduction/MiB);
+                            __func__, n_ctx_max, cparams->n_ctx, memory_reduction/MiB);
                         if (nd <= 1) {
                             LOG_TRC("%s: entire model can be fit by reducing context\n", __func__);
                             return;
@@ -355,14 +440,14 @@ static void common_params_fit_impl(
                     } else {
                         const int64_t memory_reduction = sum_projected_used - sum_projected_used_min_ctx;
                         LOG_TRC("%s: context size reduced from %" PRIu32 " to %" PRIu32 " -> need %" PRId64 " MiB less memory in total\n",
-                            __func__, hp_nct, cparams->n_ctx, memory_reduction/MiB);
+                            __func__, n_ctx_max, cparams->n_ctx, memory_reduction/MiB);
                     }
                 } else {
                     if (n_ctx_min == UINT32_MAX) {
-                        LOG_TRC("%s: user has requested full context size of %" PRIu32 " -> no change\n", __func__, hp_nct);
+                        LOG_TRC("%s: user has requested full context size of %" PRIu32 " -> no change\n", __func__, n_ctx_max);
                     } else {
                         LOG_TRC("%s: default model context size is %" PRIu32 " which is <= the min. context size of %" PRIu32 " -> no change\n",
-                            __func__, hp_nct, n_ctx_min);
+                            __func__, n_ctx_max, n_ctx_min_total);
                     }
                 }
             } else {
@@ -507,8 +592,9 @@ static void common_params_fit_impl(
         llama_model_params mparams_copy = *mparams;
         set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, mparams_copy);
 
-        const dmds_t dmd_nl = common_get_device_memory_data_impl(
+        dmds_t dmd_nl = common_get_device_memory_data_impl(
             path_model, &mparams_copy, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+        add_extra_memory(dmd_nl);
 
         LOG_TRC("%s: memory for test allocation by device:\n", func_name);
         for (size_t id = 0; id < nd; id++) {
@@ -535,8 +621,9 @@ static void common_params_fit_impl(
         mparams->tensor_buft_overrides = tensor_buft_overrides;
 
         LOG_TRC("%s: getting device memory data with all MoE tensors moved to system memory:\n", __func__);
-        const dmds_t dmds_cpu_moe = common_get_device_memory_data_impl(
+        dmds_t dmds_cpu_moe = common_get_device_memory_data_impl(
             path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+        add_extra_memory(dmds_cpu_moe);
 
         for (size_t id = 0; id < nd; id++) {
             global_surplus_cpu_moe += dmds_cpu_moe[id].free;
@@ -796,11 +883,12 @@ enum common_params_fit_status common_fit_params(
         llama_model_tensor_buft_override * tensor_buft_overrides,
         size_t * margins,
         uint32_t n_ctx_min,
+        const common_fit_extra_model * extra,
         ggml_log_level log_level) {
     const int64_t t0_us = llama_time_us();
     common_params_fit_status status = COMMON_PARAMS_FIT_STATUS_SUCCESS;
     try {
-        common_params_fit_impl(path_model, mparams, cparams, tensor_split, tensor_buft_overrides, margins, n_ctx_min, log_level);
+        common_params_fit_impl(path_model, mparams, cparams, tensor_split, tensor_buft_overrides, margins, n_ctx_min, extra, log_level);
         LOG_TRC("%s: successfully fit params to free device memory\n", __func__);
     } catch (const common_params_fit_exception & e) {
         LOG_WRN("%s: failed to fit params to free device memory: %s\n", __func__, e.what());

common/fit.h

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

common/hf-cache.cpp

@@ -4,9 +4,7 @@
 #include "common.h"
 #include "log.h"
 #include "http.h"
-
-#define JSON_ASSERT GGML_ASSERT
-#include <nlohmann/json.hpp>
+#include "json.h"
 
 #include <filesystem>
 #include <fstream>
@@ -15,8 +13,6 @@
 #include <string_view>
 #include <stdexcept>
 
-namespace nl = nlohmann;
-
 #if defined(_WIN32)
 #define WIN32_LEAN_AND_MEAN
 #ifndef NOMINMAX
@@ -195,8 +191,8 @@ static void safe_write_file(const fs::path & path, const std::string & data) {
     }
 }
 
-static nl::json api_get(const std::string & url,
-                        const std::string & token) {
+static common_json api_get(const std::string & url,
+                           const std::string & token) {
     auto [cli, parts] = common_http_client(url);
 
     httplib::Headers headers = {
@@ -214,10 +210,10 @@ static nl::json api_get(const std::string & url,
         auto body = res->body;
 
         if (res->status == 200) {
-            return nl::json::parse(res->body);
+            return common_json::parse(res->body);
         }
         try {
-            body = nl::json::parse(res->body)["error"].get<std::string>();
+            body = common_json::parse(res->body)["error"].get<std::string>();
         } catch (...) { }
 
         throw std::runtime_error("GET failed (" + std::to_string(res->status) + "): " + body);
@@ -280,7 +276,7 @@ static std::string get_repo_commit(const std::string & repo_id,
         safe_write_file(refs_path / name, commit);
         return commit;
 
-    } catch (const nl::json::exception & e) {
+    } catch (const common_json_error & e) {
         LOG_ERR("%s: JSON error: %s\n", __func__, e.what());
     } catch (const std::exception & e) {
         LOG_ERR("%s: error: %s\n", __func__, e.what());
@@ -358,7 +354,7 @@ hf_files get_repo_files(const std::string & repo_id,
 
             files.push_back(file);
         }
-    } catch (const nl::json::exception & e) {
+    } catch (const common_json_error & e) {
         LOG_ERR("%s: JSON error: %s\n", __func__, e.what());
     } catch (const std::exception & e) {
         LOG_ERR("%s: error: %s\n", __func__, e.what());

common/jinja/README.md

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

common/jinja/caps.cpp

@@ -4,14 +4,14 @@
 
 // note: the json dependency is only for defining input in a convenient way
 // we can remove it in the future when we figure out a better way to define inputs using jinja::value
-#include <nlohmann/json.hpp>
+#include "json.h"
 
 #include <functional>
 #include <sstream>
 
 #define FILENAME "jinja-caps"
 
-using json = nlohmann::ordered_json;
+using json = common_json;
 
 namespace jinja {
 

common/jinja/value.cpp

@@ -3,7 +3,7 @@
 #include "value.h"
 
 // for converting from JSON to jinja values
-#include <nlohmann/json.hpp>
+#include "json.h"
 
 #include <sstream>
 #include <string>
@@ -1355,7 +1355,7 @@ const func_builtins & value_undefined_t::get_builtins() const {
 //////////////////////////////////
 
 
-static value from_json(const nlohmann::ordered_json & j, bool mark_input) {
+static value from_json(const common_json & j, bool mark_input) {
     if (j.is_null()) {
         return mk_val<value_none>();
     } else if (j.is_boolean()) {
@@ -1452,7 +1452,7 @@ bool value_compare(const value & a, const value & b, value_compare_op op) {
 }
 
 template<>
-void global_from_json(context & ctx, const nlohmann::ordered_json & json_obj, bool mark_input) {
+void global_from_json(context & ctx, const common_json & json_obj, bool mark_input) {
     // printf("global_from_json: %s\n" , json_obj.dump(2).c_str());
     if (json_obj.is_null() || !json_obj.is_object()) {
         throw std::runtime_error("global_from_json: input JSON value must be an object");

common/jinja/value.h

@@ -86,7 +86,7 @@ struct context; // forward declaration
 // marking input can be useful for tracking data provenance
 // and preventing template injection attacks
 //
-// Note: T_JSON can be nlohmann::ordered_json
+// Note: T_JSON can be common_json
 template<typename T_JSON>
 void global_from_json(context & ctx, const T_JSON & json_obj, bool mark_input);
 

common/json-schema-to-grammar.cpp

@@ -1,9 +1,8 @@
 #include "json-schema-to-grammar.h"
 #include "common.h"
 
-#include <nlohmann/json.hpp>
-
 #include <algorithm>
+#include <limits>
 #include <map>
 #include <regex>
 #include <sstream>
@@ -12,7 +11,7 @@
 #include <unordered_set>
 #include <vector>
 
-using json = nlohmann::ordered_json;
+using json = common_json;
 
 static std::string build_repetition(const std::string & item_rule, int min_items, int max_items, const std::string & separator_rule = "") {
     auto has_max = max_items != std::numeric_limits<int>::max();
@@ -917,7 +916,11 @@ public:
             return _add_rule(rule_name, _resolve_ref(schema["$ref"]));
         }
         if (schema.contains("oneOf") || schema.contains("anyOf")) {
-            std::vector<json> alt_schemas = schema.contains("oneOf") ? schema["oneOf"].get<std::vector<json>>() : schema["anyOf"].get<std::vector<json>>();
+            const json & alts = schema.contains("oneOf") ? schema.at("oneOf") : schema.at("anyOf");
+            std::vector<json> alt_schemas;
+            for (const auto & alt : alts) {
+                alt_schemas.push_back(alt);
+            }
             return _add_rule(rule_name, _generate_union_rule(name, alt_schemas));
         }
         if (schema_type.is_array()) {
@@ -1111,7 +1114,7 @@ common_schema_info::~common_schema_info() = default;
 common_schema_info::common_schema_info(common_schema_info &&) noexcept = default;
 common_schema_info & common_schema_info::operator=(common_schema_info &&) noexcept = default;
 
-void common_schema_info::resolve_refs(nlohmann::ordered_json & schema) {
+void common_schema_info::resolve_refs(common_json & schema) {
     impl_->resolve_refs(schema, "");
 }
 
@@ -1119,7 +1122,7 @@ void common_schema_info::resolve_refs(nlohmann::ordered_json & schema) {
 // Some models emit raw string values rather than JSON-encoded strings for string parameters.
 // If any branch of the schema (via oneOf, anyOf, $ref, etc.) permits a string, this returns
 // true, allowing callers to handle the value as a raw string for simplicity.
-bool common_schema_info::resolves_to_string(const nlohmann::ordered_json & schema) {
+bool common_schema_info::resolves_to_string(const common_json & schema) {
     std::unordered_set<std::string> visited_refs;
 
     std::function<bool(const json &)> check = [&](const json & s) -> bool {
@@ -1227,7 +1230,7 @@ bool common_schema_info::resolves_to_string(const nlohmann::ordered_json & schem
     return check(schema);
 }
 
-std::string json_schema_to_grammar(const json & schema, bool force_gbnf) {
+std::string json_schema_to_grammar(const common_json & schema, bool force_gbnf) {
 #ifdef LLAMA_USE_LLGUIDANCE
     if (!force_gbnf) {
         return "%llguidance {}\nstart: %json " + schema.dump();
@@ -1248,10 +1251,10 @@ std::string build_grammar(const std::function<void(const common_grammar_builder
         /* .add_rule = */ [&](const std::string & name, const std::string & rule) {
             return converter._add_rule(name, rule);
         },
-        /* .add_schema = */ [&](const std::string & name, const nlohmann::ordered_json & schema) {
+        /* .add_schema = */ [&](const std::string & name, const common_json & schema) {
             return converter.visit(schema, name == "root" ? "" : name);
         },
-        /* .resolve_refs = */ [&](nlohmann::ordered_json & schema) {
+        /* .resolve_refs = */ [&](common_json & schema) {
             converter.resolve_refs(schema, "");
         }
     };

common/json-schema-to-grammar.h

@@ -1,12 +1,12 @@
 #pragma once
 
-#include <nlohmann/json_fwd.hpp>
+#include "json.h"
 
 #include <functional>
 #include <memory>
 #include <string>
 
-std::string json_schema_to_grammar(const nlohmann::ordered_json & schema,
+std::string json_schema_to_grammar(const common_json & schema,
                                    bool force_gbnf = false);
 
 class common_schema_converter;
@@ -24,14 +24,14 @@ class common_schema_info {
     common_schema_info(common_schema_info &&) noexcept;
     common_schema_info & operator=(common_schema_info &&) noexcept;
 
-    void resolve_refs(nlohmann::ordered_json & schema);
-    bool resolves_to_string(const nlohmann::ordered_json & schema);
+    void resolve_refs(common_json & schema);
+    bool resolves_to_string(const common_json & schema);
 };
 
 struct common_grammar_builder {
     std::function<std::string(const std::string &, const std::string &)> add_rule;
-    std::function<std::string(const std::string &, const nlohmann::ordered_json &)> add_schema;
-    std::function<void(nlohmann::ordered_json &)> resolve_refs;
+    std::function<std::string(const std::string &, const common_json &)> add_schema;
+    std::function<void(common_json &)> resolve_refs;
 };
 
 struct common_grammar_options {

common/json.cpp

@@ -0,0 +1,433 @@
+#include "json.h"
+
+#include "ggml.h"
+
+#define JSON_ASSERT GGML_ASSERT
+#include <nlohmann/json.hpp>
+
+#include <iterator>
+#include <new>
+#include <set>
+#include <unordered_map>
+#include <vector>
+
+using nlohmann::ordered_json;
+
+// a common_json is the backing value, so any value of a tree can be used as a common_json
+static_assert(sizeof(ordered_json)  <= sizeof(common_json),  "common_json storage is too small");
+static_assert(alignof(ordered_json) <= alignof(common_json), "common_json alignment is too weak");
+
+// runs fn and gives every error of the backing library as a common_json_error
+template <typename F>
+static decltype(auto) guard(F && fn) {
+    try {
+        return fn();
+    } catch (const ordered_json::exception & e) {
+        throw common_json_error(e.what());
+    }
+}
+
+static ordered_json & as_json(common_json * self) {
+    return *reinterpret_cast<ordered_json *>(self);
+}
+
+static const ordered_json & as_json(const common_json * self) {
+    return *reinterpret_cast<const ordered_json *>(self);
+}
+
+static common_json & as_common(ordered_json & json) {
+    return *reinterpret_cast<common_json *>(&json);
+}
+
+static const common_json & as_common(const ordered_json & json) {
+    return *reinterpret_cast<const common_json *>(&json);
+}
+
+static ordered_json to_json(const common_json_value & val) {
+    switch (val.type) {
+        case common_json_value::VAL_NULL:   return nullptr;
+        case common_json_value::VAL_BOOL:   return val.val_bool;
+        case common_json_value::VAL_INT:    return val.val_int;
+        case common_json_value::VAL_UINT:   return val.val_uint;
+        case common_json_value::VAL_DOUBLE: return val.val_double;
+        case common_json_value::VAL_STRING: return val.val_string;
+        case common_json_value::VAL_JSON:
+            // one owner means no one else can see this tree, so it is safe to move it out
+            // note: this makes a value single use, same as the json_ref of the backing library
+            if (val.val_json.use_count() == 1) {
+                return std::move(as_json(val.val_json.get()));
+            }
+            return as_json(val.val_json.get());
+    }
+
+    return nullptr;
+}
+
+common_json_value::common_json_value(const char * val) {
+    if (val) {
+        type       = VAL_STRING;
+        val_string = val;
+    } else {
+        type = VAL_NULL;
+    }
+}
+
+common_json_value::common_json_value(const common_json & val) :
+    type(VAL_JSON), val_json(std::make_shared<common_json>(val)) {}
+
+common_json_value::common_json_value(common_json && val) :
+    type(VAL_JSON), val_json(std::make_shared<common_json>(std::move(val))) {}
+
+// the ctors and get<T>() below are explicit specializations, giving strong symbols
+// an explicit instantiation is a weak symbol, dropped by some LTO builds (clang-cl)
+template <typename T>
+static std::shared_ptr<common_json> set_json(const std::set<T> & vals) {
+    common_json out = common_json::array();
+
+    for (const auto & val : vals) {
+        out.push_back(val);
+    }
+
+    return std::make_shared<common_json>(std::move(out));
+}
+
+// a set value is usable only for the types below
+#define COMMON_JSON_SET(...) template <> common_json_value::common_json_value(const std::set<__VA_ARGS__> & vals) : type(VAL_JSON), val_json(set_json(vals)) {}
+
+COMMON_JSON_SET(int)
+COMMON_JSON_SET(std::string)
+
+#undef COMMON_JSON_SET
+
+template <typename T>
+static std::shared_ptr<common_json> map_json(const T & vals) {
+    common_json out = common_json::object();
+
+    for (const auto & val : vals) {
+        out.set({ val.first, val.second });
+    }
+
+    return std::make_shared<common_json>(std::move(out));
+}
+
+// a map value is usable only for the types below
+#define COMMON_JSON_MAP(...) template <> common_json_value::common_json_value(const std::map<std::string, __VA_ARGS__> & vals) : type(VAL_JSON), val_json(map_json(vals)) {}
+
+COMMON_JSON_MAP(bool)
+COMMON_JSON_MAP(std::string)
+
+#undef COMMON_JSON_MAP
+
+// an unordered map value is usable only for the types below
+#define COMMON_JSON_UMAP(...) template <> common_json_value::common_json_value(const std::unordered_map<std::string, __VA_ARGS__> & vals) : type(VAL_JSON), val_json(map_json(vals)) {}
+
+COMMON_JSON_UMAP(size_t)
+
+#undef COMMON_JSON_UMAP
+
+template <typename T>
+static std::shared_ptr<common_json> vec_json(const std::vector<T> & vals) {
+    common_json out = common_json::array();
+
+    for (const auto & val : vals) {
+        out.push_back(val);
+    }
+
+    return std::make_shared<common_json>(std::move(out));
+}
+
+// a vector value is usable only for the types below
+// note: std::vector<bool> is not here, its proxy reference does not convert
+#define COMMON_JSON_VEC(...) template <> common_json_value::common_json_value(const std::vector<__VA_ARGS__> & vals) : type(VAL_JSON), val_json(vec_json(vals)) {}
+
+COMMON_JSON_VEC(int)
+COMMON_JSON_VEC(unsigned char)
+COMMON_JSON_VEC(unsigned int)
+COMMON_JSON_VEC(long)
+COMMON_JSON_VEC(unsigned long)
+COMMON_JSON_VEC(long long)
+COMMON_JSON_VEC(unsigned long long)
+COMMON_JSON_VEC(float)
+COMMON_JSON_VEC(double)
+COMMON_JSON_VEC(std::string)
+COMMON_JSON_VEC(std::vector<float>)
+COMMON_JSON_VEC(common_json)
+
+#undef COMMON_JSON_VEC
+
+common_json_value::common_json_value(std::initializer_list<common_json_item> items) :
+    type(VAL_JSON), val_json(std::make_shared<common_json>(items)) {}
+
+// null, same as the backing library
+// operator[] turns it into an object, push_back() into an array
+common_json::common_json() {
+    new (storage) ordered_json();
+}
+
+common_json::common_json(const common_json & other) {
+    new (storage) ordered_json(as_json(&other));
+}
+
+common_json::common_json(common_json && other) noexcept {
+    new (storage) ordered_json(std::move(as_json(&other)));
+}
+
+common_json::common_json(std::initializer_list<common_json_item> items) {
+    new (storage) ordered_json(ordered_json::object());
+
+    for (const auto & item : items) {
+        set(item);
+    }
+}
+
+common_json::common_json(const common_json_value & val) {
+    new (storage) ordered_json(to_json(val));
+}
+
+common_json::common_json(std::nullptr_t) {
+    new (storage) ordered_json(nullptr);
+}
+
+common_json & common_json::operator=(common_json other) noexcept {
+    as_json(this).swap(as_json(&other));
+
+    return *this;
+}
+
+common_json::~common_json() {
+    as_json(this).~basic_json();
+}
+
+common_json common_json::parse(const std::string & text) {
+    try {
+        // the assignment moves the parsed tree in, it does not copy
+        common_json out;
+        as_json(&out) = ordered_json::parse(text);
+        return out;
+    } catch (const std::exception & e) {
+        throw common_json_error(e.what());
+    }
+}
+
+common_json common_json::parse_no_throw(const std::string & text) {
+    common_json out;
+    as_json(&out) = ordered_json::parse(text, nullptr, false);
+    return out;
+}
+
+bool common_json::is_discarded() const {
+    return as_json(this).is_discarded();
+}
+
+common_json common_json::array() {
+    common_json out;
+    as_json(&out) = ordered_json::array();
+    return out;
+}
+
+common_json common_json::array(std::initializer_list<common_json_value> vals) {
+    common_json out;
+    ordered_json & arr = as_json(&out);
+    arr = ordered_json::array();
+
+    for (const auto & val : vals) {
+        arr.push_back(to_json(val));
+    }
+
+    return out;
+}
+
+common_json common_json::object() {
+    common_json out;
+    as_json(&out) = ordered_json::object();
+    return out;
+}
+
+common_json common_json::object(std::initializer_list<common_json_item> items) {
+    return common_json(items);
+}
+
+common_json common_json::make(const common_json_value & val) {
+    return common_json(val);
+}
+
+bool common_json::is_null()           const { return as_json(this).is_null(); }
+bool common_json::is_object()         const { return as_json(this).is_object(); }
+bool common_json::is_array()          const { return as_json(this).is_array(); }
+bool common_json::is_string()         const { return as_json(this).is_string(); }
+bool common_json::is_boolean()        const { return as_json(this).is_boolean(); }
+bool common_json::is_number()         const { return as_json(this).is_number(); }
+bool common_json::is_number_integer() const { return as_json(this).is_number_integer(); }
+bool common_json::is_number_float()   const { return as_json(this).is_number_float(); }
+
+bool   common_json::empty() const { return as_json(this).empty(); }
+size_t common_json::size()  const { return as_json(this).size(); }
+
+bool common_json::contains(const std::string & key) const {
+    return as_json(this).contains(key);
+}
+
+bool common_json::operator==(const common_json_value & val) const {
+    // compare a tree in place, to_json() would copy it
+    if (val.type == common_json_value::VAL_JSON) {
+        return as_json(this) == as_json(val.val_json.get());
+    }
+    return as_json(this) == to_json(val);
+}
+
+bool common_json::operator!=(const common_json_value & val) const {
+    return !(*this == val);
+}
+
+common_json       & common_json::at(const std::string & key)       { return guard([&]() -> common_json       & { return as_common(as_json(this).at(key)); }); }
+const common_json & common_json::at(const std::string & key) const { return guard([&]() -> const common_json & { return as_common(as_json(this).at(key)); }); }
+common_json       & common_json::at(size_t idx)                    { return guard([&]() -> common_json       & { return as_common(as_json(this).at(idx)); }); }
+const common_json & common_json::at(size_t idx)              const { return guard([&]() -> const common_json & { return as_common(as_json(this).at(idx)); }); }
+
+common_json       & common_json::operator[](const std::string & key)       { return guard([&]() -> common_json       & { return as_common(as_json(this)[key]); }); }
+const common_json & common_json::operator[](const std::string & key) const { return guard([&]() -> const common_json & { return as_common(as_json(this).at(key)); }); }
+common_json       & common_json::operator[](size_t idx)                    { return guard([&]() -> common_json       & { return as_common(as_json(this)[idx]); }); }
+const common_json & common_json::operator[](size_t idx)              const { return guard([&]() -> const common_json & { return as_common(as_json(this).at(idx)); }); }
+
+common_json       & common_json::front()       { return as_common(as_json(this).front()); }
+const common_json & common_json::front() const { return as_common(as_json(this).front()); }
+common_json       & common_json::back()        { return as_common(as_json(this).back()); }
+const common_json & common_json::back()  const { return as_common(as_json(this).back()); }
+
+void common_json::clear() {
+    as_json(this).clear();
+}
+
+void common_json::erase(const std::string & key) {
+    guard([&] { as_json(this).erase(key); });
+}
+
+void common_json::erase(size_t idx) {
+    guard([&] { as_json(this).erase(idx); });
+}
+
+void common_json::assign(const common_json_value & val) {
+    as_json(this) = to_json(val);
+}
+
+void common_json::set(const common_json_item & item) {
+    guard([&] { as_json(this)[item.key] = to_json(item.val); });
+}
+
+void common_json::push_back(const common_json_value & val) {
+    guard([&] { as_json(this).push_back(to_json(val)); });
+}
+
+void common_json::push_back(std::initializer_list<common_json_item> items) {
+    common_json val(items);
+
+    guard([&] { as_json(this).push_back(std::move(as_json(&val))); });
+}
+
+size_t common_json::count(const std::string & key) const {
+    return as_json(this).count(key);
+}
+
+void common_json::insert(const common_json & vals) {
+    guard([&] {
+        ordered_json & self = as_json(this);
+
+        self.insert(self.end(), as_json(&vals).begin(), as_json(&vals).end());
+    });
+}
+
+std::string common_json::dump(int indent) const {
+    return guard([&] { return as_json(this).dump(indent); });
+}
+
+std::string common_json::dump_safe(int indent) const {
+    return as_json(this).dump(indent, ' ', false, ordered_json::error_handler_t::replace);
+}
+
+// an array is indexed directly, an object needs a walk from the start
+common_json & common_json::iterator::operator*() const {
+    return guard([&]() -> common_json & {
+        ordered_json & j = as_json(node);
+
+        if (j.is_object()) {
+            return as_common(std::next(j.begin(), idx).value());
+        }
+        if (j.is_array()) {
+            return as_common(j[idx]);
+        }
+
+        // a plain value gives itself once, same as the backing library
+        return *node;
+    });
+}
+
+std::string common_json::iterator::key() const {
+    return guard([&] { return std::next(as_json(node).begin(), idx).key(); });
+}
+
+common_json::iterator common_json::begin() const {
+    return iterator(const_cast<common_json *>(this), 0);
+}
+
+common_json::iterator common_json::end() const {
+    return iterator(const_cast<common_json *>(this), size());
+}
+
+// the keys follow the backing library: the index for an array, "" for a plain value
+common_json::items_view::entry common_json::items_view::iterator::operator*() const {
+    return guard([&]() -> entry {
+        ordered_json & j = as_json(node);
+
+        if (j.is_object()) {
+            auto it = std::next(j.begin(), idx);
+
+            return { it.key(), as_common(it.value()) };
+        }
+        if (j.is_array()) {
+            return { std::to_string(idx), as_common(j[idx]) };
+        }
+
+        return { std::string(), *node };
+    });
+}
+
+common_json::items_view common_json::items() const {
+    return items_view(const_cast<common_json *>(this), size());
+}
+
+// the backing library cannot build a common_json, so this one is just a copy
+template <> common_json common_json::get<common_json>() const {
+    return *this;
+}
+
+// get<T>() is usable only for the types below
+
+#define COMMON_JSON_GET(...) template <> __VA_ARGS__ common_json::get<__VA_ARGS__>() const { return guard([&] { return as_json(this).get<__VA_ARGS__>(); }); }
+
+COMMON_JSON_GET(bool)
+COMMON_JSON_GET(int)
+COMMON_JSON_GET(unsigned int)
+COMMON_JSON_GET(long)
+COMMON_JSON_GET(unsigned long)
+COMMON_JSON_GET(long long)
+COMMON_JSON_GET(unsigned long long)
+COMMON_JSON_GET(float)
+COMMON_JSON_GET(double)
+COMMON_JSON_GET(std::string)
+COMMON_JSON_GET(std::vector<float>)
+COMMON_JSON_GET(std::vector<std::string>)
+COMMON_JSON_GET(std::set<std::string>)
+COMMON_JSON_GET(std::vector<int>)
+COMMON_JSON_GET(std::vector<size_t>)
+COMMON_JSON_GET(std::unordered_map<std::string, size_t>)
+
+#undef COMMON_JSON_GET
+
+// must stay below the get<std::string> specialization
+common_json::operator std::string() const {
+    return get<std::string>();
+}
+
+std::string common_json::value(const std::string & key, const char * def) const {
+    return contains(key) ? at(key).get<std::string>() : std::string(def);
+}

common/json.h

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

common/peg-parser.cpp

@@ -10,7 +10,6 @@
 #include <initializer_list>
 #include <map>
 #include <memory>
-#include <nlohmann/json.hpp>
 #include <regex>
 #include <set>
 #include <stdexcept>
@@ -1120,8 +1119,8 @@ common_peg_parser common_peg_parser_builder::chars(const std::string & classes,
     return wrap(arena_.add_parser(common_peg_chars_parser{classes, ranges, negated, min, max}));
 }
 
-common_peg_parser common_peg_parser_builder::schema(const common_peg_parser & p, const std::string & name, const nlohmann::ordered_json & schema, bool raw) {
-    return wrap(arena_.add_parser(common_peg_schema_parser{p.id(), name, std::make_shared<nlohmann::ordered_json>(schema), raw}));
+common_peg_parser common_peg_parser_builder::schema(const common_peg_parser & p, const std::string & name, const common_json & schema, bool raw) {
+    return wrap(arena_.add_parser(common_peg_schema_parser{p.id(), name, std::make_shared<common_json>(schema), raw}));
 }
 
 common_peg_parser common_peg_parser_builder::rule(const std::string & name, const common_peg_parser & p, bool trigger) {
@@ -1805,8 +1804,8 @@ void common_peg_arena::build_grammar(const common_grammar_builder & builder, boo
     }
 }
 
-static nlohmann::json serialize_parser_variant(const common_peg_parser_variant & variant) {
-    using json = nlohmann::json;
+static common_json serialize_parser_variant(const common_peg_parser_variant & variant) {
+    using json = common_json;
 
     return std::visit([](const auto & p) -> json {
         using T = std::decay_t<decltype(p)>;
@@ -1860,7 +1859,7 @@ static nlohmann::json serialize_parser_variant(const common_peg_parser_variant &
                 {"type", "schema"},
                 {"child", p.child},
                 {"name", p.name},
-                {"schema", p.schema ? *p.schema : nullptr},
+                {"schema", p.schema ? *p.schema : json(nullptr)},
                 {"raw", p.raw}
             };
         } else if constexpr (std::is_same_v<T, common_peg_rule_parser>) {
@@ -1888,19 +1887,19 @@ static nlohmann::json serialize_parser_variant(const common_peg_parser_variant &
     }, variant);
 }
 
-nlohmann::json common_peg_arena::to_json() const {
-    auto parsers = nlohmann::json::array();
+common_json common_peg_arena::to_json() const {
+    auto parsers = common_json::array();
     for (const auto & parser : parsers_) {
         parsers.push_back(serialize_parser_variant(parser));
     }
-    return nlohmann::json{
+    return common_json{
         {"parsers", parsers},
         {"rules", rules_},
         {"root", root_}
     };
 }
 
-static common_peg_parser_variant deserialize_parser_variant(const nlohmann::json & j) {
+static common_peg_parser_variant deserialize_parser_variant(const common_json & j) {
     if (!j.contains("type") || !j["type"].is_string()) {
         throw std::runtime_error("Parser variant JSON missing or invalid 'type' field");
     }
@@ -1969,9 +1968,9 @@ static common_peg_parser_variant deserialize_parser_variant(const nlohmann::json
         }
         common_peg_chars_parser parser;
         parser.pattern = j["pattern"];
-        parser.negated = j["negated"];
-        parser.min_count = j["min_count"];
-        parser.max_count = j["max_count"];
+        parser.negated = j["negated"].get<bool>();
+        parser.min_count = j["min_count"].get<int>();
+        parser.max_count = j["max_count"].get<int>();
         for (const auto & range_json : j["ranges"]) {
             if (!range_json.contains("start") || !range_json.contains("end")) {
                 throw std::runtime_error("char_range missing 'start' or 'end' field");
@@ -2007,7 +2006,7 @@ static common_peg_parser_variant deserialize_parser_variant(const nlohmann::json
         parser.child = j["child"].get<common_peg_parser_id>();
         parser.name = j["name"];
         if (!j["schema"].is_null()) {
-            parser.schema = std::make_shared<nlohmann::ordered_json>(j["schema"]);
+            parser.schema = std::make_shared<common_json>(j["schema"]);
         }
         parser.raw = j["raw"].get<bool>();
         return parser;
@@ -2069,7 +2068,7 @@ static common_peg_parser_variant deserialize_parser_variant(const nlohmann::json
     throw std::runtime_error("Unknown parser type: " + type);
 }
 
-common_peg_arena common_peg_arena::from_json(const nlohmann::json & j) {
+common_peg_arena common_peg_arena::from_json(const common_json & j) {
     if (!j.contains("parsers") || !j["parsers"].is_array()) {
         throw std::runtime_error("JSON missing or invalid 'parsers' array");
     }
@@ -2109,7 +2108,7 @@ std::string common_peg_arena::save() const {
 }
 
 void common_peg_arena::load(const std::string & data) {
-    *this = from_json(nlohmann::json::parse(data));
+    *this = from_json(common_json::parse(data));
 }
 
 common_peg_arena build_peg_parser(const std::function<common_peg_parser(common_peg_parser_builder & builder)> & fn) {

common/peg-parser.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <nlohmann/json_fwd.hpp>
+#include "json.h"
 
 #include <memory>
 #include <set>
@@ -245,7 +245,7 @@ struct common_peg_until_parser {
 struct common_peg_schema_parser {
     common_peg_parser_id child;
     std::string name;
-    std::shared_ptr<nlohmann::ordered_json> schema;
+    std::shared_ptr<common_json> schema;
 
     // Indicates if the GBNF should accept a raw string that matches the schema.
     bool raw;
@@ -332,8 +332,8 @@ class common_peg_arena {
 
     std::string dump(common_peg_parser_id id) const;
 
-    nlohmann::json to_json() const;
-    static common_peg_arena from_json(const nlohmann::json & j);
+    common_json to_json() const;
+    static common_peg_arena from_json(const common_json & j);
 
     std::string save() const;
     void load(const std::string & data);
@@ -490,7 +490,7 @@ class common_peg_parser_builder {
 
     // Wraps a parser with JSON schema metadata for grammar generation.
     // Used internally to convert JSON schemas to GBNF grammar rules.
-    common_peg_parser schema(const common_peg_parser & p, const std::string & name, const nlohmann::ordered_json & schema, bool raw = false);
+    common_peg_parser schema(const common_peg_parser & p, const std::string & name, const common_json & schema, bool raw = false);
 
     // Creates a named rule, stores it in the grammar, and returns a ref.
     // If trigger=true, marks this rule as an entry point for lazy grammar generation.

common/speculative.cpp

@@ -2322,6 +2322,9 @@ common_params common_base_params_to_speculative(const common_params & params) {
     const auto & params_spec = params.speculative.draft;
     common_params result = params;
 
+    result.embedding    = false;
+    result.pooling_type = LLAMA_POOLING_TYPE_UNSPECIFIED;
+
     if (has_draft) {
         result.devices               = params_spec.devices;
         result.model                 = params_spec.mparams;
@@ -2385,6 +2388,9 @@ common_speculative_init_result::common_speculative_init_result(
         cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
     }
 
+    // the draft context holds as many tokens per sequence as the target context
+    cparams.n_ctx = llama_n_ctx(ctx_tgt);
+
     // note: for small models maybe we can set this to the maximum possible draft from all speculative types
     //       the extra memory for small models is likely negligible?
     cparams.n_rs_seq  = 0;

conversion/__init__.py

@@ -58,12 +58,16 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "DSparkDraftModel": "qwen",
     "DSparkSpeculator": "qwen",
     "Lfm2DSparkDraftModel": "qwen",
+    "LingDSparkModel": "qwen",
     "DeepseekV4ForCausalLM": "deepseek",
     "DeepseekV4DSparkModel": "deepseek",
     "DistilBertForMaskedLM": "bert",
     "DistilBertForSequenceClassification": "bert",
     "DistilBertModel": "bert",
     "Dots1ForCausalLM": "dots1",
+    "Dots3NoteForCausalLM": "dots3",
+    "Dots3NoteForConditionalGeneration": "dots3",
+    "Dots3NoteTextForCausalLM": "dots3",
     "DotsOCRForCausalLM": "qwen",
     "DreamModel": "dream",
     "Ernie4_5ForCausalLM": "ernie",
@@ -279,6 +283,8 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
     "CogVLMForCausalLM": "cogvlm",
     "DeepseekOCR2ForCausalLM": "deepseek",
     "DeepseekOCRForCausalLM": "deepseek",
+    "Dots3NoteForCausalLM": "dots3",
+    "Dots3NoteForConditionalGeneration": "dots3",
     "DotsOCRForCausalLM": "dotsocr",
     "Exaone4_5_ForConditionalGeneration": "exaone",
     "Gemma3ForConditionalGeneration": "gemma",

conversion/dots3.py

@@ -0,0 +1,323 @@
+from __future__ import annotations
+
+import math
+import re
+
+import torch
+
+from typing import TYPE_CHECKING, Any, Callable, Iterable
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import MmprojModel, ModelBase, gguf
+
+from .deepseek import DeepseekV2Model
+
+
+@ModelBase.register("Dots3NoteForCausalLM", "Dots3NoteForConditionalGeneration", "Dots3NoteTextForCausalLM")
+class Dots3NoteModel(DeepseekV2Model):
+    model_arch = gguf.MODEL_ARCH.DOTS3NOTE
+    skip_mtp = False
+    supports_mtp_export = True
+
+    # trunk layer count, stashed before indexing for filter_tensors (mirrors DeepseekV32Model)
+    _n_main_layers: int | None = None
+
+    def index_tensors(self, remote_hf_model_id: str | None = None):
+        type(self)._n_main_layers = self.hparams["num_hidden_layers"]
+        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        hparams = self.hparams
+
+        # config file doesn't specify MTP block, detect it from model weight
+        self.n_nextn = 1 if "model.mtp.embed_tokens.weight" in self.model_tensors else 0
+        if self.n_nextn:
+            self.block_count += self.n_nextn
+            self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
+
+        self.layer_types = hparams["layer_types"]
+        if len(self.layer_types) < hparams["num_hidden_layers"]:
+            raise ValueError("layer_types is shorter than num_hidden_layers")
+
+        if hparams.get("use_dsa", True) is not True:
+            raise ValueError("dots3-note conversion requires use_dsa=true")
+        if hparams.get("normalization", "RMSNorm") != "RMSNorm" or hparams.get("final_norm", "RMSNorm") != "RMSNorm":
+            raise ValueError("dots3-note conversion only supports RMSNorm")
+        if hparams.get("k_rope_only_layernorm", True) is not True:
+            raise ValueError("dots3-note conversion requires k_rope_only_layernorm=true")
+        if hparams.get("topk_method", "noaux_tc") != "noaux_tc" or hparams.get("scoring_func") != "sigmoid":
+            raise ValueError("dots3-note conversion only supports noaux_tc/sigmoid expert gating")
+        if hparams.get("n_group", 1) != 1 or hparams.get("topk_group", 1) != 1:
+            raise ValueError("dots3-note conversion does not support grouped expert routing")
+        if hparams.get("use_dynamic_rsf", False) or hparams.get("moe_gating_fp32", False):
+            raise ValueError("dots3-note conversion does not support use_dynamic_rsf/moe_gating_fp32")
+        for key in ("attention_gate_type", "swa_attention_gate_type"):
+            if hparams.get(key, "headwise") != "headwise":
+                raise ValueError(f"dots3-note conversion only supports headwise attention gate, got {key}={hparams.get(key)!r}")
+        if hparams["swa_qk_nope_head_dim"] + hparams["swa_qk_rope_head_dim"] != hparams.get("swa_head_dim", 256):
+            raise ValueError("swa_head_dim must equal swa_qk_nope_head_dim + swa_qk_rope_head_dim")
+        if hparams["swa_qk_rope_head_dim"] != hparams["qk_rope_head_dim"]:
+            # both layer kinds share a single rope_dimension_count
+            raise ValueError("swa_qk_rope_head_dim must match qk_rope_head_dim")
+
+        self.apply_lora_rescale = hparams.get("apply_mla_qkv_lora_rescale", False)
+
+    def _is_swa_layer(self, bid: int) -> bool:
+        if bid >= self.hparams["num_hidden_layers"]:
+            # note: the NextN/MTP block uses the sliding-attention MLA
+            return True
+        return self.layer_types[bid] == "sliding_attention"
+
+    def set_vocab(self):
+        from transformers import AutoTokenizer
+        tokenizer = AutoTokenizer.from_pretrained(self.dir_model)
+        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
+        tokens, toktypes, tokpre = self.get_vocab_base()
+        self.gguf_writer.add_tokenizer_model("gpt2")
+        self.gguf_writer.add_tokenizer_pre(tokpre)
+        self.gguf_writer.add_token_list(tokens)
+        self.gguf_writer.add_token_types(toktypes)
+        special_vocab._set_special_token("eot", tokenizer.get_added_vocab()["<|endofassistant|>"])  # ty: ignore[unresolved-attribute]
+        special_vocab.add_to_gguf(self.gguf_writer)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        if (titem := super().filter_tensors(item)) is None:
+            return None
+        name, gen = titem
+        if name.startswith(("vision_encoder.", "audio_encoder.")):
+            return None
+
+        assert cls._n_main_layers is not None
+        is_mtp = name.startswith("model.mtp.") or \
+            ((m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers)
+
+        # --no-mtp: drop the NextN/MTP block; --mtp: keep only that block plus the shared embeddings/norm/lm_head
+        if is_mtp and cls.no_mtp:
+            return None
+        if cls.mtp_only and not is_mtp and name not in (
+            "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
+        ):
+            return None
+
+        return name, gen
+
+    def set_gguf_parameters(self):
+        hparams = self.hparams
+
+        # head_count is a per-layer array because the two layer kinds have different head counts
+        n_layer = hparams["num_hidden_layers"]
+        hparams["num_attention_heads"] = [
+            hparams["swa_num_attention_heads"] if self._is_swa_layer(il) else hparams["num_attention_heads"]
+            for il in range(self.block_count)
+        ]
+
+        # prevent the base class from emitting key/value_length from the unused head_dim
+        hparams.pop("head_dim", None)
+
+        super().set_gguf_parameters()
+
+        # MLA geometry of the sliding-window layers (rope.freq_base_swa is emitted by the base class)
+        swa_kv_lora_rank = hparams["swa_kv_lora_rank"]
+        self.gguf_writer.add_sliding_window(hparams["sliding_window_size"])
+        self.gguf_writer.add_sliding_window_pattern([self._is_swa_layer(il) for il in range(n_layer)])
+        self.gguf_writer.add_kv_lora_rank_swa(swa_kv_lora_rank)
+        self.gguf_writer.add_key_length_swa(swa_kv_lora_rank + hparams["swa_qk_rope_head_dim"])
+        self.gguf_writer.add_value_length_swa(swa_kv_lora_rank)
+        self.gguf_writer.add_key_length_mla_swa(hparams["swa_qk_nope_head_dim"] + hparams["swa_qk_rope_head_dim"])
+        self.gguf_writer.add_value_length_mla_swa(hparams["swa_v_head_dim"])
+        if hparams["swa_q_lora_rank"] != hparams["q_lora_rank"]:
+            raise ValueError("dots3-note conversion assumes a shared q_lora_rank for both layer kinds")
+
+        if self.n_nextn:
+            self.gguf_writer.add_nextn_predict_layers(self.n_nextn)
+
+        # DSA indexer (full-attention layers only)
+        self.gguf_writer.add_indexer_head_count(hparams["index_n_heads"])
+        self.gguf_writer.add_indexer_key_length(hparams["index_head_dim"])
+        self.gguf_writer.add_indexer_top_k(hparams["index_topk"])
+        self.gguf_writer.add_indexer_types([not self._is_swa_layer(il) for il in range(n_layer)])
+
+    def prepare_metadata(self, vocab_only: bool):
+        from_dir = self.fname_out.is_dir()
+        super().prepare_metadata(vocab_only=vocab_only)
+
+        if not self.mtp_only or not from_dir:
+            return
+
+        output_type: str = self.ftype.name.partition("_")[2]
+        fname_default: str = gguf.naming_convention(
+            self.metadata.name, self.metadata.basename, self.metadata.finetune,
+            self.metadata.version, size_label=None, output_type=output_type, model_type=None)
+        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # move the MTP token embedding into the NextN block so the standard nextn mapping picks it up
+        if name == "model.mtp.embed_tokens.weight":
+            name = f"model.layers.{self.hparams['num_hidden_layers']}.embed_tokens.weight"
+            bid = self.hparams["num_hidden_layers"]
+
+        # fold the activation rescale sqrt(n_embd/lora_rank) into the preceding RMSNorm weight
+        # this also covers the indexer wq_b, which reads the same rescaled q_lora activation
+        if self.apply_lora_rescale and bid is not None:
+            if name.endswith("q_a_layernorm.weight"):
+                data_torch = data_torch * math.sqrt(self.hparams["hidden_size"] / self.hparams["q_lora_rank"])
+            elif name.endswith("kv_a_layernorm.weight"):
+                rank = self.hparams["swa_kv_lora_rank"] if self._is_swa_layer(bid) else self.hparams["kv_lora_rank"]
+                data_torch = data_torch * math.sqrt(self.hparams["hidden_size"] / rank)
+
+        # MLA absorption: split kv_b_proj into k_b (transposed) and v_b, per-layer-kind geometry
+        if name.endswith("kv_b_proj.weight"):
+            assert bid is not None
+            if self._is_swa_layer(bid):
+                n_head = self.hparams["swa_num_attention_heads"]
+                qk_nope_head_dim = self.hparams["swa_qk_nope_head_dim"]
+                v_head_dim = self.hparams["swa_v_head_dim"]
+            else:
+                n_head = self.hparams["num_attention_heads"]
+                qk_nope_head_dim = self.hparams["qk_nope_head_dim"]
+                v_head_dim = self.hparams["v_head_dim"]
+            if isinstance(n_head, list):  # set_gguf_parameters turns this into a per-layer array
+                n_head = n_head[bid]
+
+            assert data_torch.shape[0] == n_head * (qk_nope_head_dim + v_head_dim)
+
+            kv_b = data_torch.view(n_head, qk_nope_head_dim + v_head_dim, data_torch.shape[-1])
+            k_b, v_b = kv_b.split([qk_nope_head_dim, v_head_dim], dim=1)
+            k_b = k_b.transpose(1, 2)
+
+            yield from ModelBase.modify_tensors(self, k_b, name.replace("kv_b_proj", "k_b_proj"), bid)
+            yield from ModelBase.modify_tensors(self, v_b, name.replace("kv_b_proj", "v_b_proj"), bid)
+            return
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+
+@ModelBase.register("Dots3NoteForCausalLM", "Dots3NoteForConditionalGeneration")
+class Dots3NoteMmprojModel(MmprojModel):
+    has_vision_encoder = True
+    has_audio_encoder = True
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert self.hparams_vision is not None
+        assert self.hparams_audio is not None
+
+        # preprocessor_config.json nests the image params under vision_config
+        self.preprocessor_config = {**self.preprocessor_config, **self.preprocessor_config.get("vision_config", {})}
+
+        vis = self.hparams_vision
+        # in this config, hidden_size is the adapter output width; embed_dim is the tower width
+        vis["hidden_size"] = vis["embed_dim"]
+        vis["image_size"] = 0  # dynamic resolution
+        self.pyramid = [max(0, n) for n in vis["pyramid_num_routed"]]
+
+        if vis.get("adapter_type") != "patch_merger" or not vis.get("pre_pixel_shuffle"):
+            raise ValueError("dots3-note vision conversion requires adapter_type=patch_merger and pre_pixel_shuffle")
+        if vis.get("router_scoring_func", "sigmoid") != "sigmoid" or vis.get("router_scale", 1.0) != 1.0:
+            raise ValueError("dots3-note vision conversion only supports sigmoid routing with router_scale=1.0")
+        if vis.get("temporal_patch_size", 1) != 1 or vis.get("use_bias") or not vis.get("use_qk_norm"):
+            raise ValueError("unsupported dots3-note vision config variant")
+
+        aud = self.hparams_audio
+        if not aud.get("use_conv2d_stem") or not aud.get("use_rope") or not aud.get("use_rms_norm") or aud.get("use_causal"):
+            raise ValueError("unsupported dots3-note audio config variant")
+        if aud["whisper_config"].get("activation_function") != "swiglu":
+            raise ValueError("dots3-note audio conversion requires the swiglu activation")
+        if aud.get("merge_factor", 1) != 1 or aud.get("chunk_seconds") != 60:
+            raise ValueError("unsupported dots3-note audio chunking config")
+        # the graph hard-codes these rope parameters
+        rope = aud.get("rope_parameters", {})
+        if rope.get("partial_rotary_factor") != 0.5 or rope.get("rope_theta") != 10000.0:
+            raise ValueError("unsupported dots3-note audio rope config")
+
+    def get_audio_config(self) -> dict[str, Any] | None:
+        cfg = self.global_config.get("audio_config")
+        if cfg is not None:
+            # aliases so MmprojModel.find_aparam() / n_block_keys can resolve them
+            whisper = cfg["whisper_config"]
+            cfg["hidden_size"] = whisper["d_model"]
+            cfg["intermediate_size"] = whisper["encoder_ffn_dim"]
+            cfg["num_attention_heads"] = whisper["encoder_attention_heads"]
+            cfg["num_hidden_layers"] = whisper["encoder_layers"]
+        return cfg
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        assert self.hparams_vision is not None
+        assert self.hparams_audio is not None
+
+        self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.DOTS3NOTE_V)
+        self.gguf_writer.add_vision_use_silu(True)
+        self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams_vision["rms_norm_eps"])
+        self.gguf_writer.add_vision_spatial_merge_size(self.hparams_vision["spatial_merge_size"])
+        self.gguf_writer.add_vision_min_pixels(self.preprocessor_config["min_pixels"])
+        self.gguf_writer.add_vision_max_pixels(self.preprocessor_config["max_pixels"])
+        # pyramid MoE: per-block routed expert count, 0 = dense block
+        self.gguf_writer.add_vision_expert_count_per_layer(self.pyramid)
+        self.gguf_writer.add_vision_expert_used_count(int(self.hparams_vision["capacity_factor"]))
+
+        self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.DOTS3NOTE_A)
+        self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["whisper_config"]["num_mel_bins"])
+        self.gguf_writer.add_audio_attention_layernorm_eps(1e-6)  # Dots3NoteAudioRMSNorm default
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, _ = item
+        if not name.startswith(("vision_encoder.", "audio_encoder.")):
+            return None
+        return super().filter_tensors(item)
+
+    _vis_experts: dict[int, dict[str, Tensor]] | None = None
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # router params have no .weight suffix in the checkpoint, but gguf tools expect one
+        if name.endswith((".gate_weight", ".router_bias")):
+            name += ".weight"
+
+        # audio fc1 fuses gate and up for swiglu; split it
+        if ".speech_encoder.layers." in name and ".fc1." in name:
+            gate, up = data_torch.chunk(2, dim=0)
+            yield from super().modify_tensors(gate, name.replace(".fc1.", ".fc1_gate."), bid)
+            yield from super().modify_tensors(up, name.replace(".fc1.", ".fc1_up."), bid)
+            return
+
+        # vision MoE: stack per-expert weights into a single 3D tensor per block
+        if ".mlp.experts." in name:
+            assert bid is not None
+            n_expert = self.pyramid[bid]
+            if self._vis_experts is None:
+                self._vis_experts = {}
+            buf = self._vis_experts.setdefault(bid, {})
+            buf[name] = data_torch
+
+            if len(buf) >= n_expert * 3:
+                for w_name in ("fc1", "fc2", "fc3"):
+                    datas: list[Tensor] = []
+                    for xid in range(n_expert):
+                        ename = f"vision_encoder.blocks.{bid}.mlp.experts.{xid}.{w_name}.weight"
+                        datas.append(buf.pop(ename))
+                    merged = torch.stack(datas, dim=0)
+                    yield from super().modify_tensors(merged, f"vision_encoder.blocks.{bid}.mlp.experts.{w_name}.weight", bid)
+            return
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+    def prepare_tensors(self):
+        super().prepare_tensors()
+        if self._vis_experts is not None:
+            leftover = [k for d in self._vis_experts.values() for k in d.keys()]
+            if leftover:
+                raise ValueError(f"unprocessed vision experts: {leftover}")
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        # FP32 routing is load-bearing for the vision MoE (near-tied expert scores)
+        if ".ffn_gate_inp." in new_name or ".exp_probs_b." in new_name:
+            return gguf.GGMLQuantizationType.F32
+        if ".conv2d" in new_name or "a.conv_out" in new_name:
+            return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)

conversion/glm.py

@@ -112,12 +112,38 @@ class GlmOCRModel(Glm4Model):
 @ModelBase.example("zai-org/GLM-4.5-Air")
 class Glm4MoeModel(TextModel):
     model_arch = gguf.MODEL_ARCH.GLM4_MOE
+    supports_mtp_export = True
+    _n_main_layers: int | None = None
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        # GLM4_MOE has num_hidden_layers + 1 actual layers (including NextN layer)
-        self.block_count = self.hparams["num_hidden_layers"] + self.hparams.get("num_nextn_predict_layers", 0)
-        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
+        if not self.no_mtp:
+            self.block_count += self.hparams.get("num_nextn_predict_layers", 0)
+            self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
+
+    def index_tensors(self, remote_hf_model_id: str | None = None):
+        hparams = {**self.hparams, **self.hparams.get("text_config", {})}
+        key = next((k for k in ["n_layers", "num_hidden_layers", "n_layer", "num_layers"] if k in hparams), None)
+        type(self)._n_main_layers = hparams.get(key)
+        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        if (titem := super().filter_tensors(item)) is None:
+            return None
+        name, gen = titem
+
+        assert cls._n_main_layers is not None
+        is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers
+
+        if is_mtp and cls.no_mtp:
+            return None
+        if cls.mtp_only and not is_mtp and name not in (
+            "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
+        ):
+            return None
+
+        return name, gen
 
     def set_vocab(self):
         return self._set_vocab_glm()
@@ -153,10 +179,22 @@ class Glm4MoeModel(TextModel):
         if (norm_topk_prob := self.hparams.get("norm_topk_prob")) is not None:
             self.gguf_writer.add_expert_weights_norm(norm_topk_prob)
 
-        # NextN/MTP prediction layers
-        if (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None:
+        if not self.no_mtp and (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None:
             self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)
 
+    def prepare_metadata(self, vocab_only: bool):
+        from_dir = self.fname_out.is_dir()
+        super().prepare_metadata(vocab_only=vocab_only)
+
+        if not self.mtp_only or not from_dir:
+            return
+
+        output_type: str = self.ftype.name.partition("_")[2]
+        fname_default: str = gguf.naming_convention(
+            self.metadata.name, self.metadata.basename, self.metadata.finetune,
+            self.metadata.version, size_label=None, output_type=output_type, model_type=None)
+        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
+
     _experts: list[dict[str, Tensor]] | None = None
 
     # note: unlike GLM4V non-MoE, we don't need to permute Q/K here since GLM4V_MOE uses Neox ordering already
@@ -348,6 +386,7 @@ class GlmMoeDsaModel(DeepseekV2Model):
 @ModelBase.example("upstage/Solar-Open-100B")
 class SolarOpenModel(Glm4MoeModel):
     model_arch = gguf.MODEL_ARCH.GLM4_MOE
+    supports_mtp_export = False
 
     def set_vocab(self):
         from transformers import AutoTokenizer

conversion/nemotron.py

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

conversion/qwen.py

@@ -709,7 +709,13 @@ class DFlashModel(Qwen3Model):
         yield from super().modify_tensors(data_torch, name, bid)
 
 
-@ModelBase.register("Qwen3DSparkModel", "DSparkDraftModel", "DSparkSpeculator", "Lfm2DSparkDraftModel")
+@ModelBase.register(
+    "Qwen3DSparkModel",
+    "DSparkDraftModel",
+    "DSparkSpeculator",
+    "Lfm2DSparkDraftModel",
+    "LingDSparkModel",
+)
 @ModelBase.example("satgeze/Qwen3.6-27B-DSpark")
 class DSparkModel(DFlashModel):
     # DSpark = DFlash + a semi-autoregressive Markov head.

docs/autoparser.md

@@ -443,21 +443,21 @@ Each returned parser is wrapped by `wrap_for_generation_prompt()`, which prepend
 |                                           | `wrap_for_generation_prompt()`, string helpers                                  |
 | `common/chat-peg-parser.h/cpp`            | `common_chat_peg_builder`, `common_chat_peg_mapper`, and helpers                |
 | `common/chat.cpp`                         | Entry point: `common_chat_templates_apply_jinja()`                              |
-| `tools/parser/debug-template-parser.cpp`  | Debug tool for template analysis                                                |
-| `tools/parser/template-analysis.cpp`      | Template analysis tool                                                          |
+| `tests/test-chat-auto-parser.cpp`         | Auto-parser unit tests; also a debug tool when given a template path            |
+| `tests/test-chat-analysis.cpp`            | Template differential analysis debug tool                                       |
 
 ## Testing & Debugging
 
 ### Debug Tools
 
-**Template Debugger**: `tools/parser/debug-template-parser.cpp`
+**Template Debugger**: `tests/test-chat-auto-parser.cpp`
 
-- Usage: `./bin/llama-debug-template-parser path/to/template.jinja`
+- Usage: `./bin/test-chat-auto-parser path/to/template.jinja` (without a path, it runs the automated tests)
 - Shows detected format, markers, generated parser, and GBNF grammar
 
-**Template Analysis**: `tools/parser/template-analysis.cpp`
+**Template Analysis**: `tests/test-chat-analysis.cpp`
 
-- Usage: `./bin/llama-template-analysis path/to/template.jinja`
+- Usage: `./bin/test-chat-analysis --template-file path/to/template.jinja` (without arguments, it runs on all templates from the test suite)
 
 **Debug Logging**: Enable with `LLAMA_ARG_LOG_VERBOSITY=2`
 
@@ -519,7 +519,7 @@ The following templates have active tests in `tests/test-chat.cpp`:
 
 To support a new template format:
 
-1. **If it follows standard patterns** — The auto-parser should detect it automatically. Run `llama-debug-template-parser` to verify markers are correctly extracted.
+1. **If it follows standard patterns** — The auto-parser should detect it automatically. Run `test-chat-auto-parser <template_path>` to verify markers are correctly extracted.
 2. **If differential analysis extracts incorrect markers** — Add a workaround lambda to the `workarounds` vector in `common/chat-diff-analyzer.cpp`. Inspect the template source for a unique identifying substring.
 3. **If it needs fundamentally different handling** — Add a dedicated handler function in `chat.cpp` before the auto-parser block (as done for GPT-OSS, Functionary v3.2, and Ministral).
 

docs/backend/snapdragon/CMakeUserPresets.json

@@ -8,7 +8,7 @@
         "toolset":      { "value": "host=x86_64", "strategy": "external" },
         "cacheVariables": {
             "ANDROID_ABI":      "arm64-v8a",
-            "ANDROID_PLATFORM": "android-31",
+            "ANDROID_PLATFORM": "android-34",
             "CMAKE_TOOLCHAIN_FILE": "$env{ANDROID_NDK_ROOT}/build/cmake/android.toolchain.cmake",
             "CMAKE_C_FLAGS":   "-march=armv8.7a+fp16+dotprod+i8mm -fvectorize -ffp-model=fast -fno-finite-math-only -flto -D_GNU_SOURCE",
             "CMAKE_CXX_FLAGS": "-march=armv8.7a+fp16+dotprod+i8mm -fvectorize -ffp-model=fast -fno-finite-math-only -flto -D_GNU_SOURCE",

docs/backend/snapdragon/README.md

@@ -2,39 +2,47 @@
 
 ## Setup
 
-### Android
+The cross-compilation toolchain images are provided by the
+[Qualcomm Snapdragon Toolchain registry](https://github.com/snapdragon-toolchain).
+These Docker images include the Android NDK, OpenCL SDK, Hexagon SDK, CMake, and the necessary cross-compilers:
 
-The easiest way to build llama.cpp for a Snapdragon-based Android device is using the toolchain Docker image (see github.com/snapdragon-toolchain).
-This image includes Android NDK, OpenCL SDK, Hexagon SDK, CMake, etc.
+* **Android toolchain**: `ghcr.io/snapdragon-toolchain/arm64-android:v0.7`
+* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.7`
 
-This method works on Linux, macOS, and Windows. macOS and Windows users should install Docker Desktop.
-
-```
-~/src/llama.cpp$ docker run -it -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-android:v0.7
-[d]/> cd /workspace
-```
-
-Note: The rest of the **Android** build process assumes that you're running inside the toolchain container.
-
-### Windows On Snapdragon
-
-Native Windows 11 arm64 builds has the following tools dependencies:
-- MS Visual Studio 2026 (Community Edition or Pro)
-  - MSVC arm64 standard and runtime libraries
-  - UCRT and Driver Kit
-- LLVM core libraries and Clang compiler (winget)
-- CMake, Git, Python (winget)
-- Hexagon SDK Community Edition 6.6 or later (see windows.md)
-- OpenCL SDK 2.3 or later (see windows.md)
-
-Note: The rest of the **Windows** build process assumes that you're running natively in Powershell.
-Adapt below build commands accordingly.
+The unified build utility (`scripts/snapdragon/build.py`) automatically pulls
+and orchestrates these containers to perform target compilation.
+You only need to ensure that Docker (or Docker Desktop on macOS/Windows) is running on your host machine.
+Specific setup, build, and installation details for Linux and Windows on Snapdragon platforms are documented in:
+* [Linux on Snapdragon guide](linux.md)
+* [Windows on Snapdragon guide](windows.md)
 
 ## How to Build
 
-Let's build llama.cpp with CPU, OpenCL, and Hexagon backends via CMake presets:
+### Using build.py script (Recommended)
 
+The easiest way to build llama.cpp is by using the `scripts/snapdragon/build.py` script. It automatically copies the CMake presets,
+launches the correct compilation Docker container, builds the libraries and tools,
+installs them, and optionally pushes them to your ADB device.
+
+Build and deploy for Android target (accepts `android` or `adb` alias):
 ```
+$ ./scripts/snapdragon/build.py --target adb --push
+```
+
+Build and deploy for Linux target (accepts `linux` or `lnx` alias):
+```
+$ ./scripts/snapdragon/build.py --target linux:user@host --push
+```
+
+### Manual CMake Build
+
+Alternatively, you can build llama.cpp manually by entering the cross-compilation Docker container and running the CMake commands:
+
+```bash
+# Start the cross-compilation container manually:
+~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-android:v0.7
+
+# Inside the container, build the project using presets:
 [d]/workspace> cp docs/backend/snapdragon/CMakeUserPresets.json .
 
 [d]/workspace> cmake --preset arm64-android-snapdragon-release -B build-snapdragon
@@ -68,19 +76,19 @@ Preset CMake variables:
 To generate an installable "package" simply use cmake --install:
 
 ```
-[d]/workspace> cmake --install build-snapdragon --prefix pkg-snapdragon/llama.cpp
+[d]/workspace> cmake --install build-snapdragon --prefix pkg-android/llama.cpp
 -- Install configuration: "Release"
--- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-cpu.so
--- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-opencl.so
--- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-hexagon.so
--- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v73.so
--- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v75.so
--- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v79.so
--- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v81.so
--- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-cpu.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-opencl.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-hexagon.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v73.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v75.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v79.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v81.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml.so
 ...
--- Installing: /workspace/pkg-snapdragon/llama.cpp/bin/llama-bench
--- Installing: /workspace/pkg-snapdragon/llama.cpp/bin/llama-cli
+-- Installing: /workspace/pkg-android/llama.cpp/bin/llama-bench
+-- Installing: /workspace/pkg-android/llama.cpp/bin/llama-cli
 ...
 ```
 
@@ -91,14 +99,14 @@ To generate an installable "package" simply use cmake --install:
 For this step, your device needs to be configured for on-device development.
 Please see https://developer.android.com/studio/debug/dev-options for details.
 
-Once ADB is enabled, use `adb push` to install `pkg-snapdragon` on the device.
+Once ADB is enabled, use `adb push` to install `pkg-android` on the device.
 **Note that the toolchain Docker image doesn't have ADB and doesn't set up the ADB bridge. Please use native ADB on the host.**
 
 ```
-~/src/llama.cpp$ adb push pkg-snapdragon/llama.cpp /data/local/tmp/
-pkg-snapdragon/llama.cpp/bin/: 67 files pushed, 0 skipped. 190.2 MB/s (919095042 bytes in 4.607s)
-pkg-snapdragon/llama.cpp/include/: 19 files pushed, 0 skipped. 20.5 MB/s (255173 bytes in 0.012s)
-pkg-snapdragon/llama.cpp/lib/: 16 files pushed, 0 skipped. 144.4 MB/s (43801382 bytes in 0.289s)
+~/src/llama.cpp$ adb push pkg-android/llama.cpp /data/local/tmp/
+pkg-android/llama.cpp/bin/: 67 files pushed, 0 skipped. 190.2 MB/s (919095042 bytes in 4.607s)
+pkg-android/llama.cpp/include/: 19 files pushed, 0 skipped. 20.5 MB/s (255173 bytes in 0.012s)
+pkg-android/llama.cpp/lib/: 16 files pushed, 0 skipped. 144.4 MB/s (43801382 bytes in 0.289s)
 102 files pushed, 0 skipped. 186.9 MB/s (963151597 bytes in 4.914s)
 ```
 
@@ -115,24 +123,44 @@ Llama-3.2-1B-Instruct-Q4_0.gguf: 1 file pushed, 0 skipped. 38.3 MB/s (773025920
 
 ### Windows
 
-All artifacts are already installed in the `pkg-snapdragon` folder.
-To run, adapt below instructions to use Powershell scripts in `scripts/snapdragon/windows`.
+All artifacts are already installed in the `pkg-wos` folder.
+To run, you can use the `scripts/snapdragon/run.py` runner script (see details below).
 
 ## How to Run
 
-The easiest way to run llama.cpp cli tools is using provided wrapper scripts that properly set up all required environment variables.
+The easiest way to run llama.cpp cli tools is using the provided `scripts/snapdragon/run.py` wrapper script. This script automatically
+maps CLI options to environment variables, resolves executable paths, and runs the command locally, via ADB, or remotely via SSH on the
+target device.
 
-llama.cpp supports three backends on Snapdragon-based devices: CPU, Adreno GPU (GPUOpenCL), and Hexagon NPU (HTP0-4).
-You can select which backend to run the model on using the `D=` variable, which maps to the `--device` option.
+llama.cpp supports three backends on Snapdragon-based devices: CPU, Adreno GPU (GPUOpenCL), and Hexagon NPU.
+You can select which backend(s) to run the model on using the `--device` option of the tool (or `--devices` option in `run.py`).
 
 Hexagon NPU behaves as a "GPU" device when it comes to `-ngl` and other offload-related options.
 
-Here are some examples of running various llama.cpp tools via ADB.
+Here are some examples of running various llama.cpp tools.
 
-Simple question for Llama-3.2-1B
+Generating a completion with Gemma on Android (relying on default `HTP0:0` device and default thread count `-t 6`):
 
 ```
-~/src/llama.cpp$ M=Llama-3.2-1B-Instruct-Q4_0.gguf D=HTP0 ./scripts/snapdragon/adb/run-completion.sh -p "what is the most popular cookie in the world?"
+~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb -- llama-completion -m models/gemma-2-2b-it-Q4_0.gguf -f prompts/sample_prompt_1024.txt --jinja -st
+...
+ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1
+ggml-hex: Hexagon Arch version v79
+ggml-hex: allocating new session: HTP0:0
+...
+load_tensors: offloading output layer to GPU
+load_tensors: offloaded 27/27 layers to GPU
+load_tensors:          CPU model buffer size =   300.00 MiB
+load_tensors:      HTP0:0 model buffer size  =  1400.26 MiB
+...
+llama_perf_context_print: prompt eval time =     320.00 ms /  1024 tokens (    0.31 ms per token,  3200.00 tokens per second)
+llama_perf_context_print:        eval time =     2100.00 ms /   100 runs   (   21.00 ms per token,    47.62 tokens per second)
+```
+
+Simple question for Llama-3.2-1B:
+
+```
+~/src/llama.cpp$ ./scripts/snapdragon/run.py --target android --devices HTP0 -- llama-cli -m Llama-3.2-1B-Instruct-Q4_0.gguf -p "what is the most popular cookie in the world?"
 ...
 ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1
 ggml-hex: Hexagon Arch version v79
@@ -142,8 +170,7 @@ ggml-hex: new session: HTP0 : session-id 0 domain-id 3 uri file:///libggml-htp-v
 load_tensors: offloading output layer to GPU
 load_tensors: offloaded 17/17 layers to GPU
 load_tensors:          CPU model buffer size =   225.49 MiB
-load_tensors:         HTP0 model buffer size =     0.26 MiB
-load_tensors:  HTP0-REPACK model buffer size =   504.00 MiB
+load_tensors:         HTP0 model buffer size =   504.26 MiB
 ...
 I hope this helps you understand the world's most popular cookies! [end of text]
 ...
@@ -156,60 +183,25 @@ llama_perf_context_print:    graphs reused =        473
 llama_memory_breakdown_print: | memory breakdown [MiB] | total   free    self   model   context   compute    unaccounted |
 llama_memory_breakdown_print: |   - HTP0 (Hexagon)     |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
 llama_memory_breakdown_print: |   - Host               |                  439 =   225 +     136 +      77                |
-llama_memory_breakdown_print: |   - HTP0-REPACK        |                  504 =   504 +       0 +       0                |
 ```
 
-Summary request for OLMoE-1B-7B. This is a large model that requires two HTP sessions/devices
+Op test for MUL_MAT:
 
 ```
-~/src/llama.cpp$ M=OLMoE-1B-7B-0125-Instruct-Q4_0.gguf NDEV=2 D=HTP0,HTP1 ./scripts/snapdragon/adb/run-completion.sh -f surfing.txt
+~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --hex-hostbuf 0 --devices HTP0:0 -- test-backend-ops -b HTP0:0 -o MUL_MAT
 ...
-ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1
-ggml-hex: Hexagon Arch version v81
-ggml-hex: allocating new session: HTP0
-ggml-hex: allocating new session: HTP1
-...
-load_tensors: offloading output layer to GPU
-load_tensors: offloaded 17/17 layers to GPU
-load_tensors:          CPU model buffer size =   143.86 MiB
-load_tensors:         HTP1 model buffer size =     0.23 MiB
-load_tensors:  HTP1-REPACK model buffer size =  1575.00 MiB
-load_tensors:         HTP0 model buffer size =     0.28 MiB
-load_tensors:  HTP0-REPACK model buffer size =  2025.00 MiB
-...
-llama_context:        CPU  output buffer size =     0.19 MiB
-llama_kv_cache:       HTP1 KV buffer size =   238.00 MiB
-llama_kv_cache:       HTP0 KV buffer size =   306.00 MiB
-llama_kv_cache: size =  544.00 MiB (  8192 cells,  16 layers,  1/1 seqs), K (q8_0):  272.00 MiB, V (q8_0):  272.00 MiB
-llama_context:       HTP0 compute buffer size =    15.00 MiB
-llama_context:       HTP1 compute buffer size =    15.00 MiB
-llama_context:        CPU compute buffer size =    24.56 MiB
-...
-llama_perf_context_print: prompt eval time =    1730.57 ms /   212 tokens (    8.16 ms per token,   122.50 tokens per second)
-llama_perf_context_print:        eval time =    5624.75 ms /   257 runs   (   21.89 ms per token,    45.69 tokens per second)
-llama_perf_context_print:       total time =    7377.33 ms /   469 tokens
-llama_perf_context_print:    graphs reused =        255
-llama_memory_breakdown_print: | memory breakdown [MiB] | total   free    self   model   context   compute    unaccounted |
-llama_memory_breakdown_print: |   - HTP0 (Hexagon)     |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - HTP1 (Hexagon)     |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - Host               |                  742 =   144 +     544 +      54                |
-llama_memory_breakdown_print: |   - HTP1-REPACK        |                 1575 =  1575 +       0 +       0                |
-llama_memory_breakdown_print: |   - HTP0-REPACK        |                 2025 =  2025 +       0 +       0                |
-```
-
-Op test for MUL_MAT
-
-```
-~/src/llama.cpp$ HB=0 ./scripts/snapdragon/adb/run-tool.sh test-backend-ops -b HTP0 -o MUL_MAT
-...
-Backend 2/3: HTP0
+Backend 2/3: HTP0:0
 Device description: Hexagon
 Device memory: 2048 MB (2048 MB free)
 MUL_MAT(type_a=q4_0,type_b=f32,m=16,n=1,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
 MUL_MAT(type_a=q4_0,type_b=f32,m=16,n=2,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
 MUL_MAT(type_a=q4_0,type_b=f32,m=16,n=3,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
+```
 
-~/src/llama.cpp-hexagon$ M=Llama-3.2-1B-Instruct-Q4_0.gguf ./scripts/snapdragon/adb/run-bench.sh -p 128 -n 64
+Llama benchmark:
+
+```
+~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --devices HTP0 -- llama-bench -p 128 -n 64 -m Llama-3.2-1B-Instruct-Q4_0.gguf
 ...
 ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1
 ggml-hex: Hexagon Arch version v79
@@ -219,15 +211,20 @@ ggml-hex: new session: HTP0 : session-id 0 domain-id 3 uri file:///libggml-htp-v
 | ---------------| ---------: | -----: | ---------- | --: | ------: | ------: | ---: | ----: | ------------: |
 | llama 1B Q4_0  | 729.75 MiB | 1.24 B | HTP        |  99 |       4 |     128 |    0 | pp128 | 169.42 ± 1.75 |
 | llama 1B Q4_0  | 729.75 MiB | 1.24 B | HTP        |  99 |       4 |     128 |    0 |  tg64 |  51.54 ± 1.13 |
-
-build: 6a8cf8914 (6733)
 ```
 
 ## Environment variables
 
-- `GGML_HEXAGON_NDEV=1`
-  Controls the number of devices/sessions to allocate. The default is 1.
-  Most quantized models under 4B fit into a single session; an 8B model needs two, and a 20B model needs four.
+- `GGML_HEXAGON_DEVICES` (default: not set, defaults to HTP0 session)
+  Controls which NPU devices and sessions to allocate. Can be configured as:
+  - A single integer `N`: Allocates `N` sessions named `HTP0`, `HTP1`, ..., `HTP<N-1>` (behaves identically to `GGML_HEXAGON_NDEV=N`).
+  - A comma-separated list of device names in `HTP<physical_idx>:<virtual_idx>` format (or legacy `HTP<idx>` format). For example, `HTP0:0,HTP0:1` creates two virtual
+    sessions on the first physical NPU (useful for memory limits). `HTP0:0,HTP1:0` allocates one session on each of the two physical NPUs
+    on a dual-NPU device.
+
+- `GGML_HEXAGON_NDEV` (deprecated)
+  Replaced by `GGML_HEXAGON_DEVICES`. Controls the number of virtual sessions to allocate on physical NPU `0`.
+  Allocates sessions named `HTP0`, `HTP1`, etc.
 
 - `GGML_HEXAGON_NHVX=0`
   Controls the number of HVX hardware threads to use. The default is all (actual number varies depending on the hardware version).
@@ -255,26 +252,17 @@ build: 6a8cf8914 (6733)
   - `2` Extended profile with per-op `usecs`, `cycles` and default PMU counter data
   - `0x1,...,0x8` Extended profile with per-op `usecs`, `cycles` and custom PMU counter data
 
-  The logging output can be either saved into a file for post-processing or it can be piped directly into the post-processing tool to generate the report.
+  The logging output can be either saved into a file for post-processing or it can be piped directly into the post-processing tool
+  to generate the report.
   Examples:
 
-      `GGML_HEXAGON_PROFILE=1 llama-completion ... |& ./scripts/snapdragon/ggml-hexagon-profile.py -`
-
-- `GGML_HEXAGON_OPSTAGE=0x0`
-  Allows enabling specific stages of the Op processing pipeline:
-
-  - `0x1` Enable Op Queue (i.e., queuing Ops into NPU)
-  - `0x2` Enable Op Compute (MUL_MAT, etc.)
-
-  Examples:
-
-      `GGML_HEXAGON_OPSTAGE=0x1 llama-completion ...` - Ops are enqueued to the NPU but dma & compute are disabled
-      `GGML_HEXAGON_OPSTAGE=0x3 llama-completion ...` - Full queuing and processing of Ops (default)
+      `GGML_HEXAGON_PROFILE=1 ./scripts/snapdragon/run.py --target adb -- llama-cli ... |& ./scripts/snapdragon/ggml-hexagon-profile.py -`
 
 - `GGML_HEXAGON_OPFILTER=regex`
   Allows filtering (disabling) Ops that match the regex pattern:
 
   Examples:
 
-      `GGML_HEXAGON_OPFILTER="FLASH_ATTN_EXT" llama-completion ...` - Disable Flash Attention on Hexagon (falls back to CPU or GPU)
-      `GGML_HEXAGON_OPFILTER="ADD\|SUB" llama-completion ...` - Disable ADD and SUB on Hexagon (fall back to CPU or GPU)
+      `GGML_HEXAGON_OPFILTER="FLASH_ATTN_EXT" ./scripts/snapdragon/run.py --target adb -- llama-cli ...` - Disable Flash Attention on Hexagon (falls back to CPU or GPU)
+      `GGML_HEXAGON_OPFILTER="ADD\|SUB" ./scripts/snapdragon/run.py --target adb -- llama-cli ...` - Disable ADD and SUB on Hexagon (fall back to CPU or GPU)
+

docs/backend/snapdragon/developer.md

@@ -39,22 +39,21 @@ the repacking.
 
 ## Large model handling
 
-Hexagon NPU session (aka Process Domain (PD) in the Hexagon docs) is limited to a memory mapping of around 3.5GB.
-In llama.cpp/GGML the Hexagon session is mapped to a single GGML backend device (HTP0, HTP1, etc).
+Hexagon NPU sessions (aka Process Domains (PD) in the Hexagon SDK) are limited to a maximum memory mapping window of around 3.5GB.
+In llama.cpp/GGML, each Hexagon session is mapped to a single GGML backend device (e.g., `HTP0:0`, `HTP0:1`, etc. when using
+`GGML_HEXAGON_DEVICES`, or `HTP0`, `HTP1` in legacy mode).
 
-In order to map models larger than 3.5GB we need to allocate multiple devices and split the model.
-For this we're taking advantage of the llama.cpp/GGML multi-GPU layer-splitting support.
-Each Hexagon device behaves like a GPU from the offload and model splitting perspective.
+To support running models larger than 3.5GB on a single device, the Hexagon backend dynamically maps and unmaps execution buffers
+during the graph execution cycle to stay within the Process Domain window. This enables large models to run successfully on a single
+NPU device.
 
-Here is an example of running GPT-OSS-20B model on a newer Snapdragon device with 16GB of DDR.
+Alternatively, users can choose to use standard llama.cpp/GGML layer-splitting mode to partition and split the model across
+multiple Hexagon devices or virtual sessions (which behave like multiple GPUs from the offload and splitting perspective).
+
+Here is an example of running GPT-OSS-20B model on a Snapdragon device using 4 virtual sessions on a single NPU (physical index 0).
 
 ```
-M=gpt-oss-20b-Q4_0.gguf NDEV=4 D=HTP0,HTP1,HTP2,HTP3 P=surfing.txt scripts/snapdragon/adb/run-completion.sh -f surfing.txt -n 32
-...
-LD_LIBRARY_PATH=/data/local/tmp/llama.cpp/lib
-ADSP_LIBRARY_PATH=/data/local/tmp/llama.cpp/lib
-GGML_HEXAGON_NDEV=4 ./bin/llama-cli --load-mode none -m /data/local/tmp/llama.cpp/../gguf/gpt-oss-20b-Q4_0.gguf
-      -t 4 --ctx-size 8192 --batch-size 128 -ctk q8_0 -ctv q8_0 -fa on -ngl 99 --device HTP0,HTP1,HTP2,HTP3 -no-cnv -f surfing.txt
+~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --devices HTP0:0,HTP0:1,HTP0:2,HTP0:3 -- llama-cli --load-mode none -m /data/local/tmp/gguf/gpt-oss-20b-Q4_0.gguf -t 4 --ctx-size 8192 --batch-size 128 -ctk q8_0 -ctv q8_0 -fa on -ngl 99 -no-cnv -f surfing.txt
 ...
 llama_model_loader: - type  f32:  289 tensors
 llama_model_loader: - type q4_0:   96 tensors
@@ -63,33 +62,29 @@ llama_model_loader: - type mxfp4:  72 tensors
 ...
 load_tensors: offloaded 25/25 layers to GPU
 load_tensors:          CPU model buffer size =  1182.09 MiB
-load_tensors:         HTP1 model buffer size =     6.64 MiB
-load_tensors:  HTP1-REPACK model buffer size =  2505.94 MiB
-load_tensors:         HTP3 model buffer size =     5.55 MiB
-load_tensors:  HTP3-REPACK model buffer size =  2088.28 MiB
-load_tensors:         HTP0 model buffer size =     7.75 MiB
-load_tensors:  HTP0-REPACK model buffer size =  2923.59 MiB
-load_tensors:         HTP2 model buffer size =     6.64 MiB
-load_tensors:  HTP2-REPACK model buffer size =  2505.94 MiB
+load_tensors:       HTP0:1 model buffer size =  2512.58 MiB
+load_tensors:       HTP0:3 model buffer size =  2093.83 MiB
+load_tensors:       HTP0:0 model buffer size =  2931.34 MiB
+load_tensors:       HTP0:2 model buffer size =  2512.58 MiB
 ...
 llama_context: n_ctx_per_seq (8192) < n_ctx_train (131072) -- the full capacity of the model will not be utilized
 llama_context:        CPU  output buffer size =     0.77 MiB
 llama_kv_cache_iswa: creating non-SWA KV cache, size = 8192 cells
-llama_kv_cache:       HTP1 KV buffer size =    25.50 MiB
-llama_kv_cache:       HTP3 KV buffer size =    25.50 MiB
-llama_kv_cache:       HTP0 KV buffer size =    25.50 MiB
-llama_kv_cache:       HTP2 KV buffer size =    25.50 MiB
+llama_kv_cache:     HTP0:1 KV buffer size =    25.50 MiB
+llama_kv_cache:     HTP0:3 KV buffer size =    25.50 MiB
+llama_kv_cache:     HTP0:0 KV buffer size =    25.50 MiB
+llama_kv_cache:     HTP0:2 KV buffer size =    25.50 MiB
 llama_kv_cache: size =  102.00 MiB (  8192 cells,  12 layers,  1/1 seqs), K (q8_0):   51.00 MiB, V (q8_0):   51.00 MiB
 llama_kv_cache_iswa: creating     SWA KV cache, size = 256 cells
-llama_kv_cache:       HTP1 KV buffer size =     0.80 MiB
-llama_kv_cache:       HTP3 KV buffer size =     0.53 MiB
-llama_kv_cache:       HTP0 KV buffer size =     1.06 MiB
-llama_kv_cache:       HTP2 KV buffer size =     0.80 MiB
+llama_kv_cache:     HTP0:1 KV buffer size =     0.80 MiB
+llama_kv_cache:     HTP0:3 KV buffer size =     0.53 MiB
+llama_kv_cache:     HTP0:0 KV buffer size =     1.06 MiB
+llama_kv_cache:     HTP0:2 KV buffer size =     0.80 MiB
 llama_kv_cache: size =    3.19 MiB (   256 cells,  12 layers,  1/1 seqs), K (q8_0):    1.59 MiB, V (q8_0):    1.59 MiB
-llama_context:       HTP0 compute buffer size =    16.06 MiB
-llama_context:       HTP1 compute buffer size =    16.06 MiB
-llama_context:       HTP2 compute buffer size =    16.06 MiB
-llama_context:       HTP3 compute buffer size =    16.06 MiB
+llama_context:     HTP0:0 compute buffer size =    16.06 MiB
+llama_context:     HTP0:1 compute buffer size =    16.06 MiB
+llama_context:     HTP0:2 compute buffer size =    16.06 MiB
+llama_context:     HTP0:3 compute buffer size =    16.06 MiB
 llama_context:        CPU compute buffer size =    98.19 MiB
 ...
 llama_perf_context_print: prompt eval time =    3843.67 ms /   197 tokens ( 19.51 ms per token, 51.25 tokens per second)
@@ -97,13 +92,9 @@ llama_perf_context_print:        eval time =    1686.13 ms /    31 runs   ( 54.3
 llama_perf_context_print:       total time =    6266.30 ms /   228 tokens
 llama_perf_context_print:    graphs reused =         30
 llama_memory_breakdown_print: | memory breakdown [MiB] | total   free    self   model   context   compute    unaccounted |
-llama_memory_breakdown_print: |   - HTP0 (Hexagon)     |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - HTP1 (Hexagon)     |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - HTP2 (Hexagon)     |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - HTP3 (Hexagon)     |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - HTP0:0 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - HTP0:1 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - HTP0:2 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - HTP0:3 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
 llama_memory_breakdown_print: |   - Host               |                 1476 =  1208 +     105 +     162                |
-llama_memory_breakdown_print: |   - HTP1-REPACK        |                 2505 =  2505 +       0 +       0                |
-llama_memory_breakdown_print: |   - HTP3-REPACK        |                 2088 =  2088 +       0 +       0                |
-llama_memory_breakdown_print: |   - HTP0-REPACK        |                 2923 =  2923 +       0 +       0                |
-llama_memory_breakdown_print: |   - HTP2-REPACK        |                 2505 =  2505 +       0 +       0                |
 ```

docs/backend/snapdragon/linux.md

@@ -1,25 +1,37 @@
 # Snapdragon-based Linux devices
 
-## Docker Setup
+The cross-compilation is performed using the Snapdragon Linux Docker toolchain image (see
+[github.com/snapdragon-toolchain](https://github.com/snapdragon-toolchain)):
 
-The easiest way to build llama.cpp for a Snapdragon-based Linux device is using the toolchain Docker image (see [github.com/snapdragon-toolchain](https://github.com/snapdragon-toolchain)).
-This image includes OpenCL SDK, Hexagon SDK, CMake, and the ARM64 Linux cross-compilation toolchain.
+* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.7`
 
-Cross-compilation is supported on **Linux X86** hosts. The resulting binaries are deployed to and run on the target **Qualcomm Snapdragon ARM64 Linux** device.
-
-```
-~/src/llama.cpp$ docker run -it -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-linux:v0.1
-[d]/> cd /workspace
-```
-
-Note: The rest of the **Linux** build process assumes that you're running inside the toolchain container.
+The unified build utility (`scripts/snapdragon/build.py`) automatically pulls
+and orchestrates this container to perform target compilation. You only need to
+ensure that Docker is running on your host machine.
 
 
 ## How to Build
 
-Let's build llama.cpp with CPU, OpenCL, and Hexagon backends via CMake presets:
+### Using build.py script (Recommended)
 
+The easiest way to build llama.cpp is by using the `scripts/snapdragon/build.py` script. It automatically copies the CMake presets,
+launches the correct compilation Docker container, builds the libraries and tools,
+installs them, and optionally pushes them to your target device.
+
+Build and deploy for a Linux target (using SSH deployment alias `lnx` or `linux`):
 ```
+$ ./scripts/snapdragon/build.py --target lnx:user@host --push
+```
+
+### Manual CMake Build
+
+Alternatively, you can build llama.cpp manually by entering the cross-compilation Docker container and running the CMake commands:
+
+```bash
+# Start the cross-compilation container manually:
+~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-linux:v0.7
+
+# Inside the container, build the project using presets:
 [d]/workspace> cp docs/backend/snapdragon/CMakeUserPresets.json .
 
 [d]/workspace> cmake --preset arm64-linux-snapdragon-release -B build-snapdragon
@@ -30,17 +42,19 @@ Let's build llama.cpp with CPU, OpenCL, and Hexagon backends via CMake presets:
 To generate an installable "package" simply use cmake --install, then zip it:
 
 ```
-[d]/workspace> cmake --install build-snapdragon --prefix pkg-snapdragon
-[d]/workspace> zip -r pkg-snapdragon.zip pkg-snapdragon
+[d]/workspace> cmake --install build-snapdragon --prefix pkg-linux
+[d]/workspace> zip -r pkg-linux.zip pkg-linux
 ```
 
 ## How to Install
 
-For this step, you will deploy the built binaries and libraries to the target Linux device. Transfer `pkg-snapdragon.zip` to the target device, then unzip it and set up the environment variables:
+For this step, you will deploy the built binaries and libraries to the target
+Linux device. Transfer `pkg-linux.zip` to the target device, then unzip it
+and set up the environment variables:
 
 ```
-$ unzip pkg-snapdragon.zip
-$ cd pkg-snapdragon
+$ unzip pkg-linux.zip
+$ cd pkg-linux
 $ export LD_LIBRARY_PATH=./lib
 $ export ADSP_LIBRARY_PATH=./lib
 ```
@@ -52,7 +66,28 @@ $ wget https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/
 ```
 
 ## How to Run
-Next, since we have setup the environment variables, we can run the llama-cli with the Hexagon backends:
+You can run locally on the Snapdragon Linux device:
+```
+$ ./scripts/snapdragon/run.py --devices HTP0 -- llama-cli -m Llama-3.2-3B-Instruct-Q4_0.gguf -ngl 99 -p "what is the most popular cookie in the world?"
+```
+
+Or run remotely from your host development machine using the SSH target option:
+```
+$ ./scripts/snapdragon/run.py --target lnx:user@host --devices HTP0 -- llama-cli -m Llama-3.2-3B-Instruct-Q4_0.gguf -ngl 99 -p "what is the most popular cookie in the world?"
+```
+
+For multi-NPU systems, you can run a tensor split completion command targeting a remote Linux system:
+```
+$ ./scripts/snapdragon/run.py --target ubuntu:maxk@192.168.1.87 --device HTP0:0,HTP1:0 -- llama-completion -m models/gemma-2b-it-Q4_0.gguf -f prompts/sample_prompt_1024.txt --jinja -st --split-mode tensor --ctx-size 8192
+```
+
+This translates to the following command being executed remotely via SSH:
+```
++ ssh maxk@192.168.1.87 "cd ~/llama.cpp && ulimit -c unlimited && LD_LIBRARY_PATH=./lib ADSP_LIBRARY_PATH=./lib GGML_HEXAGON_DEVICES=HTP0:0,HTP1:0 GGML_HEXAGON_OPPOLL=1 ./bin/llama-completion -m models/gemma-2b-it-Q4_0.gguf -f prompts/sample_prompt_1024.txt --jinja -st --split-mode tensor --ctx-size 8192 -v -n 16 --device HTP0:0,HTP1:0 -ngl 99 --ubatch-size 1024 -fa on -t 6"
+```
+
+Alternatively, you can run the binary directly on the device:
 ```
 $ ./bin/llama-cli -m Llama-3.2-3B-Instruct-Q4_0.gguf --device HTP0 -ngl 99 -p "what is the most popular cookie in the world?"
 ```
+

docs/backend/snapdragon/windows.md

@@ -1,3 +1,18 @@
+# Snapdragon-based Windows devices
+
+## Tool Dependencies
+
+Native Windows 11 arm64 builds have the following tool dependencies:
+- MS Visual Studio 2026 (Community Edition or Pro)
+  - MSVC arm64 standard and runtime libraries
+  - UCRT and Driver Kit
+- LLVM core libraries and Clang compiler (winget)
+- CMake, Git, Python (winget)
+- Hexagon SDK Community Edition 6.6 or later (see below)
+- OpenCL SDK 2.3 or later (see below)
+
+Note: The rest of the **Windows** build process assumes that you're running natively in Powershell.
+
 ## Overview
 
 The document covers procedures for installing the latest GPU and NPU drivers, and OpenCL and Hexagon SDKs.
@@ -53,7 +68,8 @@ Download the driver from
 
     https://softwarecenter.qualcomm.com/catalog/item/Qualcomm_HND
 
-After the automated installation and reboot please make sure that the Hexagon NPU device shows up in the `Device Manager` (under `Neural Processors`).
+After the automated installation and reboot please make sure that the Hexagon NPU device shows up in the `Device Manager`
+(under `Neural Processors`).
 
 If the device is not available you can try installing all components (`qcnspmcdm8380`, `qcnspmcdm8380_ext`) manually.
 The components are extracted into
@@ -130,12 +146,12 @@ However, additional settings are required for generating and signing HTP Ops lib
 
 > cmake --preset arm64-windows-snapdragon-release -B build-wos
 ...
-> cmake --install build-wos --prefix pkg-snapdragon
+> cmake --install build-wos --prefix pkg-wos
 ```
 
 Once the build is complete HTP ops libraries will be installed like this
 ```
-> dir pkg-snapdragon/lib
+> dir pkg-wos/lib
 ...
 -a----         1/22/2026   6:01 PM         187656 libggml-htp-v73.so
 -a----         1/22/2026   6:01 PM         191752 libggml-htp-v75.so
@@ -147,8 +163,8 @@ Once the build is complete HTP ops libraries will be installed like this
 The .cat file, the signature and proper certificate installation can be verified with
 
 ```
-> signtool.exe verify /v /pa .\pkg-snapdragon\lib\libggml-htp.cat
-Verifying: .\pkg-snapdragon\lib\libggml-htp.cat
+> signtool.exe verify /v /pa .\pkg-wos\lib\libggml-htp.cat
+Verifying: .\pkg-wos\lib\libggml-htp.cat
 
 Signature Index: 0 (Primary Signature)
 Hash of file (sha256): 9820C664DA59D5EAE31DBB664127FCDAEF59CDC31502496BC567544EC2F401CF
@@ -156,6 +172,6 @@ Hash of file (sha256): 9820C664DA59D5EAE31DBB664127FCDAEF59CDC31502496BC567544EC
 Signing Certificate Chain:
         Issued to: GGML.HTP.v1
 ...
-Successfully verified: .\pkg-snapdragon\lib\libggml-htp.cat
+Successfully verified: .\pkg-wos\lib\libggml-htp.cat
 ...
 ```

docs/build.md

@@ -70,18 +70,23 @@ cmake --build build --config Release
     - Tab Workload: Desktop-development with C++
     - Tab Components (select quickly via search): C++-_CMake_ Tools for Windows, _Git_ for Windows, C++-_Clang_ Compiler for Windows, MS-Build Support for LLVM-Toolset (clang)
     - Please remember to always use a Developer Command Prompt / PowerShell for VS2022 for git, build, test
-    - For Windows on ARM (arm64, WoA) build with:
-    ```bash
-    cmake --preset arm64-windows-llvm-release -D GGML_OPENMP_FETCH=ON
-    cmake --build build-arm64-windows-llvm-release
-    ```
-    `GGML_OPENMP_FETCH` downloads the official LLVM OpenMP runtime and requires Clang, 7-Zip and network access during configuration. CMake selects the runtime from the target architecture, so this also works when cross-compiling for WoA from x64. The extracted header, import library, DLL and OpenMP license are placed under `build/_deps`. The build copies `libomp.dll` and `LICENSE-LLVM-OpenMP` to the runtime output directory and installs them together. Omit the option to use CMake's normal OpenMP detection, or pass `-D GGML_OPENMP=OFF` to disable OpenMP.
-    For building with ninja generator and clang compiler as default:
-      -set path:set LIB=C:\Program Files (x86)\Windows Kits\10\Lib\10.0.22621.0\um\x64;C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Tools\MSVC\14.41.34120\lib\x64\uwp;C:\Program Files (x86)\Windows Kits\10\Lib\10.0.22621.0\ucrt\x64
+    - For Windows on ARM (arm64, WoA), build with:
       ```bash
-      cmake --preset x64-windows-llvm-release
-      cmake --build build-x64-windows-llvm-release
+      cmake --preset arm64-windows-llvm-release -D GGML_OPENMP_FETCH=ON
+      cmake --build build-arm64-windows-llvm-release
       ```
+      - Use `ARM64 Native Tools Command Prompt for VS 2022` if you are building on an ARM64 machine.
+      - `GGML_OPENMP_FETCH` downloads the official LLVM OpenMP runtime and requires Clang, 7-Zip and network access during configuration. CMake selects the runtime from the target architecture, so this also works when cross-compiling for WoA from x64. The extracted header, import library, DLL and OpenMP license are placed under `build/_deps`. The build copies `libomp.dll` and `LICENSE-LLVM-OpenMP` to the runtime output directory and installs them together. Omit the option to use CMake's normal OpenMP detection, or pass `-D GGML_OPENMP=OFF` to disable OpenMP.
+    - For building with ninja generator and clang compiler as default:
+      - Set path:
+        ```
+        set LIB=C:\Program Files (x86)\Windows Kits\10\Lib\10.0.22621.0\um\x64;C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Tools\MSVC\14.41.34120\lib\x64\uwp;C:\Program Files (x86)\Windows Kits\10\Lib\10.0.22621.0\ucrt\x64
+        ```
+      - Run:
+        ```bash
+        cmake --preset x64-windows-llvm-release
+        cmake --build build-x64-windows-llvm-release
+        ```
 - If you want HTTPS/TLS features, you may install OpenSSL development libraries. If not installed, the project will build and run without SSL support.
   - **Debian / Ubuntu:** `sudo apt-get install libssl-dev`
   - **Fedora / RHEL / Rocky / Alma:** `sudo dnf install openssl-devel`

docs/ops.md

@@ -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 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |

docs/ops/Vulkan.csv

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

ggml/CMakeLists.txt

@@ -4,7 +4,7 @@ project("ggml" C CXX ASM)
 
 ### GGML Version
 set(GGML_VERSION_MAJOR 0)
-set(GGML_VERSION_MINOR 21)
+set(GGML_VERSION_MINOR 22)
 set(GGML_VERSION_PATCH 0)
 set(GGML_VERSION_BASE "${GGML_VERSION_MAJOR}.${GGML_VERSION_MINOR}.${GGML_VERSION_PATCH}")
 
@@ -342,9 +342,6 @@ set(GGML_PUBLIC_HEADERS
     include/gguf.h)
 
 set_target_properties(ggml PROPERTIES PUBLIC_HEADER "${GGML_PUBLIC_HEADERS}")
-#if (GGML_METAL)
-#    set_target_properties(ggml PROPERTIES RESOURCE "${CMAKE_CURRENT_SOURCE_DIR}/src/ggml-metal.metal")
-#endif()
 install(TARGETS ggml LIBRARY PUBLIC_HEADER)
 install(TARGETS ggml-base LIBRARY)
 

ggml/cmake/ggml-config.cmake.in

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

ggml/include/ggml-rpc.h

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

ggml/include/ggml.h

@@ -1724,6 +1724,19 @@ extern "C" {
             struct ggml_tensor  * a,
             int                   n_past);
 
+    GGML_API struct ggml_tensor * ggml_clamp(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            float                 min,
+            float                 max);
+
+    // in-place, returns view(a)
+    GGML_API struct ggml_tensor * ggml_clamp_inplace(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            float                 min,
+            float                 max);
+
     GGML_API struct ggml_tensor * ggml_soft_max(
             struct ggml_context * ctx,
             struct ggml_tensor  * a);
@@ -1990,14 +2003,6 @@ extern "C" {
             struct ggml_tensor  * a,
             int                   n_offs);
 
-    // clamp
-    // in-place, returns view(a)
-    GGML_API struct ggml_tensor * ggml_clamp(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            float                 min,
-            float                 max);
-
     // im2col
     // converts data into a format that effectively results in a convolution when combined with matrix multiplication
     GGML_API struct ggml_tensor * ggml_im2col(

ggml/src/ggml-alloc.c

@@ -40,6 +40,7 @@ bool ggml_op_can_inplace(enum ggml_op op) {
         case GGML_OP_SILU_BACK:
         case GGML_OP_RMS_NORM:
         case GGML_OP_RMS_NORM_BACK:
+        case GGML_OP_CLAMP:
         case GGML_OP_SOFT_MAX:
         case GGML_OP_SOFT_MAX_BACK:
             return true;

ggml/src/ggml-backend-impl.h

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

ggml/src/ggml-backend-meta.cpp

@@ -592,7 +592,18 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             GGML_ASSERT(split_states_equal(src_ss[0], src_ss[1]));
             return {assume_sync ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_PARTIAL, {0}, {1}, 1};
         }
-        GGML_ABORT("fatal error");
+        if (src_ss[0].axis == src_ss[1].axis && src_ss[0].axis >= GGML_BACKEND_SPLIT_AXIS_2 &&
+                src_ss[0].axis < GGML_MAX_DIMS) {
+            GGML_ASSERT(split_states_equal(src_ss[0], src_ss[1]));
+            return src_ss[0];
+        }
+        // batched matmul with the batches split across devices and a replicated activation
+        if (src_ss[0].axis >= GGML_BACKEND_SPLIT_AXIS_2 && src_ss[0].axis < GGML_MAX_DIMS &&
+                src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) {
+            return src_ss[0];
+        }
+        GGML_ABORT("unsupported mul_mat split states: node=%s src0=%s axis=%d src1=%s axis=%d",
+            tensor->name, tensor->src[0]->name, (int) src_ss[0].axis, tensor->src[1]->name, (int) src_ss[1].axis);
         //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1};
     };
 
@@ -602,27 +613,40 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             case GGML_BACKEND_SPLIT_AXIS_1:
             case GGML_BACKEND_SPLIT_AXIS_2:
             case GGML_BACKEND_SPLIT_AXIS_3: {
-                GGML_ASSERT(src_ss[0].n_segments == 1);
-                if (src_ss[0].axis == ggml_n_dims(tensor->src[0]) - 1 && src_ss[0].nr[0] == 1) {
-                    return {ggml_backend_meta_split_axis(ggml_n_dims(tensor) - 1), {0}, {1}, 1};
-                }
-                int64_t base_ne_in = tensor->src[0]->ne[0];
-                for (int dim = 1; dim <= src_ss[0].axis; dim++) {
+                int64_t base_ne_in = 1;
+                for (int dim = 0; dim <= src_ss[0].axis; dim++) {
                     base_ne_in *= tensor->src[0]->ne[dim];
                 }
-                base_ne_in /= src_ss[0].nr[0];
+                if (src_ss[0].n_segments == 1) {
+                    base_ne_in /= src_ss[0].nr[0];
+                    if (src_ss[0].axis == ggml_n_dims(tensor->src[0]) - 1 && src_ss[0].nr[0] == 1) {
+                        return {ggml_backend_meta_split_axis(ggml_n_dims(tensor) - 1), {0}, {1}, 1};
+                    }
+                    if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && tensor->ne[0] == tensor->src[0]->ne[0] &&
+                            tensor->ne[1] == 1 && src_ss[0].nr[0] == 1) {
+                        bool complete_rows = true;
+                        for (size_t j = 0; j < n_bufs; j++) {
+                            const int64_t ne = src_ss[0].ne[j];
+                            complete_rows = complete_rows && (ne == 0 || ne == tensor->src[0]->ne[0]);
+                        }
+                        if (complete_rows) {
+                            // Move a complete dim-0 split to the following singleton dimension.
+                            return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1};
+                        }
+                    }
+                }
+                // Reshape outputs use one segment; split-state propagation merges source segments.
                 int64_t base_ne_out = 1;
                 for (int dim = 0; dim < GGML_MAX_DIMS; dim++) {
-                    const int64_t base_ne_out_next = base_ne_out *= tensor->ne[dim];
-                    if (base_ne_out_next % base_ne_in == 0) {
-                        return {ggml_backend_meta_split_axis(dim), {0}, {uint32_t(base_ne_out_next/base_ne_in)}, 1};
+                    base_ne_out *= tensor->ne[dim];
+                    if (base_ne_out % base_ne_in == 0) {
+                        return {ggml_backend_meta_split_axis(dim), {0}, {uint32_t(base_ne_out/base_ne_in)}, 1};
                     }
-                    if (base_ne_out_next > base_ne_in) {
+                    if (base_ne_out > base_ne_in) {
                         GGML_ASSERT(src_ss[0].n_segments == 1);
                         GGML_ASSERT(src_ss[0].nr[0]      == 1);
                         return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1};
                     }
-                    base_ne_out = base_ne_out_next;
                 }
                 GGML_ABORT("shape mismatch for %s", ggml_op_name(tensor->op));
             }
@@ -747,14 +771,33 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
     };
 
     auto handle_flash_attn_ext = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
-        GGML_ASSERT(                             src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2);
-        GGML_ASSERT(                             src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2);
-        GGML_ASSERT(                             src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2);
-        GGML_ASSERT(tensor->src[4] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+        GGML_ASSERT(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+
+        if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) {
+            GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+            GGML_ASSERT(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+            GGML_ASSERT(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+            return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
+        }
+
+        GGML_ASSERT(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2);
+        const bool kv_split = src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2 &&
+                src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2;
+        const bool kv_mirrored = src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED &&
+                src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED;
+        GGML_ASSERT(kv_split || kv_mirrored);
         GGML_ASSERT(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_0);
         return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1};
     };
 
+    auto handle_lightning_indexer = [&](
+            const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+        for (size_t i = 0; i < 4; i++) {
+            GGML_ASSERT(src_ss[i].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+        }
+        return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
+    };
+
     auto handle_ssm_conv = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
         if (src_ss[0].axis == src_ss[1].axis) {
             if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) {
@@ -792,7 +835,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer));
             const ggml_backend_meta_device_context * dev_ctx = (const ggml_backend_meta_device_context *) dev->context;
             ggml_backend_meta_split_state ret = dev_ctx->get_split_state(tensor, dev_ctx->get_split_state_ud);
-            if (ret.axis >= 0 && ret.axis <= GGML_MAX_DIMS) {
+            if (ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) {
                 const int64_t granularity = ret.axis == GGML_BACKEND_SPLIT_AXIS_0 ? ggml_blck_size(tensor->type) : 1;
                 int64_t ne_sum = 0;
                 for (size_t s = 0; s < ret.n_segments; s++) {
@@ -802,6 +845,9 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
                     }
                 }
                 GGML_ASSERT(ne_sum == tensor->ne[ret.axis]);
+            } else if (ret.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) {
+                GGML_ASSERT(ret.n_segments == 1);
+                GGML_ASSERT(ret.nr[0] == 1);
             }
             return ret;
         }
@@ -922,7 +968,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
                 split_state = handle_rope(src_ss);
             } break;
             case GGML_OP_ROPE_BACK: {
-                split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+                split_state = handle_rope(src_ss);
             } break;
             case GGML_OP_CLAMP: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ false);
@@ -986,6 +1032,9 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             case GGML_OP_GATED_DELTA_NET: {
                 split_state = handle_gated_delta_net(src_ss);
             } break;
+            case GGML_OP_LIGHTNING_INDEXER: {
+                split_state = handle_lightning_indexer(src_ss);
+            } break;
             case GGML_OP_DSV4_HC_COMB:
             case GGML_OP_DSV4_HC_PRE:
             case GGML_OP_DSV4_HC_POST: {
@@ -1070,13 +1119,14 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
         if (buf_ctx->debug > 0) {
             std::string srcs_info;
             for (size_t i = 0; i < GGML_MAX_SRC; i++) {
-                if (tensor->src[i] == nullptr) {
+                if (tensor->src[i] == nullptr || tensor->src[i] == tensor) {
                     continue;
                 }
                 if (!srcs_info.empty()) {
                     srcs_info += ", ";
                 }
-                const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor->src[0], true);
+                const ggml_backend_meta_split_state split_state =
+                        ggml_backend_meta_get_split_state(tensor->src[i], true);
                 GGML_ASSERT(split_state.n_segments == 1);
                 const char * axis_name = ggml_backend_meta_split_axis_name(split_state.axis);
                 std::string ne_info;
@@ -1118,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);
 }
@@ -1209,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) {
@@ -1255,6 +1311,108 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor(ggml_backend_buffer
     return ggml_backend_meta_buffer_init_tensor_impl(buf_ctx->get_simple_tensor_container(tensor), tensor);
 }
 
+static void ggml_backend_meta_buffer_memset_tensor(
+        ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) {
+    const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer);
+    const ggml_backend_meta_split_state split_state =
+            ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false);
+    GGML_ASSERT(ggml_is_contiguous(tensor) || split_state.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+
+    if (split_state.n_segments != 1 || split_state.nr[0] != 1) {
+        GGML_ASSERT(split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS);
+        GGML_ASSERT(split_state.nr[0] != 0);
+        GGML_ASSERT(tensor->ne[3] == 1);
+
+        std::vector<size_t> simple_offsets(n_bufs, 0);
+        if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_0) {
+            GGML_ASSERT(tensor->ne[2] == 1);
+
+            const size_t row_stride = tensor->nb[1];
+            GGML_ASSERT(offset % row_stride == 0);
+            GGML_ASSERT(size   % row_stride == 0);
+            const int64_t row_start = offset / row_stride;
+            const int64_t row_count = size   / row_stride;
+            GGML_ASSERT(row_start + row_count <= tensor->ne[1]);
+
+            const int64_t blck_size = ggml_blck_size(tensor->type);
+            for (size_t s = 0; s < split_state.n_segments; s++) {
+                for (size_t r = 0; r < split_state.nr[s]; r++) {
+                    for (size_t j = 0; j < n_bufs; j++) {
+                        ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j);
+                        GGML_ASSERT(split_state.ne[s*n_bufs + j] % blck_size == 0);
+                        const size_t nbytes = split_state.ne[s*n_bufs + j]/blck_size * tensor->nb[0];
+                        for (int64_t row = 0; row < row_count; row++) {
+                            ggml_backend_tensor_memset(simple_tensor, value,
+                                    simple_offsets[j] + (row_start + row)*simple_tensor->nb[1], nbytes);
+                        }
+                        simple_offsets[j] += nbytes;
+                    }
+                }
+            }
+            return;
+        }
+
+        GGML_ASSERT(split_state.axis == GGML_BACKEND_SPLIT_AXIS_1);
+
+        const size_t row_stride = tensor->nb[2];
+        GGML_ASSERT(offset % row_stride == 0);
+        GGML_ASSERT(size   % row_stride == 0);
+        const int64_t row_start = offset / row_stride;
+        const int64_t row_count = size   / row_stride;
+        GGML_ASSERT(row_start + row_count <= tensor->ne[2]);
+
+        for (size_t s = 0; s < split_state.n_segments; s++) {
+            for (size_t r = 0; r < split_state.nr[s]; r++) {
+                for (size_t j = 0; j < n_bufs; j++) {
+                    ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j);
+                    const size_t nbytes = split_state.ne[s*n_bufs + j] * tensor->nb[1];
+                    for (int64_t row = 0; row < row_count; row++) {
+                        ggml_backend_tensor_memset(simple_tensor, value,
+                                simple_offsets[j] + (row_start + row)*simple_tensor->nb[2], nbytes);
+                    }
+                    simple_offsets[j] += nbytes;
+                }
+            }
+        }
+        return;
+    }
+
+    switch (split_state.axis) {
+        case GGML_BACKEND_SPLIT_AXIS_0:
+        case GGML_BACKEND_SPLIT_AXIS_1:
+        case GGML_BACKEND_SPLIT_AXIS_2: {
+            const size_t chunk_size_full = tensor->nb[split_state.axis + 1];
+            GGML_ASSERT(offset % chunk_size_full == 0);
+            GGML_ASSERT(size   % chunk_size_full == 0);
+            const int64_t i_start =  offset        / chunk_size_full;
+            const int64_t i_stop  = (offset + size) / chunk_size_full;
+            for (size_t j = 0; j < n_bufs; j++) {
+                ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j);
+                const size_t chunk_size = simple_tensor->nb[split_state.axis + 1];
+                if (chunk_size == 0) {
+                    continue;
+                }
+                for (int64_t i = i_start; i < i_stop; i++) {
+                    ggml_backend_tensor_memset(simple_tensor, value, i*chunk_size, chunk_size);
+                }
+            }
+        } break;
+        case GGML_BACKEND_SPLIT_AXIS_PARTIAL: {
+            GGML_ASSERT(value == 0);
+            [[fallthrough]];
+        }
+        case GGML_BACKEND_SPLIT_AXIS_MIRRORED: {
+            for (size_t j = 0; j < n_bufs; j++) {
+                ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j);
+                ggml_backend_tensor_memset(simple_tensor, value, offset, size);
+            }
+        } break;
+        default: {
+            GGML_ABORT("fatal error");
+        }
+    }
+}
+
 static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
     const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer);
     const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false);
@@ -1352,15 +1510,29 @@ static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, gg
         } break;
         case GGML_BACKEND_SPLIT_AXIS_PARTIAL: {
             GGML_ASSERT(tensor->type == GGML_TYPE_F32);
-            const int64_t ne = ggml_nelements(tensor);
-            std::vector<float> tmp;
-            tmp.reserve(ne);
-            for (int64_t i = 0; i < ne; i++) {
-                tmp.push_back(((const float *) data)[i] / n_bufs);
+            GGML_ASSERT(offset % sizeof(float) == 0);
+            GGML_ASSERT(size   % sizeof(float) == 0);
+            const size_t n_values = size / sizeof(float);
+            size_t n_contributors = 0;
+            for (size_t j = 0; j < n_bufs; j++) {
+                n_contributors += split_state.ne[j] != 0;
+            }
+            const bool has_contributor_mask = n_contributors != 0;
+            if (!has_contributor_mask) {
+                n_contributors = n_bufs;
+            }
+            std::vector<float> tmp(n_values);
+            for (size_t i = 0; i < n_values; i++) {
+                tmp[i] = ((const float *) data)[i] / n_contributors;
+            }
+            std::vector<float> zero;
+            if (has_contributor_mask) {
+                zero.resize(n_values, 0.0f);
             }
             for (size_t j = 0; j < n_bufs; j++) {
                 ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j);
-                ggml_backend_tensor_set(simple_tensor, tmp.data(), offset, size);
+                const float * partial = has_contributor_mask && split_state.ne[j] == 0 ? zero.data() : tmp.data();
+                ggml_backend_tensor_set(simple_tensor, partial, offset, size);
             }
         } break;
         default: {
@@ -1488,7 +1660,7 @@ static const ggml_backend_buffer_i ggml_backend_meta_buffer_iface = {
     /* .free_buffer     = */ ggml_backend_meta_buffer_free_buffer,
     /* .get_base        = */ ggml_backend_meta_buffer_get_base,
     /* .init_tensor     = */ ggml_backend_meta_buffer_init_tensor,
-    /* .memset_tensor   = */ nullptr, // TODO implement
+    /* .memset_tensor   = */ ggml_backend_meta_buffer_memset_tensor,
     /* .set_tensor      = */ ggml_backend_meta_buffer_set_tensor,
     /* .get_tensor      = */ ggml_backend_meta_buffer_get_tensor,
     /* .set_tensor_2d   = */ nullptr,
@@ -1502,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);
 
@@ -1841,7 +2023,7 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend,
 
         {
             // For MoE models it may make sense to delay the AllReduce in order to reduce I/O:
-            auto get_i_delayed = [&](const int i) -> int {
+            auto get_i_delayed_branch = [&](const int i) -> int {
                 int id = i; // i_delayed
                 int idr = i; // i_delayed return, last safe return value
 
@@ -1941,6 +2123,62 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend,
                 return idr;
             };
 
+            // AllReduce(a) + AllReduce(b) == AllReduce(a + b) for independent partial branches.
+            auto get_i_delayed = [&](const int i) -> int {
+                const int i_delayed = get_i_delayed_branch(i);
+                ggml_tensor * node = cgraph->nodes[i_delayed];
+
+                if (ggml_node_get_use_count(cgraph, i_delayed) != 1) {
+                    return i_delayed;
+                }
+
+                for (int id = i_delayed + 1; id < cgraph->n_nodes; id++) {
+                    ggml_tensor * next = cgraph->nodes[id];
+                    if (next->view_src == node) {
+                        return i_delayed;
+                    }
+                    for (int s = 0; s < GGML_MAX_SRC; s++) {
+                        if (next->src[s] == node) {
+                            return i_delayed;
+                        }
+                    }
+
+                    if (next->view_src != nullptr && next->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(next->view_src->buffer)) {
+                        continue;
+                    }
+                    if (ggml_backend_meta_get_split_state(next, false).axis != GGML_BACKEND_SPLIT_AXIS_PARTIAL) {
+                        continue;
+                    }
+
+                    const int i_other = id;
+                    const int i_other_delayed = get_i_delayed_branch(i_other);
+                    ggml_tensor * other = cgraph->nodes[i_other_delayed];
+                    if (ggml_node_get_use_count(cgraph, i_other_delayed) != 1 || i_other_delayed + 1 >= cgraph->n_nodes) {
+                        return i_delayed;
+                    }
+
+                    ggml_tensor * sum = cgraph->nodes[i_other_delayed + 1];
+                    if (sum->op != GGML_OP_ADD ||
+                            !ggml_are_same_shape(node, other) || node->type != other->type || sum->type != node->type ||
+                            !((sum->src[0] == node && sum->src[1] == other) ||
+                              (sum->src[0] == other && sum->src[1] == node)) ||
+                            ggml_backend_meta_get_split_state(sum, false).axis != GGML_BACKEND_SPLIT_AXIS_MIRRORED) {
+                        return i_delayed;
+                    }
+
+                    for (size_t j = 0; j < n_backends; j++) {
+                        auto & bcj = backend_ctx->backend_configs[j];
+                        const bool compute       = bcj.nodes[i]->flags       & GGML_TENSOR_FLAG_COMPUTE;
+                        const bool compute_other = bcj.nodes[i_other]->flags & GGML_TENSOR_FLAG_COMPUTE;
+                        if (compute != compute_other) {
+                            return i_delayed;
+                        }
+                    }
+                    return i_other_delayed + 1;
+                }
+                return i_delayed;
+            };
+
             int i_start = 0;
             for (int i = 0; i < cgraph->n_nodes; i++) {
                 ggml_tensor * node = cgraph->nodes[i];

ggml/src/ggml-backend.cpp

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

ggml/src/ggml-cpu/CMakeLists.txt

@@ -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}")

ggml/src/ggml-cpu/kleidiai/CMakeLists.txt

@@ -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()

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

@@ -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;
 }

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

@@ -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) {

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

@@ -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) {

ggml/src/ggml-cpu/ops.cpp

@@ -1896,7 +1896,6 @@ void ggml_compute_forward_repeat_back(
 }
 
 // ggml_compute_forward_concat
-
 static void ggml_compute_forward_concat_any(
     const ggml_compute_params * params,
     ggml_tensor * dst) {
@@ -1904,8 +1903,6 @@ static void ggml_compute_forward_concat_any(
     const ggml_tensor * src0 = dst->src[0];
     const ggml_tensor * src1 = dst->src[1];
 
-    const size_t len = ggml_type_size(src0->type);
-
     const int ith = params->ith;
     const int nth = params->nth;
 
@@ -1914,31 +1911,38 @@ static void ggml_compute_forward_concat_any(
     const int32_t dim = ggml_get_op_params_i32(dst, 0);
 
     GGML_ASSERT(dim >= 0 && dim < 4);
+    GGML_ASSERT(ggml_is_contiguous_rows(src0));
+    GGML_ASSERT(ggml_is_contiguous_rows(src1));
 
     int64_t o[4] = {0, 0, 0, 0};
+
     if (dim == 0) {
+        GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0);
+        GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0);
+
         o[dim] = src0->ne[dim]/ggml_blck_size(src0->type);
     } else {
         o[dim] = src0->ne[dim];
     }
 
-    const char * x;
+    // Region 1: copy rows from src0
+    for (int i3 = 0; i3 < ne03; i3++) {
+        for (int i2 = ith; i2 < ne02; i2 += nth) {
+            for (int i1 = 0; i1 < ne01; i1++) {
+                const char * x = (const char *) src0->data + i1*nb01 + i2*nb02 + i3*nb03;
+                      char * y = (      char *) dst->data  + i1*nb1  + i2*nb2  + i3*nb3;
+                memcpy(y, x, ggml_row_size(src0->type, ne00));
+            }
+        }
+    }
 
-    // TODO: smarter multi-theading
-    for (int i3 = 0; i3 < ne3; i3++) {
-        for (int i2 = ith; i2 < ne2; i2 += nth) {
-            for (int i1 = 0; i1 < ne1; i1++) {
-                for (int i0 = 0; i0 < ne0/ggml_blck_size(dst->type); i0++) {
-                    if (i0 < ne00/ggml_blck_size(src0->type) && i1 < ne01 && i2 < ne02 && i3 < ne03) {
-                        x = (const char *)src0->data + (i0       )*nb00 + (i1       )*nb01 + (i2       )*nb02 + (i3       )*nb03;
-                    } else {
-                        x = (const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13;
-                    }
-
-                    char * y = (char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3;
-
-                    memcpy(y, x, len);
-                }
+    // Region 2: copy rows from src1, offset into dst by o[]
+    for (int i3 = 0; i3 < ne13; i3++) {
+        for (int i2 = ith; i2 < ne12; i2 += nth) {
+            for (int i1 = 0; i1 < ne11; i1++) {
+                const char * x = (const char *) src1->data + i1*nb11         + i2*nb12         + i3*nb13;
+                      char * y = (      char *)  dst->data + (i1 + o[1])*nb1 + (i2 + o[2])*nb2 + (i3 + o[3])*nb3 + o[0]*nb0;
+                memcpy(y, x, ggml_row_size(src1->type, ne10));
             }
         }
     }
@@ -2078,14 +2082,6 @@ void ggml_compute_forward_concat(
     ggml_tensor * dst) {
 
     const ggml_tensor * src0 = dst->src[0];
-    const ggml_tensor * src1 = dst->src[1];
-
-    if (ggml_is_quantized(src0->type)) {
-        GGML_ASSERT(ggml_is_contiguous_rows(src0));
-        GGML_ASSERT(ggml_is_contiguous_rows(src1));
-        GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0);
-        GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0);
-    }
 
     switch (src0->type) {
         case GGML_TYPE_F16:

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -38,6 +38,7 @@
 #include "ggml-cuda/out-prod.cuh"
 #include "ggml-cuda/pad.cuh"
 #include "ggml-cuda/pool2d.cuh"
+#include "ggml-cuda/pool1d.cuh"
 #include "ggml-cuda/quantize.cuh"
 #include "ggml-cuda/rope.cuh"
 #include "ggml-cuda/roll.cuh"
@@ -2326,6 +2327,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_POOL_2D:
             ggml_cuda_op_pool2d(ctx, dst);
             break;
+        case GGML_OP_POOL_1D:
+            ggml_cuda_op_pool1d(ctx, dst);
+            break;
         case GGML_OP_SUM:
             ggml_cuda_op_sum(ctx, dst);
             break;
@@ -4607,8 +4611,8 @@ static std::string ggml_cuda_device_description(int device) {
     const ggml_cuda_device_info & info = ggml_cuda_info();
     std::string description = prop.name;
     if (info.device_count > info.physical_device_count) {
-        description += " (physical device " + std::to_string(info.devices[device].physical_device) +
-                       ", virtual device " + std::to_string(info.devices[device].virtual_index) + ")";
+        description += " (dev p" + std::to_string(info.devices[device].physical_device) +
+                       "/v" + std::to_string(info.devices[device].virtual_index) + ")";
     }
     return description;
 }
@@ -5245,6 +5249,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
         case GGML_OP_CONV_2D_DW:
             return op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_CONV_TRANSPOSE_2D:
+        case GGML_OP_POOL_1D:
         case GGML_OP_POOL_2D:
             return true;
         case GGML_OP_ACC:

ggml/src/ggml-cuda/mmq-config-pascal-dp4a.cuh

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

ggml/src/ggml-cuda/mmq-config-pascal-older.cuh

@@ -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);
+}

ggml/src/ggml-cuda/mmq.cu

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

ggml/src/ggml-cuda/mmq.cuh

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

ggml/src/ggml-cuda/pool1d.cu

@@ -0,0 +1,85 @@
+#include "pool1d.cuh"
+
+static __global__ void pool1d_nchw_kernel(
+        const int iw, const int ow,
+        const int kw, const int sw, const int pw,
+        const int parallel_elements,
+        const float * src, float * dst, const enum ggml_op_pool op) {
+    const int idx = threadIdx.x + blockIdx.x * blockDim.x;
+    if (idx >= parallel_elements) {
+        return;
+    }
+
+    const int nc = idx / ow;
+    const int cur_ow = idx % ow;
+
+    const float * i_ptr = src + nc * iw;
+    float * o_ptr = dst + nc * ow;
+
+    const int start = cur_ow * sw - pw;
+    const int b = max(0, start);
+    const int e = min(iw, start + kw);
+
+    float res;
+    switch (op) {
+        case GGML_OP_POOL_AVG: res = 0.0f;     break;
+        case GGML_OP_POOL_MAX: res = -FLT_MAX; break;
+        default: return;
+    }
+
+    int count = 0;
+    for (int i = b; i < e; i++) {
+#if __CUDA_ARCH__ >= 350
+        float cur = __ldg(i_ptr + i);
+#else
+        float cur = i_ptr[i];
+#endif
+        switch (op) {
+            case GGML_OP_POOL_AVG: res += cur;                break;
+            case GGML_OP_POOL_MAX: res = max(res, cur);       break;
+            default: break;
+        }
+        count++;
+    }
+
+    if (op == GGML_OP_POOL_AVG) {
+        res = (count > 0) ? (res / count) : 0.0f;
+    }
+
+    o_ptr[cur_ow] = res;
+}
+
+static void pool1d_nchw_kernel_f32_f32_cuda(
+        const int iw, const int ow,
+        const int kw, const int sw, const int pw,
+        const int parallel_elements,
+        const float * src, float * dst, const enum ggml_op_pool op,
+        cudaStream_t stream) {
+    const int num_blocks = (parallel_elements + CUDA_POOL1D_BLOCK_SIZE - 1) / CUDA_POOL1D_BLOCK_SIZE;
+    dim3 block_nums(num_blocks);
+    pool1d_nchw_kernel<<<block_nums, CUDA_POOL1D_BLOCK_SIZE, 0, stream>>>(iw, ow, kw, sw, pw, parallel_elements, src, dst, op);
+}
+
+void ggml_cuda_op_pool1d(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+    const float * src0_d = (const float *)src0->data;
+    float * dst_d = (float *)dst->data;
+    cudaStream_t stream = ctx.stream();
+
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT( dst->type == GGML_TYPE_F32);
+
+    const int32_t * opts = (const int32_t *)dst->op_params;
+    enum ggml_op_pool op = static_cast<ggml_op_pool>(opts[0]);
+    const int k0 = opts[1];
+    const int s0 = opts[2];
+    const int p0 = opts[3];
+
+    const int64_t IW = src0->ne[0];
+    const int64_t OW = dst->ne[0];
+    const int64_t nr = ggml_nrows(src0);
+
+    const int parallel_elements = (int)(nr * OW);
+
+    pool1d_nchw_kernel_f32_f32_cuda(IW, OW, k0, s0, p0, parallel_elements, src0_d, dst_d, op, stream);
+}

ggml/src/ggml-cuda/pool1d.cuh

@@ -0,0 +1,5 @@
+#include "common.cuh"
+
+#define CUDA_POOL1D_BLOCK_SIZE 256
+
+void ggml_cuda_op_pool1d(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

ggml/src/ggml-hexagon/htp-opnode.h

@@ -8,60 +8,107 @@
 #include <algorithm>
 #include <string>
 #include <vector>
+#include <memory>
 #include <stdio.h>
 #include "htp-ops.h"
 #include "htp/matmul-ops.h"
 #include "htp/flash-attn-ops.h"
 #include "htp/unary-ops.h"
+#include "htp/allreduce-ops.h"
 
 struct htp_opnode {
-    ggml_tensor * node = nullptr;
+    ggml_tensor * node   { nullptr };
+    htp_op_code   opcode { HTP_OP_INVALID };
+    int32_t       kernel_params[HTP_OP_MAX_KERN_PARAMS] {0};
 
-    std::vector<ggml_tensor *> fused;
+    std::vector<ggml_tensor *>                fused;
+    std::vector<std::shared_ptr<ggml_tensor>> dummy;
 
-    htp_op_code opcode = HTP_OP_INVALID;
+    std::vector<const ggml_tensor *> inputs;
+    std::vector<const ggml_tensor *> outputs;
+    std::string                      name;
 
-    std::vector<ggml_tensor *> extra_dsts;
-
-    int32_t kernel_params[HTP_OP_MAX_KERN_PARAMS] = {0};
-
-    htp_opnode(ggml_tensor * node = nullptr, std::vector<ggml_tensor *> fused = {}, htp_op_code opcode = HTP_OP_INVALID, std::vector<ggml_tensor *> extra_dsts = {})
-        : node(node), fused(std::move(fused)), opcode(opcode), extra_dsts(std::move(extra_dsts)) {}
-
-    ggml_op op() const {
-        return node->op;
+    int n_active_src(const ggml_tensor * t) const {
+        if (!t) return 0;
+        for (int i = GGML_MAX_SRC - 1; i >= 0; i--) {
+            if (t->src[i]) {
+                return i + 1;
+            }
+        }
+        return 0;
     }
 
-    const ggml_tensor * dst() const {
-        return fused.empty() ? node : fused.back();
+    void init(ggml_tensor * node) {
+        this->node = node;
+        if (this->node) {
+            this->name = ggml_op_desc(this->node);
+
+            // Build inputs (preserving optional nullptrs)
+            int n_inputs = n_active_src(this->node);
+            this->inputs.resize(n_inputs, nullptr);
+            for (int i = 0; i < n_inputs; i++) {
+                this->inputs[i] = this->node->src[i];
+            }
+
+            // Build outputs
+            this->outputs.push_back(this->dst());
+        }
+    }
+
+    htp_opnode(htp_op_code opcode = HTP_OP_INVALID, ggml_tensor * node = nullptr) : opcode(opcode) {
+        init(node);
+    }
+
+    ggml_op             op()   const { return node->op; }
+    const ggml_tensor * src0() const { return node->src[0]; }
+    const ggml_tensor * src1() const { return node->src[1]; }
+    const ggml_tensor * dst()  const { return outputs.empty() ? node : outputs.back(); }
+
+    ggml_tensor * add_dummy(const ggml_tensor & t) {
+        dummy.push_back(std::make_shared<ggml_tensor>(t));
+        return dummy.back().get();
     }
 
     void add_fused(ggml_tensor * t, bool extra_dst = false) {
         fused.push_back(t);
-        if (extra_dst) {
-            extra_dsts.push_back(t);
-        }
-    }
 
-    std::vector<const ggml_tensor *> get_outputs() const {
-        std::vector<const ggml_tensor *> res;
-        if (extra_dsts.empty()) {
-            res.push_back(dst());
+        name += "+";
+        name += ggml_op_desc(t);
+
+        if (extra_dst) {
+            outputs.push_back(t);
         } else {
-            res.push_back(node);
-            for (const auto * x : extra_dsts) {
-                res.push_back(x);
+            outputs.clear();
+            outputs.push_back(t);
+        }
+
+        // Remove the newly fused intermediate output tensor t from inputs (if it was there)
+        inputs.erase(std::remove(inputs.begin(), inputs.end(), t), inputs.end());
+
+        // Append new inputs from t, preserving middle nullptrs
+        int n_inputs = n_active_src(t);
+        for (int i = 0; i < n_inputs; i++) {
+            const auto * src = t->src[i];
+            if (!src) {
+                inputs.push_back(nullptr);
+            } else if (src != node &&
+                       std::find(fused.begin(), fused.end(), src) == fused.end() &&
+                       std::find(inputs.begin(), inputs.end(), src) == inputs.end()) {
+                inputs.push_back(src);
             }
         }
-        return res;
     }
 
-    const ggml_tensor * src0() const {
-        return node->src[0];
+    const std::vector<const ggml_tensor *> & get_inputs() const {
+        return inputs;
     }
 
-    const ggml_tensor * src1() const {
-        return node->src[1];
+    const std::vector<const ggml_tensor *> & get_outputs() const {
+        return outputs;
+    }
+
+    std::string op_name() const {
+        return name;
     }
 
     bool is_empty() const {
@@ -81,75 +128,6 @@ struct htp_opnode {
     bool same_input(const htp_opnode& n) const {
         return n.src1() == this->src1();
     }
-
-    std::vector<const ggml_tensor *> get_inputs() const {
-        if (fused.empty()) {
-            int last_non_null = -1;
-            for (int i = 0; i < GGML_MAX_SRC; i++) {
-                if (node->src[i]) {
-                    last_non_null = i;
-                }
-            }
-            std::vector<const ggml_tensor *> inputs(last_non_null + 1, nullptr);
-            for (int i = 0; i <= last_non_null; i++) {
-                inputs[i] = node->src[i];
-            }
-            return inputs;
-        }
-
-        std::vector<const ggml_tensor *> inputs(GGML_MAX_SRC, nullptr);
-        std::vector<const ggml_tensor *> outputs;
-        outputs.push_back(node);
-        for (const auto * f : fused) {
-            outputs.push_back(f);
-        }
-
-        auto contains = [&](const std::vector<const ggml_tensor *> & vec, const ggml_tensor * t) {
-            for (const auto * x : vec) {
-                if (x == t) return true;
-            }
-            return false;
-        };
-
-        int count = 0;
-        auto add_input = [&](const ggml_tensor * t) {
-            if (t && !contains(outputs, t) && !contains(inputs, t)) {
-                if (count < (int)inputs.size()) {
-                    inputs[count++] = t;
-                } else {
-                    inputs.push_back(t);
-                }
-            }
-        };
-
-        for (int i = 0; i < GGML_MAX_SRC; i++) {
-            if (node->src[i]) {
-                add_input(node->src[i]);
-            }
-        }
-        for (const auto * f : fused) {
-            for (int i = 0; i < GGML_MAX_SRC; i++) {
-                if (f->src[i]) {
-                    add_input(f->src[i]);
-                }
-            }
-        }
-
-        inputs.resize(count);
-        return inputs;
-    }
-
-    std::string op_name() const {
-        if (fused.empty()) {
-            return ggml_op_desc(node);
-        }
-        std::string name = ggml_op_desc(node);
-        for (const auto * f : fused) {
-            name += "+";
-            name += ggml_op_desc(f);
-        }
-        return name;
-    }
 };
 
 struct htp_opformat {
@@ -337,8 +315,7 @@ struct htp_opformat {
     }
     void format_kernel_params(char * str, size_t max_size, const htp_opnode & node) {
         if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID ||
-            node.opcode == HTP_OP_MUL_MAT_QKV || node.opcode == HTP_OP_MUL_MAT_FFN ||
-            node.opcode == HTP_OP_MUL_MAT_ADD) {
+            node.opcode == HTP_OP_MUL_MAT_NX || node.opcode == HTP_OP_MUL_MAT_ADD) {
             const auto * kparams = (const struct htp_mm_kernel_params *) node.kernel_params;
             const char * path = "unknown";
             int32_t type = kparams->kernel_type;

ggml/src/ggml-hexagon/htp/CMakeLists.txt

@@ -43,6 +43,7 @@ add_library(${HTP_LIB} SHARED
     pad-ops.c
     argsort-ops.c
     im2col-ops.c
+    allreduce-ops.c
 )
 
 target_compile_definitions(${HTP_LIB} PRIVATE

ggml/src/ggml-hexagon/htp/act-ops.c

@@ -183,6 +183,53 @@ static void swiglu_oai_f32(const float * restrict src0,
 static const float GELU_COEF_A     = 0.044715f;
 static const float SQRT_2_OVER_PI  = 0.79788456080286535587989211986876f;
 
+static inline HVX_Vector hvx_vec_fast_sigmoid_f32_2it(HVX_Vector v) {
+    v = Q6_Vqf32_vmpy_VsfVsf(v, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
+    v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), Q6_V_vsplat_R(FAST_SIGMOID_C3));
+
+    HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v));
+    HVX_Vector x      = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
+    HVX_Vector xx     = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x);
+
+    HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2));
+    v1            = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
+
+    HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1));
+    v2            = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx);
+    v2            = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x);
+
+    HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
+    v3            = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
+
+    HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
+    HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
+
+    // Newton-Raphson with 2 iterations
+    HVX_Vector two_sf = hvx_vec_splat_f32(2.0f);
+    HVX_Vector i_sf   = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(0x7EEEEBB3), v5);
+    HVX_Vector r_qf   = Q6_Vqf32_vmpy_VsfVsf(
+        i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5)))));
+    r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32(
+        r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5))));
+    HVX_Vector res = Q6_Vsf_equals_Vqf32(r_qf);
+
+    res = Q6_Vqf32_vmpy_VsfVsf(v3, res);
+
+    return Q6_Vsf_equals_Vqf32(res);
+}
+
+static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard_2it(HVX_Vector v,
+                                                            HVX_Vector one,
+                                                            HVX_Vector max_exp,
+                                                            HVX_Vector min_exp) {
+    const HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(max_exp, v);
+    const HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(v, min_exp);
+
+    HVX_Vector out = hvx_vec_fast_sigmoid_f32_2it(v);
+    out            = Q6_V_vmux_QVV(pred_max, out, one);
+    return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero());
+}
+
 static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
     assert((unsigned long) dst  % 128 == 0);
     assert((unsigned long) src0 % 128 == 0);
@@ -200,20 +247,13 @@ static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * rest
 
     const HVX_Vector v_coef_a_times_sqrt = hvx_vec_splat_f32(GELU_COEF_A_TIMES_SQRT);
     const HVX_Vector v_sqrt_2_pi         = hvx_vec_splat_f32(SQRT_2_OVER_PI);
-    const HVX_Vector v_half              = hvx_vec_splat_f32(0.5f);
     const HVX_Vector v_one               = hvx_vec_splat_f32(1.0f);
-    const HVX_Vector v_two               = hvx_vec_splat_f32(2.0f);
-
-    // Hoisted fast sigmoid / inverse constants to avoid loop-internal overhead
-    const HVX_Vector v_log2f             = Q6_V_vsplat_R(FAST_SIGMOID_LOG2F);
-    const HVX_Vector v_c1                = Q6_V_vsplat_R(FAST_SIGMOID_C1);
-    const HVX_Vector v_c2                = Q6_V_vsplat_R(FAST_SIGMOID_C2);
-    const HVX_Vector v_inv_aprox         = Q6_V_vsplat_R(0x7EEEEBB3);
     const HVX_Vector v_max_exp           = hvx_vec_splat_f32(87.0f);
     const HVX_Vector v_min_exp           = hvx_vec_splat_f32(-87.0f);
 
     uint32_t i = 0;
 
+    _Pragma("unroll(4)")
     for (; i < nvec; i++) {
         HVX_Vector x = vsrc0[i];
         HVX_Vector g = vsrc1[i];
@@ -223,56 +263,13 @@ static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * rest
         coef = hvx_vec_add_f32_f32(coef, v_sqrt_2_pi);
         HVX_Vector inner = hvx_vec_mul_f32_f32(x, coef);
 
-        // y2 = 2 * inner
-        HVX_Vector y2 = hvx_vec_mul_f32_f32(inner, v_two);
+        // y2 = 2 * inner = inner + inner
+        HVX_Vector y2 = hvx_vec_add_f32_f32(inner, inner);
 
-        // Sigmoid guard check predicates
-        HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(v_max_exp, y2);
-        HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(y2, v_min_exp);
-
-        // Fast sigmoid approximation
-        HVX_Vector v = Q6_Vqf32_vmpy_VsfVsf(y2, v_log2f);
-        v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), v_half);
-
-        HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v));
-        HVX_Vector x_sig  = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
-        HVX_Vector xx_sig = Q6_Vqf32_vmpy_Vqf32Vqf32(x_sig, x_sig);
-
-        HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx_sig), v_c2);
-        v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, v_log2f);
-
-        HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x_sig), v_c1);
-        v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx_sig);
-        v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x_sig);
-
-        HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
-        v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
-
-        HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
-        HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
-
-        // Fast division (Newton-Raphson with 2 iterations)
-        HVX_Vector i_sf = Q6_Vw_vsub_VwVw(v_inv_aprox, v5);
-        HVX_Vector r_qf = Q6_Vqf32_vmpy_VsfVsf(
-            i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5)))));
-        r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32(
-            r_qf, Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5))));
-        HVX_Vector res_inv = Q6_Vsf_equals_Vqf32(r_qf);
-
-        HVX_Vector sig2y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, res_inv));
-
-        // Sigmoid guards
-        sig2y = Q6_V_vmux_QVV(pred_max, sig2y, v_one);
-        sig2y = Q6_V_vmux_QVV(pred_min, sig2y, Q6_V_vzero());
-
-        // tanh(inner) = 2 * sigmoid(2 * inner) - 1
-        HVX_Vector tanh_val = hvx_vec_mul_f32_f32(sig2y, v_two);
-        tanh_val = hvx_vec_sub_f32_f32(tanh_val, v_one);
-
-        HVX_Vector tanh_plus_one = hvx_vec_add_f32_f32(tanh_val, v_one);
-        HVX_Vector half_x = hvx_vec_mul_f32_f32(x, v_half);
-        HVX_Vector gelu_x = hvx_vec_mul_f32_f32(half_x, tanh_plus_one);
+        // Fast sigmoid approximation (2 iterations)
+        HVX_Vector sig2y = hvx_vec_fast_sigmoid_f32_guard_2it(y2, v_one, v_max_exp, v_min_exp);
 
+        HVX_Vector gelu_x = hvx_vec_mul_f32_f32(x, sig2y);
         vdst[i] = hvx_vec_mul_f32_f32(gelu_x, g);
     }
 
@@ -285,50 +282,11 @@ static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * rest
         coef = hvx_vec_add_f32_f32(coef, v_sqrt_2_pi);
         HVX_Vector inner = hvx_vec_mul_f32_f32(x, coef);
 
-        HVX_Vector y2 = hvx_vec_mul_f32_f32(inner, v_two);
+        HVX_Vector y2 = hvx_vec_add_f32_f32(inner, inner);
 
-        HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(v_max_exp, y2);
-        HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(y2, v_min_exp);
-
-        HVX_Vector v = Q6_Vqf32_vmpy_VsfVsf(y2, v_log2f);
-        v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), v_half);
-
-        HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v));
-        HVX_Vector x_sig  = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
-        HVX_Vector xx_sig = Q6_Vqf32_vmpy_Vqf32Vqf32(x_sig, x_sig);
-
-        HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx_sig), v_c2);
-        v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, v_log2f);
-
-        HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x_sig), v_c1);
-        v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx_sig);
-        v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x_sig);
-
-        HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
-        v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
-
-        HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
-        HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
-
-        HVX_Vector i_sf = Q6_Vw_vsub_VwVw(v_inv_aprox, v5);
-        HVX_Vector r_qf = Q6_Vqf32_vmpy_VsfVsf(
-            i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5)))));
-        r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32(
-            r_qf, Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5))));
-        HVX_Vector res_inv = Q6_Vsf_equals_Vqf32(r_qf);
-
-        HVX_Vector sig2y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, res_inv));
-
-        sig2y = Q6_V_vmux_QVV(pred_max, sig2y, v_one);
-        sig2y = Q6_V_vmux_QVV(pred_min, sig2y, Q6_V_vzero());
-
-        HVX_Vector tanh_val = hvx_vec_mul_f32_f32(sig2y, v_two);
-        tanh_val = hvx_vec_sub_f32_f32(tanh_val, v_one);
-
-        HVX_Vector tanh_plus_one = hvx_vec_add_f32_f32(tanh_val, v_one);
-        HVX_Vector half_x = hvx_vec_mul_f32_f32(x, v_half);
-        HVX_Vector gelu_x = hvx_vec_mul_f32_f32(half_x, tanh_plus_one);
+        HVX_Vector sig2y = hvx_vec_fast_sigmoid_f32_guard_2it(y2, v_one, v_max_exp, v_min_exp);
 
+        HVX_Vector gelu_x = hvx_vec_mul_f32_f32(x, sig2y);
         HVX_Vector res = hvx_vec_mul_f32_f32(gelu_x, g);
         hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(float), res);
     }

ggml/src/ggml-hexagon/htp/allreduce-ops.c

@@ -0,0 +1,398 @@
+#pragma clang diagnostic ignored "-Wunused-variable"
+#pragma clang diagnostic ignored "-Wunused-function"
+#pragma clang diagnostic ignored "-Wunused-but-set-variable"
+
+#include <HAP_farf.h>
+#include <HAP_perf.h>
+#include <stdatomic.h>
+#include <math.h>
+#include <string.h>
+
+#define GGML_COMMON_DECL_C
+#include "ggml-common.h"
+#include "htp-ctx.h"
+#include "htp-ops.h"
+#include "hvx-utils.h"
+#include "htp-tensor.h"
+#include "hex-dma.h"
+#include "hex-profile.h"
+#include "allreduce-ops.h"
+
+struct htp_allreduce_context {
+    struct htp_ops_context * octx;
+    uint32_t n_ranks;
+    uint32_t n_dsts;
+    uint32_t nelem;
+    uint32_t ne0;
+    uint32_t ne1;
+    uint32_t row_size_aligned;
+    uint32_t rank_elem_start;
+    uint32_t rank_nelem;
+    uint32_t elems_per_thread;
+    uint32_t block_elems;
+    uint32_t vtcm_size_per_thread;
+    bool     is_row_bcast;
+    uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];
+    uint8_t * dst_spad_base;
+    uint8_t * res_spad_base;
+};
+
+#define DEFINE_ALLREDUCE_THREAD_DMA_1D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD)                                             \
+static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                       \
+    struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data;                                      \
+    struct htp_ops_context * octx = actx->octx;                                                                       \
+                                                                                                                      \
+    const uint32_t n_ranks     = actx->n_ranks;                                                                       \
+    const uint32_t n_dsts      = actx->n_dsts;                                                                        \
+    const uint32_t block_elems = actx->block_elems;                                                                   \
+                                                                                                                      \
+    const uint32_t dr  = actx->elems_per_thread;                                                                      \
+    const uint32_t ir0 = actx->rank_elem_start + dr * ith;                                                            \
+    const uint32_t ir1 = MIN(ir0 + dr, actx->rank_elem_start + actx->rank_nelem);                                     \
+    if (ir0 >= ir1) return;                                                                                           \
+                                                                                                                      \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                            \
+    dma_queue * q = octx->ctx->dma[ith];                                                                              \
+                                                                                                                      \
+    uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];                                                                 \
+    for (uint32_t s = 0; s < n_ranks; s++) {                                                                          \
+        src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread);                               \
+    }                                                                                                                 \
+    uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread);                               \
+    uint8_t * res_spad_base = HAS_ADD ? (actx->res_spad_base + (ith * actx->vtcm_size_per_thread)) : NULL;            \
+                                                                                                                      \
+    const size_t spad_half = actx->vtcm_size_per_thread / 2;                                                          \
+    uint32_t ir_prefetch = ir0;                                                                                       \
+    int spad_idx = 0;                                                                                                 \
+                                                                                                                      \
+    for (int k = 0; k < 2 && ir_prefetch < ir1; k++) {                                                                \
+        uint32_t cur_elems = MIN(block_elems, ir1 - ir_prefetch);                                                     \
+        size_t   cur_bytes = cur_elems * sizeof(TYPE);                                                                \
+        uint8_t * d_spad = dst_spad_base + spad_idx * spad_half;                                                      \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                       \
+            uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir_prefetch * sizeof(TYPE);                           \
+            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 0);                       \
+        }                                                                                                             \
+        for (uint32_t s = 0; s < n_ranks; s++) {                                                                      \
+            uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half;                                               \
+            const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE);                \
+            dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), cur_bytes, cur_bytes, cur_bytes, 1);                       \
+        }                                                                                                             \
+        if (HAS_ADD) {                                                                                                \
+            uint8_t * r_spad = res_spad_base + spad_idx * spad_half;                                                  \
+            const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE);      \
+            dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), cur_bytes, cur_bytes, cur_bytes, 1);                       \
+        }                                                                                                             \
+        ir_prefetch += cur_elems;                                                                                     \
+        spad_idx ^= 1;                                                                                                \
+    }                                                                                                                 \
+                                                                                                                      \
+    for (uint32_t ir = ir0; ir < ir1; ) {                                                                             \
+        uint32_t cur_elems = MIN(block_elems, ir1 - ir);                                                              \
+        size_t   cur_bytes = cur_elems * sizeof(TYPE);                                                                \
+        uint8_t * d_spad = NULL;                                                                                      \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                       \
+            d_spad = (uint8_t *) dma_queue_pop(q).src;                                                                \
+        }                                                                                                             \
+        uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS];                                                                    \
+        for (uint32_t s = 0; s < n_ranks; s++) {                                                                      \
+            s_spad[s] = (uint8_t *) dma_queue_pop(q).dst;                                                             \
+        }                                                                                                             \
+        uint8_t * r_spad = HAS_ADD ? (uint8_t *) dma_queue_pop(q).dst : NULL;                                         \
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);                                             \
+        HVX_ADD_FN(d_spad, s_spad[0], s_spad[1], cur_elems);                                                          \
+        for (uint32_t s = 2; s < n_ranks; s++) {                                                                      \
+            HVX_ADD_FN(d_spad, d_spad, s_spad[s], cur_elems);                                                         \
+        }                                                                                                             \
+        if (HAS_ADD) {                                                                                                \
+            HVX_ADD_FN(d_spad, d_spad, r_spad, cur_elems);                                                            \
+        }                                                                                                             \
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);                                              \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                       \
+            uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir * sizeof(TYPE);                                    \
+            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 1);                       \
+        }                                                                                                             \
+        if (ir_prefetch < ir1) {                                                                                      \
+            uint32_t next_elems = MIN(block_elems, ir1 - ir_prefetch);                                                \
+            size_t   next_bytes = next_elems * sizeof(TYPE);                                                          \
+            for (uint32_t s = 0; s < n_ranks; s++) {                                                                  \
+                const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE);           \
+                dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), next_bytes, next_bytes, next_bytes, 1);            \
+            }                                                                                                         \
+            if (HAS_ADD) {                                                                                            \
+                const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE); \
+                dma_queue_push(q, dma_make_ptr(r_spad, r_next), next_bytes, next_bytes, next_bytes, 1);               \
+            }                                                                                                         \
+            ir_prefetch += next_elems;                                                                                \
+        }                                                                                                             \
+        ir += cur_elems;                                                                                              \
+    }                                                                                                                 \
+    dma_queue_flush(q);                                                                                               \
+}
+
+DEFINE_ALLREDUCE_THREAD_DMA_1D(f16,     __fp16, hvx_add_f16_aaa, 0)
+DEFINE_ALLREDUCE_THREAD_DMA_1D(f32,     float,  hvx_add_f32_aaa, 0)
+DEFINE_ALLREDUCE_THREAD_DMA_1D(add_f16, __fp16, hvx_add_f16_aaa, 1)
+DEFINE_ALLREDUCE_THREAD_DMA_1D(add_f32, float,  hvx_add_f32_aaa, 1)
+
+#define DEFINE_ALLREDUCE_THREAD_DMA_2D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD, IS_ROW_BCAST)                                                           \
+static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                                                   \
+    struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data;                                                                  \
+    struct htp_ops_context * octx = actx->octx;                                                                                                   \
+                                                                                                                                                  \
+    const uint32_t n_ranks          = actx->n_ranks;                                                                                              \
+    const uint32_t n_dsts           = actx->n_dsts;                                                                                               \
+    const uint32_t ne0              = actx->ne0;                                                                                                  \
+    const uint32_t block_rows       = actx->block_elems;                                                                                          \
+    const uint32_t row_size_aligned = actx->row_size_aligned;                                                                                     \
+    const uint32_t row_bytes        = ne0 * sizeof(TYPE);                                                                                         \
+                                                                                                                                                  \
+    const uint32_t dr  = actx->elems_per_thread;                                                                                                  \
+    const uint32_t r0  = actx->rank_elem_start + dr * ith;                                                                                        \
+    const uint32_t r1  = MIN(r0 + dr, actx->rank_elem_start + actx->rank_nelem);                                                                  \
+    if (r0 >= r1) return;                                                                                                                         \
+                                                                                                                                                  \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                                                        \
+    dma_queue * q = octx->ctx->dma[ith];                                                                                                          \
+                                                                                                                                                  \
+    uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];                                                                                             \
+    for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                      \
+        src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread);                                                           \
+    }                                                                                                                                             \
+    uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread);                                                           \
+    uint8_t * res_spad_base = HAS_ADD ? (IS_ROW_BCAST ? actx->res_spad_base : (actx->res_spad_base + (ith * actx->vtcm_size_per_thread))) : NULL; \
+                                                                                                                                                  \
+    const size_t spad_half = actx->vtcm_size_per_thread / 2;                                                                                      \
+    uint32_t r_prefetch = r0;                                                                                                                     \
+    int spad_idx = 0;                                                                                                                             \
+                                                                                                                                                  \
+    for (int k = 0; k < 2 && r_prefetch < r1; k++) {                                                                                              \
+        uint32_t cur_rows = MIN(block_rows, r1 - r_prefetch);                                                                                     \
+        uint8_t * d_spad = dst_spad_base + spad_idx * spad_half;                                                                                  \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                                                   \
+            uint8_t * d_ddr  = (uint8_t *) octx->dsts[d]->data + r_prefetch * octx->dsts[d]->nb[1];                                               \
+            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, 0);                                 \
+        }                                                                                                                                         \
+        for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                  \
+            uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half;                                                                           \
+            const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1];                                      \
+            dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), row_size_aligned, octx->src[s]->nb[1], row_bytes, cur_rows);                           \
+        }                                                                                                                                         \
+        if (HAS_ADD && !IS_ROW_BCAST) {                                                                                                           \
+            uint8_t * r_spad = res_spad_base + spad_idx * spad_half;                                                                              \
+            const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1];                  \
+            dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, cur_rows);                 \
+        }                                                                                                                                         \
+        r_prefetch += cur_rows;                                                                                                                   \
+        spad_idx ^= 1;                                                                                                                            \
+    }                                                                                                                                             \
+                                                                                                                                                  \
+    for (uint32_t r = r0; r < r1; ) {                                                                                                             \
+        uint32_t cur_rows = MIN(block_rows, r1 - r);                                                                                              \
+        uint8_t * d_spad = NULL;                                                                                                                  \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                                                   \
+            d_spad = (uint8_t *) dma_queue_pop(q).src;                                                                                            \
+        }                                                                                                                                         \
+        uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS];                                                                                                \
+        for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                  \
+            s_spad[s] = (uint8_t *) dma_queue_pop(q).dst;                                                                                         \
+        }                                                                                                                                         \
+        uint8_t * r_spad = (HAS_ADD && !IS_ROW_BCAST) ? (uint8_t *) dma_queue_pop(q).dst : NULL;                                                  \
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);                                                                          \
+        for (uint32_t row = 0; row < cur_rows; row++) {                                                                                           \
+            uint8_t * d_row = d_spad + row * row_size_aligned;                                                                                    \
+            const uint8_t * s0_row = s_spad[0] + row * row_size_aligned;                                                                          \
+            const uint8_t * s1_row = s_spad[1] + row * row_size_aligned;                                                                          \
+            HVX_ADD_FN(d_row, s0_row, s1_row, ne0);                                                                                               \
+            for (uint32_t s = 2; s < n_ranks; s++) {                                                                                              \
+                const uint8_t * ss_row = s_spad[s] + row * row_size_aligned;                                                                      \
+                HVX_ADD_FN(d_row, d_row, ss_row, ne0);                                                                                            \
+            }                                                                                                                                     \
+            if (HAS_ADD) {                                                                                                                        \
+                const uint8_t * res_row = IS_ROW_BCAST ? res_spad_base : (r_spad + row * row_size_aligned);                                       \
+                HVX_ADD_FN(d_row, d_row, res_row, ne0);                                                                                           \
+            }                                                                                                                                     \
+        }                                                                                                                                         \
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);                                                                           \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                                                   \
+            uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + r * octx->dsts[d]->nb[1];                                                         \
+            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, cur_rows);                          \
+        }                                                                                                                                         \
+        if (r_prefetch < r1) {                                                                                                                    \
+            uint32_t next_rows = MIN(block_rows, r1 - r_prefetch);                                                                                \
+            for (uint32_t s = 0; s < n_ranks; s++) {                                                                                              \
+                const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1];                                 \
+                dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), row_size_aligned, octx->src[s]->nb[1], row_bytes, next_rows);                  \
+            }                                                                                                                                     \
+            if (HAS_ADD && !IS_ROW_BCAST) {                                                                                                       \
+                const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1];             \
+                dma_queue_push(q, dma_make_ptr(r_spad, r_next), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, next_rows);           \
+            }                                                                                                                                     \
+            r_prefetch += next_rows;                                                                                                              \
+        }                                                                                                                                         \
+        r += cur_rows;                                                                                                                            \
+    }                                                                                                                                             \
+    dma_queue_flush(q);                                                                                                                           \
+}
+
+DEFINE_ALLREDUCE_THREAD_DMA_2D(f16,           __fp16, hvx_add_f16_aaa, 0, 0)
+DEFINE_ALLREDUCE_THREAD_DMA_2D(f32,           float,  hvx_add_f32_aaa, 0, 0)
+DEFINE_ALLREDUCE_THREAD_DMA_2D(add_f16,       __fp16, hvx_add_f16_aaa, 1, 0)
+DEFINE_ALLREDUCE_THREAD_DMA_2D(add_f32,       float,  hvx_add_f32_aaa, 1, 0)
+DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f16, __fp16, hvx_add_f16_aaa, 1, 1)
+DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f32, float,  hvx_add_f32_aaa, 1, 1)
+
+int op_allreduce(struct htp_ops_context * octx) {
+    const struct htp_allreduce_kernel_params * kparams = (const struct htp_allreduce_kernel_params *) octx->kernel_params;
+    const struct htp_tensor * dst = octx->dst;
+
+    const uint32_t rank    = (uint32_t) kparams->rank;
+    const uint32_t n_ranks = (uint32_t) kparams->n_ranks;
+
+    if (n_ranks < 2 || n_ranks > HTP_ALLREDUCE_MAX_RANKS || rank >= n_ranks) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    if (dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_F32) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    const uint32_t nelem = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3];
+    const uint32_t fence_seq_entry = (uint32_t) octx->op_params[0];
+    const uint32_t fence_seq_exit  = (uint32_t) octx->op_params[1];
+
+    // 1. Entry Barrier: Synchronize all ranks before reading
+    struct htp_thread_trace * tr0 = &octx->ctx->trace[0];
+    htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
+
+    const struct htp_tensor * my_sync = octx->src[n_ranks + rank];
+    atomic_uint * my_fence = (atomic_uint *) my_sync->data;
+
+    atomic_store(&my_fence[0], fence_seq_entry);
+    asm volatile ("syncht" : : : "memory");
+    Q6_dccleaninva_A((void *) my_fence);
+
+    for (uint32_t j = 0; j < n_ranks; j++) {
+        if (j == rank) continue;
+        const struct htp_tensor * peer_sync = octx->src[n_ranks + j];
+        atomic_uint * peer_fence = (atomic_uint *) peer_sync->data;
+        uint64_t spins = 0;
+        while (1) {
+            Q6_dccleaninva_A((void *) peer_fence);
+            uint32_t val = atomic_load(&peer_fence[0]);
+            if (val == fence_seq_entry || val == fence_seq_exit) {
+                break;
+            }
+            if (++spins > HTP_FENCE_TIMEOUT) {
+                FARF(ERROR, "ggml-hex: allreduce entry fence-wait TIMEOUT: rank %u waiting on %u (fence %p seq %u)\n", rank, j, peer_fence, fence_seq_entry);
+                return HTP_STATUS_INTERNAL_ERR;
+            }
+            hex_pause();
+        }
+    }
+    asm volatile ("syncht" : : : "memory");
+
+    htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
+
+    // 2. Multi-threaded Reduction across assigned rank chunk
+    if (nelem > 0) {
+        const uint32_t n_threads            = (uint32_t) kparams->n_threads;
+        const uint32_t block_elems          = (uint32_t) kparams->block_elems;
+        const uint32_t elems_per_thread     = (uint32_t) kparams->elems_per_thread;
+        const uint32_t vtcm_size_per_thread = (uint32_t) kparams->vtcm_size_per_thread;
+
+        const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD);
+
+        struct htp_allreduce_context actx;
+        actx.octx                 = octx;
+        actx.n_ranks              = n_ranks;
+        actx.n_dsts               = (uint32_t) kparams->n_dsts ? (uint32_t) kparams->n_dsts : n_ranks;
+        actx.nelem                = nelem;
+        actx.ne0                  = (uint32_t) kparams->ne0;
+        actx.ne1                  = (uint32_t) kparams->ne1;
+        actx.row_size_aligned     = (uint32_t) kparams->row_size_aligned;
+        actx.rank_elem_start      = (uint32_t) kparams->rank_elem_start;
+        actx.rank_nelem           = (uint32_t) kparams->rank_nelem;
+        actx.elems_per_thread     = elems_per_thread;
+        actx.block_elems          = block_elems;
+        actx.vtcm_size_per_thread = vtcm_size_per_thread;
+        actx.is_row_bcast         = (kparams->is_row_bcast != 0);
+
+        work_queue_func_t reduce_fun = NULL;
+        switch (kparams->kernel_type) {
+            case HTP_ALLREDUCE_KERNEL_DMA_1D:
+                if (has_add) {
+                    reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_1d_add_f16 : allreduce_thread_dma_1d_add_f32;
+                } else {
+                    reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_1d_f16 : allreduce_thread_dma_1d_f32;
+                }
+                break;
+            case HTP_ALLREDUCE_KERNEL_DMA_2D:
+                if (has_add) {
+                    if (kparams->is_row_bcast) {
+                        reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_add_bcast_f16 : allreduce_thread_dma_2d_add_bcast_f32;
+                    } else {
+                        reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_add_f16 : allreduce_thread_dma_2d_add_f32;
+                    }
+                } else {
+                    reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_f16 : allreduce_thread_dma_2d_f32;
+                }
+                break;
+            default:
+                return HTP_STATUS_NO_SUPPORT;
+        }
+
+        uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base;
+        for (uint32_t s = 0; s < n_ranks; s++) {
+            actx.src_spad_base[s] = vtcm_ptr;
+            vtcm_ptr += n_threads * vtcm_size_per_thread;
+        }
+        actx.dst_spad_base = vtcm_ptr;
+        vtcm_ptr += n_threads * vtcm_size_per_thread;
+        if (has_add) {
+            actx.res_spad_base = vtcm_ptr;
+            vtcm_ptr += (actx.is_row_bcast ? 1 : n_threads) * vtcm_size_per_thread;
+        }
+
+        if (has_add && actx.is_row_bcast) {
+            const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data;
+            const uint32_t row_bytes = actx.ne0 * (dst->type == HTP_TYPE_F16 ? sizeof(__fp16) : sizeof(float));
+            dma_queue * q = octx->ctx->dma[0];
+            dma_queue_push(q, dma_make_ptr(actx.res_spad_base, r_ddr), actx.row_size_aligned, 0, row_bytes, 1);
+            dma_queue_pop(q);
+        }
+
+        work_queue_run(octx->ctx->work_queue, reduce_fun, &actx, n_threads);
+    }
+
+    // 4. Exit Barrier: Synchronize all ranks after writing
+    htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
+
+    atomic_store(&my_fence[0], fence_seq_exit);
+    asm volatile ("syncht" : : : "memory");
+    Q6_dccleaninva_A((void *) my_fence);
+
+    for (uint32_t j = 0; j < n_ranks; j++) {
+        if (j == rank) continue;
+        const struct htp_tensor * peer_sync = octx->src[n_ranks + j];
+        atomic_uint * peer_fence = (atomic_uint *) peer_sync->data;
+        uint64_t spins = 0;
+        while (1) {
+            Q6_dccleaninva_A((void *) peer_fence);
+            uint32_t val = atomic_load(&peer_fence[0]);
+            if (val == fence_seq_exit) {
+                break;
+            }
+            if (++spins > HTP_FENCE_TIMEOUT) {
+                FARF(ERROR, "ggml-hex: allreduce exit fence-wait TIMEOUT: rank %u waiting on %u (fence %p seq %u)\n", rank, j, peer_fence, fence_seq_exit);
+                return HTP_STATUS_INTERNAL_ERR;
+            }
+            hex_pause();
+        }
+    }
+    asm volatile ("syncht" : : : "memory");
+
+    htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
+
+    return HTP_STATUS_OK;
+}

ggml/src/ggml-hexagon/htp/allreduce-ops.h

@@ -0,0 +1,40 @@
+#ifndef ALLREDUCE_OPS_H
+#define ALLREDUCE_OPS_H
+
+#include <stdint.h>
+
+#define HTP_ALLREDUCE_MAX_RANKS 4
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+enum htp_allreduce_kernel_type {
+    HTP_ALLREDUCE_KERNEL_UNSUPPORTED = 0,
+    HTP_ALLREDUCE_KERNEL_DMA_1D,
+    HTP_ALLREDUCE_KERNEL_DMA_2D,
+};
+
+struct htp_allreduce_kernel_params {
+    int32_t rank;
+    int32_t n_ranks;
+    int32_t n_threads;
+    int32_t block_elems;          // 1D: block_elems, 2D: block_rows
+    int32_t elems_per_thread;     // 1D: nelem_per_thread, 2D: nrows_per_thread
+    int32_t vtcm_size_per_thread;
+    int32_t vtcm_size;
+    int32_t kernel_type;
+    int32_t ne0;
+    int32_t ne1;
+    int32_t row_size_aligned;
+    int32_t rank_elem_start;
+    int32_t rank_nelem;
+    int32_t n_dsts;
+    int32_t is_row_bcast;
+};
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* ALLREDUCE_OPS_H */

ggml/src/ggml-hexagon/htp/cpy-ops.c

@@ -4,6 +4,7 @@
 
 #include <HAP_farf.h>
 #include <HAP_perf.h>
+#include <qurt_memory.h>
 
 #include <math.h>
 #include <string.h>
@@ -14,6 +15,7 @@
 #include "htp-ops.h"
 #include "htp-ops.h"
 #include "hvx-utils.h"
+#include "htp-tensor.h"
 
 struct htp_copy_context {
     struct htp_ops_context * octx;
@@ -78,7 +80,7 @@ static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, vo
     }                                                                                                          \
 }
 
-DEFINE_CPY_SAMESHAPE(f32, float, 4)
+DEFINE_CPY_SAMESHAPE(f32,  float, 4)
 DEFINE_CPY_SAMESHAPE(f16, __fp16, 2)
 
 #define DEFINE_CPY_RESHAPE(NAME, ELEM_TYPE, ELEM_SIZE)                                                         \
@@ -179,7 +181,7 @@ static void cpy_thread_##NAME##_reshape(unsigned int nth, unsigned int ith, void
     }                                                                                                          \
 }
 
-DEFINE_CPY_RESHAPE(f32, float, 4)
+DEFINE_CPY_RESHAPE(f32,  float, 4)
 DEFINE_CPY_RESHAPE(f16, __fp16, 2)
 
 static void cpy_thread_f16_f32_sameshape(unsigned int nth, unsigned int ith, void * data) {
@@ -232,6 +234,41 @@ static void cpy_thread_f32_f16_sameshape(unsigned int nth, unsigned int ith, voi
     }
 }
 
+static inline void cpy_dma_sametype_sameshape(
+    struct htp_ops_context * octx,
+    const struct htp_tensor * dst,
+    const struct htp_tensor * src0,
+    uint32_t elem_size,
+    uint32_t ne00, uint32_t ne01, uint32_t ne02, uint32_t ne03,
+    uint32_t nb01, uint32_t nb02, uint32_t nb03,
+    uint32_t  nb1, uint32_t  nb2, uint32_t nb3
+) {
+    const bool contiguous_outer =
+        (ne02 == 1 || (nb02 == ne01 * nb01 && nb2 == ne01 * nb1)) &&
+        (ne03 == 1 || (nb03 == ne02 * nb02 && nb3 == ne02 * nb2));
+
+    dma_queue * q = octx->ctx->dma[0];
+
+    if (contiguous_outer) {
+        dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data), nb1, nb01, ne00 * elem_size, ne01 * ne02 * ne03);
+        dma_queue_pop(q);
+        return;
+    }
+
+    for (uint32_t i03 = 0; i03 < ne03; i03++) {
+        for (uint32_t i02 = 0; i02 < ne02; i02++) {
+            uint8_t* dst_ptr  = (uint8_t*) dst->data  + i02*nb2  + i03*nb3;
+            uint8_t* src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03;
+            if (!dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01)) {
+                dma_queue_flush(q);
+                dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01);
+            }
+        }
+    }
+
+    dma_queue_flush(q);
+}
+
 int op_cpy(struct htp_ops_context * octx) {
     cpy_preamble;
 
@@ -264,14 +301,11 @@ int op_cpy(struct htp_ops_context * octx) {
 
     ct.src0_nrows_per_thread = (nr + n_threads - 1) / n_threads;
 
-    worker_callback_t copy_fun;
+    worker_callback_t copy_fun = NULL;
+    bool use_dma = false;
 
     if (sametype && sameshape) {
-        if (src0->type == HTP_TYPE_F32) {
-            copy_fun = cpy_thread_f32_sameshape;
-        } else {
-            copy_fun = cpy_thread_f16_sameshape;
-        }
+        use_dma = true;
     } else if (sameshape) {
         /**/ if (dst->type == HTP_TYPE_F16 && src0->type == HTP_TYPE_F32)
             copy_fun = cpy_thread_f16_f32_sameshape;
@@ -289,7 +323,28 @@ int op_cpy(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    worker_pool_run_func(octx->ctx->worker_pool, copy_fun, &ct, n_threads);
+    if (use_dma) {
+        cpy_dma_sametype_sameshape(octx, dst, src0, ct.src0_type_size, ne00, ne01, ne02, ne03, nb01, nb02, nb03, nb1, nb2, nb3);
+    } else {
+        worker_pool_run_func(octx->ctx->worker_pool, copy_fun, &ct, n_threads);
+    }
+
+    const struct htp_tensor *sync = octx->src[1];
+    if (sync) {
+        if (!use_dma) {
+            // htp_tensor_flush_all(octx->ctx, octx->dsts, 1);
+            qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
+        }
+
+        atomic_uint * sync_fence = (atomic_uint *) sync->data;
+        const uint32_t seq = (uint32_t) octx->op_params[0];
+
+        atomic_store(&sync_fence[0], seq);
+        asm volatile ("syncht" : : : "memory");
+        Q6_dccleaninva_A((void *) sync_fence);
+
+        FARF(HIGH, "ggml-hex: sync-release : fence %p seq %u\n", sync_fence, seq);
+    }
 
     return HTP_STATUS_OK;
 }

ggml/src/ggml-hexagon/htp/dma-queue.h

@@ -244,17 +244,18 @@ static inline dma_ptr dma_queue_pop(dma_queue * q) {
         return dptr;
     }
 
-    dma_descriptor_2d * desc = &r->desc[r->pop_idx];
+    dptr = r->dptr[r->pop_idx];
+
+    volatile dma_descriptor_2d * desc = &r->desc[r->pop_idx];
 
     // Wait for desc to complete
     if (!desc->done) {
+        // FARF(ALWAYS, "dma-poll: idx %u dst %p src %p", r->pop_idx, dptr.dst, dptr.src);
         while (!desc->done) {
             dmpoll();
         }
     }
 
-    dptr = r->dptr[r->pop_idx];
-
     htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);
 
     r->pop_idx = (r->pop_idx + 1) & r->idx_mask;

ggml/src/ggml-hexagon/htp/flash-attn-ops.c

@@ -30,6 +30,8 @@
 #include "ggml-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"
+#include "hvx-quant.h"
 
 #include "flash-attn-ops.h"
 #include "hvx-fa-kernels.h"
@@ -85,12 +87,17 @@ struct htp_fa_context {
     uint8_t * spad_m;
     uint8_t * spad_a;
 
+    const struct htp_tensor * k;
+    const struct htp_tensor * v;
+
     uint64_t t_start;
 };
 
 struct hmx_fa_context {
     const struct htp_ops_context * octx;
     const struct htp_tensor *      sinks;  // attention sinks (src[4]), NULL if absent
+    const struct htp_tensor *      k;
+    const struct htp_tensor *      v;
     bool         pipeline;  // true when n_kv_blocks >= FA_MIN_KV_BLOCKS && n_threads >= 2
     uint32_t     n_threads;
 
@@ -214,8 +221,8 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
     const uint32_t DV = nev0;
 
     const size_t size_q_row = DK * ((q->type == HTP_TYPE_F32) ? 4 : 2);
-    const size_t size_k_row = DK * sizeof(__fp16);
-    const size_t size_v_row = DV * sizeof(__fp16);
+    const size_t size_k_row = htp_tensor_get_row_size(k->type, DK);
+    const size_t size_v_row = htp_tensor_get_row_size(v->type, DV);
 
     // Scratchpad buffers for Q, K, V, Mask, and VKQ32 accumulator
     uint8_t * spad_q = factx->spad_q + factx->size_q_block * ith;
@@ -364,6 +371,23 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
             uint8_t * v_base = dma_queue_pop(dma).dst; // V
             __fp16  * m_base = mask ? dma_queue_pop(dma).dst : NULL; // M
 
+            if (factx->k->type == HTP_TYPE_Q8_0) {
+                htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir);
+                for (uint32_t r = 0; r < current_block_size; ++r) {
+                    __fp16 * row_k = (__fp16 *)(k_base + r * factx->size_k_row_padded);
+                    hvx_dequantize_row_q8_0_f16(row_k, row_k, DK);
+                }
+                htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir);
+            }
+            if (factx->v->type == HTP_TYPE_Q8_0) {
+                htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, ir);
+                for (uint32_t r = 0; r < current_block_size; ++r) {
+                    __fp16 * row_v = (__fp16 *)(v_base + r * factx->size_v_row_padded);
+                    hvx_dequantize_row_q8_0_f16(row_v, row_v, DV);
+                }
+                htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, ir);
+            }
+
             htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_QK, ir);
 
             // Inner loop processing the block from VTCM
@@ -625,6 +649,12 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data)
 
     struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
+    if (factx->k->type == HTP_TYPE_Q8_0) {
+        for (uint32_t r = start; r < end; ++r) {
+            __fp16 * row_k = (__fp16 *)((char *)args->curr_k + r * args->src_stride * sizeof(__fp16));
+            hvx_dequantize_row_q8_0_f16(row_k, row_k, factx->DK);
+        }
+    }
     hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles[args->buf_idx], (const __fp16 *) args->curr_k, total_rows, factx->DK,
                              args->src_stride, start, end);
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
@@ -673,6 +703,12 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data)
 
     struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start));
+    if (factx->v->type == HTP_TYPE_Q8_0) {
+        for (uint32_t r = start; r < end; ++r) {
+            __fp16 * row_v = (__fp16 *)((char *)args->v_src + r * args->src_stride * sizeof(__fp16));
+            hvx_dequantize_row_q8_0_f16(row_v, row_v, factx->DV);
+        }
+    }
     hmx_interleave_cols_to_tiles(v_tiles_dst, (const __fp16 *) args->v_src, total_rows, factx->DV,
                              args->src_stride, (uint32_t) args->n_col_tiles, start, end);
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start));
@@ -1809,6 +1845,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     memset(&factx, 0, sizeof(factx));
     factx.octx           = octx;
     factx.sinks          = octx->src[4];  // NULL if this op has no attention sinks
+    factx.k              = k;
+    factx.v              = v;
     factx.n_threads      = kparams->n_threads;
     factx.DK             = DK;
     factx.DV             = DV;
@@ -1853,10 +1891,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     // ======== VTCM allocation (GQA-aware) ========
     // K/V row sizes drive the DMA descriptors (not the VTCM layout) and are used
     // throughout the KV loop below.
-    const size_t size_k_row        = DK * sizeof(__fp16);
-    const size_t size_v_row        = DV * sizeof(__fp16);
-    const size_t size_k_row_padded = hex_round_up(size_k_row, 128);
-    const size_t size_v_row_padded = hex_round_up(size_v_row, 128);
+    const size_t size_k_row        = htp_tensor_get_row_size(k->type, DK);
+    const size_t size_v_row        = htp_tensor_get_row_size(v->type, DV);
+    const size_t size_k_row_padded = hex_round_up(DK * sizeof(__fp16), 128);
+    const size_t size_v_row_padded = hex_round_up(DV * sizeof(__fp16), 128);
 
     // Build the VTCM layout once (shared with the host estimator) and place every
     // scratch buffer at its computed offset.
@@ -2348,7 +2386,9 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
     const struct htp_tensor * dst  = octx->dst;
 
     // Check support
-    if ((q->type != HTP_TYPE_F16 && q->type != HTP_TYPE_F32) || k->type != HTP_TYPE_F16 || v->type != HTP_TYPE_F16) {
+    if ((q->type != HTP_TYPE_F16 && q->type != HTP_TYPE_F32) ||
+        (k->type != HTP_TYPE_F16 && k->type != HTP_TYPE_Q8_0) ||
+        (v->type != HTP_TYPE_F16 && v->type != HTP_TYPE_Q8_0)) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
@@ -2364,6 +2404,8 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
 
     struct htp_fa_context factx;
     factx.octx = octx;
+    factx.k = k;
+    factx.v = v;
 
     factx.t_start = HAP_perf_get_qtimer_count();