tests : disable DOTS3NOTE arch test for WebGPU (#27654)

Co-authored-by: Stanisław Szymczyk <sszymczy@gmail.com>

.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

@@ -73,6 +73,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 +119,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 +183,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 +208,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 +238,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

@@ -72,6 +72,16 @@ jobs:
             -DGGML_CUDA_CUB_3DOT2=ON
           cmake --build build
 
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: cuda-ubuntu-24.04-cuda
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
+
   hip:
     runs-on: ubuntu-22.04
     container: rocm/dev-ubuntu-22.04:6.1.2
@@ -103,6 +113,16 @@ jobs:
             -DGGML_HIP=ON
           cmake --build build --config Release -j $(nproc)
 
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: cuda-ubuntu-22.04-hip
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
+
   musa:
     runs-on: ubuntu-22.04
     container: mthreads/musa:rc4.3.0-devel-ubuntu22.04-amd64
@@ -131,3 +151,13 @@ jobs:
           cmake -B build -S . \
             -DGGML_MUSA=ON
           time cmake --build build --config Release -j $(nproc)
+
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: cuda-ubuntu-22.04-musa
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}

.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/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/release.yml

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

.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' }}

.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
 

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"

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);
             }
         }
 
@@ -2238,7 +2238,7 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
 
                 std::set<std::string> required;
                 if (params.contains("required")) {
-                    params.at("required").get_to(required);
+                    required = params.at("required").get<std::set<std::string>>();
                 }
 
                 auto schema_info = common_schema_info();
@@ -2860,7 +2860,7 @@ static common_chat_params common_chat_params_init_minimax_m3(const common_chat_t
 
                 std::set<std::string> required;
                 if (schema.contains("required")) {
-                    schema.at("required").get_to(required);
+                    required = schema.at("required").get<std::set<std::string>>();
                 }
 
                 std::vector<common_peg_parser> required_elements;
@@ -2972,10 +2972,10 @@ static void system_message_not_supported(json & messages) {
             auto & second_msg = messages[1];
             second_msg["content"] = first_msg.at("content").get<std::string>()
                 + "\n" + second_msg.at("content").get<std::string>();
-            messages.erase(messages.begin());
+            messages.erase(0);
         } else {
             LOG_WRN("Removing system prompt due to template not supporting system role\n");
-            messages.erase(messages.begin());
+            messages.erase(0);
         }
     }
 }

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/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,6 +58,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "DSparkDraftModel": "qwen",
     "DSparkSpeculator": "qwen",
     "Lfm2DSparkDraftModel": "qwen",
+    "LingDSparkModel": "qwen",
     "DeepseekV4ForCausalLM": "deepseek",
     "DeepseekV4DSparkModel": "deepseek",
     "DistilBertForMaskedLM": "bert",
@@ -282,6 +283,8 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
     "CogVLMForCausalLM": "cogvlm",
     "DeepseekOCR2ForCausalLM": "deepseek",
     "DeepseekOCRForCausalLM": "deepseek",
+    "Dots3NoteForCausalLM": "dots3",
+    "Dots3NoteForConditionalGeneration": "dots3",
     "DotsOCRForCausalLM": "dotsocr",
     "Exaone4_5_ForConditionalGeneration": "exaone",
     "Gemma3ForConditionalGeneration": "gemma",

conversion/dots3.py

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

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/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).
 

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-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;
@@ -1255,6 +1305,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 +1504,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 +1654,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,
@@ -1841,7 +2007,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 +2107,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-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/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-metal/ggml-metal-device.cpp

@@ -17,10 +17,10 @@ struct ggml_metal_device_deleter {
 
 typedef std::unique_ptr<ggml_metal_device, ggml_metal_device_deleter> ggml_metal_device_ptr;
 
-ggml_metal_device_t ggml_metal_device_get(int device) {
+ggml_metal_device_t ggml_metal_device_get(int device, int n_devices) {
     static std::vector<ggml_metal_device_ptr> devs;
 
-    devs.emplace_back(ggml_metal_device_init(device));
+    devs.emplace_back(ggml_metal_device_init(device, n_devices));
 
     return devs.back().get();
 }

ggml/src/ggml-metal/ggml-metal-device.h

@@ -259,6 +259,8 @@ enum ggml_metal_device_id {
 
 struct ggml_metal_device_props {
     int device;
+    int device_phys;
+    int device_virt;
     char name[128];
     char desc[128];
 
@@ -286,10 +288,10 @@ typedef struct ggml_metal_event * ggml_metal_event_t;
 void ggml_metal_event_encode_signal(ggml_metal_event_t ev, ggml_metal_cmd_buf_t cmd_buf);
 void ggml_metal_event_encode_wait  (ggml_metal_event_t ev, ggml_metal_cmd_buf_t cmd_buf);
 
-ggml_metal_device_t ggml_metal_device_init(int device);
+ggml_metal_device_t ggml_metal_device_init(int device, int n_devices);
 void ggml_metal_device_free(ggml_metal_device_t dev);
 
-ggml_metal_device_t ggml_metal_device_get(int device);
+ggml_metal_device_t ggml_metal_device_get(int device, int n_devices);
 
 void * ggml_metal_device_get_obj  (ggml_metal_device_t dev); // id<MTLDevice>
 void * ggml_metal_device_get_queue(ggml_metal_device_t dev); // id<MTLCommandQueue>

ggml/src/ggml-metal/ggml-metal-device.m

@@ -711,7 +711,7 @@ static enum ggml_metal_device_id ggml_metal_device_id_parse(const char * name) {
     return GGML_METAL_DEVICE_GENERIC;
 }
 
-ggml_metal_device_t ggml_metal_device_init(int device) {
+ggml_metal_device_t ggml_metal_device_init(int device, int n_devices) {
     ggml_metal_device_t dev = calloc(1, sizeof(struct ggml_metal_device));
 
     assert(dev != NULL);
@@ -728,6 +728,12 @@ ggml_metal_device_t ggml_metal_device_init(int device) {
             dev->addr_virt = 0x000000400ULL;
 
             dev->props.device = device;
+
+            // the Metal backend uses the system default device as the single physical device;
+            // additional (virtual) devices are emulated on top of it via GGML_METAL_DEVICES
+            dev->props.device_phys = 0;
+            dev->props.device_virt = device;
+
             dev->props.has_simdgroup_reduction  = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7];
             dev->props.has_simdgroup_reduction |= [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML];
 
@@ -891,7 +897,13 @@ ggml_metal_device_t ggml_metal_device_init(int device) {
             }
 
             snprintf(dev->props.name, sizeof(dev->props.name), "%s%d", "MTL", device);
-            snprintf(dev->props.desc, sizeof(dev->props.desc), "%s", [[dev->mtl_device name] UTF8String]);
+            const char * gpu_name = [[dev->mtl_device name] UTF8String];
+            if (n_devices > 1) {
+                snprintf(dev->props.desc, sizeof(dev->props.desc), "%s (dev p%d/v%d)",
+                         gpu_name, dev->props.device_phys, dev->props.device_virt);
+            } else {
+                snprintf(dev->props.desc, sizeof(dev->props.desc), "%s", gpu_name);
+            }
 
             dev->library = ggml_metal_library_init(dev);
             if (!dev->library) {

ggml/src/ggml-metal/ggml-metal.cpp

@@ -891,7 +891,7 @@ static ggml_backend_dev_t ggml_backend_metal_device_init(ggml_backend_reg_t reg,
     return new ggml_backend_device {
         /* .iface   = */ ggml_backend_metal_device_i,
         /* .reg     = */ reg,
-        /* .context = */ ggml_metal_device_get(device),
+        /* .context = */ ggml_metal_device_get(device, g_devices),
     };
 }
 

ggml/src/ggml-sycl/convert.cpp

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

ggml/src/ggml-sycl/dequantize.hpp

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

ggml/src/ggml-sycl/dmmv.cpp

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

ggml/src/ggml-sycl/esimd.hpp

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

ggml/src/ggml-sycl/ggml-sycl.cpp

@@ -3796,6 +3796,7 @@ inline bool ggml_sycl_supports_reorder_mmvq(enum ggml_type type) {
         case GGML_TYPE_Q1_0:
         case GGML_TYPE_Q4_0:
         case GGML_TYPE_Q8_0:
+        case GGML_TYPE_Q2_K:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q4_K:
         case GGML_TYPE_Q5_K:
@@ -3809,6 +3810,7 @@ inline bool ggml_sycl_supports_reorder_mmvq(enum ggml_type type) {
 static bool ggml_sycl_supports_reorder_esimd(enum ggml_type type) {
 #ifdef GGML_SYCL_DMMV_HAS_ESIMD
     switch (type) {
+        case GGML_TYPE_Q2_K:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q4_K:
         case GGML_TYPE_Q5_K:

ggml/src/ggml-sycl/mmvq.cpp

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

ggml/src/ggml-sycl/quants.hpp

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

ggml/src/ggml-sycl/vecdotq.hpp

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

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -955,6 +955,7 @@ struct vk_device_struct {
     vk_pipeline pipeline_diag[2];
     vk_pipeline pipeline_clamp[2];
     vk_pipeline pipeline_pad_f32;
+    vk_pipeline pipeline_pad_reflect_1d_f32;
     vk_pipeline pipeline_roll_f32;
     vk_pipeline pipeline_repeat_i32, pipeline_repeat_back_f32;
     vk_pipeline pipeline_repeat_i16;
@@ -5630,6 +5631,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     ggml_vk_create_pipeline(device, device->pipeline_diag[1], "diag_f16", diag_f16_len, diag_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
 
     ggml_vk_create_pipeline(device, device->pipeline_pad_f32, "pad_f32", pad_f32_len, pad_f32_data, "main", 2, sizeof(vk_op_pad_push_constants), {512, 1, 1}, {}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_pad_reflect_1d_f32, "pad_reflect_1d_f32", pad_reflect_1d_f32_len, pad_reflect_1d_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
 
     ggml_vk_create_pipeline(device, device->pipeline_roll_f32, "roll_f32", roll_f32_len, roll_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
 
@@ -11336,6 +11338,11 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
             return ctx->device->pipeline_pad_f32;
         }
         return nullptr;
+    case GGML_OP_PAD_REFLECT_1D:
+        if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+            return ctx->device->pipeline_pad_reflect_1d_f32;
+        }
+        return nullptr;
     case GGML_OP_ROLL:
         if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
             return ctx->device->pipeline_roll_f32;
@@ -12239,6 +12246,7 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
     case GGML_OP_CLAMP:
     case GGML_OP_LEAKY_RELU:
     case GGML_OP_PAD:
+    case GGML_OP_PAD_REFLECT_1D:
     case GGML_OP_ROLL:
     case GGML_OP_REPEAT:
     case GGML_OP_REPEAT_BACK:
@@ -13111,6 +13119,17 @@ static void ggml_vk_pad(ggml_backend_vk_context * ctx, vk_context& subctx, const
     ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_PAD, std::move(p));
 }
 
+static void ggml_vk_pad_reflect_1d(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
+    const uint32_t p0 = (uint32_t)dst->op_params[0];
+    const uint32_t p1 = (uint32_t)dst->op_params[1];
+
+    vk_op_unary_push_constants p = vk_op_unary_push_constants_init(src0, dst, ggml_nelements(dst));
+    memcpy(&p.param1, &p0, sizeof(float));
+    memcpy(&p.param2, &p1, sizeof(float));
+
+    ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_PAD_REFLECT_1D, std::move(p));
+}
+
 static void ggml_vk_roll(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
     const int32_t s0 = ggml_get_op_params_i32(dst, 0);
     const int32_t s1 = ggml_get_op_params_i32(dst, 1);
@@ -15520,6 +15539,10 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
     case GGML_OP_PAD:
         ggml_vk_pad(ctx, compute_ctx, src0, node);
 
+        break;
+    case GGML_OP_PAD_REFLECT_1D:
+        ggml_vk_pad_reflect_1d(ctx, compute_ctx, src0, node);
+
         break;
     case GGML_OP_ROLL:
         ggml_vk_roll(ctx, compute_ctx, src0, node);
@@ -18446,6 +18469,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
         case GGML_OP_SCALE:
             return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_PAD:
+        case GGML_OP_PAD_REFLECT_1D:
         case GGML_OP_ROLL:
             return op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_DIAG_MASK_INF:
@@ -19228,6 +19252,8 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
         } else if (tensor->op == GGML_OP_PAD) {
             tensor_clone = ggml_pad_ext(ggml_ctx, src_clone[0], tensor->op_params[0], tensor->op_params[1], tensor->op_params[2], tensor->op_params[3],
                                                                 tensor->op_params[4], tensor->op_params[5], tensor->op_params[6], tensor->op_params[7]);
+        } else if (tensor->op == GGML_OP_PAD_REFLECT_1D) {
+            tensor_clone = ggml_pad_reflect_1d(ggml_ctx, src_clone[0], tensor->op_params[0], tensor->op_params[1]);
         } else if (tensor->op == GGML_OP_REPEAT) {
             tensor_clone = ggml_repeat(ggml_ctx, src_clone[0], tensor);
         } else if (tensor->op == GGML_OP_REPEAT_BACK) {

ggml/src/ggml-vulkan/vulkan-shaders/pad_reflect_1d.comp

@@ -0,0 +1,43 @@
+#version 450
+
+#include "types.glsl"
+#include "generic_unary_head.glsl" // included to use functions like fastdiv etc.
+
+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
+
+void main() {
+
+    const uint idx = get_idx();
+
+    if (idx >= p.ne) {
+        return;
+    }
+
+    const uint p0 = floatBitsToUint(p.param1);
+    const uint p1 = floatBitsToUint(p.param2);
+
+    const uint i3 = fastdiv(idx, p.ne1_012mp, fastdiv_L(p.ne1_Ls, 0));
+    const uint i3_offset = i3 * p.ne12 * p.ne11 * p.ne10;
+
+    const uint i2 = fastdiv(idx - i3_offset, p.ne1_01mp, fastdiv_L(p.ne1_Ls, 1));
+    const uint i2_offset = i2 * p.ne11 * p.ne10;
+
+    const uint i1 = fastdiv(idx - i3_offset - i2_offset, p.ne1_0mp, fastdiv_L(p.ne1_Ls, 2));
+    const uint i0 = idx - i3_offset - i2_offset - i1 * p.ne10;
+
+    uint src_col;
+
+    if (i0 < p0) {
+        src_col = p0 - i0;    // left pad area
+    } else if (i0 < p0 + p.ne00) {
+        src_col = i0 - p0;    // center area
+    } else {
+        src_col = 2u * p.ne00 - 2u - (i0 - p0);  // right pad area
+    }
+
+    const uint src_idx = i3 * p.nb03 + i2 * p.nb02 + i1 * p.nb01 + src_col * p.nb00;
+    const uint d_idx = i3 * p.nb13 + i2 * p.nb12 + i1 * p.nb11 + i0 * p.nb10;
+
+    // copy the computed value to the destination tensor
+    data_d[get_doffset() + d_idx] = D_TYPE(data_a[get_aoffset() + src_idx]);
+}

ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp

@@ -896,6 +896,7 @@ void process_shaders() {
     string_to_spv("scale_f32", "scale.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
 
     string_to_spv("pad_f32", "pad.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
+    string_to_spv("pad_reflect_1d_f32", "pad_reflect_1d.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
 
     string_to_spv("concat_i8", "concat.comp", {{"A_TYPE", "uint8_t"}, {"B_TYPE", "uint8_t"}, {"D_TYPE", "uint8_t"}});
     string_to_spv("concat_i16", "concat.comp", {{"A_TYPE", "uint16_t"}, {"B_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}});

ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl

@@ -5,10 +5,9 @@ enable subgroups;
 enable chromium_experimental_subgroup_matrix;
 
 #define BYTE_HELPERS
-#include "common_decls.tmpl"
-
 #define FLASH_ATTN_SCALAR_KV
 #include "flash_attn_decls.tmpl"
+#include "common_decls.tmpl"
 
 // Default values
 // The actual values are defined in shader-lib.

ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_tile.wgsl

@@ -2,8 +2,8 @@ enable f16;
 enable subgroups;
 
 #define BYTE_HELPERS
-#include "common_decls.tmpl"
 #include "flash_attn_decls.tmpl"
+#include "common_decls.tmpl"
 
 // Default values
 // The actual values are defined in shader-lib.

ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_vec_split.wgsl

@@ -3,9 +3,9 @@ enable f16;
 enable subgroups;
 
 #define BYTE_HELPERS
-#include "common_decls.tmpl"
 #define FLASH_ATTN_VEC_SPLIT
 #include "flash_attn_decls.tmpl"
+#include "common_decls.tmpl"
 
 // Default values
 // The actual values are defined in shader-lib.

ggml/src/ggml.c

@@ -4042,6 +4042,41 @@ struct ggml_tensor * ggml_diag_mask_zero_inplace(
     return ggml_diag_mask_zero_impl(ctx, a, n_past, true);
 }
 
+// ggml_clamp
+
+static struct ggml_tensor * ggml_clamp_impl(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * a,
+        float                 min,
+        float                 max,
+        bool                  inplace) {
+    struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a);
+
+    float params[] = { min, max };
+    ggml_set_op_params(result, params, sizeof(params));
+
+    result->op     = GGML_OP_CLAMP;
+    result->src[0] = a;
+
+    return result;
+}
+
+struct ggml_tensor * ggml_clamp(
+    struct ggml_context * ctx,
+    struct ggml_tensor  * a,
+    float                 min,
+    float                 max) {
+    return ggml_clamp_impl(ctx, a, min, max, false);
+}
+
+struct ggml_tensor * ggml_clamp_inplace(
+    struct ggml_context * ctx,
+    struct ggml_tensor  * a,
+    float                 min,
+    float                 max) {
+    return ggml_clamp_impl(ctx, a, min, max, true);
+}
+
 // ggml_soft_max
 
 static struct ggml_tensor * ggml_soft_max_impl(
@@ -4438,25 +4473,6 @@ struct ggml_tensor * ggml_rope_set_offset(
     return a;
 }
 
-// ggml_clamp
-
-struct ggml_tensor * ggml_clamp(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        float                 min,
-        float                 max) {
-    // TODO: when implement backward, fix this:
-    struct ggml_tensor * result = ggml_view_tensor(ctx, a);
-
-    float params[] = { min, max };
-    ggml_set_op_params(result, params, sizeof(params));
-
-    result->op     = GGML_OP_CLAMP;
-    result->src[0] = a;
-
-    return result;
-}
-
 static int64_t ggml_calc_conv_output_size(int64_t ins, int64_t ks, int s, int p, int d) {
     return (ins + 2 * p - d * (ks - 1) - 1) / s + 1;
 }

gguf-py/gguf/constants.py

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

gguf-py/gguf/gguf_writer.py

@@ -1327,6 +1327,12 @@ class GGUFWriter:
     def add_vision_spatial_merge_size(self, value: int) -> None:
         self.add_uint32(Keys.ClipVision.SPATIAL_MERGE_SIZE, value)
 
+    def add_vision_expert_count_per_layer(self, value: Sequence[int]) -> None:
+        self.add_array(Keys.ClipVision.EXPERT_COUNT_PER_LAYER, value)
+
+    def add_vision_expert_used_count(self, value: int) -> None:
+        self.add_uint32(Keys.ClipVision.EXPERT_USED_COUNT, value)
+
     def add_vision_use_gelu(self, value: bool) -> None:
         self.add_bool(Keys.ClipVision.USE_GELU, value)
 

gguf-py/gguf/tensor_mapping.py

@@ -1454,6 +1454,7 @@ class TensorNameMap:
             "mlp_AR.linear_{bid}", # PaddleOCR-VL
             "merger.mlp.{bid}",
             "vision_tower.merger.mlp.{bid}", # dots.ocr
+            "vision_encoder.adapter.mlp.{bid}", # dots3note
             "vit.perceive.proj.{bid}", # HunyuanVL (proj.0 = conv1, proj.2 = conv2)
         ),
 
@@ -1504,6 +1505,7 @@ class TensorNameMap:
             "vision_model.radio_model.model.patch_generator.embedder", # Nemotron Nano v2 VL
             "model.vision_tower.patch_embedder.input_proj", # gemma4
             "vision_tower.patch_embed.patchifier.proj", # dots.ocr
+            "vision_encoder.patch_embed.proj", # dots3note
             "vision_model.conv1", # Step3-VL
             "model.vision_embedder.patch_dense", # gemma4 unified
             "model.vision_tower.patch_embedder.patch_embedding", # muse-glimmer
@@ -1512,6 +1514,7 @@ class TensorNameMap:
         MODEL_TENSOR.V_ENC_EMBD_NORM: (
             "visual.post_conv_layernorm", # glm4v
             "vision_tower.patch_embed.patchifier.norm", # dots.ocr
+            "vision_encoder.patch_embed.norm", # dots3note
         ),
 
         MODEL_TENSOR.V_ENC_EMBD_PATCH_NORM: (
@@ -1551,6 +1554,7 @@ class TensorNameMap:
         MODEL_TENSOR.V_ENC_ATTN_QKV: (
             "visual.blocks.{bid}.attn.qkv", # qwen3vl
             "vision_tower.blocks.{bid}.attn.qkv", # dots.ocr
+            "vision_encoder.blocks.{bid}.attn.qkv", # dots3note
             "model.vision.transformer.layers.{bid}.attention.query_key_value", # cogvlm
             "model.vision_model.transformer.layers.{bid}.self_attn.qkv_proj", # Deepseek-OCR CLIP
             "vision_tower.encoder.blocks.{bid}.wqkv", # Kimi-K2.5
@@ -1579,6 +1583,7 @@ class TensorNameMap:
         ),
 
         MODEL_TENSOR.V_ENC_ATTN_Q_NORM: (
+            "vision_encoder.blocks.{bid}.attn.q_norm", # dots3note
             "vision_tower.vision_model.encoder.layers.{bid}.attn.q_norm", # InternVL
             "model.vision_tower.encoder.layer.{bid}.attention.q_norm", # Intern-S1
             "visual.blocks.{bid}.attn.q_norm", # GLM-OCR
@@ -1606,6 +1611,7 @@ class TensorNameMap:
         ),
 
         MODEL_TENSOR.V_ENC_ATTN_K_NORM: (
+            "vision_encoder.blocks.{bid}.attn.k_norm", # dots3note
             "vision_tower.vision_model.encoder.layers.{bid}.attn.k_norm", # InternVL
             "model.vision_tower.encoder.layer.{bid}.attention.k_norm", # Intern-S1
             "visual.blocks.{bid}.attn.k_norm", # GLM-OCR
@@ -1651,6 +1657,7 @@ class TensorNameMap:
             "siglip2.vision_model.encoder.layers.{bid}.layer_norm1",
             "vision_model.radio_model.model.blocks.{bid}.norm1", # Nemotron Nano v2 VL
             "vision_tower.blocks.{bid}.norm1", # dots.ocr
+            "vision_encoder.blocks.{bid}.norm_1", # dots3note
             "vision_model.transformer.resblocks.{bid}.ln_1", # Step3-VL
             "model.qwen2_model.model.model.layers.{bid}.input_layernorm", # Deepseek-OCR-2 qwen2
             "model.vision_tower.layers.{bid}.norm1", # muse-glimmer
@@ -1678,6 +1685,7 @@ class TensorNameMap:
             "model.qwen2_model.model.model.layers.{bid}.self_attn.o_proj", # Deepseek-OCR-2 qwen2
             "vision_model.model.layers.{bid}.self_attn.o_proj.linear", # gemma4
             "vision_tower.blocks.{bid}.attn.proj", # dots.ocr
+            "vision_encoder.blocks.{bid}.attn.proj", # dots3note
             "vision_model.transformer.resblocks.{bid}.attn.out_proj", # Step3-VL
             "model.vision_tower.layers.{bid}.attn.proj", # muse-glimmer
         ),
@@ -1706,12 +1714,14 @@ class TensorNameMap:
             "vision_model.radio_model.model.blocks.{bid}.norm2", # Nemotron Nano v2 VL
             "vision_model.model.layers.{bid}.pre_feedforward_layernorm", # gemma4
             "vision_tower.blocks.{bid}.norm2", # dots.ocr
+            "vision_encoder.blocks.{bid}.norm_2", # dots3note
             "vision_model.transformer.resblocks.{bid}.ln_2", # Step3-VL
             "model.qwen2_model.model.model.layers.{bid}.post_attention_layernorm", # Deepseek-OCR-2 qwen2
             "model.vision_tower.layers.{bid}.norm2", # muse-glimmer
         ),
 
         MODEL_TENSOR.V_ENC_FFN_UP: (
+            "vision_encoder.blocks.{bid}.mlp.fc3", # dots3note
             "model.vision_tower.vision_model.encoder.layers.{bid}.mlp.fc1", # Granite4Vision
             "vision_tower.vision_model.encoder.layers.{bid}.mlp.fc1",
             "model.vision_tower.encoder.layers.{bid}.mlp.fc1", # minicpmv4_6
@@ -1737,6 +1747,7 @@ class TensorNameMap:
         ),
 
         MODEL_TENSOR.V_ENC_FFN_GATE: (
+            "vision_encoder.blocks.{bid}.mlp.fc1", # dots3note
             "vision_tower.transformer.layers.{bid}.feed_forward.gate_proj", # pixtral-hf
             "vision_encoder.transformer.layers.{bid}.feed_forward.w1", # pixtral
             "visual.blocks.{bid}.mlp.gate_proj", # qwen2.5vl
@@ -1745,6 +1756,7 @@ class TensorNameMap:
         ),
 
         MODEL_TENSOR.V_ENC_FFN_DOWN: (
+            "vision_encoder.blocks.{bid}.mlp.fc2", # dots3note
             "model.vision_tower.vision_model.encoder.layers.{bid}.mlp.fc2", # Granite4Vision
             "vision_tower.vision_model.encoder.layers.{bid}.mlp.fc2",
             "model.vision_tower.encoder.layers.{bid}.mlp.fc2", # minicpmv4_6
@@ -1769,6 +1781,29 @@ class TensorNameMap:
             "model.vision_tower.layers.{bid}.mlp.fc2", # muse-glimmer
         ),
 
+
+        MODEL_TENSOR.V_ENC_FFN_GATE_INP: (
+            "vision_encoder.blocks.{bid}.mlp.gate_weight", # dots3note
+        ),
+
+        MODEL_TENSOR.V_ENC_FFN_EXP_PROBS_B: (
+            "vision_encoder.blocks.{bid}.mlp.router_bias", # dots3note
+        ),
+
+        # note: expert weights are stacked into a single 3D tensor in conversion code,
+        # which emits the pseudo-names below
+        MODEL_TENSOR.V_ENC_FFN_GATE_EXPS: (
+            "vision_encoder.blocks.{bid}.mlp.experts.fc1", # dots3note
+        ),
+
+        MODEL_TENSOR.V_ENC_FFN_UP_EXPS: (
+            "vision_encoder.blocks.{bid}.mlp.experts.fc3", # dots3note
+        ),
+
+        MODEL_TENSOR.V_ENC_FFN_DOWN_EXPS: (
+            "vision_encoder.blocks.{bid}.mlp.experts.fc2", # dots3note
+        ),
+
         MODEL_TENSOR.V_ENC_ATTN_POST_NORM: (
             "vision_model.model.layers.{bid}.post_attention_layernorm", # gemma4
         ),
@@ -1800,6 +1835,7 @@ class TensorNameMap:
             "vision_model.layernorm_pre", # llama4
             "model.vision_model.pre_layrnorm", # Deepseek-OCR CLIP
             "vision_tower.patch_embed.patchifier.norm", # dots.ocr
+            "vision_encoder.patch_embed.norm", # dots3note
             "vision_model.ln_pre", # Step3-VL
             "model.vision_tower.ln_pre", # muse-glimmer
         ),
@@ -1821,6 +1857,7 @@ class TensorNameMap:
         MODEL_TENSOR.V_MM_POST_NORM: (
             "visual.merger.post_projection_norm", # glm4v
             "vision_tower.post_trunk_norm", # dots.ocr
+            "vision_encoder.post_trunk_norm", # dots3note
             "vit.perceive.after_rms", # HunyuanVL
         ),
 
@@ -1838,6 +1875,7 @@ class TensorNameMap:
             "mlp_AR.pre_norm", # PaddleOCR-VL
             "merger.ln_q",
             "vision_tower.merger.ln_q", # dots.ocr
+            "vision_encoder.adapter.ln_q", # dots3note
             "model.merger.mlp.0.pre_norm", # minicpmv4_6
         ),
 
@@ -2173,10 +2211,12 @@ class TensorNameMap:
 
         MODEL_TENSOR.A_ENC_CONV2D: (
             "audio_tower.conv2d{bid}", # qwen3omni
+            "audio_encoder.dots_encoder.speech_encoder.conv2d{bid}", # dots3note
         ),
 
         MODEL_TENSOR.A_ENC_CONV_OUT: (
             "audio_tower.conv_out", # qwen3omni
+            "audio_encoder.dots_encoder.speech_encoder.conv_out", # dots3note
             "speaker_encoder.mfa.conv", # qwen3tts speaker encoder: multi-layer feature aggregation
         ),
 
@@ -2184,12 +2224,14 @@ class TensorNameMap:
 
         MODEL_TENSOR.A_POST_NORM: (
             "audio_tower.layer_norm", # ultravox
+            "audio_encoder.dots_encoder.speech_encoder.layer_norm", # dots3note
             "audio_tower.ln_post", # qwen2omni
             "encoder.layer_norm", # mimo-audio-tokenizer
         ),
 
         MODEL_TENSOR.A_ENC_ATTN_Q: (
             "audio_tower.layers.{bid}.self_attn.q_proj", # ultravox
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn.q_proj", # dots3note
             "conformer.layers.{bid}.self_attn.linear_q", # lfm2
             "conformer.layers.{bid}.attention.attn.q_proj", # gemma3n
             "conformer.layers.{bid}.self_attn.q_proj", # gemma4
@@ -2200,6 +2242,7 @@ class TensorNameMap:
 
         MODEL_TENSOR.A_ENC_ATTN_K: (
             "audio_tower.layers.{bid}.self_attn.k_proj", # ultravox
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn.k_proj", # dots3note
             "conformer.layers.{bid}.self_attn.linear_k", # lfm2
             "conformer.layers.{bid}.attention.attn.k_proj", # gemma3n
             "conformer.layers.{bid}.self_attn.k_proj", # gemma4
@@ -2210,6 +2253,7 @@ class TensorNameMap:
 
         MODEL_TENSOR.A_ENC_ATTN_V: (
             "audio_tower.layers.{bid}.self_attn.v_proj", # ultravox
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn.v_proj", # dots3note
             "conformer.layers.{bid}.self_attn.linear_v", # lfm2
             "conformer.layers.{bid}.attention.attn.v_proj", # gemma3n
             "conformer.layers.{bid}.self_attn.v_proj", # gemma4
@@ -2241,6 +2285,7 @@ class TensorNameMap:
 
         MODEL_TENSOR.A_ENC_INPUT_NORM: (
             "audio_tower.layers.{bid}.self_attn_layer_norm", # ultravox
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn_layer_norm", # dots3note
             "conformer.layers.{bid}.norm_self_att", # lfm2
             "conformer.layers.{bid}.attention.pre_attn_norm", # gemma3n
             "sound_encoder.encoder.layers.{bid}.norm_self_att", # parakeet
@@ -2250,6 +2295,7 @@ class TensorNameMap:
 
         MODEL_TENSOR.A_ENC_OUTPUT: (
             "audio_tower.layers.{bid}.self_attn.out_proj", # ultravox
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn.out_proj", # dots3note
             "conformer.layers.{bid}.self_attn.linear_out", # lfm2
             "conformer.layers.{bid}.attention.post", # gemma3n
             "conformer.layers.{bid}.self_attn.post", # gemma4
@@ -2260,6 +2306,7 @@ class TensorNameMap:
 
         MODEL_TENSOR.A_ENC_OUTPUT_NORM: (
             "audio_tower.layers.{bid}.final_layer_norm", # ultravox
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.final_layer_norm", # dots3note
             "conformer.layers.{bid}.norm_out", # lfm2
             "conformer.layers.{bid}.attention.post_norm", # gemma3n
             "sound_encoder.encoder.layers.{bid}.norm_out", # parakeet
@@ -2285,6 +2332,7 @@ class TensorNameMap:
         ),
 
         MODEL_TENSOR.A_ENC_FFN_UP: (
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.fc1_up", # dots3note (split from fc1 in conversion code)
             "audio_tower.layers.{bid}.fc1", # ultravox
             "conformer.layers.{bid}.feed_forward1.linear1", # lfm2
             "conformer.layers.{bid}.ffw_layer_start.ffw_layer_1", # gemma3n
@@ -2294,9 +2342,12 @@ class TensorNameMap:
             "encoder.layers.{bid}.fc1", # mimo-audio-tokenizer
         ),
 
-        MODEL_TENSOR.A_ENC_FFN_GATE: (),
+        MODEL_TENSOR.A_ENC_FFN_GATE: (
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.fc1_gate", # dots3note (split from fc1 in conversion code)
+        ),
 
         MODEL_TENSOR.A_ENC_FFN_DOWN: (
+            "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.fc2", # dots3note
             "audio_tower.layers.{bid}.fc2", # ultravox
             "conformer.layers.{bid}.feed_forward1.linear2", # lfm2
             "conformer.layers.{bid}.ffw_layer_start.ffw_layer_2", # gemma3n
@@ -2380,6 +2431,7 @@ class TensorNameMap:
 
         MODEL_TENSOR.A_MMPROJ: (
             "audio.multi_modal_projector.linear_{bid}", # ultravox, meralion
+            "audio_encoder.audio_adapter.proj.{bid}", # dots3note (proj.1, proj.3)
             "audio_adapter.model.{bid}", # lfm2
             "audio_tower.proj{bid}", # qwen3omni
             "sound_projection.linear{bid}", # parakeet (linear1, linear2)
@@ -2394,6 +2446,7 @@ class TensorNameMap:
 
         MODEL_TENSOR.A_MM_NORM_PRE: (
             "audio.multi_modal_projector.ln_pre", # ultravox
+            "audio_encoder.audio_adapter.proj.0", # dots3note
             "sound_projection.norm", # parakeet
         ),
 

include/llama.h

@@ -733,7 +733,7 @@ extern "C" {
 
     // Removes all tokens that belong to the specified sequence and have positions in [p0, p1)
     // Returns false if a partial sequence cannot be removed. Removing a whole sequence never fails
-    // seq_id < 0 : match any sequence
+    // seq_id < 0 : match any sequence [TAG_LLAMA_SEQ_ID_NEG]
     // p0 < 0     : [0,  p1]
     // p1 < 0     : [p0, inf)
     LLAMA_API bool llama_memory_seq_rm(

scripts/ccache-clear.sh

@@ -0,0 +1,105 @@
+#!/bin/bash
+# Delete GitHub Actions caches matching a key prefix, oldest first.
+#
+# Usage: ccache-clear.sh --key KEY [--older DURATION] [--min N] [--dry-run]
+#   --key:       cache key prefix to match and delete (without the ccache- prefix)
+#   --older:     only delete caches created more than DURATION ago (e.g. 5m, 1h, 1d);
+#                by default all matching caches are deleted
+#   --min:       stop deleting if fewer than N caches would remain (default: 0)
+#   --dry-run:   only print the caches that would be deleted, without deleting them
+#
+# Env (when running in GitHub Actions):
+#   GH_TOKEN:          token for the gh CLI
+#   GITHUB_REPOSITORY: owner/repo of the caches to manage
+set -euo pipefail
+
+KEY=""
+OLDER=""
+MIN=0
+DRY_RUN=false
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --key)     [[ $# -ge 2 ]] || { echo "Missing value for $1" >&2; exit 1; }; KEY="$2"; shift 2 ;;
+        --older)   [[ $# -ge 2 ]] || { echo "Missing value for $1" >&2; exit 1; }; OLDER="$2"; shift 2 ;;
+        --min)     [[ $# -ge 2 ]] || { echo "Missing value for $1" >&2; exit 1; }; MIN="$2"; shift 2 ;;
+        --dry-run) DRY_RUN=true; shift ;;
+        *) echo "Unknown argument: $1"; exit 1 ;;
+    esac
+done
+
+command -v gh >/dev/null 2>&1 || { echo "Error: GitHub CLI (gh) is required" >&2; exit 1; }
+[[ -n "${GITHUB_REPOSITORY:-}" ]] || { echo "Error: GITHUB_REPOSITORY not set" >&2; exit 1; }
+[[ -n "$KEY" ]] || { echo "Error: --key is required" >&2; exit 1; }
+[[ "$MIN" =~ ^[0-9]+$ ]] || { echo "Invalid min value: $MIN" >&2; exit 1; }
+
+# Convert a duration (e.g. 90m, 1h, 1d, plain seconds) to seconds
+to_seconds() {
+  local val="$1"
+  [[ "$val" =~ ^[0-9]+$ ]] && { echo "$val"; return 0; }
+  local num="${val%?}" unit="${val: -1}" mult
+  [[ "$num" =~ ^[0-9]+$ ]] || return 1
+  case "$unit" in
+    s) mult=1 ;;
+    m) mult=60 ;;
+    h) mult=3600 ;;
+    d) mult=86400 ;;
+    *) return 1 ;;
+  esac
+  echo $((num * mult))
+}
+
+# Convert an ISO-8601 UTC timestamp (e.g. 2026-08-23T16:51:23.313693Z) to epoch seconds
+to_epoch() {
+  local val="$1" out
+  # GNU date (e.g. Linux)
+  if out=$(date -d "$val" +%s 2>/dev/null) && [[ "$out" =~ ^[0-9]+$ ]]; then
+    echo "$out"
+    return 0
+  fi
+  # BSD date (e.g. macOS); fractional seconds are not needed, TZ forces UTC
+  out=$(TZ=UTC date -j -f "%Y-%m-%dT%H:%M:%S" "${val:0:19}" +%s 2>/dev/null) || return 1
+  [[ "$out" =~ ^[0-9]+$ ]] || return 1
+  echo "$out"
+}
+
+CACHES=$(gh cache list --repo "$GITHUB_REPOSITORY" --key "ccache-$KEY" --json id,key,createdAt --jq '.[] | [.createdAt, .id, .key] | @tsv' | LC_ALL=C sort)
+if [[ -z "$CACHES" ]]; then
+    echo "No caches found with key prefix: $KEY"
+    exit 0
+fi
+
+TOTAL=$(( $(wc -l <<< "$CACHES") ))
+
+echo "Found $TOTAL cache(s) with key prefix: $KEY (oldest first):"
+while IFS=$'\t' read -r CREATED ID CACHE_KEY; do
+    printf '  %s  %s  %s\n' "$CREATED" "$ID" "$CACHE_KEY"
+done <<< "$CACHES"
+
+CUTOFF=""
+if [[ -n "$OLDER" ]]; then
+    OLDER_SECONDS=$(to_seconds "$OLDER") || { echo "Invalid older value: $OLDER (expected e.g. 90m, 1h, 1d)" >&2; exit 1; }
+    CUTOFF=$(( $(date +%s) - OLDER_SECONDS ))
+fi
+
+# Caches are sorted oldest first
+DELETED=0
+while IFS=$'\t' read -r CREATED ID CACHE_KEY; do
+    if [[ -n "$CUTOFF" ]]; then
+        CREATED_SECONDS=$(to_epoch "$CREATED") || { echo "Failed to parse date: $CREATED" >&2; exit 1; }
+        if [[ "$CREATED_SECONDS" -ge "$CUTOFF" ]]; then
+            echo "Rest are not older than $OLDER, stopping"
+            break
+        fi
+    fi
+    if (( TOTAL - DELETED - 1 < MIN )); then
+        echo "Keeping at least $MIN cache(s), stopping"
+        break
+    fi
+    if [[ "$DRY_RUN" == "true" ]]; then
+        echo "Would delete cache: $ID ($CACHE_KEY)"
+    else
+        echo "Deleting cache: $ID ($CACHE_KEY)"
+        gh cache delete --repo "$GITHUB_REPOSITORY" "$ID"
+    fi
+    DELETED=$((DELETED + 1))
+done <<< "$CACHES"

scripts/sync_vendor.py

@@ -27,7 +27,7 @@ vendor = {
     f"https://raw.githubusercontent.com/yhirose/cpp-httplib/{HTTPLIB_VERSION}/split.py":  "split.py",
     f"https://raw.githubusercontent.com/yhirose/cpp-httplib/{HTTPLIB_VERSION}/LICENSE":   "vendor/cpp-httplib/LICENSE",
 
-    "https://raw.githubusercontent.com/sheredom/subprocess.h/9ce0d701b6fb10f8f8c4445edd31e7c60a1237e3/subprocess.h": "vendor/sheredom/subprocess.h",
+    "https://raw.githubusercontent.com/sheredom/subprocess.h/0dccaa9aa176dd6d7ef8afeca3c18d6e80a32795/subprocess.h": "vendor/sheredom/subprocess.h",
 
     f"https://raw.githubusercontent.com/Cyan4973/xxHash/{XXHASH_COMMIT}/xxhash.c":      "vendor/hash/xxhash/xxhash.c",
     f"https://raw.githubusercontent.com/Cyan4973/xxHash/{XXHASH_COMMIT}/xxhash.h":      "vendor/hash/xxhash/xxhash.h",

skills/add-new-model/SKILL.md

@@ -66,7 +66,7 @@ These recur often enough in review comments on past add-model PRs that they're w
 - Optional hparams that are genuinely absent from some configs (e.g. a shared-expert count) should be read with an explicit optional/fallback accessor, not assumed present.
 - Hparams that are actually load-bearing (the model produces wrong output or crashes without them, e.g. `sliding_window_pattern`, norm-eps) must hard-error if missing, not silently fall back to a default.
 - Don't bake a default chat template into the C++ binary - inject it into the GGUF at conversion time instead, since one `llm_arch` can be reused by multiple fine-tunes with different templates, and a baked-in C++ default fails silently for those.
-- Before writing a dedicated tool-call/output parser, check whether the existing autoparser already handles the template (`llama-debug-template-parser <jinja>` shows what it detects).
+- Before writing a dedicated tool-call/output parser, check whether the existing autoparser already handles the template (`test-chat-auto-parser <jinja>` shows what it detects).
 - Marking a custom EOS/closing-tag token as `eot` at conversion time isn't always sufficient - in long/agentic generations a model can emit the closing sequence as literal text instead of the token, so generation never stops on EOG and raw text leaks past the parser. Verify this case, not just the token path.
 - If reusing or aliasing an existing pre-tokenizer for convenience, justify and test that choice explicitly - silent reuse is an easy source of subtle tokenizer bugs.
 - Watch for excessive graph splits caused by building per-layer view/index tensors inside the layer loop - hoist tensors that don't vary per layer out of the loop (relevant if you hit `GGML_SCHED_MAX_SPLIT_INPUTS`).

src/llama-arch.cpp

@@ -1038,6 +1038,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
         case LLM_ARCH_NEMOTRON_H_MOE:
         case LLM_ARCH_LFM2:
         case LLM_ARCH_LFM2MOE:
+        case LLM_ARCH_BAILINGMOE3:
             return true;
         default:
             return false;
@@ -1059,7 +1060,6 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
         case LLM_ARCH_OLMOE:
         case LLM_ARCH_DEEPSEEK2:
         case LLM_ARCH_DEEPSEEK32:
-        case LLM_ARCH_DEEPSEEK4:
         case LLM_ARCH_DOTS3NOTE:
         case LLM_ARCH_GLM_DSA:
         case LLM_ARCH_BITNET:

src/llama-context.cpp

@@ -3218,8 +3218,6 @@ size_t llama_context::state_read_data(llama_io_read_i & io) {
 }
 
 size_t llama_context::state_seq_write_data(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
-    GGML_UNUSED(seq_id);
-
     if (memory) {
         memory->state_write(io, seq_id, flags);
     }
@@ -3228,8 +3226,6 @@ size_t llama_context::state_seq_write_data(llama_io_write_i & io, llama_seq_id s
 }
 
 size_t llama_context::state_seq_read_data(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
-    GGML_UNUSED(seq_id);
-
     if (memory) {
         memory->state_read(io, seq_id, flags);
     }

src/llama-kv-cache-dsv4.cpp

@@ -599,6 +599,33 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
         }
     }
 
+    if (ratio == DSV4_HCA_RATIO && !plan.state_pos.empty() && plan.state_write_idxs.empty()) {
+        assert(kv_size > 0);
+        // the last slot must not be live, or the dummy write would corrupt it;
+        // a full stream implies a completed block, which implies real writes
+        assert(plan.n_kv < (int64_t) kv_size);
+
+        // Keep the compress/write ops in the graph when no HCA block completes
+        // in this ubatch. The dummy block writes to the last cache slot and is
+        // masked out.
+        uint32_t i = 0;
+        while (i < ubatch.n_tokens && ubatch.pos[i] < 0) {
+            ++i;
+        }
+        assert(i < ubatch.n_tokens);
+
+        const llama_seq_id seq_id = ubatch.seq_id[i][0];
+        const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size);
+        const int32_t source_idx = state_source_idx(seq_id, ubatch.pos[i]);
+
+        plan.state_write_idxs.push_back(cache_off + kv_size - 1);
+        plan.state_write_pos .push_back(0);
+
+        for (uint32_t j = 0; j < ratio; ++j) {
+            plan.state_read_idxs.push_back(source_idx);
+        }
+    }
+
     if (overlap) {
         // [ all blocks' prev-window indices | all blocks' cur-window indices ]
         plan.state_read_idxs.reserve(overlap_prev_reads.size() + overlap_cur_reads.size());
@@ -608,7 +635,10 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
                 overlap_cur_reads.begin(), overlap_cur_reads.end());
     }
 
-    plan.n_kv = GGML_PAD(plan.n_kv, 256u);
+    // Keep the mask (and with it the compressed-attention branch) present even
+    // before the first block is visible, so the graph topology never changes.
+    // Padded slots are masked out; comp cache buffers are zero-initialized.
+    plan.n_kv = std::max<int64_t>(GGML_PAD(plan.n_kv, 256u), 256);
 
     std::sort(persist_rows.begin(), persist_rows.end(),
             [](const persist_row & a, const persist_row & b) {
@@ -620,16 +650,26 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
         plan.state_persist_dst_idxs.push_back(row.dst);
     }
 
-
     if (n_rs_seq > 0) {
-        for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
-            const llama_seq_id seq_id = ubatch.seq_id_unq[s];
-            if (seq_id < 0 || (uint32_t) seq_id >= n_stream) {
-                continue;
+        // Emit restore/snapshot entries for all layout streams so that the
+        // graph tensor sizes do not depend on the ubatch's sequence count.
+        // Streams not present in the ubatch get no-op entries.
+        for (uint32_t stream = 0; stream < n_stream; ++stream) {
+            llama_seq_id seq_id = -1;
+            if (n_stream == 1) {
+                // a unified stream serves any single sequence
+                seq_id = ubatch.n_seqs_unq > 0 ? ubatch.seq_id_unq[0] : -1;
+            } else {
+                for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
+                    if (ubatch.seq_id_unq[s] == (llama_seq_id) stream) {
+                        seq_id = ubatch.seq_id_unq[s];
+                        break;
+                    }
+                }
             }
 
-            const int64_t stream_off = dsv4_stream_offset(n_stream, seq_id, state_size);
-            const uint32_t rollback = (uint32_t) seq_id < rs_idx.size() ? rs_idx[seq_id] : 0;
+            const int64_t stream_off = (int64_t) stream*state_size;
+            const uint32_t rollback = seq_id >= 0 && (uint32_t) seq_id < rs_idx.size() ? rs_idx[seq_id] : 0;
             // Keep the restore graph fixed-width when no rollback is pending.
             const int64_t src_plane = rollback > 0 && rollback <= n_rs_seq ? (int64_t) rollback*state_rows : 0;
             for (uint32_t r = 0; r < state_size; ++r) {
@@ -639,35 +679,33 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
 
             std::vector<uint32_t> token_idxs;
             token_idxs.reserve(ubatch.n_tokens);
-            for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
-                if (dsv4_token_has_seq(ubatch, i, seq_id)) {
-                    token_idxs.push_back(i);
+            if (seq_id >= 0) {
+                for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
+                    if (dsv4_token_has_seq(ubatch, i, seq_id)) {
+                        token_idxs.push_back(i);
+                    }
                 }
             }
-            if (token_idxs.empty()) {
-                continue;
-            }
 
             const uint32_t n_seq_tokens = (uint32_t) token_idxs.size();
             const int64_t scratch_off = (int64_t) state_rows*(1 + n_rs_seq);
             for (uint32_t d = 1; d <= n_rs_seq; ++d) {
                 const int64_t dst_plane = (int64_t) d*state_rows;
+                const uint32_t prefix = d <= n_seq_tokens ? n_seq_tokens - d : 0;
 
                 for (uint32_t r = 0; r < state_size; ++r) {
-                    int32_t src;
-                    if (d <= n_seq_tokens) {
-                        const uint32_t prefix = n_seq_tokens - d;
-                        src = (int32_t) (stream_off + r);
+                    int32_t src = (int32_t) (stream_off + r);
 
-                        for (uint32_t j = 0; j < prefix; ++j) {
-                            const uint32_t i_tok = token_idxs[j];
-                            if (ubatch.pos[i_tok] >= 0 && (uint32_t) (ubatch.pos[i_tok]%state_size) == r) {
-                                src = (int32_t) (scratch_off + i_tok);
-                            }
+                    for (uint32_t j = 0; j < prefix; ++j) {
+                        const uint32_t i_tok = token_idxs[j];
+                        if (ubatch.pos[i_tok] >= 0 && (uint32_t) (ubatch.pos[i_tok]%state_size) == r) {
+                            src = (int32_t) (scratch_off + i_tok);
                         }
-                    } else {
-                        const int64_t src_plane = (int64_t) (d - n_seq_tokens)*state_rows;
-                        src = (int32_t) (src_plane + stream_off + r);
+                    }
+
+                    if (n_seq_tokens == 0) {
+                        // no-op: copy the snapshot plane onto itself
+                        src = (int32_t) (dst_plane + stream_off + r);
                     }
 
                     plan.state_snapshot_src_idxs.push_back(src);
@@ -683,10 +721,16 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
     }();
 
     if (debug) {
-        LLAMA_LOG_INFO("%s: ratio=%u, n_tokens=%u, state_persist_dst=%s, state_write_pos=%s\n",
-                __func__, ratio, ubatch.n_tokens,
+        LLAMA_LOG_DEBUG("%s: ratio=%u, n_tokens=%u, n_seqs_unq=%u, state_persist_dst=%s, state_write_pos=%s\n",
+                __func__, ratio, ubatch.n_tokens, ubatch.n_seqs_unq,
                 dsv4_plan_positions(plan.state_persist_dst_idxs).c_str(),
                 dsv4_plan_positions(plan.state_write_pos).c_str());
+        for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
+            const llama_seq_id seq_id = ubatch.seq_id_unq[s];
+            const uint32_t rollback = seq_id >= 0 && (uint32_t) seq_id < rs_idx.size() ? rs_idx[seq_id] : 0;
+            LLAMA_LOG_DEBUG("%s:   seq %d pos [%d, %d] rollback=%u\n", __func__, seq_id,
+                    ubatch.pos[0], ubatch.pos[ubatch.n_tokens - 1], rollback);
+        }
     }
 
     return plan;
@@ -704,8 +748,17 @@ static std::vector<llama_kv_cache_dsv4_context::comp_plan> dsv4_build_comp_plans
     std::vector<llama_kv_cache_dsv4_context::comp_plan> plans;
     plans.reserve(ubatches.size());
 
+    // the first ubatch touching a seq consumes its rollback restore
+    std::vector<uint32_t> rs(rs_idx);
     for (const llama_ubatch & ubatch : ubatches) {
-        plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream, n_rs_seq, rs_idx));
+        plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream, n_rs_seq, rs));
+
+        for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
+            const llama_seq_id seq_id = ubatch.seq_id_unq[s];
+            if (seq_id >= 0 && (size_t) seq_id < rs.size()) {
+                rs[seq_id] = 0;
+            }
+        }
     }
 
     return plans;
@@ -803,16 +856,15 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan(
         return plan;
     }
 
-    const uint32_t n_seqs       = std::max<uint32_t>(1, ubatch.n_seqs);
-    const uint32_t n_seq_tokens = std::max<uint32_t>(1, ubatch.n_seq_tokens);
-    const uint64_t n_blocks_u64 = (uint64_t) n_seqs*((n_seq_tokens + ratio - 1)/ratio);
-    const size_t n_blocks = (size_t) std::max<uint64_t>(1, n_blocks_u64);
-    GGML_ASSERT((uint64_t) n_blocks == std::max<uint64_t>(1, n_blocks_u64));
+    // worst case over every seq split: sum of per-seq ceil(tokens/ratio) is at
+    // most floor(n_tokens/ratio) + n_seqs
+    const uint32_t n_seqs = std::max<uint32_t>(1, ubatch.n_seqs);
+    const size_t n_blocks = (size_t) ubatch.n_tokens/ratio + n_seqs;
 
     const uint64_t state_rows = (uint64_t) state_size*n_stream;
     const size_t n_persist = (size_t) std::min<uint64_t>(ubatch.n_tokens, state_rows);
-    const size_t n_restore = n_rs_seq > 0 ? (size_t) state_size*std::max<uint32_t>(1, ubatch.n_seqs_unq) : 0;
-    const size_t n_snapshot = (size_t) n_rs_seq*state_size*std::max<uint32_t>(1, ubatch.n_seqs_unq);
+    const size_t n_restore = n_rs_seq > 0 ? (size_t) state_size*n_stream : 0;
+    const size_t n_snapshot = (size_t) n_rs_seq*state_size*n_stream;
 
     plan.state_pos .resize(ubatch.n_tokens);
     plan.state_persist_src_idxs.resize(n_persist);
@@ -1356,7 +1408,9 @@ llama_memory_context_ptr llama_kv_cache_dsv4::init_batch(
             if (has_coupled) {
                 ubatch = balloc.split_seq(n_ubatch);
             } else {
-                ubatch = balloc.split_equal(n_ubatch, raw_per_seq || comp_per_seq, 0);
+                // [TAG_RECURRENT_ROLLBACK_SPLITS]
+                // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch
+                ubatch = balloc.split_equal(n_ubatch, raw_per_seq || comp_per_seq, n_rs_seq > 0 ? n_rs_seq + 1 : 0);
             }
 
             if (ubatch.n_tokens == 0) {
@@ -1433,6 +1487,11 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
             return false;
         }
 
+        // pending rollback is single-use: stacked partial removals don't compose
+        if (rs_idx[seq_id] != 0) {
+            return false;
+        }
+
         const bool res = kv_raw->seq_rm(seq_id, p0, p1);
         if (res) {
             rs_idx[seq_id] = (uint32_t) rollback;
@@ -1594,9 +1653,7 @@ void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id,
     kv_raw->state_read(io, seq_id, flags);
 
     if (!partial_only) {
-        kv_csa->clear(true);
-        kv_hca->clear(true);
-        kv_lid->clear(true);
+        clear_compressed(seq_id, true);
 
         dsv4_state_read_k_cache(io, kv_csa.get(), seq_id, flags);
         dsv4_state_read_k_cache(io, kv_hca.get(), seq_id, flags);
@@ -1680,6 +1737,7 @@ void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) {
             kv->seq_rm(seq_id, -1, -1);
 
             if (data) {
+                //TODO: do not clear the kv-cache during `seq_rm`, ref: https://github.com/ggml-org/llama.cpp/pull/26490#discussion_r3798143663
                 for (uint32_t il : kv->get_layer_ids()) {
                     dsv4_clear_tensor_stream(kv->get_k_storage(il), (uint32_t) seq_id);
                 }

src/llama-kv-cache.cpp

@@ -383,6 +383,7 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
         return true;
     }
 
+    // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG]
     GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()));
 
     if (p0 < 0) {
@@ -2043,6 +2044,7 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama
 
     GGML_UNUSED(flags);
 
+    // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG]
     GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()));
 
     uint32_t n_stream_cur;

src/llama-memory-recurrent.cpp

@@ -158,13 +158,14 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos
         p1 = std::numeric_limits<llama_pos>::max();
     }
 
+    if ((uint32_t) seq_id >= this->n_seq_max) {
+        LLAMA_LOG_ERROR("%s: invalid seq_id (%d) - larger than n_seq_max (%d)\n", __func__, seq_id, this->n_seq_max);
+        return false;
+    }
+
     const bool rm_all = p0 == 0 && p1 == std::numeric_limits<llama_pos>::max();
     if (rm_all) {
-        if (seq_id >= 0) {
-            set_rs_idx(seq_id, 0);
-        } else {
-            std::fill(rs_idx.begin(), rs_idx.end(), 0);
-        }
+        set_rs_idx(seq_id, 0);
     }
 
     // models like Mamba or RWKV can't have a state partially erased at the end
@@ -181,7 +182,9 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos
             // partial rollback via per-token snapshot index (bounded by n_rs_seq)
             if (0 < p0 && p0 <= cell.pos && p1 > cell.pos) {
                 const llama_pos rollback = cell.pos - (p0 - 1);
-                if (rollback >= 1 && rollback <= (llama_pos) n_rs_seq) {
+                // pending rollback is single-use
+                const bool pending = rs_idx[seq_id] != 0;
+                if (!pending && rollback >= 1 && rollback <= (llama_pos) n_rs_seq) {
                     set_rs_idx(seq_id, (uint32_t) rollback);
                     cell.pos = p0 - 1;
                     return true;
@@ -390,10 +393,17 @@ llama_pos llama_memory_recurrent::seq_pos_max(llama_seq_id seq_id) const {
 }
 
 void llama_memory_recurrent::set_rs_idx(llama_seq_id seq_id, uint32_t idx) {
-    if (seq_id < 0 || (size_t) seq_id >= rs_idx.size()) {
+    if (seq_id < 0) {
+        std::fill(rs_idx.begin(), rs_idx.end(), 0);
         return;
     }
-    rs_idx[seq_id] = (idx > n_rs_seq) ? n_rs_seq : idx;
+
+    assert(n_seq_max == rs_idx.size());
+
+    GGML_ASSERT((uint32_t) seq_id < n_seq_max);
+    GGML_ASSERT(idx <= n_rs_seq);
+
+    rs_idx[seq_id] = idx;
 }
 
 std::map<ggml_backend_buffer_type_t, size_t> llama_memory_recurrent::memory_breakdown() const {
@@ -742,6 +752,7 @@ void llama_memory_recurrent::state_write(llama_io_write_i & io, llama_seq_id seq
     uint32_t cell_range_begin = size;
     for (uint32_t i = 0; i < size; ++i) {
         const auto & cell = cells[i];
+        // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG]
         if ((seq_id == -1 && !cell.is_empty()) || cell.has_seq_id(seq_id)) {
             ++cell_count;
             uint32_t rs_idx_cur = 0;
@@ -827,6 +838,7 @@ void llama_memory_recurrent::state_read(llama_io_read_i & io, llama_seq_id seq_i
     }
 
     if (!res) {
+        // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG]
         if (seq_id == -1) {
             clear(true);
         } else {
@@ -836,11 +848,7 @@ void llama_memory_recurrent::state_read(llama_io_read_i & io, llama_seq_id seq_i
     }
 
     if (n_rs_seq != 0) {
-        if (seq_id == -1) {
-            std::fill(rs_idx.begin(), rs_idx.end(), 0);
-        } else {
-            set_rs_idx(seq_id, 0);
-        }
+        set_rs_idx(seq_id, 0);
     }
 }
 

src/llama-model-saver.cpp

@@ -293,6 +293,21 @@ void llama_model_saver::add_kv_from_model() {
     add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS,    hparams.indexer_local_blocks);
     add_kv(LLM_KV_ATTENTION_INDEXER_TYPES,           hparams.is_indexer_full_impl, true);
     add_kv(LLM_KV_ATTENTION_RECURRENT_LAYERS,        hparams.is_recr_impl, true);
+    add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,      hparams.dsv4_o_group_count);
+    add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK,        hparams.dsv4_o_lora_rank);
+    add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, hparams.dsv4_compress_rope_base);
+    if (model->arch == LLM_ARCH_DEEPSEEK4 || hparams.dsv4_hc_mult > 0) {
+        // the loader requires one compress ratio per layer, including nextn layers
+        const std::vector<uint32_t> compress_ratios(
+                hparams.dsv4_compress_ratios.begin(), hparams.dsv4_compress_ratios.begin() + hparams.n_layer_all);
+        add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS,         compress_ratios);
+    } else {
+        add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS,         hparams.dsv4_compress_ratios, true);
+    }
+    add_kv(LLM_KV_HYPER_CONNECTION_COUNT,               hparams.dsv4_hc_mult);
+    add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters);
+    add_kv(LLM_KV_HYPER_CONNECTION_EPSILON,             hparams.dsv4_hc_eps);
+    add_kv(LLM_KV_HASH_LAYER_COUNT,                     hparams.dsv4_hash_layer_count);
 
     const float rope_scaling_factor = hparams.rope_freq_scale_train == 1.0f ? 0.0f : 1.0f/hparams.rope_freq_scale_train;
 
@@ -417,11 +432,16 @@ void llama_model_saver::add_tensors_from_model() {
     add_tensor(model->output_s);
     add_tensor(model->output_in_s);
     add_tensor(model->output_res_score);
+    add_tensor(model->nextn_proj_pre);
+    add_tensor(model->nextn_proj_post);
     add_tensor(model->cls);
     add_tensor(model->cls_b);
     add_tensor(model->cls_out);
     add_tensor(model->cls_out_b);
     add_tensor(model->cls_norm);
+    add_tensor(model->hc_head_fn);
+    add_tensor(model->hc_head_base);
+    add_tensor(model->hc_head_scale);
 
     for (const struct llama_layer & layer : model->layers) {
         for (size_t i = 0; i < sizeof(layer)/sizeof(struct ggml_tensor *); ++i) {

src/llama-model.cpp

@@ -365,6 +365,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
     const llama_meta_device_get_split_state_userdata * ud = (const llama_meta_device_get_split_state_userdata *) userdata;
     const llama_hparams & hparams = ud->model->hparams;
     const std::string tensor_name = tensor->name;
+    const bool is_dsv4 = ud->model->arch == LLM_ARCH_DEEPSEEK4 ||
+        (ud->model->arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0);
 
     static const std::regex pattern_q_weight        ("blk\\.\\d*\\.attn_q.weight");
     static const std::regex pattern_kv_weight       ("blk\\.\\d*\\.attn_(k|v).weight");
@@ -374,9 +376,13 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
     static const std::regex pattern_qkv_bias        ("blk\\.\\d*\\.attn_qkv.bias");
     static const std::regex pattern_qk_norm         ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight");
     static const std::regex pattern_kv_cache        ("cache_(k|v)_l\\d*");
+    static const std::regex pattern_dsv4_state      ("dsv4_(csa|hca|lid)_state_(kv|score)_l\\d*");
     static const std::regex pattern_attn_sinks      ("blk\\.\\d*\\.attn_sinks.weight");
     static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight");
     static const std::regex pattern_attn_out_bias   ("blk\\.\\d*\\.attn_output.bias");
+    static const std::regex pattern_attn_out_a_weight("blk\\.\\d*\\.attn_output_a\\.weight");
+    static const std::regex pattern_attn_out_b_weight("blk\\.\\d*\\.attn_output_b\\.weight");
+    static const std::regex pattern_attn_q_b_weight ("blk\\.\\d*\\.attn_q_b\\.weight");
     static const std::regex pattern_attn_gate_weight("blk\\.\\d*\\.attn_gate.weight");
 
     static const std::regex pattern_ssm_dt          ("blk\\.\\d*\\.ssm_dt.bias");
@@ -395,8 +401,11 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
     static const std::regex pattern_ffn_gate_bias     ("blk\\.\\d*\\.ffn_gate(_exps)?.bias");
     static const std::regex pattern_ffn_gate_up_weight("blk\\.\\d*\\.ffn_gate_up(_exps)?.weight");
     static const std::regex pattern_ffn_down_weight   ("blk\\.\\d*\\.ffn_down(_exps)?.weight");
-    static const std::regex pattern_ffn_down_bias     ("blk\\.\\d*\\.ffn_down.bias");
-    static const std::regex pattern_ffn_down_exps_bias("blk\\.\\d*\\.ffn_down_exps.bias");
+    static const std::regex pattern_ffn_down_bias         ("blk\\.\\d*\\.ffn_down.bias");
+    static const std::regex pattern_ffn_down_exps_bias    ("blk\\.\\d*\\.ffn_down_exps.bias");
+    static const std::regex pattern_ffn_up_shexp_weight   ("blk\\.\\d*\\.ffn_up_shexp.weight");
+    static const std::regex pattern_ffn_gate_shexp_weight ("blk\\.\\d*\\.ffn_gate_shexp.weight");
+    static const std::regex pattern_ffn_down_shexp_weight ("blk\\.\\d*\\.ffn_down_shexp.weight");
 
     static const std::regex pattern_output_weight("output\\.weight");
     static const std::regex pattern_output_bias  ("output\\.bias");
@@ -453,6 +462,32 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
     };
 
     auto get_tensor_config = [&]() -> tensor_config {
+        if (is_dsv4) {
+            if (std::regex_match(tensor_name, pattern_kv_cache) ||
+                    std::regex_match(tensor_name, pattern_dsv4_state)) {
+                return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+            }
+            if (std::regex_match(tensor_name, pattern_attn_sinks)) {
+                return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output_a.weight");
+            }
+            if (std::regex_match(tensor_name, pattern_attn_q_b_weight)) {
+                return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output_a.weight");
+            }
+            if (std::regex_match(tensor_name, pattern_attn_out_a_weight)) {
+                return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_2);
+            }
+            if (std::regex_match(tensor_name, pattern_attn_out_b_weight)) {
+                return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0);
+            }
+            if (std::regex_match(tensor_name, pattern_ffn_up_shexp_weight) ||
+                    std::regex_match(tensor_name, pattern_ffn_gate_shexp_weight)) {
+                return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "ffn_down_shexp.weight");
+            }
+            if (std::regex_match(tensor_name, pattern_ffn_down_shexp_weight)) {
+                return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "ffn_down_shexp.weight");
+            }
+        }
+
         // standard attention
         if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_kv_weight)) {
             return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight", "ssm_out.weight");
@@ -520,11 +555,14 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
             return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED);
         }
         if (std::regex_match(tensor_name, pattern_ffn_down_exps_bias)) {
-            return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_PARTIAL);
+            return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_PARTIAL, "ffn_down_exps.weight");
         }
 
         // output
         if (std::regex_match(tensor_name, pattern_output_weight)) {
+            if (is_dsv4) {
+                return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+            }
             return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1);
         }
         if (std::regex_match(tensor_name, pattern_output_bias)) {
@@ -554,6 +592,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
                     GGML_ASSERT(tensor->ne[axis] == 2*key_dim + value_dim);
                     return {{key_dim, 2}, {value_dim, 1}};
                 }
+                if (std::regex_match(tensor_name, pattern_r_cache)) {
+                    return {{key_dim * (hparams.ssm_d_conv - 1), 2}, {value_dim * (hparams.ssm_d_conv - 1), 1}};
+                }
             } else {
                 const int64_t head_ratio = n_v_heads / n_k_heads;
                 if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_ssm_conv1d)) {
@@ -642,12 +683,34 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
                 blck_size_perf *= 2;
             }
 
+            const int64_t granularity_q    = std::lcm(n_embd_q, blck_size_perf);
+            const int64_t granularity_head = granularity_q / hparams.n_embd_head_k(il); // for tensors with one value per head
             if (std::regex_match(tensor_name, pattern_attn_sinks)) {
                 GGML_ASSERT(segments.size() == 1);
-                return {std::lcm(n_embd_q, blck_size_perf)/n_embd_q * n_gqa};
+                if (is_dsv4) {
+                    return {hparams.n_head(il) / hparams.dsv4_o_group_count};
+                }
+                return {granularity_head};
             }
 
-            const int64_t granularity_q = std::lcm(n_embd_q, blck_size_perf);
+            if (is_dsv4) {
+                if (std::regex_match(tensor_name, pattern_attn_q_b_weight)) {
+                    GGML_ASSERT(segments.size() == 1);
+                    // the grouped output projection requires each device to hold whole groups of heads
+                    const int64_t n_head_group = hparams.n_head(il) / hparams.dsv4_o_group_count;
+                    return {n_head_group * hparams.n_embd_head_k(il)};
+                }
+                if (std::regex_match(tensor_name, pattern_attn_out_a_weight)) {
+                    GGML_ASSERT(segments.size() == 1);
+                    return {1};
+                }
+                if (std::regex_match(tensor_name, pattern_attn_out_b_weight)) {
+                    GGML_ASSERT(segments.size() == 1);
+                    // the boundaries must align with wo_a's per-group split, so quant blocks must not straddle groups
+                    GGML_ASSERT(hparams.dsv4_o_lora_rank % blck_size == 0);
+                    return {hparams.dsv4_o_lora_rank};
+                }
+            }
             if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_q_bias)) {
                 GGML_ASSERT(segments.size() == 1);
                 // some models have Q gate tensors, for those cases the granularity needs to be doubled:
@@ -660,6 +723,13 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
                 GGML_ASSERT(segments.size() == 1);
                 return {granularity_q};
             }
+            if (std::regex_match(tensor_name, pattern_attn_gate_weight)) {
+                GGML_ASSERT(segments.size() == 1);
+                if (tensor->ne[1] == hparams.n_head(il)) {
+                    return {granularity_head};
+                }
+                return {granularity_q};
+            }
 
             const int64_t granularity_kv = granularity_q / n_gqa;
             if (std::regex_match(tensor_name, pattern_kv_weight) ||
@@ -677,7 +747,11 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
         // FFN
         if (std::regex_match(tensor_name, pattern_ffn_up_weight) || std::regex_match(tensor_name, pattern_ffn_up_bias) ||
                 std::regex_match(tensor_name, pattern_ffn_gate_weight) || std::regex_match(tensor_name, pattern_ffn_gate_bias) ||
-                std::regex_match(tensor_name, pattern_ffn_gate_up_weight) || std::regex_match(tensor_name, pattern_ffn_down_weight)) {
+                std::regex_match(tensor_name, pattern_ffn_gate_up_weight) ||
+                std::regex_match(tensor_name, pattern_ffn_down_weight) ||
+                std::regex_match(tensor_name, pattern_ffn_up_shexp_weight) ||
+                std::regex_match(tensor_name, pattern_ffn_gate_shexp_weight) ||
+                std::regex_match(tensor_name, pattern_ffn_down_shexp_weight)) {
             const int64_t blck_size_perf = std::lcm(blck_size, 128);
             GGML_ASSERT(segments.size() == 1);
             return {blck_size_perf};
@@ -728,6 +802,16 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
         memset(split_state.ne, 0, sizeof(split_state.ne));
         split_state.nr[0] = 1;
         split_state.n_segments = 1;
+        if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) {
+            GGML_ASSERT(tc.tensor_axis_0 != tensor);
+            const ggml_backend_meta_split_state source_split_state = llama_meta_device_get_split_state(tc.tensor_axis_0, userdata);
+            GGML_ASSERT(source_split_state.axis >= 0 && source_split_state.axis < GGML_MAX_DIMS);
+            for (size_t j = 0; j < ud->n_devices; j++) {
+                for (size_t is = 0; is < source_split_state.n_segments; is++) {
+                    split_state.ne[j] += source_split_state.ne[is*ud->n_devices + j] * source_split_state.nr[is];
+                }
+            }
+        }
     }
     return split_state;
     GGML_UNUSED(userdata);

src/models/bailingmoe3.cpp

@@ -1,6 +1,8 @@
 #include "models.h"
 #include "llama-memory-recurrent.h"
 
+#include <algorithm>
+
 void llama_model_bailingmoe3::load_arch_hparams(llama_model_loader & ml) {
     ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS,      hparams.f_norm_rms_eps);
     ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA,         hparams.n_embd_head_k_mla_impl);
@@ -179,7 +181,9 @@ static ggml_tensor * bailingmoe3_causal_conv1d(
         int64_t n_seq_tokens,
         int64_t n_seqs,
         int64_t n_tokens,
-        int64_t cache_head) {
+        int64_t cache_head,
+        uint32_t mem_size,
+        uint32_t n_rs_seq) {
     const int64_t d_inner = head_dim * n_head;
     const int64_t conv_state_size = (d_conv - 1) * d_inner;
     const int64_t total_state_size = 3 * conv_state_size;
@@ -193,13 +197,18 @@ static ggml_tensor * bailingmoe3_causal_conv1d(
     x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
     ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0);
 
-    ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
-            conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]);
-    ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv_x,
-            ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
-                (d_conv - 1) * ggml_element_size(conv_states_all),
-                total_state_size * ggml_element_size(conv_states_all),
-                (cache_head * total_state_size + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
+    const int64_t K = (int64_t) n_rs_seq + 1;
+    const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
+
+    for (int64_t slot = 0; slot < n_written; ++slot) {
+        ggml_tensor * conv_snap = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
+                conv_x->nb[1], conv_x->nb[2], (conv_x->ne[0] - (d_conv - 1) - slot) * conv_x->nb[0]);
+        ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_snap,
+                ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
+                    (d_conv - 1) * ggml_element_size(conv_states_all),
+                    total_state_size * ggml_element_size(conv_states_all),
+                    ((slot * mem_size + cache_head) * total_state_size + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
+    }
 
     ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
     ggml_tensor * out = ggml_ssm_conv(ctx0, conv_x, conv_weight);
@@ -237,6 +246,8 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
     GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);
 
     for (int il = 0; il < n_layer; ++il) {
+        res->t_layer_inp[il] = inpL;
+
         const auto & layer = model.layers[il];
         ggml_tensor * inpSA = inpL;
         ggml_tensor * cur = build_norm(inpL, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
@@ -245,18 +256,19 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
         if (hparams.is_recr(il)) {
             const auto * mctx_cur = inp_rs->mctx;
             const auto cache_head = mctx_cur->get_head();
+            const auto mem_size = mctx_cur->get_size();
             ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
             ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
 
             ggml_tensor * q = bailingmoe3_causal_conv1d(
                     gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv,
-                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head);
+                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
             ggml_tensor * k = bailingmoe3_causal_conv1d(
                     gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv,
-                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head);
+                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
             ggml_tensor * v = bailingmoe3_causal_conv1d(
                     gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv,
-                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head);
+                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
 
             ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_f_a, cur);
             gate = ggml_add(ctx0, gate, layer.ssm_dt_b);
@@ -276,11 +288,8 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
             ggml_tensor * state = build_rs(inp_rs, states_all, hparams.n_embd_s(), n_seqs);
             state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head, n_seqs);
 
-            auto result = build_delta_net(q, k, v, gate, beta, state, il);
-            ggml_tensor * out = ggml_cont(ctx0, result.first);
-            ggml_build_forward_expand(gf, ggml_cpy(ctx0, result.second,
-                    ggml_view_1d(ctx0, states_all, hparams.n_embd_s() * n_seqs,
-                        cache_head * hparams.n_embd_s() * ggml_element_size(states_all))));
+            ggml_tensor * out = ggml_cont(ctx0, build_recurrent_attn(
+                    inp_rs, states_all, q, k, v, gate, beta, state, il));
 
             ggml_tensor * out_gate = ggml_mul_mat(ctx0, layer.ssm_g_a, cur);
             out_gate = ggml_reshape_3d(ctx0, out_gate, head_dim, n_head, n_tokens);

src/models/deepseek4.cpp

@@ -1,3 +1,4 @@
+#include "llama-hparams.h"
 #include "models.h"
 
 #include "llama-kv-cache-dsv4.h"
@@ -58,6 +59,7 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {
     if (n_compress_ratios < hparams.n_layer_all) {
         throw std::runtime_error("DeepSeek-V4 compress_ratios is shorter than block_count");
     }
+    GGML_ASSERT(n_compress_ratios <= LLAMA_MAX_LAYERS);
     ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios);
 
     ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);

src/models/dflash.cpp

@@ -117,6 +117,10 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
     output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)
     output_norm     = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM,    "weight"), { n_embd }, 0); // decoder final norm
 
+    // optional: reduced-vocab drafts ship their own lm head, full-vocab drafts can share the target's via ctx_other
+    // a draft with its own embeddings + head references no target tensors and can run on devices the target does not use (e.g. -devd with a tensor-split target)
+    output   = create_tensor(tn(LLM_TENSOR_OUTPUT,     "weight"), { n_embd, n_vocab_draft }, TENSOR_NOT_REQUIRED);
+
     if (hparams.dsv4_hc_mult > 0) {
         const int64_t q_lora_rank     = hparams.n_lora_q;
         const int64_t n_ff_exp        = hparams.n_ff_exp;
@@ -167,9 +171,6 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
         return;
     }
 
-    // optional: reduced-vocab drafts ship their own, full-vocab drafts share the target's via ctx_other
-    output   = create_tensor(tn(LLM_TENSOR_OUTPUT,     "weight"), { n_embd, n_vocab_draft }, TENSOR_NOT_REQUIRED);
-
     for (int i = 0; i < n_layer; ++i) {
         auto & layer = layers[i];
 

src/models/glm4-moe.cpp

@@ -29,10 +29,19 @@ void llama_model_glm4_moe::load_arch_hparams(llama_model_loader & ml) {
     }
 }
 
-void llama_model_glm4_moe::load_arch_tensors(llama_model_loader &) {
+void llama_model_glm4_moe::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
     const int64_t n_expert_shared = hparams.n_expert_shared;
 
+    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
+    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
+    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
+    const int trunk_flags = mtp_only  ? TENSOR_NOT_REQUIRED : 0;
+    int       mtp_flags   = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+
+    if (!ml.load_mtp) {
+        mtp_flags |= TENSOR_SKIP;
+    }
 
     GGML_ASSERT(hparams.n_expert > 0 && "n_expert must be > 0 for GLM4_MOE MoE layers");
     GGML_ASSERT(hparams.n_expert_used > 0 && "n_expert_used must be > 0 for GLM4_MOE MoE layers");
@@ -47,16 +56,9 @@ void llama_model_glm4_moe::load_arch_tensors(llama_model_loader &) {
         output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);
     }
 
-    // Load ALL tensors including NextN layer to satisfy total tensor count
-    // but only PROCESS up to last layer (skipping final NextN layer) in forward pass
     for (int i = 0; i < n_layer_all; ++i) {
-        int flags = 0;
-        if (i >= n_layer) {
-            // skip all tensors in the NextN layers
-            flags |= TENSOR_SKIP;
-        }
-
         auto & layer = layers[i];
+        const int flags = i < n_layer ? trunk_flags : mtp_flags;
 
         layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, flags);
 
@@ -110,24 +112,186 @@ void llama_model_glm4_moe::load_arch_tensors(llama_model_loader &) {
             layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), { n_embd, n_ff }, flags);
         }
 
-        // NextN/MTP tensors (preserved but unused) - conditionally load for last nextn_predict_layers
+        // NextN/MTP tensors
         if (i >= n_layer) {
             layer.nextn.eh_proj          = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, flags);
             layer.nextn.enorm            = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags);
             layer.nextn.hnorm            = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), { n_embd }, flags);
 
             // Optional tensors
-            layer.nextn.embed_tokens     = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), { n_embd, n_vocab }, flags | TENSOR_NOT_REQUIRED);
-            layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), { n_embd, n_vocab }, flags | TENSOR_NOT_REQUIRED);
-            layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), { n_embd }, flags | TENSOR_NOT_REQUIRED);
+            layer.nextn.embed_tokens     = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED | flags);
+            layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED | flags);
+            layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), { n_embd }, TENSOR_NOT_REQUIRED | flags);
         }
     }
 }
 
 std::unique_ptr<llm_graph_context> llama_model_glm4_moe::build_arch_graph(const llm_graph_params & params) const {
+    if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
+        return std::make_unique<graph_mtp>(*this, params);
+    }
     return std::make_unique<graph>(*this, params);
 }
 
+llama_model_glm4_moe::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
+    : llm_graph_context(params) {
+    GGML_ASSERT(hparams.n_layer_nextn > 0 && "GLM4_MOE MTP requires n_layer_nextn > 0");
+    GGML_ASSERT(hparams.n_layer_nextn == 1 && "GLM4_MOE MTP currently only supports a single MTP block");
+
+    const int64_t n_embd_head = hparams.n_embd_head_v();
+    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+
+    const int il = hparams.n_layer() + cparams.nextn_layer_offset;
+    GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
+                cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
+                "nextn_layer_offset out of range [0, n_layer_nextn)");
+
+    const auto & layer = model.layers[il];
+
+    GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");
+    GGML_ASSERT(layer.nextn.enorm   && "MTP block missing nextn.enorm");
+    GGML_ASSERT(layer.nextn.hnorm   && "MTP block missing nextn.hnorm");
+    GGML_ASSERT(layer.ffn_gate_inp  && "MTP block missing ffn_gate_inp");
+
+    auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);
+
+    inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+    ggml_set_input(inp->tokens);
+
+    inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
+    ggml_set_input(inp->embd);
+
+    ggml_tensor * tok_embd;
+    if (ubatch.token) {
+        ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
+        tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
+    } else {
+        tok_embd = inp->embd;
+    }
+    cb(tok_embd, "mtp_tok_embd", il);
+
+    inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
+    ggml_set_input(inp->h);
+    ggml_set_name(inp->h, "mtp_h_input");
+
+    ggml_tensor * h_embd = inp->h;
+
+    res->add_input(std::move(inp));
+
+    ggml_tensor * inp_pos     = build_inp_pos();
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+    auto * inp_attn = build_attn_inp_kv();
+
+    ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
+    cb(h_norm, "mtp_hnorm", il);
+
+    ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
+    cb(e_norm, "mtp_enorm", il);
+
+    ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, 0);
+    cb(concat, "mtp_concat", il);
+
+    ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
+    cb(cur, "mtp_eh_proj", il);
+
+    ggml_tensor * inpSA = cur;
+
+    cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+    cb(cur, "mtp_attn_norm", il);
+
+    auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur,
+            n_embd_head, n_head, n_head_kv, il);
+
+    if (layer.attn_q_norm) {
+        Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il);
+        cb(Qcur, "mtp_Qcur_normed", il);
+    }
+    if (layer.attn_k_norm) {
+        Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il);
+        cb(Kcur, "mtp_Kcur_normed", il);
+    }
+
+    Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot,
+            rope_type, n_ctx_orig, freq_base, freq_scale,
+            ext_factor, attn_factor, beta_fast, beta_slow);
+
+    Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot,
+            rope_type, n_ctx_orig, freq_base, freq_scale,
+            ext_factor, attn_factor, beta_fast, beta_slow);
+
+    cb(Qcur, "mtp_Qcur", il);
+    cb(Kcur, "mtp_Kcur", il);
+    cb(Vcur, "mtp_Vcur", il);
+
+    cur = build_attn(inp_attn,
+            layer.wo, nullptr, layer.wo_s,
+            Qcur, Kcur, Vcur, nullptr, nullptr, nullptr,
+            1.0f / sqrtf(float(n_embd_head)), il);
+    cb(cur, "mtp_attn_out", il);
+
+    ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+    cb(ffn_inp, "mtp_ffn_inp", il);
+
+    cur = build_norm(ffn_inp, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il);
+    cb(cur, "mtp_post_attn_norm", il);
+
+    ggml_tensor * routed_out = build_moe_ffn(cur,
+            layer.ffn_gate_inp,
+            layer.ffn_up_exps,
+            layer.ffn_gate_exps,
+            layer.ffn_down_exps,
+            layer.ffn_exp_probs_b,
+            n_expert, n_expert_used,
+            LLM_FFN_SILU, hparams.expert_weights_norm,
+            hparams.expert_weights_scale,
+            (llama_expert_gating_func_type) hparams.expert_gating_func,
+            il);
+    cb(routed_out, "mtp_ffn_moe_out", il);
+
+    ggml_tensor * shared_out = build_ffn(cur,
+            layer.ffn_up_shexp,   nullptr, nullptr,
+            layer.ffn_gate_shexp, nullptr, nullptr,
+            layer.ffn_down_shexp, nullptr, nullptr,
+            nullptr,
+            LLM_FFN_SILU, LLM_FFN_PAR, il);
+    cb(shared_out, "mtp_ffn_shexp_out", il);
+
+    cur = ggml_add(ctx0, routed_out, shared_out);
+    cb(cur, "mtp_ffn_out", il);
+
+    cur = ggml_add(ctx0, cur, ffn_inp);
+    cb(cur, "mtp_post_ffn", il);
+
+    ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
+            ? layer.nextn.shared_head_norm
+            : model.output_norm;
+    GGML_ASSERT(head_norm_w && "GLM4_MOE MTP: missing both nextn.shared_head_norm and output_norm");
+
+    cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);
+    cb(cur, "h_nextn", -1);
+    res->t_h_nextn = cur;
+
+    if (inp_out_ids) {
+        cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+    }
+    cb(cur, "mtp_shared_head_norm", -1);
+
+    ggml_tensor * head_w = layer.nextn.shared_head_head
+            ? layer.nextn.shared_head_head
+            : model.output;
+    ggml_tensor * head_s = layer.nextn.shared_head_head
+            ? layer.nextn.shared_head_head_s
+            : model.output_s;
+    GGML_ASSERT(head_w && "GLM4_MOE MTP: missing LM head (nextn.shared_head_head or model.output)");
+
+    cur = build_lora_mm(head_w, cur, head_s);
+    cb(cur, "result_output", -1);
+
+    res->t_logits = cur;
+    ggml_build_forward_expand(gf, cur);
+}
+
 llama_model_glm4_moe::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
     const int64_t n_embd_head = hparams.n_embd_head_v();
 
@@ -154,8 +318,7 @@ llama_model_glm4_moe::graph::graph(const llama_model & model, const llm_graph_pa
 
     ggml_tensor * inp_out_ids = build_inp_out_ids();
 
-    // Only process up to last layer (skip final NextN layer)
-    // Final layer tensors are loaded but not processed in forward pass
+    // NextN layers are processed by graph_mtp.
     for (int il = 0; il < n_layer; ++il) {
         ggml_tensor * inpSA = inpL;
 
@@ -205,7 +368,7 @@ llama_model_glm4_moe::graph::graph(const llama_model & model, const llm_graph_pa
                     model.layers[il].wo, NULL, model.layers[il].wo_s,
                     Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
         }
-        if (il == n_layer - 1 && inp_out_ids) {
+        if (il == n_layer - 1 && inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) {
             cur   = ggml_get_rows(ctx0, cur, inp_out_ids);
             inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
         }
@@ -265,6 +428,13 @@ llama_model_glm4_moe::graph::graph(const llama_model & model, const llm_graph_pa
     cur = inpL;
     cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
 
+    cb(cur, "h_nextn", -1);
+    res->t_h_nextn = cur;
+
+    if (cparams.embeddings_nextn && !cparams.embeddings_nextn_masked && inp_out_ids) {
+        cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+    }
+
     cb(cur, "result_norm", -1);
     res->t_embd = cur;
 

src/models/mamba-base.cpp

@@ -182,13 +182,14 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
     ggml_tensor * conv = build_rs(inp, conv_states_all, hparams.n_embd_r(), n_seqs);
     conv               = ggml_reshape_3d(ctx0, conv, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs);
 
-    // {n_embd, n_tokens} => {n_embd, n_seq_tokens, n_seqs}
-    cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs);
-
     // d_in_proj = 2 * self.d_inner + 2 * self.ngroups * self.d_state + self.nheads
 
-    // {n_embd, d_in_proj} @ {n_embd, n_seq_tokens, n_seqs} => {d_in_proj, n_seq_tokens, n_seqs}
+    // Keep the projection 2D: with a {n_embd, 1, n_seqs} batch the CUDA backend
+    // dispatches a column-batched GEMV for what is a large dense GEMM.
+    // {n_embd, d_in_proj} @ {n_embd, n_tokens} => {d_in_proj, n_tokens}
     ggml_tensor * zxBCdt = build_lora_mm(model.layers[il].ssm_in, cur, model.layers[il].ssm_in_s);
+    // {d_in_proj, n_tokens} => {d_in_proj, n_seq_tokens, n_seqs}
+    zxBCdt = ggml_reshape_3d(ctx0, zxBCdt, zxBCdt->ne[0], n_seq_tokens, n_seqs);
 
     // split the above in three
     ggml_tensor * z   = ggml_view_4d(ctx0, zxBCdt, head_dim, n_head, n_seq_tokens, n_seqs, head_dim * zxBCdt->nb[0],
@@ -290,15 +291,12 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
             y = build_norm(y, model.layers[il].ssm_norm, NULL, LLM_NORM_RMS, il);
         }
 
-        y = ggml_reshape_3d(ctx0, y, d_inner, n_seq_tokens, n_seqs);
+        y = ggml_reshape_2d(ctx0, y, d_inner, n_seq_tokens * n_seqs);
 
-        // {d_inner, n_embd} @ {d_inner, n_seq_tokens, n_seqs} => {n_embd, n_seq_tokens, n_seqs}
+        // {d_inner, n_embd} @ {d_inner, n_tokens} => {n_embd, n_tokens}
         cur = build_lora_mm(model.layers[il].ssm_out, y, model.layers[il].ssm_out_s);
     }
 
-    // {n_embd, n_seq_tokens, n_seqs} => {n_embd, n_tokens}
-    cur = ggml_reshape_2d(ctx0, cur, cur->ne[0], n_seq_tokens * n_seqs);
     cb(cur, "mamba_out", il);
-
     return cur;
 }

src/models/models.h

@@ -1412,6 +1412,10 @@ struct llama_model_glm4_moe : public llama_model_base {
         graph(const llama_model & model, const llm_graph_params & params);
     };
 
+    struct graph_mtp : public llm_graph_context {
+        graph_mtp(const llama_model & model, const llm_graph_params & params);
+    };
+
     std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
 };
 

tests/CMakeLists.txt

@@ -228,6 +228,15 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
     set_tests_properties(test-recurrent-state-rollback-nemotron-h PROPERTIES
         FIXTURES_REQUIRED generate-models
     )
+    llama_test(
+        test-recurrent-state-rollback
+        NAME test-recurrent-state-rollback-dsv4
+        LABEL main
+        ARGS -m "${MODEL_DIR}/deepseek4-moe.gguf"
+    )
+    set_tests_properties(test-recurrent-state-rollback-dsv4 PROPERTIES
+        FIXTURES_REQUIRED generate-models
+    )
 endif()
 
 llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp)
@@ -235,6 +244,8 @@ llama_build_and_test(test-jinja.cpp)
 llama_test(test-jinja NAME test-jinja-py ARGS -py LABEL python)
 llama_build_and_test(test-chat-auto-parser.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR})
 llama_build_and_test(test-chat-template.cpp)
+# debug tool for chat template differential analysis (not registered as a test, run it manually)
+llama_build(test-chat-analysis.cpp)
 llama_build_and_test(test-log.cpp)
 llama_build_and_test(
     test-peg-parser.cpp

tests/peg-parser/test-json-serialization.cpp

@@ -8,7 +8,7 @@ void test_json_serialization(testing &t) {
     auto json_serialized = original.to_json().dump();
 
     t.test("compare before/after", [&](testing &t) {
-        auto deserialized = common_peg_arena::from_json(nlohmann::json::parse(json_serialized));
+        auto deserialized = common_peg_arena::from_json(common_json::parse(json_serialized));
 
         // Test complex JSON
         std::string input = R"({"name": "test", "values": [1, 2, 3], "nested": {"a": true}})";
@@ -23,6 +23,6 @@ void test_json_serialization(testing &t) {
     });
 
     t.bench("deserialize", [&]() {
-        auto deserialized = common_peg_arena::from_json(nlohmann::json::parse(json_serialized));
+        auto deserialized = common_peg_arena::from_json(common_json::parse(json_serialized));
     }, 100);
 }

tests/peg-parser/tests.h

@@ -1,7 +1,7 @@
 #pragma once
 
 // Common includes for all test files
-#include <nlohmann/json.hpp>
+#include "json.h"
 #include <string>
 #include <vector>
 
@@ -11,9 +11,9 @@
 #include "simple-tokenize.h"
 
 struct bench_tool_call {
-    std::string            id;
-    std::string            name;
-    nlohmann::ordered_json args;
+    std::string id;
+    std::string name;
+    common_json args;
 };
 
 // Test function declarations

tests/test-chat-analysis.cpp

@@ -11,9 +11,9 @@
 #include <vector>
 #include <algorithm>
 
-#include "nlohmann/json.hpp"
+#include "json.h"
 
-using json = nlohmann::ordered_json;
+using json = common_json;
 
 // ANSI color codes - using 256-color palette for brighter colors (all bold)
 #define ANSI_RESET       "\033[0m"
@@ -84,11 +84,12 @@ static std::string read_file(const std::string & path) {
 }
 
 static void print_usage(const char * program_name) {
-    LOG_ERR("Usage: %s [options]\n", program_name);
+    LOG_ERR("Debug the auto-parser's differential analysis: render a template with/without tools, reasoning, etc. and show the diffs.\n");
+    LOG_ERR("\nUsage: %s [options]\n", program_name);
     LOG_ERR("\nOptions:\n");
     LOG_ERR("  --template <name>       Analyze specific template from test suite (e.g., 'deepseek' or 'DeepSeek-V3.1')\n");
     LOG_ERR("  --template-file <path>  Analyze custom template file\n");
-    LOG_ERR("  --all                   Analyze all templates from test suite\n");
+    LOG_ERR("  --all                   Analyze all templates from test suite (default when no arguments are given)\n");
     LOG_ERR("\nExamples:\n");
     LOG_ERR("  %s --all\n", program_name);
     LOG_ERR("  %s --template deepseek\n", program_name);
@@ -97,14 +98,17 @@ static void print_usage(const char * program_name) {
 
 static bool parse_options(int argc, char ** argv, analysis_options & opts) {
     if (argc < 2) {
-        print_usage(argv[0]);
-        return false;
+        // default mode: analyze all templates from the test suite
+        opts.analyze_all = true;
     }
 
     for (int i = 1; i < argc; ++i) {
         std::string arg = argv[i];
 
-        if (arg == "--all") {
+        if (arg == "-h" || arg == "--help") {
+            print_usage(argv[0]);
+            return false;
+        } else if (arg == "--all") {
             opts.analyze_all = true;
         } else if (arg == "--template") {
             if (i + 1 >= argc) {

tests/test-chat-auto-parser.cpp

@@ -2,11 +2,18 @@
 #include "chat-auto-parser.h"
 #include "chat-peg-parser.h"
 #include "chat.h"
+#include "gguf.h"
+#include "jinja/runtime.h"
+#include "log.h"
 #include "peg-parser.h"
 #include "testing.h"
 
+#include <cstdlib>
+#include <filesystem>
 #include <fstream>
 #include <iostream>
+#include <iterator>
+#include <optional>
 #include <sstream>
 #include <string>
 
@@ -94,11 +101,447 @@ static void test_bailing_v3_tool_format(testing & t);
 
 static void test_role_markers_all_templates(testing & t);
 
+static json build_tools_definition();
+
+//
+// debug mode: analyze a single template and dump the generated parser and grammar
+//
+
+enum class output_mode {
+    ANALYSIS,  // Only output analysis results (default)
+    TEMPLATE,  // Only output rendered template
+    BOTH       // Output both
+};
+
+enum class input_message_type {
+    NONE,                    // Don't render any message scenarios (only analysis)
+    CONTENT_ONLY,            // Simple assistant message with content
+    REASONING_CONTENT,       // Message with reasoning_content + content
+    TOOL_CALL_ONLY,          // Message with tool_calls only
+    CONTENT_TOOL_CALL,       // Message with content + tool_calls
+    REASONING_TOOL_CALL,     // Message with reasoning_content + tool_calls
+    CONTENT_FAKE_TOOL_CALL,  // Message with content but no actual tool_calls (for testing)
+    ALL                      // Render all scenarios
+};
+
+struct debug_options {
+    std::string        template_path;
+    bool               with_tools        = true;
+    bool               generation_prompt = true;
+    bool               enable_reasoning  = true;
+    bool               debug_jinja       = false;
+    bool               force_tool_call   = false;
+    bool               parallel_tool_calls = true;
+    output_mode        mode              = output_mode::BOTH;
+    input_message_type input_message     = input_message_type::NONE;
+};
+
+static std::string read_file(const std::string & path) {
+    std::ifstream fin(path, std::ios::binary);
+    if (!fin.is_open()) {
+        throw std::runtime_error("Could not open file: " + path);
+    }
+    std::ostringstream buf;
+    buf << fin.rdbuf();
+    return buf.str();
+}
+
+static std::string read_gguf_chat_template(const std::string & path) {
+    struct gguf_init_params params = { /*no_alloc =*/true,  // We only need metadata, not tensor data
+                                       /*ctx=*/nullptr };
+
+    struct gguf_context * ctx = gguf_init_from_file(path.c_str(), params);
+    if (ctx == nullptr) {
+        throw std::runtime_error("Could not open GGUF file: " + path);
+    }
+
+    const char * key    = "tokenizer.chat_template";
+    int64_t      key_id = gguf_find_key(ctx, key);
+
+    if (key_id == -1) {
+        gguf_free(ctx);
+        throw std::runtime_error("GGUF file does not contain chat template key: " + std::string(key));
+    }
+
+    const char * template_str = gguf_get_val_str(ctx, key_id);
+    if (template_str == nullptr) {
+        gguf_free(ctx);
+        throw std::runtime_error("GGUF file contains chat template key but value is null");
+    }
+
+    std::string result = template_str;
+    gguf_free(ctx);
+    return result;
+}
+
+static void print_usage(const char * program_name) {
+    LOG_ERR("Test the chat template auto-parser; also usable as a debug tool that shows the generated PEG parser, GBNF grammar and triggers for a given template.\n");
+    LOG_ERR("\nUsage: %s [filter_regex]                       run the automated tests (default)\n", program_name);
+    LOG_ERR("       %s <template_or_gguf_path> [options]     debug a single template\n", program_name);
+    LOG_ERR("\nDebug mode options:\n");
+    LOG_ERR("  --no-tools              Disable tool definitions\n");
+    LOG_ERR("  --force-tool-call       Set tool calls to forced\n");
+    LOG_ERR("  --parallel-tool-calls=0|1 Set parallel_tool_calls (default: 1)\n");
+    LOG_ERR("  --generation-prompt=0|1 Set add_generation_prompt (default: 1)\n");
+    LOG_ERR("  --enable-reasoning=0|1  Enable reasoning parsing (default: 1)\n");
+    LOG_ERR("  --output=MODE           Output mode: analysis, template, both (default: both)\n");
+    LOG_ERR("  --debug-jinja           Enable Jinja fine-grained debug\n");
+    LOG_ERR("  --input-message=TYPE    Message type to render:\n");
+    LOG_ERR("                          content_only, reasoning_content, tool_call_only,\n");
+    LOG_ERR("                          content_tool_call, reasoning_tool_call,\n");
+    LOG_ERR("                          content_fake_tool_call, all\n");
+    LOG_ERR("\nExamples:\n");
+    LOG_ERR("  %s template.jinja --input-message=all --generation-prompt=1\n", program_name);
+    LOG_ERR("  %s template.jinja --output=template --input-message=tool_call_only\n", program_name);
+}
+
+static bool parse_bool_option(const std::string & value) {
+    return value == "1" || value == "true" || value == "yes";
+}
+
+static bool parse_debug_options(int argc, char ** argv, debug_options & opts) {
+    opts.template_path = argv[1];
+
+    for (int i = 2; i < argc; ++i) {
+        std::string arg = argv[i];
+
+        if (arg == "--force-tool-call") {
+            opts.force_tool_call = true;
+        } else if (arg == "--debug-jinja") {
+            opts.debug_jinja = true;
+        } else if (arg == "--no-tools") {
+            opts.with_tools = false;
+        } else if (arg.rfind("--parallel-tool-calls=", 0) == 0) {
+            opts.parallel_tool_calls = parse_bool_option(arg.substr(22));
+        } else if (arg.rfind("--generation-prompt=", 0) == 0) {
+            opts.generation_prompt = parse_bool_option(arg.substr(20));
+        } else if (arg.rfind("--enable-reasoning=", 0) == 0) {
+            opts.enable_reasoning = parse_bool_option(arg.substr(19));
+        } else if (arg.rfind("--output=", 0) == 0) {
+            std::string mode = arg.substr(9);
+            if (mode == "analysis") {
+                opts.mode = output_mode::ANALYSIS;
+            } else if (mode == "template") {
+                opts.mode = output_mode::TEMPLATE;
+            } else if (mode == "both") {
+                opts.mode = output_mode::BOTH;
+            } else {
+                LOG_ERR("Unknown output mode: %s\n", mode.c_str());
+                return false;
+            }
+        } else if (arg.rfind("--input-message=", 0) == 0) {
+            std::string type = arg.substr(16);
+            if (type == "content_only") {
+                opts.input_message = input_message_type::CONTENT_ONLY;
+            } else if (type == "reasoning_content") {
+                opts.input_message = input_message_type::REASONING_CONTENT;
+            } else if (type == "tool_call_only") {
+                opts.input_message = input_message_type::TOOL_CALL_ONLY;
+            } else if (type == "content_tool_call") {
+                opts.input_message = input_message_type::CONTENT_TOOL_CALL;
+            } else if (type == "reasoning_tool_call") {
+                opts.input_message = input_message_type::REASONING_TOOL_CALL;
+            } else if (type == "content_fake_tool_call") {
+                opts.input_message = input_message_type::CONTENT_FAKE_TOOL_CALL;
+            } else if (type == "all") {
+                opts.input_message = input_message_type::ALL;
+            } else {
+                LOG_ERR("Unknown input message type: %s\n", type.c_str());
+                return false;
+            }
+        } else {
+            LOG_ERR("Unknown option: %s\n", arg.c_str());
+            print_usage(argv[0]);
+            return false;
+        }
+    }
+
+    return true;
+}
+
+static json build_debug_user_message() {
+    return json{
+        { "role",    "user"                               },
+        { "content", "Hello, please help me with a task." }
+    };
+}
+
+static json build_content_only_message() {
+    return json{
+        { "role",    "assistant"                                   },
+        { "content", "Hello! I'm here to help you with your task." }
+    };
+}
+
+static json build_reasoning_content_message() {
+    return json{
+        { "role",              "assistant"                                                               },
+        { "content",           "Hello! I'm here to help you with your task."                             },
+        { "reasoning_content", "The user is greeting me and asking for help. I should respond politely." }
+    };
+}
+
+static json build_tool_call_only_message() {
+    return json{
+        { "role",       "assistant"      },
+        { "content",    nullptr          },
+        { "tool_calls",
+         json::array({ json{
+              { "type", "function" },
+              { "function", json{ { "name", "test_function_name" },
+                                  { "arguments", json::object({ { "param1", "value1" }, { "param2", "value2" } }) } } },
+              { "id", "123456789" } } }) }
+    };
+}
+
+static json build_content_tool_call_message() {
+    return json{
+        { "role",       "assistant"                                                                              },
+        { "content",    "I'll help you by calling a function."                                                   },
+        { "tool_calls",
+         json::array({ json{
+              { "type", "function" },
+              { "function",
+                json{ { "name", "test_function_name" },
+                      { "arguments", json::object({ { "param1", "value1" }, { "param2", "value2" } }) } } } } }) }
+    };
+}
+
+static json build_reasoning_tool_call_message() {
+    return json{
+        { "role",              "assistant"                                                                       },
+        { "content",           nullptr                                                                           },
+        { "reasoning_content", "I need to call a function to help with this task."                               },
+        { "tool_calls",
+         json::array({ json{
+              { "type", "function" },
+              { "function",
+                json{ { "name", "test_function_name" },
+                      { "arguments", json::object({ { "param1", "value1" }, { "param2", "value2" } }) } } } } }) }
+    };
+}
+
+static json build_content_fake_tool_call_message() {
+    // This message has content but NO tool_calls field
+    // It's used to test if a template renders tool definitions but not tool calls
+    return json{
+        { "role",    "assistant"                            },
+        { "content", "I'll help you by calling a function." }
+    };
+}
+
+static void render_scenario(const common_chat_template & tmpl,
+                            const std::string &          scenario_name,
+                            const json &                 messages,
+                            const json &                 tools,
+                            bool                         add_generation_prompt,
+                            bool                         enable_thinking) {
+    LOG_ERR("\n=== Scenario: %s ===\n", scenario_name.c_str());
+    LOG_ERR("add_generation_prompt: %s, enable_thinking: %s\n", add_generation_prompt ? "true" : "false",
+            enable_thinking ? "true" : "false");
+
+    // When add_generation_prompt is true, add a trailing user message to trigger the prompt
+    json final_messages = messages;
+    if (add_generation_prompt && !messages.empty() && messages.back().value("role", "") == "assistant") {
+        final_messages.push_back(json{
+            { "role",    "user"                                       },
+            { "content", "Now please continue with another response." }
+        });
+    }
+
+    LOG_ERR("Messages:\n%s\n", final_messages.dump(2).c_str());
+
+    try {
+        generation_params inputs;
+        inputs.messages                         = final_messages;
+        inputs.add_generation_prompt            = add_generation_prompt;
+        inputs.extra_context["enable_thinking"] = enable_thinking;
+
+        if (!tools.is_null() && tools.is_array() && !tools.empty()) {
+            inputs.tools = tools;
+        }
+
+        std::string output = common_chat_template_direct_apply(tmpl, inputs);
+
+        LOG_ERR("\n--- Rendered Output ---\n");
+        LOG_ERR("%s\n", output.c_str());
+        LOG_ERR("--- End Output (length: %zu) ---\n", output.length());
+    } catch (const std::exception & e) {
+        LOG_ERR("Rendering failed: %s\n", e.what());
+    }
+}
+
+static void render_all_scenarios(const common_chat_template & tmpl,
+                                 const json &                 tools,
+                                 bool                         add_generation_prompt,
+                                 bool                         enable_thinking,
+                                 input_message_type           message_type) {
+    json user_msg = build_debug_user_message();
+
+    auto render_if = [&](input_message_type type, const std::string & name, const json & assistant_msg) {
+        if (message_type == input_message_type::ALL || message_type == type) {
+            json messages = json::array({ user_msg, assistant_msg });
+            render_scenario(tmpl, name, messages, tools, add_generation_prompt, enable_thinking);
+        }
+    };
+
+    render_if(input_message_type::CONTENT_ONLY, "content_only", build_content_only_message());
+    render_if(input_message_type::REASONING_CONTENT, "reasoning_content", build_reasoning_content_message());
+    render_if(input_message_type::TOOL_CALL_ONLY, "tool_call_only", build_tool_call_only_message());
+    render_if(input_message_type::CONTENT_TOOL_CALL, "content_tool_call", build_content_tool_call_message());
+    render_if(input_message_type::REASONING_TOOL_CALL, "reasoning_tool_call", build_reasoning_tool_call_message());
+    render_if(input_message_type::CONTENT_FAKE_TOOL_CALL, "content_fake_tool_call",
+              build_content_fake_tool_call_message());
+
+    // Also render with add_generation_prompt=true to show the prompt ending
+    if (message_type == input_message_type::ALL) {
+        LOG_ERR("\n\n=== Generation Prompt Scenarios (add_generation_prompt=true) ===\n");
+
+        json prompt_messages = json::array({ user_msg });
+        render_scenario(tmpl, "generation_prompt_only", prompt_messages, tools, true, enable_thinking);
+
+        // With enable_thinking toggled
+        render_scenario(tmpl, "generation_prompt_thinking_disabled", prompt_messages, tools, true, false);
+    }
+}
+
+static generation_params prepare_debug_params(const debug_options & opts, const json & tools) {
+    generation_params params;
+    params.messages         = json::array({ build_debug_user_message() });
+    params.reasoning_format = opts.enable_reasoning ? COMMON_REASONING_FORMAT_DEEPSEEK : COMMON_REASONING_FORMAT_NONE;
+    params.enable_thinking  = opts.enable_reasoning;
+    params.add_generation_prompt = opts.generation_prompt;
+
+    if (opts.with_tools) {
+        params.tools       = tools;
+        params.tool_choice = opts.force_tool_call ? COMMON_CHAT_TOOL_CHOICE_REQUIRED : COMMON_CHAT_TOOL_CHOICE_AUTO;
+    } else {
+        params.tools       = json();
+        params.tool_choice = COMMON_CHAT_TOOL_CHOICE_NONE;
+    }
+    params.parallel_tool_calls = opts.parallel_tool_calls;
+    return params;
+}
+
+static int debug_single_template(const debug_options & opts) {
+    std::string template_source;
+    try {
+        // Check if the file is a GGUF file
+        if (opts.template_path.size() >= 5 &&
+            opts.template_path.compare(opts.template_path.size() - 5, 5, ".gguf") == 0) {
+            template_source = read_gguf_chat_template(opts.template_path);
+        } else {
+            template_source = read_file(opts.template_path);
+        }
+    } catch (const std::exception & e) {
+        LOG_ERR("Error reading template: %s\n", e.what());
+        return 1;
+    }
+
+    LOG_ERR("Analyzing template: %s\n", opts.template_path.c_str());
+    LOG_ERR("Options: with_tools=%s, generation_prompt=%s, enable_reasoning=%s\n", opts.with_tools ? "true" : "false",
+            opts.generation_prompt ? "true" : "false", opts.enable_reasoning ? "true" : "false");
+
+    try {
+        common_chat_template chat_template(template_source, "", "");
+
+        json tools = opts.with_tools ? build_tools_definition() : json();
+
+        generation_params params = prepare_debug_params(opts, tools);
+        common_chat_params            parser_data;
+        if (std::optional<common_chat_params> spec_tmpl =
+                common_chat_try_specialized_template(chat_template, template_source, params)) {
+            LOG_ERR("\n");
+            LOG_ERR("This template uses a specialized parser, analysis results will not be available.\n");
+            parser_data = *spec_tmpl;
+        } else {
+            // Render template scenarios if requested
+            if (opts.input_message != input_message_type::NONE &&
+                (opts.mode == output_mode::TEMPLATE || opts.mode == output_mode::BOTH)) {
+                LOG_ERR("\n");
+                LOG_ERR("================================================================================\n");
+                LOG_ERR("                         TEMPLATE RENDERING OUTPUT\n");
+                LOG_ERR("================================================================================\n");
+
+                render_all_scenarios(chat_template, tools, opts.generation_prompt, opts.enable_reasoning,
+                                     opts.input_message);
+            }
+
+            // Output analysis if requested
+            if (opts.mode == output_mode::ANALYSIS || opts.mode == output_mode::BOTH) {
+                LOG_ERR("\n");
+                LOG_ERR("================================================================================\n");
+                LOG_ERR("                           TEMPLATE ANALYSIS\n");
+                LOG_ERR("================================================================================\n");
+
+                struct autoparser analysis;
+                analysis.analyze_template(chat_template);
+
+                // Generate Parser
+                parser_data = peg_generator::generate_parser(chat_template, params, analysis);
+            }
+        }
+
+        if (!std::empty(parser_data.parser)) {
+            LOG_ERR("\n=== Generated Parser ===\n");
+            common_peg_arena arena;
+            arena.load(parser_data.parser);
+            LOG_ERR("%s\n", arena.dump(arena.root()).c_str());
+
+            LOG_ERR("\n=== Generated Grammar ===\n");
+            LOG_ERR("%s\n", parser_data.grammar.c_str());
+
+            LOG_ERR("\n=== Generated Lazy Grammar ===\n");
+            LOG_ERR("%d\n", parser_data.grammar_lazy);
+
+            LOG_ERR("\n=== Generated Grammar Triggers ===\n");
+            for (const common_grammar_trigger & cgt : parser_data.grammar_triggers) {
+                LOG_ERR("Token: %d | Type: %d | Value: %s\n", cgt.token, cgt.type, cgt.value.c_str());
+            }
+
+            LOG_ERR("\n=== Preserved Tokens ===\n");
+            for (const std::string & token : parser_data.preserved_tokens) {
+                LOG_ERR("  '%s'\n", token.c_str());
+            }
+        }
+    } catch (const std::exception & e) {
+        LOG_ERR("Analysis failed: %s\n", e.what());
+        return 1;
+    }
+
+    return 0;
+}
+
 int main(int argc, char * argv[]) {
+    if (argc > 1) {
+        std::string arg = argv[1];
+        if (arg == "-h" || arg == "--help") {
+            common_log_set_verbosity_thold(99);
+            print_usage(argv[0]);
+            return 0;
+        }
+
+        // debug mode: if the first argument is an existing file, analyze that template instead of running the automated tests
+        if (std::filesystem::is_regular_file(arg)) {
+            common_log_set_verbosity_thold(99);
+
+            debug_options opts;
+            if (!parse_debug_options(argc, argv, opts)) {
+                return 1;
+            }
+
+            if (opts.debug_jinja || std::getenv("LLAMA_DEBUG_JINJA") != nullptr) {
+                jinja::enable_debug(true);
+            }
+
+            return debug_single_template(opts);
+        }
+    }
+
     testing t(std::cout);
     t.verbose = true;
 
-    // usage: test-chat-auto-parser-helpers [filter_regex]
+    // usage: test-chat-auto-parser [filter_regex]
 
     if (argc > 1) {
         t.set_filter(argv[1]);