llama.cpp : bump version to 0.4.1 (#28900)

* tests : add README for updating the per-backend fusion baselines

Assisted-by: pi:llama.cpp/Qwen3.8-27B

* ci : trigger fusion on changes to test-llama-archs.cpp and src/models

the dummy models and their architectures drive the fusion baselines, so a
change to either can alter the per-fusion counters and should re-run the
fusion job.

Assisted-by: pi:llama.cpp/Qwen3.8-27B

* tests : merge the fusion build commands in the README

assisted-by: pi:llama.cpp/Qwen3.8-27B

* pi : require explicit permission before posting PR/issue comments

assisted-by: pi:llama.cpp/Qwen3.8-27B

.github/workflows/build-cpu.yml

@@ -221,7 +221,6 @@ jobs:
       #     7z x "-o${env:RUNNER_TEMP}" $env:RUNNER_TEMP/sde.tar
       #     $sde = $(join-path $env:RUNNER_TEMP sde-external-${env:SDE_VERSION}-win/sde.exe)
       #     cd build
-      #     $env:LLAMA_SKIP_TESTS_SLOW_ON_EMULATOR = 1
       #     & $sde -future -- ctest -L main -C Release --verbose --timeout 900
 
       - name: ccache-clear

.github/workflows/build-ibm.yml

@@ -34,10 +34,15 @@ env:
   LLAMA_ARG_LOG_TIMESTAMPS: 1
 
 jobs:
-
   ubuntu-24-s390x:
+    name: ubuntu-24-s390x (VXE ${{ matrix.vxe }})
     runs-on: ubuntu-24.04-s390x
 
+    strategy:
+      fail-fast: false
+      matrix:
+        vxe: ["ON", "OFF"]  # `-DGGML_VXE=ON/OFF`
+
     steps:
       - name: Clone
         id: checkout
@@ -77,7 +82,8 @@ jobs:
         run: |
           cmake -B build \
             -DLLAMA_FATAL_WARNINGS=ON \
-            -DGGML_RPC=ON
+            -DGGML_RPC=ON \
+            -DGGML_VXE=${{ matrix.vxe }}
           time cmake --build build --config Release -j $(nproc)
 
       - name: Test

.github/workflows/build-openvino.yml

@@ -33,7 +33,7 @@ env:
   LLAMA_ARG_LOG_PREFIX: 1
   LLAMA_ARG_LOG_TIMESTAMPS: 1
   # TODO: fix failing tests on OpenVINO backend
-  CTEST_EXCLUDE: "test-llama-archs|^test-recurrent-state-|test-backend-ops|test-save-load-state"
+  CTEST_EXCLUDE: "test-llama-archs|^test-recurrent-state-|test-save-load-state"
 
 jobs:
   ubuntu-24-openvino:

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

@@ -7,6 +7,7 @@ on:
       - master
     paths: [
       '.github/workflows/build-self-hosted.yml',
+      'ci/run.sh',
       '**/CMakeLists.txt',
       '**/.cmake',
       '**/*.h',
@@ -27,6 +28,7 @@ on:
     types: [opened, synchronize, reopened]
     paths: [
       '.github/workflows/build-self-hosted.yml',
+      'ci/run.sh',
       '**/CMakeLists.txt',
       '**/.cmake',
       '**/*.h',
@@ -160,10 +162,10 @@ jobs:
       - name: Dawn Dependency
         id: dawn-depends
         run: |
-          DAWN_VERSION="v20260317.182325"
+          DAWN_VERSION="v20260908.214631"
           DAWN_OWNER="google"
           DAWN_REPO="dawn"
-          DAWN_ASSET_NAME="Dawn-18eb229ef5f707c1464cc581252e7603c73a3ef0-ubuntu-latest-Release"
+          DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-ubuntu-latest-Release"
           echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
           curl -L -o artifact.tar.gz \
             "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
@@ -246,10 +248,10 @@ jobs:
       - name: Dawn Dependency
         id: dawn-depends
         run: |
-          DAWN_VERSION="v20260317.182325"
+          DAWN_VERSION="v20260908.214631"
           DAWN_OWNER="google"
           DAWN_REPO="dawn"
-          DAWN_ASSET_NAME="Dawn-18eb229ef5f707c1464cc581252e7603c73a3ef0-macos-latest-Release"
+          DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-macos-latest-Release"
           echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
           curl -L -o artifact.tar.gz \
             "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
@@ -395,7 +397,11 @@ jobs:
       - name: Test
         id: ggml-ci
         run: |
-          LLAMA_ARG_THREADS=$(nproc) GG_BUILD_HIGH_PERF=1 GG_BUILD_NO_BF16=1 GG_BUILD_EXTRA_TESTS_0=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
+          LLAMA_ARG_THREADS=$(nproc) \
+          GG_BUILD_HIGH_PERF=1 \
+          GG_BUILD_NO_BF16=1 \
+          GG_BUILD_EXTRA_TESTS_0=1 \
+          bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
 
   cpu-arm64-graviton4-kleidiai:
     runs-on: ah-ubuntu_22_04-c8g_8x
@@ -434,6 +440,8 @@ jobs:
       - name: Test
         id: ggml-ci
         run: |
+          LLAMA_ARG_THREADS=$(nproc) \
           GG_BUILD_KLEIDIAI=1 \
           GG_BUILD_EXTRA_TESTS_0=1 \
-          bash ./ci/run.sh ./tmp/results ./tmp/mnt
+          GG_BUILD_HIGH_PERF=1 \
+          bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp

.github/workflows/build-vulkan.yml

@@ -164,9 +164,7 @@ jobs:
           export GGML_VK_VISIBLE_DEVICES=0
           export GGML_VK_DISABLE_F16=1
           export GGML_VK_DISABLE_COOPMAT=1
-          # 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
+          ctest -L main --verbose --timeout 900
 
   windows:
     runs-on: windows-2025

.github/workflows/build-wasm.yml

@@ -68,7 +68,7 @@ jobs:
 
       - name: Fetch emdawnwebgpu
         run: |
-          DAWN_TAG="v20260317.182325"
+          DAWN_TAG="v20260908.214631"
           EMDAWN_PKG="emdawnwebgpu_pkg-${DAWN_TAG}.zip"
           echo "Downloading ${EMDAWN_PKG}"
           curl -L -o emdawn.zip \

.github/workflows/build-webgpu.yml

@@ -77,10 +77,10 @@ jobs:
       - name: Dawn Dependency
         id: dawn-depends
         run: |
-          DAWN_VERSION="v20260317.182325"
+          DAWN_VERSION="v20260908.214631"
           DAWN_OWNER="google"
           DAWN_REPO="dawn"
-          DAWN_ASSET_NAME="Dawn-18eb229ef5f707c1464cc581252e7603c73a3ef0-macos-latest-Release"
+          DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-macos-latest-Release"
           echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
           curl -L -o artifact.tar.gz \
             "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
@@ -147,10 +147,10 @@ jobs:
         id: dawn-depends
         run: |
           sudo apt-get install -y libxrandr-dev libxinerama-dev libxcursor-dev mesa-common-dev libx11-xcb-dev libxi-dev
-          DAWN_VERSION="v20260317.182325"
+          DAWN_VERSION="v20260908.214631"
           DAWN_OWNER="google"
           DAWN_REPO="dawn"
-          DAWN_ASSET_NAME="Dawn-18eb229ef5f707c1464cc581252e7603c73a3ef0-ubuntu-latest-Release"
+          DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-ubuntu-latest-Release"
           echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
           curl -L -o artifact.tar.gz \
             "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
@@ -190,6 +190,4 @@ jobs:
         id: cmake_test
         run: |
           cd build
-          # 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
+          ctest -L main --verbose --timeout 900

.github/workflows/code-style.yml

@@ -15,7 +15,7 @@ concurrency:
 
 jobs:
   model-naming:
-    runs-on: [self-hosted, fast]
+    runs-on: ubuntu-slim
     steps:
       - uses: actions/checkout@v6
       - name: Check model naming conventions

.github/workflows/editorconfig.yml

@@ -15,7 +15,7 @@ concurrency:
 
 jobs:
   editorconfig:
-    runs-on: [self-hosted, fast]
+    runs-on: ubuntu-slim
     steps:
       - uses: actions/checkout@v6
       - uses: editorconfig-checker/action-editorconfig-checker@840e866d93b8e032123c23bac69dece044d4d84c # v2.2.0

.github/workflows/fusion.yml

@@ -9,7 +9,9 @@ on:
       '.github/workflows/fusion.yml',
       'ggml/**',
       'tests/fusion/**',
-      'tests/test-fusion.cpp'
+      'tests/test-fusion.cpp',
+      'tests/test-llama-archs.cpp',
+      'src/models/**'
     ]
 
   pull_request:
@@ -18,7 +20,9 @@ on:
       '.github/workflows/fusion.yml',
       'ggml/**',
       'tests/fusion/**',
-      'tests/test-fusion.cpp'
+      'tests/test-fusion.cpp',
+      'tests/test-llama-archs.cpp',
+      'src/models/**'
     ]
 
 concurrency:

.github/workflows/server-sanitize.yml

@@ -116,7 +116,6 @@ jobs:
         run: |
           source .venv/bin/activate
           cd tools/server/tests
-          export ${{ matrix.extra_args }}
           PYTEST_WORKERS=1 ./tests.sh
 
       - name: Slow tests
@@ -125,5 +124,4 @@ jobs:
         run: |
           source .venv/bin/activate
           cd tools/server/tests
-          export ${{ matrix.extra_args }}
           PYTEST_WORKERS=1 SLOW_TESTS=1 ./tests.sh

.pi/gg/SYSTEM.md

@@ -6,6 +6,7 @@ General:
 - PR and commit titles format: `<module> : <title>`. Lookup recents for examples
 - Don't try to build or run the code unless you are explicitly asked to do so
 - Use the `gh` CLI tool when querying PRs, issues, or other GitHub resources
+- When [MODEL] is needed, first try to get it from the `PI_MODEL_NAME` env var before asking the user
 
 Coding:
 - When in doubt, always refer to the CONTRIBUTING.md file of the project
@@ -20,8 +21,9 @@ Pull requests (PRs):
 - 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]
+- If `PI_MODEL_NAME` env var is not set, ask the user to tell you what model was used and write it in place of [MODEL]
 - Always create the pull requests in draft mode
+- Never reply to review comments or post comments on issues/PRs without explicit permission from the user
 
 Commits:
 - On every commit that you make, include a "Assisted-by: pi:llama.cpp/[MODEL]" tag

CMakeLists.txt

@@ -5,7 +5,7 @@ include(CheckIncludeFileCXX)
 ### llama.cpp version
 set(LLAMA_VERSION_MAJOR 0)
 set(LLAMA_VERSION_MINOR 4)
-set(LLAMA_VERSION_PATCH 0)
+set(LLAMA_VERSION_PATCH 1)
 set(LLAMA_VERSION_BASE "${LLAMA_VERSION_MAJOR}.${LLAMA_VERSION_MINOR}.${LLAMA_VERSION_PATCH}")
 
 # whether this is a development/nightly build
@@ -197,6 +197,16 @@ llama_option_depr(WARNING     LLAMA_CURL)
 include("cmake/license.cmake")
 license_add_file("llama.cpp" "LICENSE")
 
+#
+# compile options
+#
+
+# clang stores the modification time of the precompiled header sources inside the
+# header and rejects it when they differ, so the timestamp is left out of it
+add_compile_options(
+    "$<$<COMPILE_LANG_AND_ID:C,Clang,IntelLLVM>:SHELL:-Xclang -fno-pch-timestamp>"
+    "$<$<COMPILE_LANG_AND_ID:CXX,Clang,IntelLLVM>:SHELL:-Xclang -fno-pch-timestamp>")
+
 #
 # 3rd-party
 #

ci/run.sh

@@ -58,8 +58,6 @@ if [ -n "${GG_BUILD_ROCM}" ] && [ -n "${GITHUB_RUN_ID}" ]; then
 fi
 
 rm -f $OUT/*.log
-rm -f $OUT/*.exit
-rm -f $OUT/*.md
 
 sd=`dirname $0`
 cd $sd/../
@@ -190,7 +188,7 @@ if [ ! -z ${GG_BUILD_OPENVINO} ]; then
     CMAKE_EXTRA="${CMAKE_EXTRA} -DGGML_OPENVINO=ON"
 
     # TODO: fix failing tests on OpenVINO backend
-    CTEST_EXTRA="-E test-llama-archs|^test-recurrent-state-|test-backend-ops|test-save-load-state"
+    CTEST_EXTRA="-E test-llama-archs|^test-recurrent-state-|test-save-load-state"
 fi
 
 ## helpers
@@ -211,10 +209,6 @@ function gg_wget {
     cd $cwd
 }
 
-function gg_printf {
-    printf -- "$@" >> $OUT/README.md
-}
-
 function gg_run {
     ci=$1
 
@@ -223,13 +217,10 @@ function gg_run {
 
     gg_run_$ci | tee $OUT/$ci.log
     cur=$?
-    echo "$cur" > $OUT/$ci.exit
 
     set +x
     set +o pipefail
 
-    gg_sum_$ci
-
     ret=$((ret | cur))
 }
 
@@ -250,22 +241,11 @@ function gg_run_ctest_debug {
     (cmake -G "${CMAKE_GENERATOR}" -DCMAKE_BUILD_TYPE=Debug ${CMAKE_EXTRA} .. ) 2>&1 | tee -a $OUT/${ci}-cmake.log
     (time cmake --build . --config Debug -j$(nproc)) 2>&1 | tee -a $OUT/${ci}-make.log
 
-    (time ctest -C Debug --output-on-failure -L main -E "test-opt|test-backend-ops|test-llama-archs" ${CTEST_EXTRA}) 2>&1 | tee -a $OUT/${ci}-ctest.log
+    (time ctest -C Debug --output-on-failure -L main -E "test-opt|test-llama-archs" ${CTEST_EXTRA}) 2>&1 | tee -a $OUT/${ci}-ctest.log
 
     set +e
 }
 
-function gg_sum_ctest_debug {
-    gg_printf '### %s\n\n' "${ci}"
-
-    gg_printf 'Runs ctest in debug mode\n'
-    gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)"
-    gg_printf '```\n'
-    gg_printf '%s\n' "$(cat $OUT/${ci}-ctest.log)"
-    gg_printf '```\n'
-    gg_printf '\n'
-}
-
 # ctest_release
 
 function gg_run_ctest_release {
@@ -290,16 +270,6 @@ function gg_run_ctest_release {
     set +e
 }
 
-function gg_sum_ctest_release {
-    gg_printf '### %s\n\n' "${ci}"
-
-    gg_printf 'Runs ctest in release mode\n'
-    gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)"
-    gg_printf '```\n'
-    gg_printf '%s\n' "$(cat $OUT/${ci}-ctest.log)"
-    gg_printf '```\n'
-}
-
 # test_llama_archs_tensor_split
 
 function gg_run_test_llama_archs_tensor_split {
@@ -324,16 +294,6 @@ function gg_run_test_llama_archs_tensor_split {
     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_llama_archs_models
 
 function gg_run_test_llama_archs_models {
@@ -353,16 +313,6 @@ function gg_run_test_llama_archs_models {
     set +e
 }
 
-function gg_sum_test_llama_archs_models {
-    gg_printf '### %s\n\n' "${ci}"
-
-    gg_printf 'Generates the dummy models used by the model-dependent tests\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 {
@@ -376,17 +326,6 @@ function gg_run_test_scripts {
     set +e
 }
 
-function gg_sum_test_scripts {
-    gg_printf '### %s\n\n' "${ci}"
-
-    gg_printf 'Runs test scripts\n'
-    gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)"
-    gg_printf '```\n'
-    gg_printf '%s\n' "$(cat $OUT/${ci}-scripts.log)"
-    gg_printf '```\n'
-    gg_printf '\n'
-}
-
 function gg_get_model {
     #local gguf_0="$MNT/models/qwen3/0.6B/ggml-model-f16.gguf"
     local gguf_0="$MNT/models/qwen3/0.6B/ggml-model-q4_0.gguf"
@@ -430,26 +369,6 @@ function gg_run_ctest_with_model_release {
     cd ..
 }
 
-function gg_sum_ctest_with_model_debug {
-    gg_printf '### %s\n\n' "${ci}"
-
-    gg_printf 'Runs ctest with model files in debug mode\n'
-    gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)"
-    gg_printf '```\n'
-    gg_printf '%s\n' "$(cat $OUT/${ci}-ctest.log)"
-    gg_printf '```\n'
-}
-
-function gg_sum_ctest_with_model_release {
-    gg_printf '### %s\n\n' "${ci}"
-
-    gg_printf 'Runs ctest with model files in release mode\n'
-    gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)"
-    gg_printf '```\n'
-    gg_printf '%s\n' "$(cat $OUT/${ci}-ctest.log)"
-    gg_printf '```\n'
-}
-
 # qwen3_0_6b
 
 function gg_run_qwen3_0_6b {
@@ -554,50 +473,24 @@ function gg_run_qwen3_0_6b {
         return 0
     }
 
-    check_ppl "f16"  "$(cat $OUT/${ci}-tg-f16.log  | grep "^\[1\]")" | tee -a $OUT/${ci}-ppl.log
+    check_ppl "f16"  "$(cat $OUT/${ci}-tg-f16.log  | grep "^\[1\]")"
     if [ -z ${GG_BUILD_NO_BF16} ]; then
-        check_ppl "bf16" "$(cat $OUT/${ci}-tg-bf16.log | grep "^\[1\]")" | tee -a $OUT/${ci}-ppl.log
+        check_ppl "bf16" "$(cat $OUT/${ci}-tg-bf16.log | grep "^\[1\]")"
     fi
-    check_ppl "q8_0" "$(cat $OUT/${ci}-tg-q8_0.log | grep "^\[1\]")" | tee -a $OUT/${ci}-ppl.log
-    check_ppl "q4_0" "$(cat $OUT/${ci}-tg-q4_0.log | grep "^\[1\]")" | tee -a $OUT/${ci}-ppl.log
-    check_ppl "q4_1" "$(cat $OUT/${ci}-tg-q4_1.log | grep "^\[1\]")" | tee -a $OUT/${ci}-ppl.log
-    check_ppl "q5_0" "$(cat $OUT/${ci}-tg-q5_0.log | grep "^\[1\]")" | tee -a $OUT/${ci}-ppl.log
-    check_ppl "q5_1" "$(cat $OUT/${ci}-tg-q5_1.log | grep "^\[1\]")" | tee -a $OUT/${ci}-ppl.log
-   #check_ppl "q2_k" "$(cat $OUT/${ci}-tg-q2_k.log | grep "^\[1\]")" | tee -a $OUT/${ci}-ppl.log # note: ppl > 20.0 for this quant and model
-    check_ppl "q3_k" "$(cat $OUT/${ci}-tg-q3_k.log | grep "^\[1\]")" | tee -a $OUT/${ci}-ppl.log
-    check_ppl "q4_k" "$(cat $OUT/${ci}-tg-q4_k.log | grep "^\[1\]")" | tee -a $OUT/${ci}-ppl.log
-    check_ppl "q5_k" "$(cat $OUT/${ci}-tg-q5_k.log | grep "^\[1\]")" | tee -a $OUT/${ci}-ppl.log
-    check_ppl "q6_k" "$(cat $OUT/${ci}-tg-q6_k.log | grep "^\[1\]")" | tee -a $OUT/${ci}-ppl.log
-
-    cat $OUT/${ci}-imatrix.log | grep "Final" >> $OUT/${ci}-imatrix-sum.log
+    check_ppl "q8_0" "$(cat $OUT/${ci}-tg-q8_0.log | grep "^\[1\]")"
+    check_ppl "q4_0" "$(cat $OUT/${ci}-tg-q4_0.log | grep "^\[1\]")"
+    check_ppl "q4_1" "$(cat $OUT/${ci}-tg-q4_1.log | grep "^\[1\]")"
+    check_ppl "q5_0" "$(cat $OUT/${ci}-tg-q5_0.log | grep "^\[1\]")"
+    check_ppl "q5_1" "$(cat $OUT/${ci}-tg-q5_1.log | grep "^\[1\]")"
+   #check_ppl "q2_k" "$(cat $OUT/${ci}-tg-q2_k.log | grep "^\[1\]")" # note: ppl > 20.0 for this quant and model
+    check_ppl "q3_k" "$(cat $OUT/${ci}-tg-q3_k.log | grep "^\[1\]")"
+    check_ppl "q4_k" "$(cat $OUT/${ci}-tg-q4_k.log | grep "^\[1\]")"
+    check_ppl "q5_k" "$(cat $OUT/${ci}-tg-q5_k.log | grep "^\[1\]")"
+    check_ppl "q6_k" "$(cat $OUT/${ci}-tg-q6_k.log | grep "^\[1\]")"
 
     set +e
 }
 
-function gg_sum_qwen3_0_6b {
-    gg_printf '### %s\n\n' "${ci}"
-
-    gg_printf 'Qwen3 0.6B:\n'
-    gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)"
-    gg_printf '- perplexity:\n%s\n' "$(cat $OUT/${ci}-ppl.log)"
-    gg_printf '- imatrix:\n```\n%s\n```\n' "$(cat $OUT/${ci}-imatrix-sum.log)"
-    gg_printf '- f16:\n```\n%s\n```\n'  "$(cat $OUT/${ci}-tg-f16.log)"
-    if [ -z ${GG_BUILD_NO_BF16} ]; then
-        gg_printf '- bf16:\n```\n%s\n```\n' "$(cat $OUT/${ci}-tg-bf16.log)"
-    fi
-    gg_printf '- q8_0:\n```\n%s\n```\n' "$(cat $OUT/${ci}-tg-q8_0.log)"
-    gg_printf '- q4_0:\n```\n%s\n```\n' "$(cat $OUT/${ci}-tg-q4_0.log)"
-    gg_printf '- q4_1:\n```\n%s\n```\n' "$(cat $OUT/${ci}-tg-q4_1.log)"
-    gg_printf '- q5_0:\n```\n%s\n```\n' "$(cat $OUT/${ci}-tg-q5_0.log)"
-    gg_printf '- q5_1:\n```\n%s\n```\n' "$(cat $OUT/${ci}-tg-q5_1.log)"
-    gg_printf '- q2_k:\n```\n%s\n```\n' "$(cat $OUT/${ci}-tg-q2_k.log)"
-    gg_printf '- q3_k:\n```\n%s\n```\n' "$(cat $OUT/${ci}-tg-q3_k.log)"
-    gg_printf '- q4_k:\n```\n%s\n```\n' "$(cat $OUT/${ci}-tg-q4_k.log)"
-    gg_printf '- q5_k:\n```\n%s\n```\n' "$(cat $OUT/${ci}-tg-q5_k.log)"
-    gg_printf '- q6_k:\n```\n%s\n```\n' "$(cat $OUT/${ci}-tg-q6_k.log)"
-    gg_printf '- save-load-state: \n```\n%s\n```\n' "$(cat $OUT/${ci}-save-load-state.log)"
-}
-
 # bge-small
 
 function gg_run_embd_bge_small {
@@ -639,15 +532,6 @@ function gg_run_embd_bge_small {
     set +e
 }
 
-function gg_sum_embd_bge_small {
-    gg_printf '### %s\n\n' "${ci}"
-
-    gg_printf 'BGE Small (BERT):\n'
-    gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)"
-    gg_printf '- f16: \n```\n%s\n```\n' "$(cat $OUT/${ci}-tg-f16.log)"
-    gg_printf '- q8_0:\n```\n%s\n```\n' "$(cat $OUT/${ci}-tg-q8_0.log)"
-}
-
 # rerank_tiny
 
 function gg_run_rerank_tiny {
@@ -704,93 +588,97 @@ function gg_run_rerank_tiny {
     set +e
 }
 
-function gg_sum_rerank_tiny {
-    gg_printf '### %s\n\n' "${ci}"
-
-    gg_printf 'Rerank Tiny (Jina):\n'
-    gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)"
-    gg_printf '- f16: \n```\n%s\n```\n' "$(cat $OUT/${ci}-rk-f16.log)"
-}
-
 function gg_check_build_requirements {
     if ! command -v git &> /dev/null; then
-        gg_printf 'git not found, please install\n'
+        echo 'git not found, please install'
         exit 1
     fi
 
     if ! command -v git-lfs &> /dev/null; then
-        gg_printf 'git-lfs not found, please install\n'
+        echo 'git-lfs not found, please install'
         exit 1
     fi
 
     if ! git config --get filter.lfs.clean &> /dev/null; then
-        gg_printf 'git-lfs not initialized, please run `git lfs install`\n'
+        echo 'git-lfs not initialized, please run `git lfs install`'
         exit 1
     fi
 
     if ! command -v wget &> /dev/null; then
-        gg_printf 'wget not found, please install\n'
+        echo 'wget not found, please install'
         exit 1
     fi
 
     if ! command -v python3 &> /dev/null; then
-        gg_printf 'python3 not found, please install\n'
+        echo 'python3 not found, please install'
         exit 1
     fi
 
     if ! command -v pip3 &> /dev/null; then
-        gg_printf 'pip3 not found, please install\n'
+        echo 'pip3 not found, please install'
         exit 1
     fi
 
     if ! python3 -m ensurepip --help &> /dev/null; then
-        gg_printf 'ensurepip not found, please install python3-venv package\n'
+        echo 'ensurepip not found, please install python3-venv package'
         exit 1
     fi
 
     if ! command -v cmake &> /dev/null; then
-        gg_printf 'cmake not found, please install\n'
+        echo 'cmake not found, please install'
         exit 1
     fi
 
     if ! command -v ccache &> /dev/null; then
-        gg_printf 'ccache not found, please consider installing for faster builds\n'
+        echo 'ccache not found, please consider installing for faster builds'
     fi
 
     if ! command -v ctest &> /dev/null; then
-        gg_printf 'ctest not found, please install\n'
+        echo 'ctest not found, please install'
         exit 1
     fi
 
     if ! command -v unzip &> /dev/null; then
-        gg_printf 'unzip not found, please install\n'
+        echo 'unzip not found, please install'
         exit 1
     fi
 }
 
-function gg_run_test_backend_ops_cpu {
+function gg_run_test_backend_ops {
     cd ${SRC}
 
     cd build-ci-release
 
     set -e
 
-    (time ./bin/test-backend-ops -b CPU ) 2>&1 | tee -a $OUT/${ci}-test-backend-ops-cpu.log
+    local n_jobs=$(nproc)
+    if [ "${n_jobs}" -gt 2 ]; then
+        n_jobs=2
+    fi
+    local args_extra="-j ${n_jobs}"
+
+    # TODO: fix multi-threaded for ROCm
+    #       https://github.com/ggml-org/llama.cpp/actions/runs/34576278519/job/103297889044?pr=28740#step:3:4865
+    if [ ! -z ${GG_BUILD_ROCM} ]; then
+        args_extra=""
+    fi
+
+    # TODO: MoltenVK bug?
+    #       https://github.com/ggml-org/llama.cpp/actions/runs/34611260059/job/103302413736?pr=28740#step:3:5897
+    if [ ! -z "${GG_BUILD_VULKAN}" ] && [ "$(uname -s)" = "Darwin" ]; then
+        args_extra=""
+    fi
+
+    # TODO: reduce the test-backend-ops timeout to 1800s
+    if [ ! -z ${GG_BUILD_HIGH_PERF} ]; then
+        (time timeout 3600 ./bin/test-backend-ops ${args_extra} -b CPU) 2>&1 | tee -a $OUT/${ci}-test-backend-ops.log
+    else
+        (time timeout 3600 ./bin/test-backend-ops ${args_extra}       ) 2>&1 | tee -a $OUT/${ci}-test-backend-ops.log
+    fi
 
     set +e
 }
 
-function gg_sum_test_backend_ops_cpu {
-    gg_printf '### %s\n\n' "${ci}"
-
-    gg_printf 'Runs test-backend-ops for CPU backend\n'
-    gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)"
-    gg_printf '```\n'
-    gg_printf '%s\n' "$(cat $OUT/${ci}-test-backend-ops-cpu.log)"
-    gg_printf '```\n'
-    gg_printf '\n'
-}
-
 ## main
 
 export LLAMA_ARG_LOG_PREFIX=1
@@ -819,13 +707,11 @@ 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_backend_ops
+
 test $ret -eq 0 && gg_run test_llama_archs_models
 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
-
 if [ -z ${GG_BUILD_LOW_PERF} ]; then
     test $ret -eq 0 && gg_run embd_bge_small
     test $ret -eq 0 && gg_run rerank_tiny
@@ -840,6 +726,4 @@ if [ -z ${GG_BUILD_LOW_PERF} ]; then
     test $ret -eq 0 && gg_run ctest_with_model_release
 fi
 
-cat $OUT/README.md
-
 exit $ret

common/CMakeLists.txt

@@ -84,6 +84,8 @@ add_library(${TARGET}
     imatrix-loader.cpp
     imatrix-loader.h
     json-schema-to-grammar.cpp
+    json-schema.cpp
+    json-schema.h
     json.cpp
     json.h
     llguidance.cpp
@@ -134,6 +136,8 @@ set_target_properties(${TARGET} PROPERTIES
 target_include_directories(${TARGET} PUBLIC .)
 target_link_libraries     (${TARGET} PUBLIC vendor::nlohmann vendor::sheredom)
 target_compile_features   (${TARGET} PUBLIC cxx_std_17)
+target_precompile_headers (${TARGET} PRIVATE common.h)
+target_precompile_headers (${TARGET} PRIVATE chat.h)
 
 if (LLAMA_SUBPROCESS)
     target_compile_definitions(${TARGET} PUBLIC LLAMA_SUBPROCESS)

common/arg.cpp

@@ -2277,14 +2277,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
     ).set_sampling());
     add_opt(common_arg(
         {"-j", "--json-schema"}, "SCHEMA",
-        "JSON schema to constrain generations (https://json-schema.org/), e.g. `{}` for any JSON object\nFor schemas w/ external $refs, use --grammar + example/json_schema_to_grammar.py instead",
+        "JSON schema to constrain generations (https://json-schema.org/), e.g. `{\"type\": \"object\"}` for any JSON object",
         [](common_params & params, const std::string & value) {
             params.sampling.grammar = {COMMON_GRAMMAR_TYPE_OUTPUT_FORMAT, json_schema_to_grammar(json::parse(value))};
         }
     ).set_sampling());
     add_opt(common_arg(
         {"-jf", "--json-schema-file"}, "FILE",
-        "File containing a JSON schema to constrain generations (https://json-schema.org/), e.g. `{}` for any JSON object\nFor schemas w/ external $refs, use --grammar + example/json_schema_to_grammar.py instead",
+        "File containing a JSON schema to constrain generations (https://json-schema.org/), e.g. `{\"type\": \"object\"}` for any JSON object",
         [](common_params & params, const std::string & value) {
             std::ifstream file(value);
             if (!file) {
@@ -3875,7 +3875,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         {"--no-log-jsonl"},
         "Log as JSONL (one JSON object per line) to stdout, this also disables colored logging (default: disabled)",
         [](common_params &, bool value) {
-            common_log_set_jsonl(common_log_main(), value);
+            common_log_set_jsonl(value);
         }
     ).set_env("LLAMA_ARG_LOG_JSONL"));
     add_opt(common_arg(

common/chat-auto-parser-generator.cpp

@@ -5,6 +5,7 @@
 #include "common.h"
 #include "json-schema-to-grammar.h"
 #include "log.h"
+#include "parsers/parsers.h"
 #include "peg-parser.h"
 
 #include <stdexcept>
@@ -12,16 +13,6 @@
 
 using json = common_json;
 
-// Helper to iterate over tools/functions
-static void foreach_function(const json & tools, const std::function<void(const json &)> & fn) {
-    for (const auto & tool : tools) {
-        if (!tool.contains("type") || tool.at("type") != "function" || !tool.contains("function")) {
-            continue;
-        }
-        fn(tool);
-    }
-}
-
 namespace autoparser {
 
 parser_build_context::parser_build_context(common_chat_peg_builder & p, const generation_params & inputs) :
@@ -87,15 +78,6 @@ common_chat_params peg_generator::generate_parser(const common_chat_template &
     if (include_grammar) {
         data.grammar_lazy = !has_response_format && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
         data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                auto         schema   = function.contains("parameters") ? function.at("parameters") : json::object();
-                builder.resolve_refs(schema);
-            });
-            if (has_response_format) {
-                auto schema = inputs.json_schema;
-                builder.resolve_refs(schema);
-            }
             parser.build_grammar(builder, data.grammar_lazy);
         });
 
@@ -312,7 +294,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_json(parser_build_context
     foreach_function(inputs.tools, [&](const json & tool) {
         const auto & func   = tool.at("function");
         std::string  name   = func.at("name");
-        const auto & schema = func.contains("parameters") ? func.at("parameters") : json::object();
+        const auto   schema = common_chat_tool_parameters(func);
 
         // Build call_id parser based on position (if supported)
         bool have_call_id = false;
@@ -383,43 +365,31 @@ common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_conte
     common_peg_parser tool_choice = p.choice();
 
     foreach_function(inputs.tools, [&](const json & tool) {
-        const auto &          func       = tool.at("function");
-        std::string           name       = func.at("name");
-        auto                  params     = func.contains("parameters") ? func.at("parameters") : json::object();
-        const auto &          properties = params.contains("properties") ? params.at("properties") : json::object();
-
-        std::set<std::string> required;
-        if (params.contains("required")) {
-            required = params.at("required").get<std::set<std::string>>();
-        }
-
-        auto schema_info = common_schema_info();
-        schema_info.resolve_refs(params);
+        const auto & func = tool.at("function");
+        std::string  name = func.at("name");
 
         // Build parser for each argument, separating required and optional
         std::vector<common_peg_parser> required_parsers;
         std::vector<common_peg_parser> optional_parsers;
-        for (const auto & [param_name, param_schema] : properties.items()) {
-            bool is_required = required.find(param_name) != required.end();
-
+        foreach_parameter(func, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
             auto arg =
-                p.tool_arg(p.tool_arg_open(arguments.name_prefix + p.tool_arg_name(p.literal(param_name)) +
+                p.tool_arg(p.tool_arg_open(arguments.name_prefix + p.tool_arg_name(p.literal(param.name)) +
                                            arguments.name_suffix) +
                            arguments.value_prefix +
-                           (schema_info.resolves_to_string(param_schema) ?
+                           (param.schema->may_be_string() ?
                                 p.ac(p.tool_arg_string_value(until_suffix) +
                                     p.tool_arg_close(p.literal(arguments.value_suffix)), arguments.value_suffix) :
                                 (p.tool_arg_json_value(p.schema(
-                                    p.json(), "tool-" + name + "-arg-" + param_name + "-schema", param_schema, false)) +
+                                    p.json(), "tool-" + name + "-arg-" + param.name + "-schema", doc, *param.schema)) +
                                     p.tool_arg_close(p.literal(arguments.value_suffix)))));
 
-            auto named_arg = p.rule("tool-" + name + "-arg-" + param_name, arg);
-            if (is_required) {
+            auto named_arg = p.rule("tool-" + name + "-arg-" + param.name, arg);
+            if (param.required) {
                 required_parsers.push_back(named_arg);
             } else {
                 optional_parsers.push_back(named_arg);
             }
-        }
+        });
 
         // Build required arg sequence in definition order
         common_peg_parser args_seq = p.eps();

common/chat-peg-parser.cpp

@@ -488,7 +488,7 @@ common_peg_parser common_chat_peg_builder::standard_constructed_tools(
         }
         const auto &   function = tool_def.at("function");
         std::string    name     = function.at("name");
-        ordered_json   params   = function.contains("parameters") ? function.at("parameters") : ordered_json::object();
+        ordered_json   params   = common_chat_tool_parameters(function);
 
         // Build argument parsers
         auto args = eps();
@@ -565,7 +565,7 @@ common_peg_parser common_chat_peg_builder::python_style_tool_calls(
         }
         const auto &   function = tool_def.at("function");
         std::string    name     = function.at("name");
-        ordered_json   params   = function.contains("parameters") ? function.at("parameters") : ordered_json::object();
+        ordered_json   params   = common_chat_tool_parameters(function);
 
         auto args = eps();
         if (params.contains("properties") && !params["properties"].empty()) {
@@ -640,7 +640,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_function_is_key(
         }
         const auto &   function = tool_def.at("function");
         std::string    name     = function.at("name");
-        ordered_json   params   = function.contains("parameters") ? function.at("parameters") : ordered_json::object();
+        ordered_json   params   = common_chat_tool_parameters(function);
 
         // Build inner object fields
         std::vector<common_peg_parser> inner_fields;
@@ -726,7 +726,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_nested_keys(
         }
         const auto &   function = tool_def.at("function");
         std::string    name     = function.at("name");
-        ordered_json   params   = function.contains("parameters") ? function.at("parameters") : ordered_json::object();
+        ordered_json   params   = common_chat_tool_parameters(function);
 
         auto nested_name = literal("\"" + nested_name_field + "\"") + space() + literal(":") + space() +
                           atomic(literal("\"") + tool_name(literal(name)) + literal("\""));
@@ -795,7 +795,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys(
         }
         const auto &   function = tool_def.at("function");
         std::string    name     = function.at("name");
-        ordered_json   params   = function.contains("parameters") ? function.at("parameters") : ordered_json::object();
+        ordered_json   params   = common_chat_tool_parameters(function);
 
         auto tool_name_ = name_key_parser + space() + literal(":") + space() +
                          atomic(literal("\"") + tool_name(literal(name)) + literal("\""));

common/chat.cpp

@@ -574,6 +574,16 @@ json common_chat_tools_to_json_oaicompat(const std::vector<common_chat_tool> & t
     return result;
 }
 
+json common_chat_tool_parameters(const json & function) {
+    if (function.contains("parameters")) {
+        const auto & params = function.at("parameters");
+        if (!params.is_null() && !(params.is_object() && params.empty())) {
+            return params;
+        }
+    }
+    return json{{"type", "object"}, {"properties", json::object()}};
+}
+
 std::vector<common_chat_tool> common_chat_tools_parse_oaicompat(const json & tools) {
     std::vector<common_chat_tool> result;
 

common/chat.h

@@ -360,6 +360,9 @@ common_json common_chat_msgs_to_json_oaicompat(const std::vector<common_chat_msg
 
 common_json common_chat_tools_to_json_oaicompat(const std::vector<common_chat_tool> & tools);
 
+// The parameters schema of a function tool. A tool without parameters, or with an empty {}, takes zero arguments.
+common_json common_chat_tool_parameters(const common_json & function);
+
 // 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);
 

common/common.cpp

@@ -1586,6 +1586,11 @@ common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx) {
         return COMMON_CONTEXT_SEQ_RM_TYPE_NO;
     }
 
+    if (llama_n_rs_seq(ctx) > 0) {
+        COM_TRC("%s", "the context supports bounded partial sequence removal\n");
+        return COMMON_CONTEXT_SEQ_RM_TYPE_RS;
+    }
+
     common_context_seq_rm_type res = COMMON_CONTEXT_SEQ_RM_TYPE_PART;
 
     llama_memory_clear(mem, true);
@@ -1602,12 +1607,6 @@ common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx) {
         goto done;
     }
 
-    if (llama_n_rs_seq(ctx) > 0) {
-        COM_TRC("%s", "the context supports bounded partial sequence removal\n");
-        res = COMMON_CONTEXT_SEQ_RM_TYPE_RS;
-        goto done;
-    }
-
     // try to remove the last tokens
     if (!llama_memory_seq_rm(mem, 0, 1, -1)) {
         COM_TRC("%s", "the context does not support partial sequence removal\n");

common/fit.cpp

@@ -1,5 +1,6 @@
 #include "fit.h"
 
+#include "json.h"
 #include "log.h"
 
 #include "../src/llama-ext.h"
@@ -915,6 +916,9 @@ void common_memory_breakdown_print(const struct llama_context * ctx) {
 
     std::vector<std::array<std::string, 9>> table_data;
     table_data.reserve(devices.size());
+
+    // same data as the table below, for --log-jsonl consumers
+    common_json rows = common_json::array();
     const std::string template_header = "%s: | %s | %s   %s    %s   %s   %s   %s    %s |\n";
     const std::string template_gpu    = "%s: | %s | %s = %s + (%s = %s + %s + %s) + %s |\n";
     const std::string template_other  = "%s: | %s | %s   %s    %s = %s + %s + %s    %s |\n";
@@ -989,6 +993,19 @@ void common_memory_breakdown_print(const struct llama_context * ctx) {
             std::to_string(mb.context / MiB),
             std::to_string(mb.compute / MiB),
             std::to_string(unaccounted / static_cast<int64_t>(MiB))});
+
+        rows.push_back({
+            {"kind",        "device"},
+            {"name",        name},
+            {"description", desc},
+            {"total",       total / MiB},
+            {"free",        free / MiB},
+            {"self",        self / MiB},
+            {"model",       mb.model / MiB},
+            {"context",     mb.context / MiB},
+            {"compute",     mb.compute / MiB},
+            {"unaccounted", unaccounted / static_cast<int64_t>(MiB)},
+        });
     }
 
     // print memory breakdown for host:
@@ -1004,6 +1021,15 @@ void common_memory_breakdown_print(const struct llama_context * ctx) {
             std::to_string(mb_host.context / MiB),
             std::to_string(mb_host.compute / MiB),
             ""}); // unaccounted
+
+        rows.push_back({
+            {"kind",    "host"},
+            {"name",    "Host"},
+            {"self",    self / MiB},
+            {"model",   mb_host.model / MiB},
+            {"context", mb_host.context / MiB},
+            {"compute", mb_host.compute / MiB},
+        });
     }
 
     // print memory breakdown for all remaining buffer types:
@@ -1025,6 +1051,16 @@ void common_memory_breakdown_print(const struct llama_context * ctx) {
             std::to_string(mb.context / MiB),
             std::to_string(mb.compute / MiB),
             ""}); // unaccounted
+
+        rows.push_back({
+            {"kind",    "buffer_type"},
+            {"name",    name},
+            {"self",    self / MiB},
+            {"model",   mb.model / MiB},
+            {"context", mb.context / MiB},
+            {"compute", mb.compute / MiB},
+        });
+
         seen_buffer_types.insert(buft);
     }
 
@@ -1042,6 +1078,11 @@ void common_memory_breakdown_print(const struct llama_context * ctx) {
             __func__, td[1].c_str(), td[2].c_str(), td[3].c_str(), td[4].c_str(), td[5].c_str(),
             td[6].c_str(), td[7].c_str(), td[8].c_str());
     }
+
+    LOG_JSON("fit_memory_breakdown", common_json({
+        {"unit", "MiB"},
+        {"rows", rows},
+    }));
 }
 
 void common_fit_print(

common/jinja/runtime.cpp

@@ -842,6 +842,12 @@ value member_expression::execute_impl(context & ctx) {
         } else {
             property = this->property->execute(ctx);
         }
+    } else if (is_stmt<integer_literal>(this->property)) {
+        // syntax: obj.index
+        property = mk_val<value_int>(cast_stmt<integer_literal>(this->property)->val);
+        if (property->as_int() < 0) {
+            throw std::runtime_error("Static member property cannot be negative");
+        }
     } else {
         // syntax: obj.prop
         if (!is_stmt<identifier>(this->property)) {

common/json-schema-to-grammar.cpp

@@ -1,5 +1,7 @@
 #include "json-schema-to-grammar.h"
 #include "common.h"
+#include "trie.h"
+#include "unicode.h"
 
 #include <algorithm>
 #include <limits>
@@ -336,18 +338,20 @@ static size_t gbnf_escape_length(const std::string & pattern, size_t pos) {
     return 2 + n_hex;
 }
 
-class common_schema_converter {
+class common_chat_schema_converter {
 private:
-    friend class common_schema_info;
     friend std::string build_grammar(const std::function<void(const common_grammar_builder &)> & cb, const common_grammar_options & options);
-    std::function<json(const std::string &)> _fetch_json;
     bool _dotall;
     std::map<std::string, std::string> _rules;
-    std::unordered_map<std::string, json> _refs;
     std::unordered_set<std::string> _refs_being_resolved;
     std::vector<std::string> _errors;
     std::vector<std::string> _warnings;
 
+    template <typename T>
+    static const T & as(const common_chat_schema & node) {
+        return static_cast<const T &>(node);
+    }
+
     std::string _add_rule(const std::string & name, const std::string & rule) {
         std::string esc_name = regex_replace(name, INVALID_RULE_CHARS_RE, "-");
         if (_rules.find(esc_name) == _rules.end() || _rules[esc_name] == rule) {
@@ -363,11 +367,11 @@ private:
         return key;
     }
 
-    std::string _generate_union_rule(const std::string & name, const std::vector<json> & alt_schemas) {
+    std::string _generate_union_rule(const std::string & name, const std::vector<common_chat_schema_ptr> & alt_schemas) {
         std::vector<std::string> rules;
         rules.reserve(alt_schemas.size());
         for (size_t i = 0; i < alt_schemas.size(); i++) {
-            rules.push_back(visit(alt_schemas[i], name + (name.empty() ? "alternative-" : "-") + std::to_string(i)));
+            rules.push_back(visit(*alt_schemas[i], name + (name.empty() ? "alternative-" : "-") + std::to_string(i)));
         }
         return string_join(rules, " | ");
     }
@@ -634,85 +638,68 @@ private:
             -> ["] ( [a] ([l] ([s] ([o] char+ | [^"o] char*) | [^"s] char*) | [n] ([d] char+ | [^"d] char*) | [^"ln] char*) | [^"a] char* )? ["]
     */
     std::string _not_strings(const std::vector<std::string> & strings) {
-
-        struct TrieNode {
-            std::map<char, TrieNode> children;
-            bool is_end_of_string;
-
-            TrieNode() : is_end_of_string(false) {}
-
-            void insert(const std::string & string) {
-                auto *node = this;
-                for (char c : string) {
-                    node = &node->children[c];
-                }
-                node->is_end_of_string = true;
-            }
-        };
-
-        TrieNode trie;
-        for (const auto & s : strings) {
-            trie.insert(s);
-        }
+        common_trie trie(strings);
 
         std::string char_rule = _add_primitive("char", PRIMITIVE_RULES.at("char"));
         std::ostringstream out;
         out << "[\"] ( ";
-        std::function<void(const TrieNode &)> visit = [&](const TrieNode & node) {
-            std::ostringstream rejects;
+        std::function<void(size_t)> visit = [&](size_t idx) {
+            const auto & node = trie.nodes[idx];
+            std::string rejects;
             auto first = true;
-            for (const auto & kv : node.children) {
-                rejects << kv.first;
+            for (const auto & [cpt, child] : node.children) {
+                std::string c = common_unicode_cpt_to_utf8(cpt);
+                rejects += c;
                 if (first) {
                     first = false;
                 } else {
                     out << " | ";
                 }
-                out << "[" << kv.first << "]";
-                if (!kv.second.children.empty()) {
+                out << "[" << c << "]";
+                if (!trie.nodes[child].children.empty()) {
                     out << " (";
-                    visit(kv.second);
+                    visit(child);
                     out << ")";
-                } else if (kv.second.is_end_of_string) {
+                } else {
                     out << " " << char_rule << "+";
                 }
             }
             if (!node.children.empty()) {
-                if (!first) {
-                    out << " | ";
-                }
-                out << "[^\"" << rejects.str() << "] " << char_rule << "*";
+                out << " | [^\"" << rejects << "] " << char_rule << "*";
             }
         };
-        visit(trie);
+        visit(0);
 
         out << " )";
-        if (!trie.is_end_of_string) {
+        if (trie.nodes[0].pattern < 0) {
             out << "?";
         }
         out << " [\"]";
         return out.str();
     }
 
-    std::string _resolve_ref(const std::string & ref) {
-        auto it = ref.find('#');
-        std::string ref_fragment = it != std::string::npos ? ref.substr(it + 1) : ref;
+    std::string _resolve_ref(const common_chat_schema_ref & schema) {
+        auto it = schema.ref.find('#');
+        std::string ref_fragment = it != std::string::npos ? schema.ref.substr(it + 1) : schema.ref;
         static const std::regex nonalphanumeric_regex(R"([^a-zA-Z0-9-]+)");
         std::string ref_name = "ref" + std::regex_replace(ref_fragment, nonalphanumeric_regex, "-");
-        if (_rules.find(ref_name) == _rules.end() && _refs_being_resolved.find(ref) == _refs_being_resolved.end()) {
-            _refs_being_resolved.insert(ref);
-            json resolved = _refs[ref];
-            ref_name = visit(resolved, ref_name);
-            _refs_being_resolved.erase(ref);
+        if (_rules.find(ref_name) == _rules.end() && _refs_being_resolved.find(schema.ref) == _refs_being_resolved.end()) {
+            if (!schema.target) {
+                _errors.push_back("Unresolved $ref " + schema.ref);
+                return "";
+            }
+            _refs_being_resolved.insert(schema.ref);
+            ref_name = visit(*schema.target, ref_name);
+            _refs_being_resolved.erase(schema.ref);
         }
         return ref_name;
     }
 
     std::string _build_object_rule(
-        const std::vector<std::pair<std::string, json>> & properties,
+        const std::vector<std::pair<std::string, const common_chat_schema *>> & properties,
         const std::unordered_set<std::string> & required,
         const std::string & name,
-        const json & additional_properties)
+        const common_chat_schema * additional_properties)
     {
         std::vector<std::string> required_props;
         std::vector<std::string> optional_props;
@@ -722,7 +709,7 @@ private:
             const auto &prop_name = kv.first;
             const auto &prop_schema = kv.second;
 
-            std::string prop_rule_name = visit(prop_schema, name + (name.empty() ? "" : "-") + prop_name);
+            std::string prop_rule_name = visit(*prop_schema, name + (name.empty() ? "" : "-") + prop_name);
             prop_kv_rule_names[prop_name] = _add_rule(
                 name + (name.empty() ? "" : "-") + prop_name + "-kv",
                 format_literal(json(prop_name).dump()) + " space \":\" space " + prop_rule_name
@@ -734,10 +721,10 @@ private:
             }
             prop_names.push_back(prop_name);
         }
-        if ((additional_properties.is_boolean() && additional_properties.get<bool>()) || additional_properties.is_object()) {
+        if (additional_properties) {
             std::string sub_name = name + (name.empty() ? "" : "-") + "additional";
             std::string value_rule =
-                additional_properties.is_object() ? visit(additional_properties, sub_name + "-value")
+                additional_properties->kind() != common_chat_schema::KIND_ANY ? visit(*additional_properties, sub_name + "-value")
                 : _add_primitive("value", PRIMITIVE_RULES.at("value"));
 
             auto key_rule =
@@ -825,267 +812,163 @@ private:
     }
 
 public:
-    common_schema_converter(
-        const std::function<json(const std::string &)> & fetch_json,
-        bool dotall)
-          : _fetch_json(fetch_json), _dotall(dotall)
-    {
+    explicit common_chat_schema_converter(bool dotall) : _dotall(dotall) {
         _rules["space"] = SPACE_RULE;
     }
 
-    void resolve_refs(json & schema, const std::string & url) {
-        /*
-        * Resolves all $ref fields in the given schema, fetching any remote schemas,
-        * replacing each $ref with absolute reference URL and populates _refs with the
-        * respective referenced (sub)schema dictionaries.
-        */
-        std::function<void(json &)> visit_refs = [&](json & n) {
-            if (n.is_array()) {
-                for (auto & x : n) {
-                    visit_refs(x);
-                }
-            } else if (n.is_object()) {
-                if (n.contains("$ref")) {
-                    std::string ref = n["$ref"];
-                    if (_refs.find(ref) == _refs.end()) {
-                        json target;
-                        if (ref.find("https://") == 0) {
-                            std::string base_url = ref.substr(0, ref.find('#'));
-                            auto it = _refs.find(base_url);
-                            if (it != _refs.end()) {
-                                target = it->second;
-                            } else {
-                                // Fetch the referenced schema and resolve its refs
-                                auto referenced = _fetch_json(ref);
-                                resolve_refs(referenced, base_url);
-                                _refs[base_url] = referenced;
-                            }
-                            if (ref.find('#') == std::string::npos || ref.substr(ref.find('#') + 1).empty()) {
-                                return;
-                            }
-                        } else if (ref.find("#/") == 0) {
-                            target = schema;
-                            n["$ref"] = url + ref;
-                            ref = url + ref;
-                        } else {
-                            _errors.push_back("Unsupported ref: " + ref);
-                            return;
-                        }
-                        std::string pointer = ref.substr(ref.find('#') + 1);
-                        std::vector<std::string> tokens = string_split(pointer, "/");
-                        for (size_t i = 1; i < tokens.size(); ++i) {
-                            const std::string& sel = tokens[i];
-                            if (target.is_object() && target.contains(sel)) {
-                                target = target[sel];
-                            } else if (target.is_array()) {
-                                size_t sel_index;
-                                try {
-                                    sel_index = std::stoull(sel);
-                                } catch (const std::invalid_argument & e) {
-                                    sel_index = target.size();
-                                }
-                                if (sel_index >= target.size()) {
-                                    _errors.push_back("Error resolving ref " + ref + ": " + sel + " not in " + target.dump());
-                                    return;
-                                }
-                                target = target[sel_index];
-                            } else {
-                                _errors.push_back("Error resolving ref " + ref + ": " + sel + " not in " + target.dump());
-                                return;
-                            }
-                        }
-                        _refs[ref] = target;
-                    }
-                } else {
-                    for (const auto & kv : n.items()) {
-                        visit_refs(kv.value());
-                    }
-                }
-            }
-        };
-
-        visit_refs(schema);
+    std::string add_schema(const std::string & name, const common_chat_schema & schema) {
+        return visit(schema, name);
     }
 
     static std::string _generate_constant_rule(const json & value) {
         return format_literal(value.dump());
     }
 
-    std::string visit(const json & schema, const std::string & name) {
-        json schema_type = schema.contains("type") ? schema["type"] : json();
-        std::string schema_format = schema.contains("format") ? schema["format"].get<std::string>() : "";
-        std::string rule_name = is_reserved_name(name) ? name + "-" : name.empty() ? "root" : name;
+    std::string _visit_primitive(const std::string & rule_name, const std::string & type) {
+        return _add_primitive(rule_name == "root" ? "root" : type, PRIMITIVE_RULES.at(type));
+    }
 
-        if (schema.contains("$ref")) {
-            return _add_rule(rule_name, _resolve_ref(schema["$ref"]));
-        }
-        if (schema.contains("oneOf") || schema.contains("anyOf")) {
-            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()) {
-            std::vector<json> schema_types;
-            for (const auto & t : schema_type) {
-                json schema_copy(schema);
-                schema_copy["type"] = t;
-                schema_types.push_back(schema_copy);
-            }
-            return _add_rule(rule_name, _generate_union_rule(name, schema_types));
-        }
-        if (schema.contains("const")) {
-            return _add_rule(rule_name, _generate_constant_rule(schema["const"]));
-        }
-        if (schema.contains("enum")) {
-            std::vector<std::string> enum_values;
-            for (const auto & v : schema["enum"]) {
-                enum_values.push_back(_generate_constant_rule(v));
-            }
-            return _add_rule(rule_name, "(" + string_join(enum_values, " | ") + ")");
-        }
-        if ((schema_type.is_null() || schema_type == "object")
-                && (schema.contains("properties") ||
-                    (schema.contains("additionalProperties") && schema["additionalProperties"] != true))) {
-            std::unordered_set<std::string> required;
-            if (schema.contains("required") && schema["required"].is_array()) {
-                for (const auto & item : schema["required"]) {
-                    if (item.is_string()) {
-                        required.insert(item.get<std::string>());
+    std::string _visit_all_of(const common_chat_schema_all_of & schema, const std::string & name, const std::string & rule_name) {
+        std::unordered_set<std::string> required;
+        std::vector<std::pair<std::string, const common_chat_schema *>> properties;
+        std::map<std::string, size_t> enum_values;
+        std::function<void(const common_chat_schema &, bool)> add_component = [&](const common_chat_schema & comp, bool is_required) {
+            if (comp.kind() == common_chat_schema::KIND_REF) {
+                if (const auto * target = as<common_chat_schema_ref>(comp).target) {
+                    add_component(*target, is_required);
+                }
+            } else if (comp.kind() == common_chat_schema::KIND_OBJECT) {
+                for (const auto & prop : as<common_chat_schema_object>(comp).properties) {
+                    properties.emplace_back(prop.name, prop.schema.get());
+                    if (is_required) {
+                        required.insert(prop.name);
                     }
                 }
-            }
-            std::vector<std::pair<std::string, json>> properties;
-            if (schema.contains("properties")) {
-                for (const auto & prop : schema["properties"].items()) {
-                    properties.emplace_back(prop.key(), prop.value());
+            } else if (comp.kind() == common_chat_schema::KIND_ENUM) {
+                for (const auto & v : as<common_chat_schema_enum>(comp).values) {
+                    enum_values[_generate_constant_rule(v)] += 1;
                 }
             }
-            return _add_rule(rule_name,
-                _build_object_rule(
-                    properties, required, name,
-                    schema.contains("additionalProperties") ? schema["additionalProperties"] : json()));
+        };
+        for (const auto & child : schema.children) {
+            if (child->kind() == common_chat_schema::KIND_ANY_OF) {
+                for (const auto & alt : as<common_chat_schema_any_of>(*child).children) {
+                    add_component(*alt, false);
+                }
+            } else {
+                add_component(*child, true);
+            }
         }
-        if ((schema_type.is_null() || schema_type == "object" || schema_type == "string") && schema.contains("allOf")) {
-            std::unordered_set<std::string> required;
-            std::vector<std::pair<std::string, json>> properties;
-            std::map<std::string, size_t> enum_values;
-            const std::string& hybrid_name = name;
-            std::function<void(const json &, bool)> add_component = [&](const json & comp_schema, bool is_required) {
-                if (comp_schema.contains("$ref")) {
-                    add_component(_refs[comp_schema["$ref"]], is_required);
-                } else if (comp_schema.contains("properties")) {
-                    for (const auto & prop : comp_schema["properties"].items()) {
-                        properties.emplace_back(prop.key(), prop.value());
-                        if (is_required) {
-                            required.insert(prop.key());
-                        }
-                    }
-                } else if (comp_schema.contains("enum")) {
-                    for (const auto & v : comp_schema["enum"]) {
-                        const auto rule = _generate_constant_rule(v);
-                        if (enum_values.find(rule) == enum_values.end()) {
-                            enum_values[rule] = 0;
-                        }
-                        enum_values[rule] += 1;
-                    }
-                } else {
-                  // todo warning
-                }
-            };
-            for (const auto & t : schema["allOf"]) {
-                if (t.contains("anyOf")) {
-                    for (const auto & tt : t["anyOf"]) {
-                        add_component(tt, false);
-                    }
-                } else {
-                    add_component(t, true);
+        if (!enum_values.empty()) {
+            std::vector<std::string> enum_intersection;
+            for (const auto & p : enum_values) {
+                if (p.second == schema.children.size()) {
+                    enum_intersection.push_back(p.first);
                 }
             }
-            if (!enum_values.empty()) {
-                std::vector<std::string> enum_intersection;
-                for (const auto & p : enum_values) {
-                    if (p.second == schema["allOf"].size()) {
-                        enum_intersection.push_back(p.first);
-                    }
-                }
-                if (!enum_intersection.empty()) {
-                    return _add_rule(rule_name, "(" + string_join(enum_intersection, " | ") + ")");
-                }
+            if (!enum_intersection.empty()) {
+                return _add_rule(rule_name, "(" + string_join(enum_intersection, " | ") + ")");
             }
-            return _add_rule(rule_name, _build_object_rule(properties, required, hybrid_name, json()));
         }
-        if ((schema_type.is_null() || schema_type == "array") && (schema.contains("items") || schema.contains("prefixItems"))) {
-            json items = schema.contains("items") ? schema["items"] : schema["prefixItems"];
-            if (items.is_array()) {
+        return _add_rule(rule_name, _build_object_rule(properties, required, name, nullptr));
+    }
+
+    std::string visit(const common_chat_schema & schema, const std::string & name) {
+        std::string rule_name = is_reserved_name(name) ? name + "-" : name.empty() ? "root" : name;
+        std::string sub_name  = name + (name.empty() ? "" : "-");
+
+        switch (schema.kind()) {
+            case common_chat_schema::KIND_REF:
+                return _add_rule(rule_name, _resolve_ref(as<common_chat_schema_ref>(schema)));
+            case common_chat_schema::KIND_ANY_OF:
+                return _add_rule(rule_name, _generate_union_rule(name, as<common_chat_schema_any_of>(schema).children));
+            case common_chat_schema::KIND_ALL_OF:
+                return _visit_all_of(as<common_chat_schema_all_of>(schema), name, rule_name);
+            case common_chat_schema::KIND_CONST:
+                return _add_rule(rule_name, _generate_constant_rule(as<common_chat_schema_const>(schema).value));
+            case common_chat_schema::KIND_ENUM: {
+                std::vector<std::string> enum_values;
+                for (const auto & v : as<common_chat_schema_enum>(schema).values) {
+                    enum_values.push_back(_generate_constant_rule(v));
+                }
+                return _add_rule(rule_name, "(" + string_join(enum_values, " | ") + ")");
+            }
+            case common_chat_schema::KIND_OBJECT: {
+                const auto & obj = as<common_chat_schema_object>(schema);
+                if (obj.properties.empty() && obj.additional_properties && obj.additional_properties->kind() == common_chat_schema::KIND_ANY) {
+                    return _add_rule(rule_name, _add_primitive("object", PRIMITIVE_RULES.at("object")));
+                }
+                std::vector<std::pair<std::string, const common_chat_schema *>> properties;
+                std::unordered_set<std::string> required;
+                for (const auto & prop : obj.properties) {
+                    properties.emplace_back(prop.name, prop.schema.get());
+                    if (prop.required) {
+                        required.insert(prop.name);
+                    }
+                }
+                return _add_rule(rule_name, _build_object_rule(properties, required, name, obj.additional_properties.get()));
+            }
+            case common_chat_schema::KIND_TUPLE: {
+                const auto & items = as<common_chat_schema_tuple>(schema).items;
                 std::string rule = "\"[\" space ";
                 for (size_t i = 0; i < items.size(); i++) {
                     if (i > 0) {
                         rule += " \",\" space ";
                     }
-                    rule += visit(items[i], name + (name.empty() ? "" : "-") + "tuple-" + std::to_string(i));
+                    rule += visit(*items[i], sub_name + "tuple-" + std::to_string(i));
                 }
                 rule += " space \"]\"";
                 return _add_rule(rule_name, rule);
             }
-            std::string item_rule_name = visit(items, name + (name.empty() ? "" : "-") + "item");
-            int min_items = schema.contains("minItems") ? schema["minItems"].get<int>() : 0;
-            json max_items_json = schema.contains("maxItems") ? schema["maxItems"] : json();
-            int max_items = max_items_json.is_number_integer() ? max_items_json.get<int>() : std::numeric_limits<int>::max();
-
-            return _add_rule(rule_name, "\"[\" space " + build_repetition(item_rule_name, min_items, max_items, "\",\" space") + " space \"]\"");
-        }
-        if ((schema_type.is_null() || schema_type == "string") && schema.contains("pattern")) {
-            return _visit_pattern(schema["pattern"], rule_name);
-        }
-        if ((schema_type.is_null() || schema_type == "string") && std::regex_match(schema_format, std::regex("^uuid[1-5]?$"))) {
-            return _add_primitive(rule_name == "root" ? "root" : schema_format, PRIMITIVE_RULES.at("uuid"));
-        }
-        if ((schema_type.is_null() || schema_type == "string") && STRING_FORMAT_RULES.find(schema_format + "-string") != STRING_FORMAT_RULES.end()) {
-            auto prim_name = schema_format + "-string";
-            return _add_rule(rule_name, _add_primitive(prim_name, STRING_FORMAT_RULES.at(prim_name)));
-        }
-        if (schema_type == "string" && (schema.contains("minLength") || schema.contains("maxLength"))) {
-            std::string char_rule = _add_primitive("char", PRIMITIVE_RULES.at("char"));
-            int min_len = schema.contains("minLength") ? schema["minLength"].get<int>() : 0;
-            int max_len = schema.contains("maxLength") ? schema["maxLength"].get<int>() : std::numeric_limits<int>::max();
-            return _add_rule(rule_name, "\"\\\"\" " + build_repetition(char_rule, min_len, max_len) + " \"\\\"\"");
-        }
-        if (schema_type == "integer" && (schema.contains("minimum") || schema.contains("exclusiveMinimum") || schema.contains("maximum") || schema.contains("exclusiveMaximum"))) {
-            int64_t min_value = std::numeric_limits<int64_t>::min();
-            int64_t max_value = std::numeric_limits<int64_t>::max();
-            if (schema.contains("minimum")) {
-                min_value = schema["minimum"].get<int64_t>();
-            } else if (schema.contains("exclusiveMinimum")) {
-                min_value = schema["exclusiveMinimum"].get<int64_t>() + 1;
+            case common_chat_schema::KIND_ARRAY: {
+                const auto & arr = as<common_chat_schema_array>(schema);
+                if (arr.items->kind() == common_chat_schema::KIND_ANY && arr.min_items == 0 && arr.max_items < 0) {
+                    return _visit_primitive(rule_name, "array");
+                }
+                std::string item_rule_name = visit(*arr.items, sub_name + "item");
+                int max_items = arr.max_items < 0 ? std::numeric_limits<int>::max() : arr.max_items;
+                return _add_rule(rule_name, "\"[\" space " + build_repetition(item_rule_name, arr.min_items, max_items, "\",\" space") + " space \"]\"");
             }
-            if (schema.contains("maximum")) {
-                max_value = schema["maximum"].get<int64_t>();
-            } else if (schema.contains("exclusiveMaximum")) {
-                max_value = schema["exclusiveMaximum"].get<int64_t>() - 1;
+            case common_chat_schema::KIND_STRING: {
+                const auto & str = as<common_chat_schema_string>(schema);
+                if (!str.pattern.empty()) {
+                    return _visit_pattern(str.pattern, rule_name);
+                }
+                if (str.format == common_chat_schema::FORMAT_UUID) {
+                    return _visit_primitive(rule_name, "uuid");
+                }
+                if (str.format != common_chat_schema::FORMAT_NONE) {
+                    std::string prim_name = std::string(str.format == common_chat_schema::FORMAT_DATE ? "date" : str.format == common_chat_schema::FORMAT_TIME ? "time" : "date-time") + "-string";
+                    return _add_rule(rule_name, _add_primitive(prim_name, STRING_FORMAT_RULES.at(prim_name)));
+                }
+                if (str.min_length > 0 || str.max_length >= 0) {
+                    std::string char_rule = _add_primitive("char", PRIMITIVE_RULES.at("char"));
+                    int max_len = str.max_length < 0 ? std::numeric_limits<int>::max() : str.max_length;
+                    return _add_rule(rule_name, "\"\\\"\" " + build_repetition(char_rule, str.min_length, max_len) + " \"\\\"\"");
+                }
+                return _visit_primitive(rule_name, "string");
             }
-            std::stringstream out;
-            out << "(";
-            build_min_max_int(min_value, max_value, out);
-            out << ")";
-            return _add_rule(rule_name, out.str());
+            case common_chat_schema::KIND_INTEGER: {
+                const auto & i = as<common_chat_schema_integer>(schema);
+                if (i.minimum == std::numeric_limits<int64_t>::min() && i.maximum == std::numeric_limits<int64_t>::max()) {
+                    return _visit_primitive(rule_name, "integer");
+                }
+                std::stringstream out;
+                out << "(";
+                build_min_max_int(i.minimum, i.maximum, out);
+                out << ")";
+                return _add_rule(rule_name, out.str());
+            }
+            case common_chat_schema::KIND_NUMBER:
+                return _visit_primitive(rule_name, "number");
+            case common_chat_schema::KIND_BOOLEAN:
+                return _visit_primitive(rule_name, "boolean");
+            case common_chat_schema::KIND_NULL:
+                return _visit_primitive(rule_name, "null");
+            case common_chat_schema::KIND_ANY:
+                return _add_rule(rule_name, _add_primitive("value", PRIMITIVE_RULES.at("value")));
         }
-        if (schema.empty() || schema_type == "object") {
-            return _add_rule(rule_name, _add_primitive("object", PRIMITIVE_RULES.at("object")));
-        }
-        if (schema_type.is_null() && schema.is_object()) {
-            // No type constraint and no recognized structural keywords (e.g. {"description": "..."}).
-            // Per JSON Schema semantics this is equivalent to {} and accepts any value.
-            return _add_rule(rule_name, _add_primitive("value", PRIMITIVE_RULES.at("value")));
-        }
-        if (!schema_type.is_string() || PRIMITIVE_RULES.find(schema_type.get<std::string>()) == PRIMITIVE_RULES.end()) {
-            _errors.push_back("Unrecognized schema: " + schema.dump());
-            return "";
-        }
-        // TODO: support minimum, maximum, exclusiveMinimum, exclusiveMaximum at least for zero
-        return _add_primitive(rule_name == "root" ? "root" : schema_type.get<std::string>(), PRIMITIVE_RULES.at(schema_type.get<std::string>()));
+        return "";
     }
 
     void check_errors() {
@@ -1106,134 +989,6 @@ public:
     }
 };
 
-// common_schema_info implementation (pimpl)
-
-common_schema_info::common_schema_info()
-    : impl_(std::make_unique<common_schema_converter>(
-        [](const std::string &) { return json(); },
-        false)) {}
-
-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(common_json & schema) {
-    impl_->resolve_refs(schema, "");
-}
-
-// Determines if a JSON schema can resolve to a string type through any path.
-// 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 common_json & schema) {
-    std::unordered_set<std::string> visited_refs;
-
-    std::function<bool(const json &)> check = [&](const json & s) -> bool {
-        if (!s.is_object()) {
-            return false;
-        }
-
-        // Handle $ref
-        if (s.contains("$ref")) {
-            const std::string & ref = s["$ref"];
-            if (visited_refs.find(ref) != visited_refs.end()) {
-                // Circular reference, assume not a string to be safe
-                return false;
-            }
-            visited_refs.insert(ref);
-            auto it = impl_->_refs.find(ref);
-            if (it != impl_->_refs.end()) {
-                return check(it->second);
-            }
-            return false;
-        }
-
-        // Check type field
-        if (s.contains("type")) {
-            const json & schema_type = s["type"];
-            if (schema_type.is_string()) {
-                if (schema_type == "string") {
-                    return true;
-                }
-            } else if (schema_type.is_array()) {
-                // Type can be an array like ["string", "null"]
-                for (const auto & t : schema_type) {
-                    if (t == "string") {
-                        return true;
-                    }
-                }
-            }
-        }
-
-        // Check oneOf/anyOf - if any alternative can be a string
-        if (s.contains("oneOf")) {
-            for (const auto & alt : s["oneOf"]) {
-                if (check(alt)) {
-                    return true;
-                }
-            }
-        }
-        if (s.contains("anyOf")) {
-            for (const auto & alt : s["anyOf"]) {
-                if (check(alt)) {
-                    return true;
-                }
-            }
-        }
-
-        // Check allOf - all components must be compatible with string type
-        if (s.contains("allOf")) {
-            bool all_string = true;
-            for (const auto & component : s["allOf"]) {
-                if (!check(component)) {
-                    all_string = false;
-                    break;
-                }
-            }
-            if (all_string) {
-                return true;
-            }
-        }
-
-        // Check const - if the constant value is a string
-        if (s.contains("const")) {
-            if (s["const"].is_string()) {
-                return true;
-            }
-        }
-
-        // Check enum - if any enum value is a string
-        if (s.contains("enum")) {
-            for (const auto & val : s["enum"]) {
-                if (val.is_string()) {
-                    return true;
-                }
-            }
-        }
-
-        // String-specific keywords imply string type
-        if (s.contains("pattern") || s.contains("minLength") || s.contains("maxLength")) {
-            return true;
-        }
-
-        // Check format - many formats imply string
-        if (s.contains("format")) {
-            const std::string & fmt = s["format"];
-            if (fmt == "date" || fmt == "time" || fmt == "date-time" ||
-                fmt == "uri" || fmt == "email" || fmt == "hostname" ||
-                fmt == "ipv4" || fmt == "ipv6" || fmt == "uuid" ||
-                fmt.find("uuid") == 0) {
-                return true;
-            }
-        }
-
-        return false;
-    };
-
-    return check(schema);
-}
-
 std::string json_schema_to_grammar(const common_json & schema, bool force_gbnf) {
 #ifdef LLAMA_USE_LLGUIDANCE
     if (!force_gbnf) {
@@ -1242,25 +997,29 @@ std::string json_schema_to_grammar(const common_json & schema, bool force_gbnf)
 #else
     (void)force_gbnf;
 #endif // LLAMA_USE_LLGUIDANCE
-    return build_grammar([&](const common_grammar_builder & callbacks) {
-        auto copy = schema;
-        callbacks.resolve_refs(copy);
-        callbacks.add_schema("", copy);
-    });
+    try {
+        return json_schema_to_grammar(common_chat_schema_from_json(schema));
+    } catch (const std::runtime_error & e) {
+        throw std::invalid_argument(std::string("JSON schema conversion failed:\n") + e.what());
+    }
+}
+
+std::string json_schema_to_grammar(const common_chat_schema_document & schema) {
+    common_chat_schema_converter converter(false);
+    converter.visit(*schema.root, "");
+    converter.check_errors();
+    return converter.format_grammar();
 }
 
 std::string build_grammar(const std::function<void(const common_grammar_builder &)> & cb, const common_grammar_options & options) {
-    common_schema_converter converter([&](const std::string &) { return json(); }, options.dotall);
+    common_chat_schema_converter converter(options.dotall);
     common_grammar_builder builder {
         /* .add_rule = */ [&](const std::string & name, const std::string & rule) {
             return converter._add_rule(name, rule);
         },
-        /* .add_schema = */ [&](const std::string & name, const common_json & schema) {
-            return converter.visit(schema, name == "root" ? "" : name);
+        /* .add_schema = */ [&](const std::string & name, const common_chat_schema & schema) {
+            return converter.add_schema(name == "root" ? "" : name, schema);
         },
-        /* .resolve_refs = */ [&](common_json & schema) {
-            converter.resolve_refs(schema, "");
-        }
     };
     cb(builder);
     converter.check_errors();

common/json-schema-to-grammar.h

@@ -1,37 +1,17 @@
 #pragma once
 
+#include "json-schema.h"
 #include "json.h"
 
 #include <functional>
-#include <memory>
 #include <string>
 
-std::string json_schema_to_grammar(const common_json & schema,
-                                   bool force_gbnf = false);
-
-class common_schema_converter;
-
-// Probes a JSON schema to extract information about its structure and type constraints.
-class common_schema_info {
-    std::unique_ptr<common_schema_converter> impl_;
-
-  public:
-    common_schema_info();
-    ~common_schema_info();
-
-    common_schema_info(const common_schema_info &) = delete;
-    common_schema_info & operator=(const common_schema_info &) = delete;
-    common_schema_info(common_schema_info &&) noexcept;
-    common_schema_info & operator=(common_schema_info &&) noexcept;
-
-    void resolve_refs(common_json & schema);
-    bool resolves_to_string(const common_json & schema);
-};
+std::string json_schema_to_grammar(const common_json & schema, bool force_gbnf = false);
+std::string json_schema_to_grammar(const common_chat_schema_document & 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 common_json &)> add_schema;
-    std::function<void(common_json &)> resolve_refs;
+    std::function<std::string(const std::string &, const std::string &)>    add_rule;
+    std::function<std::string(const std::string &, const common_chat_schema &)> add_schema;
 };
 
 struct common_grammar_options {

common/json-schema.cpp

@@ -0,0 +1,514 @@
+#include "json-schema.h"
+#include "common.h"
+
+#include <cmath>
+#include <map>
+#include <stdexcept>
+#include <string>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+class common_chat_schema_builder {
+    const common_json &      root_;
+    common_chat_schema_document & doc_;
+
+    // the targets built here, moved into doc_ once the whole schema is built
+    std::map<std::string, common_chat_schema_ptr> refs_;
+
+    // ref nodes get their target once every $ref is built, a cycle would otherwise need it too early
+    std::vector<common_chat_schema_ref *> pending_;
+
+    [[noreturn]] static void fail(const std::string & path, const std::string & msg) {
+        throw std::runtime_error("JSON schema error at " + path + ": " + msg);
+    }
+
+    static int get_count(const common_json & schema, const std::string & key, const std::string & path, int def) {
+        if (!schema.contains(key)) {
+            return def;
+        }
+        const common_json & value = schema.at(key);
+        if (!value.is_number_integer() || value.get<int>() < 0) {
+            fail(path, key + " must be a non-negative integer");
+        }
+        return value.get<int>();
+    }
+
+    // a fractional bound is rounded inwards, towards the integers it still admits
+    static int64_t get_bound(const common_json & schema, const std::string & key, const std::string & path, bool round_up) {
+        const common_json & value = schema.at(key);
+        if (value.is_number_integer()) {
+            return value.get<int64_t>();
+        }
+        if (!value.is_number()) {
+            fail(path, key + " must be a number");
+        }
+        double d = value.get<double>();
+        return (int64_t) (round_up ? std::ceil(d) : std::floor(d));
+    }
+
+    static common_chat_schema::string_format get_format(const common_json & schema, const std::string & path) {
+        if (!schema.contains("format")) {
+            return common_chat_schema::FORMAT_NONE;
+        }
+        const common_json & value = schema.at("format");
+        if (!value.is_string()) {
+            fail(path, "format must be a string");
+        }
+        std::string format = value.get<std::string>();
+        if (format == "date") {
+            return common_chat_schema::FORMAT_DATE;
+        }
+        if (format == "time") {
+            return common_chat_schema::FORMAT_TIME;
+        }
+        if (format == "date-time") {
+            return common_chat_schema::FORMAT_DATE_TIME;
+        }
+        if (format == "uuid" || (format.size() == 5 && format.compare(0, 4, "uuid") == 0 && format[4] >= '1' && format[4] <= '5')) {
+            return common_chat_schema::FORMAT_UUID;
+        }
+        return common_chat_schema::FORMAT_NONE;
+    }
+
+    const common_json & resolve_ref(const std::string & ref, const std::string & path) {
+        const common_json * target = &root_;
+        auto tokens = string_split(ref.substr(1), "/");
+        for (size_t i = 1; i < tokens.size(); i++) {
+            const std::string & sel = tokens[i];
+            if (target->is_object() && target->contains(sel)) {
+                target = &target->at(sel);
+            } else if (target->is_array()) {
+                size_t idx;
+                try {
+                    idx = std::stoull(sel);
+                } catch (const std::logic_error &) {
+                    idx = target->size();
+                }
+                if (idx >= target->size()) {
+                    fail(path, "cannot resolve $ref " + ref + ", " + sel + " is out of range");
+                }
+                target = &target->at(idx);
+            } else {
+                fail(path, "cannot resolve $ref " + ref + ", " + sel + " not found");
+            }
+        }
+        return *target;
+    }
+
+    common_chat_schema_ptr build_ref(const common_json & value, const std::string & path) {
+        if (!value.is_string()) {
+            fail(path, "$ref must be a string");
+        }
+        std::string ref = value.get<std::string>();
+        if (ref.compare(0, 2, "#/") != 0) {
+            fail(path, "unsupported $ref " + ref + ", only references into the same document are supported");
+        }
+        if (refs_.find(ref) == refs_.end()) {
+            // reserve the key first, so that a cycle back to this $ref stops here
+            refs_[ref] = nullptr;
+            refs_[ref] = build_node(resolve_ref(ref, path), ref);
+        }
+        auto node = std::make_unique<common_chat_schema_ref>(ref);
+        pending_.push_back(node.get());
+        return node;
+    }
+
+    template <typename T>
+    common_chat_schema_ptr build_alternatives(const common_json & alts, const std::string & path) {
+        if (!alts.is_array()) {
+            fail(path, "must be an array of schemas");
+        }
+        if (alts.empty()) {
+            fail(path, "must not be empty");
+        }
+        auto node = std::make_unique<T>();
+        size_t i = 0;
+        for (const auto & alt : alts) {
+            node->children.push_back(build_node(alt, path + "/" + std::to_string(i++)));
+        }
+        return node;
+    }
+
+    common_chat_schema_ptr build_object(const common_json & schema, const std::string & path) {
+        auto node = std::make_unique<common_chat_schema_object>();
+
+        std::unordered_set<std::string> required;
+        if (schema.contains("required") && schema.at("required").is_array()) {
+            for (const auto & name : schema.at("required")) {
+                if (name.is_string()) {
+                    required.insert(name.get<std::string>());
+                }
+            }
+        }
+
+        if (schema.contains("properties")) {
+            const common_json & properties = schema.at("properties");
+            if (!properties.is_object()) {
+                fail(path, "properties must be an object");
+            }
+            for (const auto & [name, prop] : properties.items()) {
+                node->properties.push_back({name, build_node(prop, path + "/properties/" + name), required.count(name) > 0});
+            }
+        }
+
+        if (schema.contains("additionalProperties")) {
+            const common_json & additional = schema.at("additionalProperties");
+            if (additional.is_boolean()) {
+                if (additional.get<bool>()) {
+                    node->additional_properties = std::make_unique<common_chat_schema_any>();
+                }
+            } else if (additional.is_object()) {
+                node->additional_properties = build_node(additional, path + "/additionalProperties");
+            } else {
+                fail(path, "additionalProperties must be a boolean or a schema");
+            }
+        } else if (!schema.contains("properties")) {
+            // {"type": "object"} on its own accepts any object
+            node->additional_properties = std::make_unique<common_chat_schema_any>();
+        }
+
+        return node;
+    }
+
+    common_chat_schema_ptr build_array(const common_json & schema, const std::string & path) {
+        auto node = std::make_unique<common_chat_schema_array>();
+        if (schema.contains("items") || schema.contains("prefixItems")) {
+            // "items" wins when both are present; as in the converter, a schema instead of an array is the item schema
+            const std::string key = schema.contains("items") ? "items" : "prefixItems";
+            const common_json & items = schema.at(key);
+            if (items.is_array()) {
+                auto tuple = std::make_unique<common_chat_schema_tuple>();
+                size_t i = 0;
+                for (const auto & item : items) {
+                    tuple->items.push_back(build_node(item, path + "/" + key + "/" + std::to_string(i++)));
+                }
+                return tuple;
+            }
+            node->items = build_node(items, path + "/" + key);
+        } else {
+            node->items = std::make_unique<common_chat_schema_any>();
+        }
+        node->min_items = get_count(schema, "minItems", path, 0);
+        node->max_items = get_count(schema, "maxItems", path, -1);
+        return node;
+    }
+
+    common_chat_schema_ptr build_string(const common_json & schema, const std::string & path) {
+        auto node = std::make_unique<common_chat_schema_string>();
+        if (schema.contains("pattern")) {
+            const common_json & pattern = schema.at("pattern");
+            if (!pattern.is_string()) {
+                fail(path, "pattern must be a string");
+            }
+            node->pattern = pattern.get<std::string>();
+        }
+        node->format     = get_format(schema, path);
+        node->min_length = get_count(schema, "minLength", path, 0);
+        node->max_length = get_count(schema, "maxLength", path, -1);
+        return node;
+    }
+
+    common_chat_schema_ptr build_integer(const common_json & schema, const std::string & path) {
+        auto node = std::make_unique<common_chat_schema_integer>();
+        if (schema.contains("minimum")) {
+            node->minimum = get_bound(schema, "minimum", path, /* round_up */ true);
+        } else if (schema.contains("exclusiveMinimum")) {
+            node->minimum = get_bound(schema, "exclusiveMinimum", path, /* round_up */ false) + 1;
+        }
+        if (schema.contains("maximum")) {
+            node->maximum = get_bound(schema, "maximum", path, /* round_up */ false);
+        } else if (schema.contains("exclusiveMaximum")) {
+            node->maximum = get_bound(schema, "exclusiveMaximum", path, /* round_up */ true) - 1;
+        }
+        return node;
+    }
+
+    common_chat_schema_ptr build_node(const common_json & schema, const std::string & path) {
+        if (!schema.is_object()) {
+            fail(path, "schema must be an object");
+        }
+        if (schema.contains("$ref")) {
+            return build_ref(schema.at("$ref"), path);
+        }
+        if (schema.contains("oneOf") || schema.contains("anyOf")) {
+            const std::string key = schema.contains("oneOf") ? "oneOf" : "anyOf";
+            return build_alternatives<common_chat_schema_any_of>(schema.at(key), path + "/" + key);
+        }
+
+        common_json type;
+        if (schema.contains("type")) {
+            type = schema.at("type");
+        }
+        if (type.is_array()) {
+            // {"type": ["a", "b"], ...} is {"anyOf": [{"type": "a", ...}, {"type": "b", ...}]}
+            if (type.empty()) {
+                fail(path, "type must not be empty");
+            }
+            auto node = std::make_unique<common_chat_schema_any_of>();
+            size_t i = 0;
+            for (const auto & t : type) {
+                common_json alt = schema;
+                alt["type"] = t;
+                node->children.push_back(build_node(alt, path + "/type/" + std::to_string(i++)));
+            }
+            return node;
+        }
+        if (schema.contains("const")) {
+            return std::make_unique<common_chat_schema_const>(schema.at("const"));
+        }
+        if (schema.contains("enum")) {
+            const common_json & values = schema.at("enum");
+            if (!values.is_array() || values.empty()) {
+                fail(path, "enum must be a non-empty array");
+            }
+            auto node = std::make_unique<common_chat_schema_enum>();
+            for (const auto & value : values) {
+                node->values.push_back(value);
+            }
+            return node;
+        }
+        if (!type.is_null() && !type.is_string()) {
+            fail(path, "type must be a string or an array of strings");
+        }
+
+        const std::string type_name = type.is_string() ? type.get<std::string>() : "";
+        const bool has_properties = schema.contains("properties") ||
+            (schema.contains("additionalProperties") && schema.at("additionalProperties") != true);
+
+        if (type_name.empty()) {
+            // without a type the structural keywords decide, in the same order as the converter
+            if (has_properties) {
+                return build_object(schema, path);
+            }
+            if (schema.contains("allOf")) {
+                return build_alternatives<common_chat_schema_all_of>(schema.at("allOf"), path + "/allOf");
+            }
+            if (schema.contains("items") || schema.contains("prefixItems")) {
+                return build_array(schema, path);
+            }
+            if (schema.contains("pattern") || schema.contains("minLength") || schema.contains("maxLength") || get_format(schema, path) != common_chat_schema::FORMAT_NONE) {
+                return build_string(schema, path);
+            }
+            return std::make_unique<common_chat_schema_any>();
+        }
+        if (type_name == "object") {
+            if (!has_properties && schema.contains("allOf")) {
+                return build_alternatives<common_chat_schema_all_of>(schema.at("allOf"), path + "/allOf");
+            }
+            return build_object(schema, path);
+        }
+        if (type_name == "string") {
+            if (schema.contains("allOf")) {
+                return build_alternatives<common_chat_schema_all_of>(schema.at("allOf"), path + "/allOf");
+            }
+            return build_string(schema, path);
+        }
+        if (type_name == "array") {
+            return build_array(schema, path);
+        }
+        if (type_name == "integer") {
+            return build_integer(schema, path);
+        }
+        if (type_name == "number") {
+            return std::make_unique<common_chat_schema_number>();
+        }
+        if (type_name == "boolean") {
+            return std::make_unique<common_chat_schema_boolean>();
+        }
+        if (type_name == "null") {
+            return std::make_unique<common_chat_schema_null>();
+        }
+        fail(path, "unrecognized type " + type_name);
+    }
+
+  public:
+    common_chat_schema_builder(const common_json & root, common_chat_schema_document & doc) : root_(root), doc_(doc) {}
+
+    common_chat_schema_ptr build() {
+        auto node = build_node(root_, "#");
+        for (auto & entry : refs_) {
+            doc_.refs[entry.first] = std::move(entry.second);
+        }
+        for (auto * ref : pending_) {
+            ref->target = doc_.refs.at(ref->ref).get();
+        }
+        return node;
+    }
+};
+
+common_chat_schema_document common_chat_schema_from_json(const common_json & schema) {
+    common_chat_schema_document doc;
+    doc.root = common_chat_schema_builder(schema, doc).build();
+    return doc;
+}
+
+static common_chat_schema::value_type json_type(const common_json & value) {
+    if (value.is_null()) {
+        return common_chat_schema::TYPE_NULL;
+    }
+    if (value.is_boolean()) {
+        return common_chat_schema::TYPE_BOOLEAN;
+    }
+    if (value.is_number_integer()) {
+        return common_chat_schema::TYPE_INTEGER;
+    }
+    if (value.is_number()) {
+        return common_chat_schema::TYPE_NUMBER;
+    }
+    if (value.is_string()) {
+        return common_chat_schema::TYPE_STRING;
+    }
+    if (value.is_array()) {
+        return common_chat_schema::TYPE_ARRAY;
+    }
+    return common_chat_schema::TYPE_OBJECT;
+}
+
+static common_chat_schema::type_set value_types_impl(const common_chat_schema & s, std::unordered_set<const common_chat_schema *> & visited) {
+    switch (s.kind()) {
+        case common_chat_schema::KIND_ANY:
+            return common_chat_schema::type_set::all();
+        case common_chat_schema::KIND_NULL:
+            return { common_chat_schema::TYPE_NULL };
+        case common_chat_schema::KIND_BOOLEAN:
+            return { common_chat_schema::TYPE_BOOLEAN };
+        case common_chat_schema::KIND_NUMBER:
+            return { common_chat_schema::TYPE_NUMBER, common_chat_schema::TYPE_INTEGER };
+        case common_chat_schema::KIND_INTEGER:
+            return { common_chat_schema::TYPE_INTEGER };
+        case common_chat_schema::KIND_STRING:
+            return { common_chat_schema::TYPE_STRING };
+        case common_chat_schema::KIND_ARRAY:
+        case common_chat_schema::KIND_TUPLE:
+            return { common_chat_schema::TYPE_ARRAY };
+        case common_chat_schema::KIND_OBJECT:
+            return { common_chat_schema::TYPE_OBJECT };
+        case common_chat_schema::KIND_CONST:
+            return { json_type(static_cast<const common_chat_schema_const &>(s).value) };
+        case common_chat_schema::KIND_ENUM: {
+            common_chat_schema::type_set types;
+            for (const auto & value : static_cast<const common_chat_schema_enum &>(s).values) {
+                types.add(json_type(value));
+            }
+            return types;
+        }
+        case common_chat_schema::KIND_REF: {
+            const auto * target = static_cast<const common_chat_schema_ref &>(s).target;
+            if (!target || !visited.insert(target).second) {
+                // a cycle contributes no type, to be safe
+                return {};
+            }
+            auto types = value_types_impl(*target, visited);
+            visited.erase(target);
+            return types;
+        }
+        case common_chat_schema::KIND_ANY_OF: {
+            common_chat_schema::type_set types;
+            for (const auto & child : static_cast<const common_chat_schema_any_of &>(s).children) {
+                types |= value_types_impl(*child, visited);
+            }
+            return types;
+        }
+        case common_chat_schema::KIND_ALL_OF: {
+            auto types = common_chat_schema::type_set::all();
+            for (const auto & child : static_cast<const common_chat_schema_all_of &>(s).children) {
+                types &= value_types_impl(*child, visited);
+            }
+            return types;
+        }
+    }
+    return {};
+}
+
+common_chat_schema::type_set common_chat_schema::value_types() const {
+    std::unordered_set<const common_chat_schema *> visited;
+    return value_types_impl(*this, visited);
+}
+
+static bool may_be_string_impl(const common_chat_schema & s, std::unordered_set<const common_chat_schema *> & visited) {
+    switch (s.kind()) {
+        case common_chat_schema::KIND_STRING:
+            return true;
+        case common_chat_schema::KIND_CONST:
+            return static_cast<const common_chat_schema_const &>(s).value.is_string();
+        case common_chat_schema::KIND_ENUM:
+            for (const auto & v : static_cast<const common_chat_schema_enum &>(s).values) {
+                if (v.is_string()) {
+                    return true;
+                }
+            }
+            return false;
+        case common_chat_schema::KIND_REF: {
+            // a cycle is taken as not a string, to be safe
+            const auto * target = static_cast<const common_chat_schema_ref &>(s).target;
+            if (!target || !visited.insert(target).second) {
+                return false;
+            }
+            bool result = may_be_string_impl(*target, visited);
+            visited.erase(target);
+            return result;
+        }
+        case common_chat_schema::KIND_ANY_OF:
+            for (const auto & child : static_cast<const common_chat_schema_any_of &>(s).children) {
+                if (may_be_string_impl(*child, visited)) {
+                    return true;
+                }
+            }
+            return false;
+        case common_chat_schema::KIND_ALL_OF: {
+            // every child must allow a string, an any child constrains nothing
+            bool any_string = false;
+            for (const auto & child : static_cast<const common_chat_schema_all_of &>(s).children) {
+                if (child->kind() == common_chat_schema::KIND_ANY) {
+                    continue;
+                }
+                if (!may_be_string_impl(*child, visited)) {
+                    return false;
+                }
+                any_string = true;
+            }
+            return any_string;
+        }
+        default:
+            return false;
+    }
+}
+
+bool common_chat_schema::may_be_string() const {
+    std::unordered_set<const common_chat_schema *> visited;
+    return may_be_string_impl(*this, visited);
+}
+
+const char * common_chat_schema::kind_name(node_kind kind) {
+    switch (kind) {
+        case KIND_ANY:     return "any";
+        case KIND_REF:     return "ref";
+        case KIND_ANY_OF:  return "anyOf";
+        case KIND_ALL_OF:  return "allOf";
+        case KIND_CONST:   return "const";
+        case KIND_ENUM:    return "enum";
+        case KIND_NULL:    return "null";
+        case KIND_BOOLEAN: return "boolean";
+        case KIND_NUMBER:  return "number";
+        case KIND_INTEGER: return "integer";
+        case KIND_STRING:  return "string";
+        case KIND_ARRAY:   return "array";
+        case KIND_TUPLE:   return "tuple";
+        case KIND_OBJECT:  return "object";
+    }
+    return "?";
+}
+
+const char * common_chat_schema::type_name(value_type type) {
+    switch (type) {
+        case TYPE_NULL:    return "null";
+        case TYPE_BOOLEAN: return "boolean";
+        case TYPE_NUMBER:  return "number";
+        case TYPE_INTEGER: return "integer";
+        case TYPE_STRING:  return "string";
+        case TYPE_ARRAY:   return "array";
+        case TYPE_OBJECT:  return "object";
+    }
+    return "?";
+}

common/json-schema.h

@@ -0,0 +1,198 @@
+#pragma once
+
+#include "json.h"
+
+#include <cstdint>
+#include <initializer_list>
+#include <map>
+#include <memory>
+#include <string>
+#include <vector>
+
+// JSON schema, covering the subset that json_schema_to_grammar() can convert.
+
+struct common_chat_schema {
+    enum node_kind {
+        KIND_ANY,
+        KIND_REF,
+        KIND_ANY_OF,
+        KIND_ALL_OF,
+        KIND_CONST,
+        KIND_ENUM,
+        KIND_NULL,
+        KIND_BOOLEAN,
+        KIND_NUMBER,
+        KIND_INTEGER,
+        KIND_STRING,
+        KIND_ARRAY,
+        KIND_TUPLE,
+        KIND_OBJECT,
+    };
+
+    enum value_type {
+        TYPE_NULL,
+        TYPE_BOOLEAN,
+        TYPE_NUMBER,
+        TYPE_INTEGER,
+        TYPE_STRING,
+        TYPE_ARRAY,
+        TYPE_OBJECT,
+    };
+
+    enum string_format {
+        FORMAT_NONE,
+        FORMAT_UUID,  // uuid, uuid1 .. uuid5
+        FORMAT_DATE,
+        FORMAT_TIME,
+        FORMAT_DATE_TIME,
+    };
+
+    class type_set {
+        uint32_t mask_ = 0;
+
+      public:
+        type_set() = default;
+        type_set(std::initializer_list<value_type> types) {
+            for (auto type : types) {
+                add(type);
+            }
+        }
+
+        static type_set all() {
+            return { TYPE_NULL, TYPE_BOOLEAN, TYPE_NUMBER, TYPE_INTEGER, TYPE_STRING, TYPE_ARRAY, TYPE_OBJECT };
+        }
+
+        void add(value_type type) { mask_ |= 1u << type; }
+
+        bool has(value_type type) const { return (mask_ & (1u << type)) != 0; }
+        bool is_only(value_type type) const { return mask_ == (1u << type); }
+        bool empty() const { return mask_ == 0; }
+
+        type_set & operator|=(const type_set & other) { mask_ |= other.mask_; return *this; }
+        type_set & operator&=(const type_set & other) { mask_ &= other.mask_; return *this; }
+
+        bool operator==(const type_set & other) const { return mask_ == other.mask_; }
+        bool operator!=(const type_set & other) const { return mask_ != other.mask_; }
+    };
+
+    virtual ~common_chat_schema() = default;
+    virtual node_kind kind() const = 0;
+
+    type_set value_types() const;
+
+    // Whether a value matching the schema may be a string, through any branch of it.
+    bool may_be_string() const;
+
+    static const char * kind_name(node_kind kind);
+    static const char * type_name(value_type type);
+};
+
+using common_chat_schema_ptr = std::unique_ptr<common_chat_schema>;
+
+struct common_chat_schema_any : common_chat_schema {
+    node_kind kind() const override { return KIND_ANY; }
+};
+
+// {"$ref": "#/..."}, only references into the same document are supported
+struct common_chat_schema_ref : common_chat_schema {
+    std::string                ref;
+    const common_chat_schema * target = nullptr;  // owned by common_chat_schema_document::refs
+
+    explicit common_chat_schema_ref(std::string ref) : ref(std::move(ref)) {}
+
+    node_kind kind() const override { return KIND_REF; }
+};
+
+// oneOf / anyOf, or a "type" array expanded to one alternative per type
+struct common_chat_schema_any_of : common_chat_schema {
+    std::vector<common_chat_schema_ptr> children;
+
+    node_kind kind() const override { return KIND_ANY_OF; }
+};
+
+struct common_chat_schema_all_of : common_chat_schema {
+    std::vector<common_chat_schema_ptr> children;
+
+    node_kind kind() const override { return KIND_ALL_OF; }
+};
+
+struct common_chat_schema_const : common_chat_schema {
+    common_json value;
+
+    explicit common_chat_schema_const(common_json value) : value(std::move(value)) {}
+
+    node_kind kind() const override { return KIND_CONST; }
+};
+
+struct common_chat_schema_enum : common_chat_schema {
+    std::vector<common_json> values;
+
+    node_kind kind() const override { return KIND_ENUM; }
+};
+
+struct common_chat_schema_null : common_chat_schema {
+    node_kind kind() const override { return KIND_NULL; }
+};
+
+struct common_chat_schema_boolean : common_chat_schema {
+    node_kind kind() const override { return KIND_BOOLEAN; }
+};
+
+struct common_chat_schema_number : common_chat_schema {
+    node_kind kind() const override { return KIND_NUMBER; }
+};
+
+// bounds are inclusive, exclusiveMinimum / exclusiveMaximum are folded in
+struct common_chat_schema_integer : common_chat_schema {
+    int64_t minimum = INT64_MIN;  // INT64_MIN for unbounded
+    int64_t maximum = INT64_MAX;  // INT64_MAX for unbounded
+
+    node_kind kind() const override { return KIND_INTEGER; }
+};
+
+struct common_chat_schema_string : common_chat_schema {
+    std::string   pattern;  // empty when absent
+    string_format format     = FORMAT_NONE;
+    int           min_length = 0;
+    int           max_length = -1;  // -1 for unbounded
+
+    node_kind kind() const override { return KIND_STRING; }
+};
+
+struct common_chat_schema_array : common_chat_schema {
+    common_chat_schema_ptr items;  // a common_chat_schema_any when "items" is absent
+    int                    min_items = 0;
+    int                    max_items = -1;  // -1 for unbounded
+
+    node_kind kind() const override { return KIND_ARRAY; }
+};
+
+struct common_chat_schema_tuple : common_chat_schema {
+    std::vector<common_chat_schema_ptr> items;
+
+    node_kind kind() const override { return KIND_TUPLE; }
+};
+
+struct common_chat_schema_property {
+    std::string            name;
+    common_chat_schema_ptr schema;
+    bool                   required = false;
+};
+
+struct common_chat_schema_object : common_chat_schema {
+    std::vector<common_chat_schema_property> properties;             // in schema order
+    common_chat_schema_ptr                   additional_properties;  // null when not allowed
+
+    node_kind kind() const override { return KIND_OBJECT; }
+};
+
+struct common_chat_schema_document {
+    common_chat_schema_ptr                        root;
+    std::map<std::string, common_chat_schema_ptr> refs;
+};
+
+// A document shared by the PEG parsers built from its nodes, which it keeps alive
+using common_chat_schema_document_ptr = std::shared_ptr<const common_chat_schema_document>;
+
+// Throws std::runtime_error when the schema falls outside the supported subset.
+common_chat_schema_document common_chat_schema_from_json(const common_json & schema);

common/log.cpp

@@ -37,6 +37,16 @@ void common_log_set_verbosity_thold(int verbosity) {
     common_log_verbosity_thold = verbosity;
 }
 
+static bool common_log_jsonl = false;
+
+bool common_log_get_jsonl(void) {
+    return common_log_jsonl;
+}
+
+void common_log_set_jsonl(bool jsonl) {
+    common_log_jsonl = jsonl;
+}
+
 static int64_t t_us() {
     return std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::system_clock::now().time_since_epoch()).count();
 }
@@ -87,6 +97,7 @@ struct common_log_entry {
     bool is_end       { false }; // signals the worker thread to stop
     bool prefix       { false };
     bool jsonl        { false };
+    bool is_json      { false }; // msg already holds a serialized JSON object
 
     common_log_entry(size_t size = 256) : msg(size) { }
 
@@ -107,6 +118,12 @@ struct common_log_entry {
         }
 
         if (jsonl) {
+            if (is_json) {
+                fprintf(fcur, "%s\n", msg.data());
+                fflush(fcur);
+                return;
+            }
+
             common_json obj = {
                 {"type",  "log"},
                 {"time",  timestamp},
@@ -156,7 +173,6 @@ struct common_log {
         file       = nullptr;
         prefix     = false;
         timestamps = false;
-        jsonl      = false;
         running    = false;
         t_start    = t_us();
 
@@ -184,7 +200,6 @@ private:
 
     bool prefix;
     bool timestamps;
-    bool jsonl;
     bool running;
 
     int64_t t_start;
@@ -273,7 +288,8 @@ public:
         entry.is_end    = false;
         entry.level     = level;
         entry.prefix    = prefix;
-        entry.jsonl     = jsonl;
+        entry.jsonl     = common_log_jsonl;
+        entry.is_json   = false;
         entry.timestamp = 0;
         if (timestamps) {
             entry.timestamp = t_us() - t_start;
@@ -283,6 +299,42 @@ public:
         cv_new.notify_one();
     }
 
+    void add_json(const char * type, const common_json & obj) {
+        const common_json full = {
+            {"type", type},
+            {"data", obj},
+        };
+
+        const std::string text = full.dump_safe();
+
+        std::unique_lock<std::mutex> lock(mtx);
+
+        // block if the queue is full
+        cv_full.wait(lock, [this]() { return !running || !is_full(); });
+
+        if (!running) {
+            // discard messages while the worker thread is paused
+            return;
+        }
+
+        auto & entry = queue[tail];
+
+        if (entry.msg.size() < text.size() + 1) {
+            entry.msg.resize(text.size() + 1);
+        }
+        memcpy(entry.msg.data(), text.c_str(), text.size() + 1);
+
+        entry.is_end    = false;
+        entry.level     = GGML_LOG_LEVEL_NONE;
+        entry.prefix    = false;
+        entry.jsonl     = true;
+        entry.is_json   = true;
+        entry.timestamp = 0;
+
+        tail = (tail + 1) % queue.size();
+        cv_new.notify_one();
+    }
+
     void resume() {
         std::lock_guard<std::mutex> lock(mtx);
 
@@ -388,12 +440,6 @@ public:
 
         this->timestamps = timestamps;
     }
-
-    void set_jsonl(bool jsonl) {
-        std::lock_guard<std::mutex> lock(mtx);
-
-        this->jsonl = jsonl;
-    }
 };
 
 //
@@ -440,6 +486,14 @@ void common_log_add(struct common_log * log, enum ggml_log_level level, const ch
     va_end(args);
 }
 
+void common_log_add_json(struct common_log * log, const char * type, const common_json & obj) {
+    if (!common_log_jsonl) {
+        return;
+    }
+
+    log->add_json(type, obj);
+}
+
 void common_log_set_file(struct common_log * log, const char * file) {
     log->set_file(file);
 }
@@ -467,10 +521,6 @@ void common_log_set_timestamps(struct common_log * log, bool timestamps) {
     log->set_timestamps(timestamps);
 }
 
-void common_log_set_jsonl(struct common_log * log, bool jsonl) {
-    log->set_jsonl(jsonl);
-}
-
 void common_log_flush(struct common_log * log) {
     log->pause();
     log->resume();

common/log.h

@@ -43,6 +43,10 @@ int  common_log_get_verbosity_thold(void);
 
 void common_log_set_verbosity_thold(int verbosity); // not thread-safe
 
+bool common_log_get_jsonl(void);
+
+void common_log_set_jsonl(bool jsonl); // not thread-safe
+
 int common_log_get_verbosity(enum ggml_log_level level);
 
 void common_log_default_callback(enum ggml_log_level level, const char * text, void * user_data);
@@ -91,7 +95,6 @@ void common_log_set_file      (struct common_log * log, const char * file); // n
 void common_log_set_colors    (struct common_log * log, log_colors colors); // not thread-safe
 void common_log_set_prefix    (struct common_log * log, bool prefix);       // whether to output prefix to each log
 void common_log_set_timestamps(struct common_log * log, bool timestamps);   // whether to output timestamps in the prefix
-void common_log_set_jsonl     (struct common_log * log, bool jsonl);        // print each log as a JSON object on one line, not thread-safe
 void common_log_flush         (struct common_log * log);                    // flush all pending log messages
 
 // helper macros for logging
@@ -127,3 +130,17 @@ void common_log_flush         (struct common_log * log);                    // f
 #define LOG_WRNV(verbosity, ...) LOG_TMPL(GGML_LOG_LEVEL_WARN,  verbosity, __VA_ARGS__)
 #define LOG_ERRV(verbosity, ...) LOG_TMPL(GGML_LOG_LEVEL_ERROR, verbosity, __VA_ARGS__)
 #define LOG_CNTV(verbosity, ...) LOG_TMPL(GGML_LOG_LEVEL_CONT,  verbosity, __VA_ARGS__)
+
+class common_json; // defined in common/json.h
+
+// helper allows different types of json output
+// no-op if --log-jsonl is not set
+void common_log_add_json(struct common_log * log, const char * type, const common_json & data);
+
+// will only print if --log-jsonl is set
+#define LOG_JSON(type, data) \
+    do { \
+        if (common_log_get_jsonl()) { \
+            common_log_add_json(common_log_main(), type, data); \
+        } \
+    } while (0)

common/parsers/cohere2moe.cpp

@@ -129,15 +129,6 @@ common_chat_params common_chat_params_init_cohere2moe(const common_chat_template
     if (include_grammar) {
         data.grammar_lazy = !has_response_format && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
         data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                auto         schema   = function.at("parameters");
-                builder.resolve_refs(schema);
-            });
-            if (has_response_format) {
-                auto schema = inputs.json_schema;
-                builder.resolve_refs(schema);
-            }
             parser.build_grammar(builder, data.grammar_lazy);
         });
 

common/parsers/deepseek.cpp

@@ -149,39 +149,28 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
             foreach_function(inputs.tools, [&](const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
-                auto         params   = function.contains("parameters") ? function.at("parameters") : json::object();
-                const auto & props    = params.contains("properties") ? params.at("properties") : json::object();
-
-                std::set<std::string> required;
-                if (params.contains("required")) {
-                    required = params.at("required").get<std::set<std::string>>();
-                }
-
-                auto schema_info = common_schema_info();
-                schema_info.resolve_refs(params);
 
                 std::vector<common_peg_parser> required_parsers;
                 std::vector<common_peg_parser> optional_parsers;
-                for (const auto & [param_name, param_schema] : props.items()) {
-                    bool is_required = required.find(param_name) != required.end();
-                    bool is_string   = schema_info.resolves_to_string(param_schema);
+                foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
+                    bool is_string = param.schema->may_be_string();
 
                     auto arg = p.tool_arg(
-                        p.tool_arg_open(p.literal(PARAM_START + " name=\"") + p.tool_arg_name(p.literal(param_name)) +
+                        p.tool_arg_open(p.literal(PARAM_START + " name=\"") + p.tool_arg_name(p.literal(param.name)) +
                                         p.literal("\" string=\"" + std::string(is_string ? "true" : "false") + "\">")) +
                         (is_string ?
                              p.tool_arg_string_value(p.until(PARAM_END)) :
-                             p.tool_arg_json_value(p.schema(p.json(), "tool-" + name + "-arg-" + param_name + "-schema",
-                                                            param_schema, false))) +
+                             p.tool_arg_json_value(p.schema(p.json(), "tool-" + name + "-arg-" + param.name + "-schema",
+                                                            doc, *param.schema))) +
                         p.tool_arg_close(p.literal(PARAM_END)));
 
-                    auto named_arg = p.rule("tool-" + name + "-arg-" + param_name, arg);
-                    if (is_required) {
+                    auto named_arg = p.rule("tool-" + name + "-arg-" + param.name, arg);
+                    if (param.required) {
                         required_parsers.push_back(named_arg);
                     } else {
                         optional_parsers.push_back(named_arg);
                     }
-                }
+                });
 
                 common_peg_parser args_seq = p.eps();
                 for (size_t i = 0; i < required_parsers.size(); i++) {
@@ -266,15 +255,6 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
     if (include_grammar) {
         data.grammar_lazy = has_tools && !require_tools;
         data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                auto         schema   = function.contains("parameters") ? function.at("parameters") : json::object();
-                builder.resolve_refs(schema);
-            });
-            if (has_response_format) {
-                auto schema = inputs.json_schema;
-                builder.resolve_refs(schema);
-            }
             parser.build_grammar(builder, data.grammar_lazy);
         });
 

common/parsers/functionary-v3-2.cpp

@@ -45,7 +45,7 @@ common_chat_params common_chat_params_init_functionary_v3_2(const common_chat_te
         foreach_function(inputs.tools, [&](const json & tool) {
             const auto & function = tool.at("function");
             std::string  name     = function.at("name");
-            const auto & schema   = function.at("parameters");
+            const auto   schema   = common_chat_tool_parameters(function);
 
             // Tool format: >>>function_name\n{json_args}
             auto tool_parser = p.tool(
@@ -82,11 +82,6 @@ common_chat_params common_chat_params_init_functionary_v3_2(const common_chat_te
         data.grammar_lazy = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
 
         data.grammar = build_grammar([&](const common_grammar_builder & builder) {
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                auto         schema   = function.at("parameters");
-                builder.resolve_refs(schema);
-            });
             parser.build_grammar(builder, data.grammar_lazy);
         });
 

common/parsers/gemma4.cpp

@@ -291,15 +291,6 @@ common_chat_params common_chat_params_init_gemma4(const common_chat_template &
     if (include_grammar) {
         data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED));
         data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                auto         schema   = function.at("parameters");
-                builder.resolve_refs(schema);
-            });
-            if (has_response_format) {
-                auto schema = inputs.json_schema;
-                builder.resolve_refs(schema);
-            }
             parser.build_grammar(builder, data.grammar_lazy);
         });
 

common/parsers/gigachat-v3.cpp

@@ -33,7 +33,7 @@ common_chat_params common_chat_params_init_gigachat_v3(
             for (const auto & tool : inputs.tools) {
                 const auto & function = tool.at("function");
                 std::string name = function.at("name");
-                const auto & schema = function.at("parameters");
+                const auto  schema = common_chat_tool_parameters(function);
 
                 auto tool_name = p.json_member("name", "\"" + p.tool_name(p.literal(name)) + "\"");
                 auto tool_args = p.json_member("arguments", p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)));
@@ -65,11 +65,6 @@ common_chat_params common_chat_params_init_gigachat_v3(
         data.grammar_lazy = has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
 
         data.grammar = build_grammar([&](const common_grammar_builder & builder) {
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                auto schema = function.at("parameters");
-                builder.resolve_refs(schema);
-            });
             parser.build_grammar(builder, data.grammar_lazy);
         });
 

common/parsers/gpt-oss.cpp

@@ -109,7 +109,7 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
             foreach_function(inputs.tools, [&](const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
-                const auto & params   = function.at("parameters");
+                const auto   params   = common_chat_tool_parameters(function);
 
                 auto func_name  = p.literal(" to=functions.") + p.tool_name(p.literal(name));
                 auto constraint = p.optional(p.space() + p.optional(p.literal("<|constrain|>")) + constrain_type);
@@ -143,15 +143,6 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
     if (include_grammar) {
         data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED));
         data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                auto         schema   = function.at("parameters");
-                builder.resolve_refs(schema);
-            });
-            if (has_response_format) {
-                auto schema = inputs.json_schema;
-                builder.resolve_refs(schema);
-            }
             parser.build_grammar(builder, data.grammar_lazy);
         });
 

common/parsers/kimi-k2.cpp

@@ -82,7 +82,7 @@ common_chat_params common_chat_params_init_kimi_k2(const common_chat_template &
         foreach_function(inputs.tools, [&](const json & tool) {
             const auto & function = tool.at("function");
             std::string  name     = function.at("name");
-            const auto & schema   = function.at("parameters");
+            const auto   schema   = common_chat_tool_parameters(function);
 
             // Match: functions.<name>:<digits>
             // Capture the full call id (functions.<name>:<digits>) using tool_id tag
@@ -116,11 +116,6 @@ common_chat_params common_chat_params_init_kimi_k2(const common_chat_template &
     if (include_grammar) {
         data.grammar_lazy = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
         data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                auto         schema   = function.at("parameters");
-                builder.resolve_refs(schema);
-            });
             parser.build_grammar(builder, data.grammar_lazy);
         });
 

common/parsers/kimi-k3.cpp

@@ -98,7 +98,7 @@ common_chat_params common_chat_params_init_kimi_k3(const common_chat_template &
         foreach_function(inputs.tools, [&](const json & tool) {
             const auto & function = tool.at("function");
             std::string  name     = function.at("name");
-            const json   schema   = function.contains("parameters") ? function.at("parameters") : json::object();
+            const json   schema   = common_chat_tool_parameters(function);
 
             // arguments come one tag per key, with the JSON type in a type="..."
             // attribute. the type is taken from the tool schema instead, as it tells
@@ -155,13 +155,6 @@ common_chat_params common_chat_params_init_kimi_k3(const common_chat_template &
     if (include_grammar) {
         data.grammar_lazy = inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_REQUIRED;
         data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                if (function.contains("parameters")) {
-                    auto schema = function.at("parameters");
-                    builder.resolve_refs(schema);
-                }
-            });
             parser.build_grammar(builder, data.grammar_lazy);
         });
 

common/parsers/lfm2.cpp

@@ -98,15 +98,6 @@ common_chat_params common_chat_params_init_lfm2(const common_chat_template &
     if (include_grammar) {
         data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED));
         data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                auto         schema   = function.at("parameters");
-                builder.resolve_refs(schema);
-            });
-            if (has_response_format) {
-                auto schema = inputs.json_schema;
-                builder.resolve_refs(schema);
-            }
             parser.build_grammar(builder, data.grammar_lazy);
         });
 

common/parsers/minicpm5.cpp

@@ -71,32 +71,27 @@ common_chat_params common_chat_params_init_minicpm5(const common_chat_template &
             foreach_function(inputs.tools, [&](const json & tool) {
                 const auto &      function = tool.at("function");
                 const std::string name     = function.at("name");
-                auto              params   = function.contains("parameters") ? function.at("parameters") : json::object();
+
+                std::vector<common_peg_parser> arg_rules;
+                foreach_parameter(function, [&](const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) {
+                    auto value_parser = p.eps();
+                    if (prop.schema->may_be_string()) {
+                        value_parser = string_value;
+                    } else {
+                        value_parser = p.tool_arg_json_value(
+                                p.schema(p.json(), "tool-" + name + "-arg-" + prop.name + "-schema", doc, *prop.schema)
+                            ) + p.tool_arg_close(p.literal("</param>"));
+                    }
+
+                    arg_rules.push_back(p.tool_arg(
+                        p.tool_arg_open(p.literal("<param name=\"") + p.tool_arg_name(p.literal(prop.name)) + p.literal("\">")) +
+                        value_parser
+                    ));
+                });
 
                 auto args = p.eps();
-                if (params.contains("properties") && params.at("properties").is_object() && !params.at("properties").empty()) {
-                    auto schema_info = common_schema_info();
-                    schema_info.resolve_refs(params);
-
-                    auto arg_choice = p.choice();
-                    for (const auto & [prop_name, prop_schema] : params.at("properties").items()) {
-                        auto value_parser = p.eps();
-                        if (schema_info.resolves_to_string(prop_schema)) {
-                            value_parser = string_value;
-                        } else {
-                            value_parser = p.tool_arg_json_value(
-                                    p.schema(p.json(), "tool-" + name + "-arg-" + prop_name + "-schema", prop_schema, false)
-                                ) + p.tool_arg_close(p.literal("</param>"));
-                        }
-
-                        auto arg_rule = p.tool_arg(
-                            p.tool_arg_open(p.literal("<param name=\"") + p.tool_arg_name(p.literal(prop_name)) + p.literal("\">")) +
-                            value_parser
-                        );
-
-                        arg_choice |= arg_rule;
-                    }
-                    args = p.zero_or_more(arg_choice + p.space());
+                if (!arg_rules.empty()) {
+                    args = p.zero_or_more(p.choice(arg_rules) + p.space());
                 }
 
                 auto tool_parser = p.tool(
@@ -123,15 +118,6 @@ common_chat_params common_chat_params_init_minicpm5(const common_chat_template &
     if (include_grammar) {
         data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED));
         data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                auto         schema   = function.contains("parameters") ? function.at("parameters") : json::object();
-                builder.resolve_refs(schema);
-            });
-            if (has_response_format) {
-                auto schema = inputs.json_schema;
-                builder.resolve_refs(schema);
-            }
             parser.build_grammar(builder, data.grammar_lazy);
         });
 

common/parsers/minimax-m3.cpp

@@ -84,29 +84,18 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
             return generation_prompt + reasoning + p.content(p.rest()) + end;
         }
 
-        auto alternatives_of = [](const json & schema) -> std::optional<json> {
-            for (const auto * keyword : { "oneOf", "anyOf" }) {
-                if (schema.contains(keyword) && schema.at(keyword).is_array() && !schema.at(keyword).empty()) {
-                    return schema.at(keyword);
-                }
-            }
-            return std::nullopt;
-        };
-
         auto tool_choice = p.choice();
         foreach_function(inputs.tools, [&](const json & tool) {
             const auto & function = tool.at("function");
             std::string  name     = function.at("name");
-            auto         params   = function.contains("parameters") ? function.at("parameters") : json::object();
-
-            auto schema_info = common_schema_info();
-            schema_info.resolve_refs(params);
+            auto         params   = common_chat_tool_parameters(function);
+            auto         doc      = std::make_shared<const common_chat_schema_document>(common_chat_schema_from_json(params));
 
             // The template expands argument values recursively in XML (see the to_xml() macro)
-            std::function<common_peg_parser(const json &, const std::string &, const std::string &)> value_of;
-            std::function<common_peg_parser(const json &, const std::string &)>                      members_of;
+            std::function<common_peg_parser(const common_chat_schema &, const std::string &, const std::string &)> value_of;
+            std::function<common_peg_parser(const common_chat_schema_object &, const std::string &)>              members_of;
 
-            auto element_of = [&](const std::string & tag, const json & schema, const std::string & rule_name) {
+            auto element_of = [&](const std::string & tag, const common_chat_schema & schema, const std::string & rule_name) {
                 const std::string close = NS + "</" + tag + ">";
                 return p.rule(rule_name,
                     p.tool_arg(
@@ -117,69 +106,57 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
                         value_of(schema, rule_name, close)));
             };
 
-            value_of = [&](const json & schema,
+            value_of = [&](const common_chat_schema & schema,
                            const std::string & rule_name,
                            const std::string & close) -> common_peg_parser {
                 auto close_tag = p.tool_arg_close(p.literal(close));
 
                 // A string accepts anything, so a union with a string alternative is a string
-                if (schema_info.resolves_to_string(schema)) {
+                if (schema.may_be_string()) {
                     return p.ac(p.tool_arg_string_value(p.until(close)) + close_tag, close);
                 }
 
-                if (auto alternatives = alternatives_of(schema)) {
+                if (schema.kind() == common_chat_schema::KIND_ANY_OF) {
                     std::vector<common_peg_parser> choices;
 
                     size_t index = 0;
-                    for (const auto & alternative : *alternatives) {
+                    for (const auto & alternative : static_cast<const common_chat_schema_any_of &>(schema).children) {
                         const std::string alt_name = rule_name + "-" + std::to_string(index++);
 
                         // There is a risk that this breaks streaming deltas, but that's a risk we
                         // assume to provide tool arg streaming.
-                        choices.push_back(value_of(alternative, alt_name, close));
+                        choices.push_back(value_of(*alternative, alt_name, close));
                     }
 
                     return p.choice(choices);
                 }
 
-                const std::string type = schema.contains("type") && schema.at("type").is_string()
-                                             ? schema.at("type").get<std::string>()
-                                             : "";
-
-                if (type == "object" && schema.contains("properties")) {
-                    return p.tag(mm3::TOOL_ARG_OBJECT, members_of(schema, rule_name)) + p.space() + close_tag;
+                if (schema.kind() == common_chat_schema::KIND_OBJECT) {
+                    const auto & object = static_cast<const common_chat_schema_object &>(schema);
+                    if (!object.properties.empty()) {
+                        return p.tag(mm3::TOOL_ARG_OBJECT, members_of(object, rule_name)) + p.space() + close_tag;
+                    }
                 }
 
-                if (type == "array" && schema.contains("items")) {
+                if (schema.kind() == common_chat_schema::KIND_ARRAY) {
                     const std::string item_close = NS + "</item>";
                     auto item = p.rule(rule_name + "-item",
                         p.tag(mm3::TOOL_ARG_ITEM,
                               p.literal(NS + "<item>") +
-                                  value_of(schema.at("items"), rule_name + "-item", item_close)));
+                                  value_of(*static_cast<const common_chat_schema_array &>(schema).items, rule_name + "-item", item_close)));
                     return p.tag(mm3::TOOL_ARG_ARRAY, p.repeat(p.space() + item, 0, -1)) + p.space() + close_tag;
                 }
 
-                return p.tool_arg_json_value(p.schema(p.json(), rule_name + "-schema", schema, false)) + close_tag;
+                return p.tool_arg_json_value(p.schema(p.json(), rule_name + "-schema", doc, schema)) + close_tag;
             };
 
             // Required properties in schema order, then any number of optional ones in any order.
-            members_of = [&](const json & schema, const std::string & rule_prefix) -> common_peg_parser {
-                const auto & props = schema.at("properties");
-
-                std::set<std::string> required;
-                if (schema.contains("required")) {
-                    required = schema.at("required").get<std::set<std::string>>();
-                }
-
+            members_of = [&](const common_chat_schema_object & object, const std::string & rule_prefix) -> common_peg_parser {
                 std::vector<common_peg_parser> required_elements;
                 std::vector<common_peg_parser> optional_elements;
-                for (const auto & [key, key_schema] : props.items()) {
-                    auto element = element_of(key, key_schema, rule_prefix + "-" + key);
-                    if (required.find(key) != required.end()) {
-                        required_elements.push_back(element);
-                    } else {
-                        optional_elements.push_back(element);
-                    }
+                for (const auto & prop : object.properties) {
+                    auto element = element_of(prop.name, *prop.schema, rule_prefix + "-" + prop.name);
+                    (prop.required ? required_elements : optional_elements).push_back(element);
                 }
 
                 common_peg_parser members = p.eps();
@@ -201,8 +178,10 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
                 return members;
             };
 
-            common_peg_parser invoke_body =
-                params.contains("properties") ? members_of(params, "tool-" + name + "-arg") : p.eps();
+            common_peg_parser invoke_body = p.eps();
+            if (doc->root->kind() == common_chat_schema::KIND_OBJECT) {
+                invoke_body = members_of(static_cast<const common_chat_schema_object &>(*doc->root), "tool-" + name + "-arg");
+            }
 
             auto func_parser = p.tool(
                 p.tool_open(p.literal(NS + "<invoke name=\"") +
@@ -238,15 +217,6 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
     if (include_grammar) {
         data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED));
         data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                auto         schema   = function.contains("parameters") ? function.at("parameters") : json::object();
-                builder.resolve_refs(schema);
-            });
-            if (has_response_format) {
-                auto schema = inputs.json_schema;
-                builder.resolve_refs(schema);
-            }
             parser.build_grammar(builder, data.grammar_lazy);
         });
 

common/parsers/ministral3.cpp

@@ -89,7 +89,7 @@ common_chat_params common_chat_params_init_ministral_3(const common_chat_templat
             foreach_function(inputs.tools, [&](const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
-                const auto & schema   = function.at("parameters");
+                const auto   schema   = common_chat_tool_parameters(function);
 
                 tool_choice |=
                     p.rule("tool-" + name, p.tool_open(p.tool_name(p.literal(name)) + "[ARGS]") +
@@ -114,15 +114,6 @@ common_chat_params common_chat_params_init_ministral_3(const common_chat_templat
         data.grammar_lazy = has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
 
         data.grammar = build_grammar([&](const common_grammar_builder & builder) {
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                auto         schema   = function.at("parameters");
-                builder.resolve_refs(schema);
-            });
-            if (has_response_format) {
-                auto schema = inputs.json_schema;
-                builder.resolve_refs(schema);
-            }
             parser.build_grammar(builder, data.grammar_lazy);
         });
 

common/parsers/muse-glimmer.cpp

@@ -74,31 +74,26 @@ common_chat_params common_chat_params_init_muse_glimmer(const common_chat_templa
             foreach_function(inputs.tools, [&](const json & tool) {
                 const auto &      function = tool.at("function");
                 const std::string name     = function.at("name");
-                auto              params   = function.contains("parameters") ? function.at("parameters") : json::object();
+
+                std::vector<common_peg_parser> arg_rules;
+                foreach_parameter(function, [&](const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) {
+                    auto value_parser = p.eps();
+                    if (prop.schema->may_be_string()) {
+                        value_parser = string_value;
+                    } else {
+                        value_parser = p.tool_arg_json_value(
+                                p.schema(p.json(), "tool-" + name + "-arg-" + prop.name + "-schema", doc, *prop.schema))
+                            + p.tool_arg_close(p.literal("</atem:parameter>"));
+                    }
+
+                    arg_rules.push_back(p.tool_arg(
+                        p.tool_arg_open(p.literal("<atem:parameter name=\"") + p.tool_arg_name(p.literal(prop.name)) + p.literal("\">")) +
+                        value_parser));
+                });
 
                 auto args = p.eps();
-                if (params.contains("properties") && params.at("properties").is_object() && !params.at("properties").empty()) {
-                    auto schema_info = common_schema_info();
-                    schema_info.resolve_refs(params);
-
-                    auto arg_choice = p.choice();
-                    for (const auto & [prop_name, prop_schema] : params.at("properties").items()) {
-                        auto value_parser = p.eps();
-                        if (schema_info.resolves_to_string(prop_schema)) {
-                            value_parser = string_value;
-                        } else {
-                            value_parser = p.tool_arg_json_value(
-                                    p.schema(p.json(), "tool-" + name + "-arg-" + prop_name + "-schema", prop_schema, false))
-                                + p.tool_arg_close(p.literal("</atem:parameter>"));
-                        }
-
-                        auto arg_rule = p.tool_arg(
-                            p.tool_arg_open(p.literal("<atem:parameter name=\"") + p.tool_arg_name(p.literal(prop_name)) + p.literal("\">")) +
-                            value_parser);
-
-                        arg_choice |= arg_rule;
-                    }
-                    args = p.zero_or_more(arg_choice + p.space());
+                if (!arg_rules.empty()) {
+                    args = p.zero_or_more(p.choice(arg_rules) + p.space());
                 }
 
                 auto tool_parser = p.tool(
@@ -131,11 +126,6 @@ common_chat_params common_chat_params_init_muse_glimmer(const common_chat_templa
     if (include_grammar) {
         data.grammar_lazy = inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_REQUIRED;
         data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                auto         schema   = function.contains("parameters") ? function.at("parameters") : json::object();
-                builder.resolve_refs(schema);
-            });
             parser.build_grammar(builder, data.grammar_lazy);
         });
         data.grammar_triggers = {

common/parsers/parsers.cpp

@@ -2,8 +2,6 @@
 
 #include "log.h"
 
-#include <set>
-
 void foreach_function(const json & tools, const std::function<void(const json &)> & fn) {
     for (const auto & tool : tools) {
         if (!tool.contains("type") || tool.at("type") != "function" || !tool.contains("function")) {
@@ -14,21 +12,14 @@ void foreach_function(const json & tools, const std::function<void(const json &)
     }
 }
 
-void foreach_parameter(const json & function, const std::function<void(const std::string &, const json &, bool)> & fn) {
-    if (!function.contains("parameters") || !function.at("parameters").is_object()) {
+void foreach_parameter(const json & function, const std::function<void(const common_chat_schema_property &, const common_chat_schema_document_ptr &)> & fn) {
+    auto         params = common_chat_tool_parameters(function);
+    auto         doc    = std::make_shared<const common_chat_schema_document>(common_chat_schema_from_json(params));
+    const auto * object = dynamic_cast<const common_chat_schema_object *>(doc->root.get());
+    if (!object) {
         return;
     }
-    const auto & params = function.at("parameters");
-    if (!params.contains("properties") || !params.at("properties").is_object()) {
-        return;
-    }
-    const auto & props = params.at("properties");
-    std::set<std::string> required;
-    if (params.contains("required") && params.at("required").is_array()) {
-        required = params.at("required").get<std::set<std::string>>();
-    }
-    for (const auto & [name, prop] : props.items()) {
-        bool is_required = (required.find(name) != required.end());
-        fn(name, prop, is_required);
+    for (const auto & prop : object->properties) {
+        fn(prop, doc);
     }
 }

common/parsers/parsers.h

@@ -20,8 +20,8 @@ using json = common_json;
 // iterate over the function tools of an OpenAI-style tools array
 void foreach_function(const json & tools, const std::function<void(const json &)> & fn);
 
-// iterate over the parameters of a function tool, flagging the ones listed as required
-void foreach_parameter(const json & function, const std::function<void(const std::string &, const json &, bool)> & fn);
+// iterate over the parameters of a function tool, with the document that owns them
+void foreach_parameter(const json & function, const std::function<void(const common_chat_schema_property &, const common_chat_schema_document_ptr &)> & fn);
 
 // render a template; the override arguments let a parser feed in messages, tools or context it has rewritten
 std::string common_chat_template_direct_apply_impl(

common/parsers/qwen3-coder.cpp

@@ -93,28 +93,49 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
 
             auto tool_choice = p.choice();
             foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function   = tool.at("function");
-                std::string  name       = function.at("name");
-                auto         parameters = function.contains("parameters") ? function.at("parameters") : json::object();
-
-                auto schema_info = common_schema_info();
-                schema_info.resolve_refs(parameters);
+                const auto & function = tool.at("function");
+                std::string  name     = function.at("name");
 
                 std::vector<common_peg_parser> required_args;
                 std::vector<common_peg_parser> optional_args;
 
-                foreach_parameter(function, [&](const std::string & param_name, const json & param_schema, bool is_required) {
-                    auto rule_name = "tool-" + name + "-arg-" + param_name;
+                foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
+                    auto rule_name = "tool-" + name + "-arg-" + param.name;
 
-                    auto arg_open = p.tool_arg_open("<parameter=" + p.tool_arg_name(p.literal(param_name)) + ">\n");
+                    auto arg_open = p.tool_arg_open("<parameter=" + p.tool_arg_name(p.literal(param.name)) + ">\n");
 
-                    auto arg_value = schema_info.resolves_to_string(param_schema) ?
-                        arg_string :
-                        p.tool_arg_json_value(p.schema(p.json(), rule_name + "-schema", param_schema)) + arg_close;
+                    auto types = param.schema->value_types();
+
+                    auto arg_value = p.eps();
+                    if (!types.has(common_chat_schema::TYPE_STRING)) {
+                        arg_value = p.tool_arg_json_value(p.schema(p.json(), rule_name + "-schema", doc, *param.schema)) + arg_close;
+                    } else if (types.is_only(common_chat_schema::TYPE_STRING)) {
+                        arg_value = arg_string;
+                    } else {
+                        // The string alternative accepts any text, so the grammar only keeps the raw string
+                        // rule. The parser still tries the JSON alternatives first to type the value.
+                        auto json_value = p.choice();
+                        if (types.has(common_chat_schema::TYPE_OBJECT)) {
+                            json_value |= p.json_object();
+                        }
+                        if (types.has(common_chat_schema::TYPE_ARRAY)) {
+                            json_value |= p.json_array();
+                        }
+                        if (types.has(common_chat_schema::TYPE_NUMBER) || types.has(common_chat_schema::TYPE_INTEGER)) {
+                            json_value |= p.json_number();
+                        }
+                        if (types.has(common_chat_schema::TYPE_BOOLEAN)) {
+                            json_value |= p.json_bool();
+                        }
+                        if (types.has(common_chat_schema::TYPE_NULL)) {
+                            json_value |= p.json_null();
+                        }
+                        arg_value = p.gbnf(p.atomic(p.tool_arg_json_value(json_value) + arg_close) | arg_string, "xml-arg-string");
+                    }
 
                     auto arg_rule = p.rule(rule_name, p.tool_arg(arg_open + arg_value));
 
-                    (is_required ? required_args : optional_args).push_back(arg_rule);
+                    (param.required ? required_args : optional_args).push_back(arg_rule);
                 });
 
                 // Accept required arguments in any order, as Qwen does not always adhere to the
@@ -158,15 +179,6 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
         data.grammar_lazy = has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
 
         data.grammar = build_grammar([&](const common_grammar_builder & builder) {
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                auto         schema   = function.contains("parameters") ? function.at("parameters") : json::object();
-                builder.resolve_refs(schema);
-            });
-            if (has_response_format) {
-                auto schema = inputs.json_schema;
-                builder.resolve_refs(schema);
-            }
             parser.build_grammar(builder, data.grammar_lazy);
         });
 

common/peg-parser.cpp

@@ -953,7 +953,7 @@ std::string common_peg_arena::dump_impl(common_peg_parser_id
         } else if constexpr (std::is_same_v<T, common_peg_until_parser>) {
             return "Until(" + string_join(p.delimiters, " | ") + ")";
         } else if constexpr (std::is_same_v<T, common_peg_schema_parser>) {
-            return "Schema(" + dump_impl(p.child, visited) + ", " + (p.schema ? p.schema->dump() : "null") + ")";
+            return "Schema(" + dump_impl(p.child, visited) + ", " + (p.node ? common_chat_schema::kind_name(p.node->kind()) : "null") + ")";
         } else if constexpr (std::is_same_v<T, common_peg_rule_parser>) {
             return "Rule(" + p.name + ", " + dump_impl(p.child, visited) + ")";
         } else if constexpr (std::is_same_v<T, common_peg_ref_parser>) {
@@ -1119,8 +1119,13 @@ 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, common_chat_schema_document_ptr doc, const common_chat_schema & node, bool raw) {
+    return wrap(arena_.add_parser(common_peg_schema_parser{p.id(), name, std::move(doc), &node, 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}));
+    auto doc = std::make_shared<const common_chat_schema_document>(common_chat_schema_from_json(schema));
+    return this->schema(p, name, doc, *doc->root, raw);
 }
 
 common_peg_parser common_peg_parser_builder::rule(const std::string & name, const common_peg_parser & p, bool trigger) {
@@ -1573,30 +1578,9 @@ static std::set<std::string> collect_reachable_rules(
 
 // GBNF generation implementation
 void common_peg_arena::build_grammar(const common_grammar_builder & builder, bool lazy) const {
+    // A raw string value is parsed by the child rather than constrained by the schema
     auto schema_delegates = [](const common_peg_schema_parser & s) -> bool {
-        if (!s.schema) {
-            return true;
-        }
-        if (s.raw && s.schema->contains("type")) {
-            const auto & type_val = s.schema->at("type");
-            if (type_val.is_string() && type_val == "string") {
-                return true;
-            }
-            // Handle nullable types like ["string", "null"] - delegate when the
-            // non-null type is string, since the tagged format uses raw text
-            if (type_val.is_array()) {
-                for (const auto & t : type_val) {
-                    if (t.is_string() && t.get<std::string>() != "null") {
-                        return t.get<std::string>() == "string";
-                    }
-                }
-            }
-        }
-        // Delegate for enum schemas in raw mode - enum values are literal strings
-        if (s.raw && !s.schema->contains("type") && s.schema->contains("enum")) {
-            return true;
-        }
-        return false;
+        return !s.node || (s.raw && s.node->may_be_string());
     };
 
     // Unwrap the parser so we can properly check if it's a sequence or choice
@@ -1731,7 +1715,7 @@ void common_peg_arena::build_grammar(const common_grammar_builder & builder, boo
                 if (schema_delegates(p)) {
                     return to_gbnf(p.child);
                 }
-                return builder.add_schema(p.name, *p.schema);
+                return builder.add_schema(p.name, *p.node);
             } else if constexpr (std::is_same_v<T, common_peg_rule_parser>) {
                 return p.name;
             } else if constexpr (std::is_same_v<T, common_peg_ref_parser>) {
@@ -1859,7 +1843,6 @@ static common_json serialize_parser_variant(const common_peg_parser_variant & va
                 {"type", "schema"},
                 {"child", p.child},
                 {"name", p.name},
-                {"schema", p.schema ? *p.schema : json(nullptr)},
                 {"raw", p.raw}
             };
         } else if constexpr (std::is_same_v<T, common_peg_rule_parser>) {
@@ -1999,15 +1982,12 @@ static common_peg_parser_variant deserialize_parser_variant(const common_json &
         return common_peg_until_parser{j["delimiters"].get<std::vector<std::string>>()};
     }
     if (type == "schema") {
-        if (!j.contains("child") || !j.contains("name") || !j.contains("schema") || !j.contains("raw")) {
+        if (!j.contains("child") || !j.contains("name") || !j.contains("raw")) {
             throw std::runtime_error("schema parser missing required fields");
         }
         common_peg_schema_parser parser;
         parser.child = j["child"].get<common_peg_parser_id>();
         parser.name = j["name"];
-        if (!j["schema"].is_null()) {
-            parser.schema = std::make_shared<common_json>(j["schema"]);
-        }
         parser.raw = j["raw"].get<bool>();
         return parser;
     }

common/peg-parser.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "json-schema.h"
 #include "json.h"
 
 #include <memory>
@@ -245,7 +246,8 @@ struct common_peg_until_parser {
 struct common_peg_schema_parser {
     common_peg_parser_id child;
     std::string name;
-    std::shared_ptr<common_json> schema;
+    common_chat_schema_document_ptr doc;  // owns node
+    const common_chat_schema * node = nullptr;
 
     // Indicates if the GBNF should accept a raw string that matches the schema.
     bool raw;
@@ -488,8 +490,10 @@ class common_peg_parser_builder {
     // A marker, i.e. text delimited by a pair of <> or []
     common_peg_parser marker();
 
-    // Wraps a parser with JSON schema metadata for grammar generation.
-    // Used internally to convert JSON schemas to GBNF grammar rules.
+    // Wraps a parser with the schema its GBNF is generated from, a node of the document that owns it
+    common_peg_parser schema(const common_peg_parser & p, const std::string & name, common_chat_schema_document_ptr doc, const common_chat_schema & node, bool raw = false);
+
+    // Parses the JSON schema into a document of its own
     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.

common/speculative.cpp

@@ -1493,7 +1493,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
 
         const int32_t n_tokens = batch_in.n_tokens;
 
-        // remember the frist and last batch index for each sequence
+        // remember the first and last batch index for each sequence
         std::fill(i_batch_beg.begin(), i_batch_beg.end(), -1);
         std::fill(i_batch_end.begin(), i_batch_end.end(), -1);
 

conversion/__init__.py

@@ -168,6 +168,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "Mamba2ForCausalLM": "mamba",
     "MambaForCausalLM": "mamba",
     "MambaLMHeadModel": "mamba",
+    "MapleForCausalLM": "maple",
     "MellumForCausalLM": "mellum",
     "MiMoV2FlashForCausalLM": "mimo",
     "MiMoV2ForCausalLM": "mimo",

conversion/maple.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+from typing import Iterable, TYPE_CHECKING, cast
+
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import LazyTorchTensor, ModelBase, TextModel, gguf
+
+
+@ModelBase.register("MapleForCausalLM")
+@ModelBase.example("deepgrove/maple-preview")
+class MapleModel(TextModel):
+    model_arch = gguf.MODEL_ARCH.MAPLE
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        hparams = self.hparams
+
+        assert hparams["hidden_act"] == "silu"
+        assert hparams.get("num_shared_experts", 0) == 0
+        assert hparams.get("norm_topk_prob", True)
+        assert hparams.get("nope_on_global_attention", False)
+
+        head_dim = hparams.get("head_dim", hparams["hidden_size"] // hparams["num_attention_heads"])
+        partial_rotary_factor = self.rope_parameters.get("partial_rotary_factor", 1.0)
+
+        self.gguf_writer.add_vocab_size(hparams["vocab_size"])
+        self.gguf_writer.add_rope_dimension_count(int(head_dim * partial_rotary_factor))
+        self.gguf_writer.add_sliding_window(hparams["sliding_window"])
+        self.gguf_writer.add_sliding_window_pattern([layer_type == "sliding_attention" for layer_type in hparams["layer_types"]])
+        self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"])
+        # the reference clamps the MoE SwiGLU gate/up at 7.0 (modeling_maple.py)
+        self.gguf_writer.add_swiglu_clamp_exp([7.0] * self.block_count)
+
+    _experts: list[dict[str, Tensor]] | None = None
+
+    @staticmethod
+    def _stack_experts(tensors: list[Tensor]) -> Tensor:
+        shape = (len(tensors), *tensors[0].shape)
+        dtype = tensors[0].dtype
+        meta = LazyTorchTensor.meta_with_dtype_and_shape(dtype, shape)
+
+        # tensors goes through args, not the closure, so that `func` matches
+        # LazyBase's single-argument shape
+        def stack(ts: list[Tensor]) -> Tensor:
+            result = torch.empty(shape, dtype=dtype)
+            for expert_id, tensor in enumerate(ts):
+                result[expert_id].copy_(LazyTorchTensor.to_eager(tensor))
+            ts.clear()
+            return result
+
+        return cast(torch.Tensor, LazyTorchTensor(meta=meta, args=(tensors,), func=stack))
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if "mlp.experts" in name:
+            n_experts = self.hparams["num_experts"]
+            assert bid is not None
+
+            if self._experts is None:
+                self._experts = [{} for _ in range(self.block_count)]
+
+            self._experts[bid][name] = data_torch
+
+            if len(self._experts[bid]) >= n_experts * 3:
+                for weight_name in ("down_proj", "gate_proj", "up_proj"):
+                    tensors = []
+
+                    for expert_id in range(n_experts):
+                        expert_name = f"model.layers.{bid}.mlp.experts.{expert_id}.{weight_name}.weight"
+                        tensors.append(self._experts[bid].pop(expert_name))
+
+                    merged_name = f"model.layers.{bid}.mlp.experts.{weight_name}.weight"
+                    yield from super().modify_tensors(self._stack_experts(tensors), merged_name, bid)
+            return
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+    def prepare_tensors(self):
+        super().prepare_tensors()
+
+        if self._experts is not None:
+            experts = [name for layer in self._experts for name in layer]
+            if experts:
+                raise ValueError(f"Unprocessed experts: {experts}")

docs/backend/SYCL.md

@@ -790,14 +790,15 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 
 | Name              | Value            | Function                                                                                                                  |
 |-------------------|------------------|---------------------------------------------------------------------------------------------------------------------------|
-| GGML_SYCL_DEBUG   | 0 (default) or 1 | Enable log function by macro: GGML_SYCL_DEBUG                                                                             |
+| GGML_SYCL_DEBUG   | 0 (default) or 1 | Enable log function: GGML_SYCL_DEBUG() for common debug. |
+| GGML_SYCL_DEV_DEBUG   | 0 (default) or 1 | Enable log function: GGML_SYCL_DEV_DEBUG() for developmental purposes by replacing GGML_SYCL_DEBUG() in special codes. Restore to GGML_SYCL_DEBUG() before committing code.|
 | GGML_SYCL_DEV2DEV_MEMCPY | 0 (default), 1, 2 | Choose the method of dev2dev memory copy.<br>Value: <br>*  0: SYCL API (default), only support dGPUs.<br>* 1: L0 API -- Better performance, only support dGPUs, found to lead to abnormal crash in some case. <br>* 2: Host Forward -- Most stable method for all cases (including iGPU + dGPU*N), but with lower performance (-2% to -5%).<br>SYCL & L0 API are easy to be impacted by Intel GPU driver issue. When you meet the garbled output or crash issues in multiple GPUs case, try with this debug flag to work around or check the issue.|
 | GGML_SYCL_ENABLE_FLASH_ATTN | 1 (default) or 0| Enable Flash-Attention. It can reduce memory usage. The performance impact depends on the LLM.|
 | GGML_SYCL_ENABLE_OPT | 0 or 1 (default)| Enable optimize features for Intel GPUs. (Recommended to 0 for Intel devices older than Gen 10) |
 | GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still on development, no better performance. |
 | GGML_SYCL_ENABLE_HOST_PINNED_MEM | 0 or 1 (default) | Enable host pinned memory to speed up copy data from host to device. When disable it, host memory will common malloc() on CPU. Disable it when use `--load-model mlock`.|
 | GGML_SYCL_HOST_PINNED_MEM_2G | 0 (default) or 1 | Limit the max memory allocation to be no more than 2GB when enable host pinned memory. USM allocations above 2 GiB take the relaxed/large-allocation path, which serializes H2D copies with compute and prevents copy/compute overlap. It will impact the startup time. Need more test. Depend on `GGML_SYCL_ENABLE_HOST_PINNED_MEM=1`.|
-| GGML_SYCL_GET_MEM_API | 0 (default) or 1  | Set to get memory info (free, total) by Level Zero or SYCL API:<br>0 - Level Zero API: support more GPUs, only run on Level Zero running time. When there is an error, fallback to call SYCL API. Depend on GGML_SYCL_SUPPORT_LEVEL_ZERO_API.<br>1 - SYCL API: legacy, support more running time, it can't get the free size of some GPUs (like Arc770). In such case, return total size for free size.|
+| GGML_SYCL_GET_MEM_API | 0 (default) or 1  | Set to get memory info (free, total) by Level Zero or SYCL API:<br>0 - Level Zero API: support more GPUs, only run on Level Zero running time. When there is an error, fallback to call SYCL API. Depend on GGML_SYCL_SUPPORT_LEVEL_ZERO_API.<br>1 - SYCL API: legacy, support more running time, it can't get the free size of some GPUs (like Arc770). In such case, return the free size as value of total size.|
 | GGML_SYCL_USE_LEVEL_ZERO_API | 1 (default) or 0 | Use Level Zero API for device memory allocation instead of SYCL. Reduces system RAM usage on Intel dGPUs by avoiding DMA-buf/TTM host memory staging. Requires GGML_SYCL_SUPPORT_LEVEL_ZERO_API=ON at build time. SYCL backend always runs on Level Zero running time even if it's set as OFF (The SYCL api will be usage for memory allocation).|
 | GGML_SYCL_ENABLE_DNN | 0 or 1 (default)| Enable running computations through oneDNN and always use oneMKL. |
 | GGML_SYCL_FA_ONEDNN | 1 (default) or 0 | Enable the oneDNN fused SDPA (flash-attention) path on supported GPUs. Set to 0 to always use the native SYCL flash-attention kernel. |

docs/backend/snapdragon/README.md

@@ -188,7 +188,7 @@ llama_memory_breakdown_print: |   - Host               |                  439 =
 Op test for MUL_MAT:
 
 ```
-~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --hex-hostbuf 0 --devices HTP0:0 -- test-backend-ops -b HTP0:0 -o MUL_MAT
+~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --devices HTP0:0 -- test-backend-ops -b HTP0:0 -o MUL_MAT
 ...
 Backend 2/3: HTP0:0
 Device description: Hexagon
@@ -213,14 +213,109 @@ ggml-hex: new session: HTP0 : session-id 0 domain-id 3 uri file:///libggml-htp-v
 | llama 1B Q4_0  | 729.75 MiB | 1.24 B | HTP        |  99 |       4 |     128 |    0 |  tg64 |  51.54 ± 1.13 |
 ```
 
+## Multi-Device Execution Modes
+
+The Hexagon backend supports multiple execution and partitioning modes to accommodate different model sizes, memory
+constraints, and single- or multi-NPU hardware topologies:
+
+### 1. Single-Device Mode with Dynamic Buffer Mapping
+
+Runs the model on a single NPU session (e.g. `HTP0` or `HTP0:0`).
+
+A single NPU session provides ~3.5GB of available virtual address space. For models larger than 3.5GB, the backend
+automatically maps and unmaps weight buffers during graph execution. This allows large models to run on a single NPU
+without manual configuration:
+
+```bash
+./scripts/snapdragon/run.py --target adb --devices HTP0:0 -- \
+    llama-cli -m models/Llama-3.2-3B-Instruct-Q4_0.gguf -ngl 99 -p "Hello"
+```
+
+### 2. Layer-Split Mode across Virtual Sessions (`HTP0,HTP1,...` or `HTP0:0,HTP0:1,...`)
+
+Partitions model layers at load time across multiple virtual sessions hosted on a single physical NPU.
+
+Each virtual session acts as an independent backend device from llama.cpp's perspective (similar to multiple GPUs).
+Because layers are permanently distributed across sessions, each session's allocated weights remain within its private 3.5GB
+address space window, eliminating runtime buffer re-mapping overhead.
+
+Here is an example of running the GPT-OSS-20B model on a Snapdragon device using 4 virtual sessions on a single NPU:
+
+```bash
+./scripts/snapdragon/run.py --target adb \
+    --devices HTP0:0,HTP0:1,HTP0:2,HTP0:3 -- \
+    llama-cli --load-mode none -m /data/local/tmp/gguf/gpt-oss-20b-Q4_0.gguf -t 4 \
+    --ctx-size 8192 --batch-size 128 -ctk q8_0 -ctv q8_0 -fa on -ngl 99 -no-cnv -f surfing.txt
+```
+
+Log output snippet:
+
+```
+...
+llama_model_loader: - type  f32:  289 tensors
+llama_model_loader: - type q4_0:   96 tensors
+llama_model_loader: - type q8_0:    2 tensors
+llama_model_loader: - type mxfp4:  72 tensors
+...
+load_tensors: offloaded 25/25 layers to GPU
+load_tensors:          CPU model buffer size =  1182.09 MiB
+load_tensors:       HTP0:1 model buffer size =  2512.58 MiB
+load_tensors:       HTP0:3 model buffer size =  2093.83 MiB
+load_tensors:       HTP0:0 model buffer size =  2931.34 MiB
+load_tensors:       HTP0:2 model buffer size =  2512.58 MiB
+...
+llama_perf_context_print: prompt eval time =    3843.67 ms /   197 tokens ( 19.51 ms per token, 51.25 tokens per second)
+llama_perf_context_print:        eval time =    1686.13 ms /    31 runs   ( 54.39 ms per token, 18.39 tokens per second)
+llama_perf_context_print:       total time =    6266.30 ms /   228 tokens
+llama_memory_breakdown_print: | memory breakdown [MiB] | total   free    self   model   context   compute    unaccounted |
+llama_memory_breakdown_print: |   - HTP0:0 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - HTP0:1 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - HTP0:2 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - HTP0:3 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - Host               |                 1476 =  1208 +     105 +     162                |
+```
+
+### 3. Tensor-Split Mode across Physical Devices (`HTP0:0,HTP1:0,...`)
+
+Distributes model tensors across distinct physical NPU hardware cores using llama.cpp's tensor parallelism
+(`--split-mode tensor`).
+
+Tensors are partitioned across physical NPUs for parallel execution (proportions are distributed equally by default without
+needing an explicit `--tensor-split` option):
+
+```bash
+./scripts/snapdragon/run.py --target adb \
+    --devices HTP0:0,HTP1:0 -- \
+    llama-cli -m models/Llama-3.2-3B-Instruct-Q4_0.gguf --split-mode tensor -ngl 99 -p "Hello"
+```
+
+### 4. Row-Split Multi-Device Mode via Device Grouping (`HTP0[0-1]`)
+
+Groups multiple physical NPU cores into a single logical device using bracket notation (`HTP0[0-1]` or `HTP0[0,1]`).
+
+Unlike host-level tensor-splitting, row-splitting is executed entirely inside the Hexagon backend:
+
+```bash
+./scripts/snapdragon/run.py --target adb \
+    --devices 'HTP0[0-1]' -- \
+    llama-cli -m models/Llama-3.2-3B-Instruct-Q4_0.gguf -ngl 99 -p "Hello"
+```
+
+You can also combine row-splitting with layer-splitting across multiple grouped devices (e.g. `--devices 'HTP0[0-1],HTP1[2-3]'`
+on 4 physical NPUs, or `--devices 'HTP0[0-1:0],HTP1[0-1:1]'` on 2 physical NPUs using virtual sessions 0 and 1).
+
 ## Environment variables
 
 - `GGML_HEXAGON_DEVICES` (default: not set, defaults to HTP0 session)
-  Controls which NPU devices and sessions to allocate. Can be configured as:
-  - A single integer `N`: Allocates `N` sessions named `HTP0`, `HTP1`, ..., `HTP<N-1>` (behaves identically to `GGML_HEXAGON_NDEV=N`).
-  - A comma-separated list of device names in `HTP<physical_idx>:<virtual_idx>` format (or legacy `HTP<idx>` format). For example, `HTP0:0,HTP0:1` creates two virtual
-    sessions on the first physical NPU (useful for memory limits). `HTP0:0,HTP1:0` allocates one session on each of the two physical NPUs
-    on a dual-NPU device.
+  Controls which NPU devices and sessions to allocate. Configurable via `--devices` in `run.py`:
+  - `N` (single integer): Allocates `N` virtual sessions named `HTP0`, `HTP1`, ..., `HTP<N-1>` on physical NPU 0.
+  - `HTP<phys>:<virt>,...`: Comma-separated list of individual devices specifying physical and virtual index:
+    - `HTP0:0,HTP0:1`: Two virtual sessions on physical NPU 0 (layer-split on single NPU).
+    - `HTP0:0,HTP1:0`: One session on physical NPU 0 and one on physical NPU 1 (tensor-split across physical cores).
+  - `HTP<name>[<phys_spec>]`: Device grouping syntax for row-split multi-device execution:
+    - `HTP0[0-1]`: A single logical device `HTP0` that groups physical cores 0 and 1.
+    - `HTP0[0-1],HTP1[2-3]`: Two layer-split devices across 4 physical NPUs (cores 0-1 and 2-3).
+    - `HTP0[0-1:0],HTP1[0-1:1]`: Two layer-split devices across 2 physical NPUs using virtual sessions 0 and 1.
 
 - `GGML_HEXAGON_NDEV` (deprecated)
   Replaced by `GGML_HEXAGON_DEVICES`. Controls the number of virtual sessions to allocate on physical NPU `0`.
@@ -229,9 +324,8 @@ ggml-hex: new session: HTP0 : session-id 0 domain-id 3 uri file:///libggml-htp-v
 - `GGML_HEXAGON_NHVX=0`
   Controls the number of HVX hardware threads to use. The default is all (actual number varies depending on the hardware version).
 
-- `GGML_HEXAGON_HOSTBUF=1`
-  Controls whether the Hexagon backend allocates host buffers. By default, all buffers except for REPACK are host buffers.
-  This option is required for testing Ops that require REPACK buffers (MUL_MAT and MUL_MAT_ID).
+- `GGML_HEXAGON_HOSTBUF=1` (default: 0, disabled)
+  Enables allocating host buffers for debugging. By default, host buffers are disabled.
 
 - `GGML_HEXAGON_VERBOSE=1`
   Enables verbose logging of Ops from the backend. Example output:
@@ -246,23 +340,26 @@ ggml-hex: new session: HTP0 : session-id 0 domain-id 3 uri file:///libggml-htp-v
   ```
 
 - `GGML_HEXAGON_PROFILE=1`
-  Enables Op profiling:
+  Enables Op profiling (configurable via `--hex-profile` in `run.py`):
 
-  - `1` Basic profile with per-op `usecs` and `cycles` counters
-  - `2` Extended profile with per-op `usecs`, `cycles` and default PMU counter data
-  - `0x1,...,0x8` Extended profile with per-op `usecs`, `cycles` and custom PMU counter data
+  - `1`: Basic profile with per-op `usecs` and `cycles` counters
+  - `2`: Extended profile with per-op `usecs`, `cycles` and default PMU counter data
+  - `0x1,...,0x8`: Extended profile with per-op `usecs`, `cycles` and custom PMU counter data
 
-  The logging output can be either saved into a file for post-processing or it can be piped directly into the post-processing tool
-  to generate the report.
-  Examples:
+  The logging output can be saved to a file or piped directly into the post-processing script:
 
-      `GGML_HEXAGON_PROFILE=1 ./scripts/snapdragon/run.py --target adb -- llama-cli ... |& ./scripts/snapdragon/ggml-hexagon-profile.py -`
+  ```bash
+  ./scripts/snapdragon/run.py --target adb --hex-profile 1 -- llama-cli ... |& \
+      ./scripts/snapdragon/ggml-hexagon-profile.py -
+  ```
 
 - `GGML_HEXAGON_OPFILTER=regex`
-  Allows filtering (disabling) Ops that match the regex pattern:
+  Filters (disables) Ops matching the regex pattern (configurable via `--hex-opfilter` in `run.py`):
 
-  Examples:
-
-      `GGML_HEXAGON_OPFILTER="FLASH_ATTN_EXT" ./scripts/snapdragon/run.py --target adb -- llama-cli ...` - Disable Flash Attention on Hexagon (falls back to CPU or GPU)
-      `GGML_HEXAGON_OPFILTER="ADD\|SUB" ./scripts/snapdragon/run.py --target adb -- llama-cli ...` - Disable ADD and SUB on Hexagon (fall back to CPU or GPU)
+  ```bash
+  # Disable Flash Attention on Hexagon (falls back to CPU or GPU)
+  ./scripts/snapdragon/run.py --target adb --hex-opfilter "FLASH_ATTN_EXT" -- llama-cli ...
 
+  # Disable ADD and SUB on Hexagon (fall back to CPU or GPU)
+  ./scripts/snapdragon/run.py --target adb --hex-opfilter "ADD|SUB" -- llama-cli ...
+  ```

docs/backend/snapdragon/developer.md

@@ -2,16 +2,16 @@
 
 ## Backend libraries
 
-The Hexagon backend consist of two parts:
+The Hexagon backend consists of two parts:
 
   - `libggml-hexagon`
-    This is the regular CPU-side GGML backend library, either shared or statically linked
+    This is the regular CPU-side GGML backend library, either shared or statically linked.
 
   - `libggml-htp-vNN`
     This is the NPU-side (HTP stands for Hexagon Tensor Processor) shared library that contains the Op dispatcher and kernels.
     The correct library is selected automatically at runtime based on the HW version.
 
-Here is an example of the build artifacts
+Here is an example of the build artifacts:
 
 ```
 ~/src/llama.cpp$ ls -l pkg-adb/llama.cpp/lib/libggml*
@@ -26,75 +26,307 @@ pkg-adb/llama.cpp/lib/libggml-htp-v81.so
 
 ## Memory buffers
 
-Hexagon NPU backend takes advantage of the Snapdragon's unified memory model where all buffers are fully accessible by the CPU and GPU.
-The NPU does have a dedicated tightly-coupled memory called VTCM but that memory is used only for intermediate data (e.g. dynamically
-quantized tensors) or temporary data (chunks of the weight tensors fetched via DMA).
-
-Please note that currently the Hexagon backend does not implement SET/GET_ROWS Ops because there is no advantage in offloading those
-to the NPU at this point.
-
-The backend does allocates non-host buffers for the tensors with datatypes that require repacking: Q4_0, Q8_0, MXFP4.
-From the MMU perspective these buffers are still regular buffers (normal access by the CPU) they are marked as non-host simply to force
-the repacking.
+The Hexagon NPU backend takes advantage of Snapdragon unified memory where all DDR buffers are accessible by CPU, GPU, and NPU.
+The NPU has dedicated tightly-coupled memory called VTCM (Vector Tightly-Coupled Memory). VTCM is used for intermediate data (such as
+dynamically quantized activations) and streaming buffers (chunks of weight and activation tensors fetched via DMA).
 
 ## Large model handling
 
-Hexagon NPU sessions (aka Process Domains (PD) in the Hexagon SDK) are limited to a maximum memory mapping window of around 3.5GB.
+Hexagon NPU sessions have a 32-bit virtual address space window of around 3.5GB.
 In llama.cpp/GGML, each Hexagon session is mapped to a single GGML backend device (e.g., `HTP0:0`, `HTP0:1`, etc. when using
 `GGML_HEXAGON_DEVICES`, or `HTP0`, `HTP1` in legacy mode).
 
-To support running models larger than 3.5GB on a single device, the Hexagon backend dynamically maps and unmaps execution buffers
-during the graph execution cycle to stay within the Process Domain window. This enables large models to run successfully on a single
-NPU device.
+To support running models larger than 3.5GB on a single device, the Hexagon backend dynamically maps and unmaps buffers:
+- Buffers are allocated in shared DDR (RPCMEM) via file descriptors (`fastrpc_mmap` using `FASTRPC_MAP_FD_DELAYED`).
+- Pinned buffers (such as KV cache and active compute buffers) remain mapped throughout execution.
+- Inactive weight buffers are dynamically mapped into the NPU session via `HAP_mmap()` during batch buffer preparation
+  (`prep_op_bufs()` in `htp/main.c`) and unmapped via `htp_iface_munmap()` when no longer needed by the active batch.
+- This dynamic sliding window allows a single NPU session to execute models that exceed the 3.5GB window.
 
-Alternatively, users can choose to use standard llama.cpp/GGML layer-splitting mode to partition and split the model across
-multiple Hexagon devices or virtual sessions (which behave like multiple GPUs from the offload and splitting perspective).
+Alternatively, users can partition and split the model across multiple virtual sessions or physical NPUs using layer-splitting,
+tensor-splitting, or row-splitting modes. For user-facing execution modes and examples, see the
+[Snapdragon user guide](README.md#multi-device-execution-modes).
 
-Here is an example of running GPT-OSS-20B model on a Snapdragon device using 4 virtual sessions on a single NPU (physical index 0).
+## Op and Kernel Development Guidelines
+
+Writing high-performance operators for Hexagon requires following specific guidelines.
+
+### DDR -> DMA -> VTCM Execution Pipeline
+
+- Strongly prefer the `DDR -> DMA -> VTCM -> compute (HVX/HMX) -> VTCM -> DMA -> DDR` data flow.
+- Direct HVX reads/writes from/to DDR are less efficient and should only be used as a fallback.
+- The DMA queue is a strict FIFO where operations must be pushed and popped in strict order.
+- Follow the pipelined multi-buffering sequence properly (typically 2x to 16x buffering) so every push has a corresponding pop:
+
+  1. In the prologue, push initial DDR -> VTCM transfers to prime the pipeline.
+  2. In the loop body, wait for buffer N via DMA pop, launch HVX/HMX compute on buffer N, push VTCM -> DDR writeback of result N,
+     and push DDR -> VTCM prefetch of buffer N+2.
+  3. In the epilogue, pop all remaining in-flight transfers to drain the pipeline.
+
+- Because every push must be matched by a pop, `dma_queue_flush()` is not required when the pipeline sequence is followed
+  properly. Flushing is only used in rare exceptions where a batch of operations is pushed without individual pops.
+- Use the DMA queue interface from [`dma-queue.h`](../../../ggml/src/ggml-hexagon/htp/dma-queue.h)
+  (`dma_queue_push_ddr_to_vtcm`, `dma_queue_pop`, `dma_queue_push_vtcm_to_ddr`).
+  See [`cumsum-ops.c`](../../../ggml/src/ggml-hexagon/htp/cumsum-ops.c) and
+  [`act-ops.c`](../../../ggml/src/ggml-hexagon/htp/act-ops.c) for reference implementations.
+
+### Avoid Scalar Reads and Writes to VTCM
+
+- Access VTCM data using DMA transfers or HVX/HMX vector instructions rather than scalar reads and writes.
+
+### Avoid Scalar Division in Inner Loops
+
+- Hexagon cores do not have hardware division instructions.
+- For recurring divisions across iterations or threads, use `fastdiv` from
+  [`hex-fastdiv.h`](../../../ggml/src/ggml-hexagon/htp/hex-fastdiv.h) with precomputed divisors (such as
+  `octx->ctx->mdev.count_div` or `octx->n_threads_div`).
+- Do not call `init_fastdiv_values()` for single-use divisions; use standard compiler division (`/`) instead.
+
+### Host-Side Precomputation via `kernel_params`
+
+- Precompute tensor shapes, strides, scale conversions, tiling layouts, and validation checks on the host CPU during graph
+  preparation in [`ggml-hexagon.cpp`](../../../ggml/src/ggml-hexagon/ggml-hexagon.cpp).
+- Pack precomputed parameters into the operator's fixed `kernel_params` structure in `htp_op_node` (such as
+  `htp_mm_kernel_params`, `htp_unary_kernel_params`, `htp_fa_kernel_params`, `htp_get_rows_kernel_params`).
+- The NPU executes directly using `octx->kernel_params` without redundant runtime metadata extraction or validation.
+- **Strict Host-Kernel Alignment**:
+  - Verify that parameters calculated by the host CPU are strictly honored by the NPU kernel.
+  - Ensure the kernel does not ignore host-computed fields (for example, falling back to `octx->n_threads` instead of
+    using `kparams->n_threads`, or ignoring precomputed `tasks_per_thread` and chunk counts).
+  - Both human developers and coding agents must audit both sides of the interface: ensure fields populated in `kernel_params`
+    in [`ggml-hexagon.cpp`](../../../ggml/src/ggml-hexagon/ggml-hexagon.cpp) are actively and consistently utilized by the
+    corresponding operator entry point and worker threads in `htp/*-ops.c`.
+
+### Tracing Instrumentation
+
+- All kernels must include trace events for performance profiling and timeline visualization in Perfetto
+  ([`hex-profile.h`](../../../ggml/src/ggml-hexagon/htp/hex-profile.h)).
+- Surround compute sections with `htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) info)` and
+  `htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) info)`.
+- Use specific event types for major phases:
+  - `HTP_TRACE_EVT_HVX_COMP`: Vector compute execution.
+  - `HTP_TRACE_EVT_DMA`: DMA transfer wait or poll cycles.
+  - `HTP_TRACE_EVT_FENCE`: Multi-device fence barrier synchronization.
+  - `HTP_TRACE_EVT_L2FLUSH`: L2 cache cleaning operations.
+- Pass meaningful progress metrics (such as row index, chunk index, or token index) in the 16-bit `info` parameter.
+
+### Work Queue and Threading
+
+- Distribute parallel work across NPU worker threads using the thread pool work queue:
+
+  ```c
+  work_queue_run(ctx->work_queue, worker_func, &op_ctx, n_threads);
+  ```
+
+- Keep worker functions independent and re-entrant. Worker threads should only operate on their designated chunk of rows or elements.
+
+### Avoid Redundant Defensive NULL Checks
+
+- Do not add defensive NULL checks or assertions for internal framework pointers or required graph operands and outputs.
+  Internal pointers include `ctx`, `octx`, local context structs like `*ctx`, `kparams`, and worker callback `data`.
+- These pointers are architectural invariants during kernel execution and host-side graph preparation.
+  Graph compute receives allocated nodes with valid required `node->src[N]` and `node->data` pointers.
+- Do not turn an invariant violation into an unsupported operation or missed fusion.
+  Checks such as `if (!octx || !octx->ctx)` clutter the code, obscure intent, and hide upstream errors.
+- **Distinction**: `octx->src[N]` pointers *can* be NULL by design and must be checked when optional.
+  Examples include attention masks, optional bias or weights in fused kernels, and frequency factors.
+
+### Multiline Macro Formatting
+
+- Keep trailing backslashes in multiline `#define` macros cleanly aligned to a consistent column.
+- Avoid trailing whitespace after macro backslashes.
+- Use [`scripts/snapdragon/ggml-hexagon-align-macros.py`](../../../scripts/snapdragon/ggml-hexagon-align-macros.py) to inspect, diff,
+  or automatically align macro definitions across Hexagon kernel sources:
+
+  ```bash
+  # Check for misaligned macros
+  python3 scripts/snapdragon/ggml-hexagon-align-macros.py ggml/src/ggml-hexagon/htp/
+
+  # Fix misaligned macros in-place
+  python3 scripts/snapdragon/ggml-hexagon-align-macros.py --fix ggml/src/ggml-hexagon/htp/
+  ```
+
+## Multi-Device Partitioning (mdev)
+
+Multi-device (mdev) mode enables row-level tensor parallel execution across multiple physical NPU cores or virtual NPU
+sessions.
+
+### 128-Byte Cache Line Alignment
+
+- Shared tensor buffers reside in DDR (RPCMEM) with a 128-byte cache line granularity
+  (`HEX_L2_LINE_SIZE` = 128 bytes, `HTP_TENSOR_MDEV_LINE_SIZE`).
+- **Rule**: Multi-device work partitions must align destination write regions to 128-byte cache line boundaries so distinct
+  devices never share or overwrite the same cache line.
+
+### Partitioning Helpers in `htp-tensor.h`
+
+Common partitioning logic is factored into reusable inline helpers in
+[`htp-tensor.h`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h):
+
+1. [`htp_tensor_mdev_rows_per_chunk`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L67):
+   Determines the minimum number of rows per chunk so that the chunk byte size is a multiple of 128 bytes:
+
+   ```
+   rows_per_chunk = 128 / hex_gcd_u32(row_size, 128)
+   ```
+
+   If row stride `nb[1]` is already a multiple of 128 bytes, `rows_per_chunk = 1`.
+   Returns `false` if the tensor cannot be safely row-partitioned (such as unaligned base pointer, permuted layout,
+   or non-128-byte aligned outer strides).
+
+2. [`htp_tensor_mdev_partition`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L94):
+   Calculates the per-device work range `struct htp_tensor_mdev_range { uint32_t start; uint32_t count; }` given
+   `total_units`, `units_per_chunk`, `mdev_idx`, `mdev_count`, and the precomputed `mdev_count_div`.
+   Handles chunk distribution across devices, assigns remainder units to the last device, and automatically triggers
+   single-device fallback when partitioning is unsafe.
+
+### Row-Partitioned Operators
+
+For row-wise operators
+(such as activations in [`act-ops.c`](../../../ggml/src/ggml-hexagon/htp/act-ops.c),
+binary ops in [`binary-ops.c`](../../../ggml/src/ggml-hexagon/htp/binary-ops.c),
+unary ops in [`unary-ops.c`](../../../ggml/src/ggml-hexagon/htp/unary-ops.c), and
+sameshape copies in [`cpy-ops.c`](../../../ggml/src/ggml-hexagon/htp/cpy-ops.c)):
+
+```c
+const uint32_t total_rows   = ne01 * ne02 * ne03;
+const size_t   dst_row_size = dst->ne[0] * elem_size;
+
+uint32_t row_start = 0;
+uint32_t nrows     = total_rows;
+
+if (octx->ctx->mdev.count > 1) {
+    uint32_t rows_per_chunk = 0;
+    htp_tensor_mdev_rows_per_chunk(dst, elem_size, (uint32_t) dst_row_size, &rows_per_chunk);
+    const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
+        total_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+    row_start = range.start;
+    nrows     = range.count;
+}
+
+if (nrows == 0) {
+    return HTP_STATUS_OK;
+}
+```
+
+### Element-Partitioned Operators
+
+For flat element-wise operations (such as reshape copies in
+[`cpy-ops.c`](../../../ggml/src/ggml-hexagon/htp/cpy-ops.c)):
+- Partition total linear elements N = ne0 * ne1 * ne2 * ne3 in 128-byte cache line chunks (`elems_per_line = (elem_size == 4) ? 32 : 64`).
+- Requires strict 1D contiguity:
+  [`htp_tensor_is_contiguous(dst, elem_size)`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L28)
+  and 128-byte aligned destination pointer
+  [`htp_tensor_mdev_data_aligned(dst)`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L47).
+- If contiguous and aligned, pass `elems_per_line` to
+  [`htp_tensor_mdev_partition`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L94);
+  otherwise pass 0 to trigger Device 0 fallback.
+
+### Single-Device Fallback (Device 0)
+
+- Fallback to Device 0 (`mdev.idx == 0`) when partitioning would cause cache line tearing or when work cannot be evenly distributed.
+- Triggers:
+  1. Destination tensor cannot be safely partitioned (`rows_per_chunk == 0` or non-contiguous/unaligned buffer).
+  2. Total aligned chunks < `mdev_count`.
+- Device 0 processes the entire tensor `[0, total_units)`.
+- Devices 1 ... N-1 receive `count = 0` and return `HTP_STATUS_OK` immediately.
+
+### Flatten Outer Dimensions Globally
+
+- **Never partition solely on `ne01` (dimension 1).**
+- Partitioning only on `ne01` repeats the device boundary across every 2D slice (`ne02`, `ne03`). If each 2D slice is small,
+  false sharing occurs repeatedly throughout the tensor.
+- Always flatten outer dimensions globally: `total_rows = ne01 * ne02 * ne03` and partition once across the combined row space.
+
+### Stateless Starting Coordinates
+
+- Do not use incremental state variables across slices that assume the thread or device starts at index 0.
+- Precompute starting multidimensional coordinates at `r = row_start` (or `e = elem_start`) once using `fastdiv`.
+- In inner loops, step base pointers directly (`ptr += stride`) or reset/wrap coordinates explicitly (`if (++i01 == ne01) { ... }`).
+
+### Clean Range Encapsulation
+
+- Initialize single-device default ranges at declaration:
+
+  ```c
+  uint32_t row_start = 0;
+  uint32_t nrows     = total_rows;
+  ```
+
+- Encapsulate all multi-device logic inside `if (octx->ctx->mdev.count > 1)`. If the block is omitted or compiled out,
+  the operator runs standard single-device execution untouched.
+- Do not propagate `mdev_` prefixes to worker functions or context structs. Worker threads are device-agnostic and
+  should only receive standard range parameters (`ctx.row_start`, `ctx.nrows`).
+- In worker threads, calculate row intervals using standard arithmetic:
+
+  ```c
+  const uint32_t ir0 = ctx->row_start + dr * ith;
+  const uint32_t ir1 = MIN(ir0 + dr, ctx->row_start + ctx->nrows);
+  ```
+
+  In single-device mode (`row_start == 0`), this naturally simplifies to `dr * ith` and `MIN(ir0 + dr, ctx->nrows)` with zero overhead.
+
+## Multi-Device Synchronization
+
+Multi-device execution synchronizes worker sessions across devices using explicit barriers and tensor cache flushing.
+
+### Synchronization Fence Protocol
+
+Multi-device execution synchronizes worker sessions through atomic fence slots and barriers defined in
+[`htp-fence.h`](../../../ggml/src/ggml-hexagon/htp/htp-fence.h):
 
 ```
-~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --devices HTP0:0,HTP0:1,HTP0:2,HTP0:3 -- llama-cli --load-mode none -m /data/local/tmp/gguf/gpt-oss-20b-Q4_0.gguf -t 4 --ctx-size 8192 --batch-size 128 -ctk q8_0 -ctv q8_0 -fa on -ngl 99 -no-cnv -f surfing.txt
-...
-llama_model_loader: - type  f32:  289 tensors
-llama_model_loader: - type q4_0:   96 tensors
-llama_model_loader: - type q8_0:    2 tensors
-llama_model_loader: - type mxfp4:  72 tensors
-...
-load_tensors: offloaded 25/25 layers to GPU
-load_tensors:          CPU model buffer size =  1182.09 MiB
-load_tensors:       HTP0:1 model buffer size =  2512.58 MiB
-load_tensors:       HTP0:3 model buffer size =  2093.83 MiB
-load_tensors:       HTP0:0 model buffer size =  2931.34 MiB
-load_tensors:       HTP0:2 model buffer size =  2512.58 MiB
-...
-llama_context: n_ctx_per_seq (8192) < n_ctx_train (131072) -- the full capacity of the model will not be utilized
-llama_context:        CPU  output buffer size =     0.77 MiB
-llama_kv_cache_iswa: creating non-SWA KV cache, size = 8192 cells
-llama_kv_cache:     HTP0:1 KV buffer size =    25.50 MiB
-llama_kv_cache:     HTP0:3 KV buffer size =    25.50 MiB
-llama_kv_cache:     HTP0:0 KV buffer size =    25.50 MiB
-llama_kv_cache:     HTP0:2 KV buffer size =    25.50 MiB
-llama_kv_cache: size =  102.00 MiB (  8192 cells,  12 layers,  1/1 seqs), K (q8_0):   51.00 MiB, V (q8_0):   51.00 MiB
-llama_kv_cache_iswa: creating     SWA KV cache, size = 256 cells
-llama_kv_cache:     HTP0:1 KV buffer size =     0.80 MiB
-llama_kv_cache:     HTP0:3 KV buffer size =     0.53 MiB
-llama_kv_cache:     HTP0:0 KV buffer size =     1.06 MiB
-llama_kv_cache:     HTP0:2 KV buffer size =     0.80 MiB
-llama_kv_cache: size =    3.19 MiB (   256 cells,  12 layers,  1/1 seqs), K (q8_0):    1.59 MiB, V (q8_0):    1.59 MiB
-llama_context:     HTP0:0 compute buffer size =    16.06 MiB
-llama_context:     HTP0:1 compute buffer size =    16.06 MiB
-llama_context:     HTP0:2 compute buffer size =    16.06 MiB
-llama_context:     HTP0:3 compute buffer size =    16.06 MiB
-llama_context:        CPU compute buffer size =    98.19 MiB
-...
-llama_perf_context_print: prompt eval time =    3843.67 ms /   197 tokens ( 19.51 ms per token, 51.25 tokens per second)
-llama_perf_context_print:        eval time =    1686.13 ms /    31 runs   ( 54.39 ms per token, 18.39 tokens per second)
-llama_perf_context_print:       total time =    6266.30 ms /   228 tokens
-llama_perf_context_print:    graphs reused =         30
-llama_memory_breakdown_print: | memory breakdown [MiB] | total   free    self   model   context   compute    unaccounted |
-llama_memory_breakdown_print: |   - HTP0:0 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - HTP0:1 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - HTP0:2 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - HTP0:3 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - Host               |                 1476 =  1208 +     105 +     162                |
+[NPU Session 0]                              [NPU Session 1]
+       |                                            |
+  (Input Prep)                                 (Input Prep)
+       |                                            |
+  Pre-Op Barrier ----------------------------- Pre-Op Barrier
+  (mdev_sync_fence)                            (mdev_sync_fence)
+       |                                            |
+  Kernel Execution                             Kernel Execution
+  (Output Slice 0)                             (Output Slice 1)
+       |                                            |
+  Tensor Cache Flush                           Tensor Cache Flush
+  (htp_tensor_flush_all)                       (htp_tensor_flush_all)
+       |                                            |
+  Post-Op/Batch Barrier ---------------------- Post-Op/Batch Barrier
+  (htp_mdev_group_barrier)                     (htp_mdev_group_barrier)
+       |                                            |
+  Return Response to Host                      Return Response to Host
 ```
+
+### Atomic Fence Slots and Cache Invalidation
+
+- Fence synchronization operates on dedicated RPCMEM shared memory mapped across all participating sessions (`ctx->mdev.fence_base`).
+- Each device owns a dedicated 128-byte cache-line aligned fence slot:
+
+  ```c
+  atomic_uint * my_fence = htp_mdev_fence_slot(fence_base, mdev_idx);
+  ```
+
+- **Writing to fence ([`htp_fence_write`](../../../ggml/src/ggml-hexagon/htp/htp-fence.h#L18))**:
+  Stores `seq` and `status`, issues a `syncht` thread synchronization barrier, and flushes/invalidates the line
+  using `Q6_dccleaninva_A(fence)`.
+- **Reading from peer fence ([`htp_fence_read`](../../../ggml/src/ggml-hexagon/htp/htp-fence.h#L26))**:
+  Executes `Q6_dccleaninva_A(fence)` and `syncht` before reading atomic values to ensure fresh data from DDR.
+
+### Deterministic Monotonic Sequence Numbers
+
+- Barrier fences use monotonically increasing sequence numbers:
+
+  ```c
+  const uint32_t seq = ++ctx->mdev.fence_seq;
+  ```
+
+- Comparing sequence numbers with signed arithmetic `(int32_t)(peer_seq - seq) >= 0` prevents race conditions or
+  misaligned barrier arrivals across iterations.
+- If any peer reports an error status (`peer_status > HTP_STATUS_OK`), the barrier propagates the error and unblocks immediately.
+
+### Tensor Cache Flush and Pipeline Completion
+
+- In the kernel, ensure all pushed DMA operations have been popped in strict FIFO order to drain the queue.
+- Use [`htp_tensor_flush_all()`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h) to flush specific dirty tensors back to DDR:
+  - [`htp_tensor_flush_all()`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h) flushes only modified tensor address ranges,
+    ensuring peer devices and the host CPU observe consistent data in DDR.
+- Never signal completion before all DMA transfers are drained and dirty tensor flushes have completed.
+

docs/build-profiling.md

@@ -0,0 +1,122 @@
+## Build profiling
+This page is a working document for analyzing the current build and try to
+identify ways to improve the build time.
+
+### Requirements
+The profiling script requires clang to be used as the compiler tool chain and
+also requires that ClangBuildAnalyzer is installed.
+
+Mac:
+```console
+brew install clang-build-analyzer
+```
+
+Linux:
+```console
+git clone https://github.com/aras-p/ClangBuildAnalyzer.git
+cd ClangBuildAnalyzer
+cmake -B build -DCMAKE_BUILD_TYPE=Release
+cmake --build build -j$(nproc)
+sudo cp build/ClangBuildAnalyzer /usr/local/bin/
+```
+
+Windows: install LLVM/clang and Ninja (e.g. via the
+[LLVM releases page](https://github.com/llvm/llvm-project/releases) and
+`winget install Ninja-build.Ninja`), then build ClangBuildAnalyzer the same
+way as on Linux:
+```console
+git clone https://github.com/aras-p/ClangBuildAnalyzer.git
+cd ClangBuildAnalyzer
+cmake -B build -G Ninja -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_BUILD_TYPE=Release
+cmake --build build --config Release
+```
+Then add `ClangBuildAnalyzer\build` to `PATH`.
+
+### Usage
+Mac/Linux:
+```console
+$ ./scripts/build-profile.sh
+```
+
+Windows:
+```console
+> .\scripts\build-profile.ps1
+```
+
+Both accept `--full`/`-Full` (include Server, Tools, and Tests) and a jobs
+override (`-jN` / `-Jobs N`).
+
+Note: on Windows, `cmake` defaults to the Visual Studio generator, which
+ignores `CMAKE_C_COMPILER`/`CMAKE_CXX_COMPILER` and silently falls back to
+MSVC. `build-profile.ps1` passes `-G Ninja` so clang is actually used, this
+is required on ARM64.
+
+### Linux (Ubuntu 24.04)
+
+Environment:
+- Clang:     18.1.3 (Ubuntu clang version 18.1.3 (1ubuntu1))
+- libstdc++: GCC 13.3.0 (Ubuntu 13.3.0-6ubuntu2~24.04.1)
+- Target:    x86_64-pc-linux-gnu
+
+```console
++------------------------+-----+------------+------------+------------+
+| Build                  | TUs | Frontend   | Backend    | Total      |
++------------------------+-----+------------+------------+------------+
+| Minimal, master        | 249 |   468.2 s  |   270.3 s  |   738.5 s  |
+| Minimal, with PCH      | 253 |   177.1 s  |   265.8 s  |   442.9 s  |
+| Full,    master        | 396 |   811.0 s  |   692.2 s  | 1,503.2 s  |
+| Full,    with PCH      | 405 |   380.0 s  |   664.7 s  | 1,044.7 s  |
+| Full,    with PCH + UB | 264 |   357.7 s  |   635.7 s  |   993.4 s  |
++------------------------+-----+------------+------------+------------+
+
+PCH  = precompiled header.
+Full = includes building Server, Tools, and Tests.
+UB   = unity build for models
+```
+Note that the number of translation units (TUs) increases when using precompiled
+headers — each PCH target adds one extra TU for the precompilation step itself.
+
+### Mac (Apple M3)
+
+Environment:
+- Clang:  Apple clang version 17.0.0 (clang-1700.3.19.1)
+- libc++: ships with Apple clang 17.0.0 (Xcode toolchain)
+- Target: arm64-apple-macosx15.6
+
+```console
++------------------------+-----+------------+------------+------------+
+| Build                  | TUs | Frontend   | Backend    | Total      |
++------------------------+-----+------------+------------+------------+
+| Minimal, master        | 256 |   154.5 s  |    94.8 s  |   249.3 s  |
+| Minimal, with PCH      | 261 |    65.9 s  |    90.0 s  |   155.9 s  |
+| Full,    master        | 407 |   265.7 s  |   209.7 s  |   475.4 s  |
+| Full,    with PCH      | 414 |   154.6 s  |   197.5 s  |   352.1 s  |
+| Full,    with PCH + UB | 274 |   143.0 s  |   192.2 s  |   335.2 s  |
++------------------------+-----+------------+------------+------------+
+
+PCH = precompiled header.
+Full = includes building Server, Tools, and Tests.
+UB   = unity build for models
+```
+
+### Windows (ARM64)
+
+Environment:
+- Clang:  clang version 22.1.8 (LLVM, `C:\Program Files\LLVM`)
+- STL:    MSVC STL (Visual Studio 2022 Build Tools 14.44.35207)
+- Target: aarch64-pc-windows-msvc
+
+```console
++------------------------+-----+------------+------------+------------+
+| Build                  | TUs | Frontend   | Backend    | Total      |
++------------------------+-----+------------+------------+------------+
+| Minimal, master        | 249 |   159.4 s  |    82.2 s  |   241.6 s  |
+| Full,    master        | 373 |   337.2 s  |   167.4 s  |   504.6 s  |
+| Minimal, with PCH + UB | 113 |    62.3 s  |    82.4 s  |   144.7 s  |
+| Full,    with PCH + UB | 240 |   233.0 s  |   185.1 s  |   418.1 s  |
++------------------------+-----+------------+------------+------------+
+
+PCH = precompiled header.
+Full = includes building Server, Tools, and Tests.
+UB   = unity build for models
+```

docs/build.md

@@ -806,7 +806,7 @@ To read documentation for how to build on Android, [click here](./android.md)
 
 ## WebGPU
 
-The WebGPU backend relies on [Dawn](https://dawn.googlesource.com/dawn). Follow the instructions [here](https://dawn.googlesource.com/dawn/+/refs/heads/main/docs/quickstart-cmake.md) to install Dawn locally so that llama.cpp can find it using CMake. The current implementation is up-to-date with Dawn commit `18eb229`.
+The WebGPU backend relies on [Dawn](https://dawn.googlesource.com/dawn). Follow the instructions [here](https://dawn.googlesource.com/dawn/+/refs/heads/main/docs/quickstart-cmake.md) to install Dawn locally so that llama.cpp can find it using CMake. The current implementation is up-to-date with Dawn commit `94c3c9c`.
 
 In the llama.cpp directory, build with CMake:
 

docs/development/parsing.md

@@ -28,7 +28,7 @@ auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
     for (const auto & tool : tools) {
         const auto & function = tool.at("function");
         std::string name = function.at("name");
-        const auto & schema = function.at("parameters");
+        const auto   schema = common_chat_tool_parameters(function);
 
         auto tool_name = p.json_member("name", "\"" + p.literal(name) + "\"");
         auto tool_args = p.json_member("arguments", p.schema(p.json(), "tool-" + name + "-schema", schema));
@@ -108,6 +108,7 @@ For a more complete example, see `test_example_native()` in
 - **`rule(name, p, trigger)`** - Creates a named rule and returns a reference
 - **`trigger_rule(name, p)`** - Creates a trigger rule (entry point for lazy grammar generation)
 - **`schema(p, name, schema, raw)`** - Wraps parser with JSON schema metadata for grammar generation
+- **`schema(p, name, doc, node, raw)`** - Same, for a node of a `common_chat_schema_document` built earlier, e.g. one tool parameter
 
 ### AST Control
 
@@ -121,9 +122,6 @@ some exceptions.
 
 ```cpp
 data.grammar = build_grammar([&](const common_grammar_builder & builder) {
-    foreach_function(params.tools, [&](const json & fn) {
-        builder.resolve_refs(fn.at("parameters"));
-    });
     parser.build_grammar(builder, data.grammar_lazy);
 });
 ```
@@ -151,7 +149,8 @@ implementation to generate the grammar instead of the underlying parser.
 
 The `raw` option emits a grammar suitable for a raw string instead of a JSON
 string. In other words, it won't be wrapped in quotes or require escaping
-quotes. It should only be used when `type == "string"`.
+quotes. It only takes effect when the schema may be a string, as reported by
+`common_chat_schema::may_be_string()`, otherwise the JSON grammar is used.
 
 The downside is that it can potentially lead to ambiguous grammars. For
 example, if a user provides the pattern `^.*$`, the following grammar may be

examples/json_schema_to_grammar.py

@@ -1,842 +0,0 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import argparse
-import itertools
-import json
-import re
-import sys
-from typing import Any, List, Optional, Set, Tuple, Union
-
-def _build_repetition(item_rule, min_items, max_items, separator_rule=None):
-
-    if max_items == 0:
-        return ""
-
-    if min_items == 0 and max_items == 1:
-        return f'{item_rule}?'
-
-    if not separator_rule:
-        if min_items == 1 and max_items is None:
-            return f'{item_rule}+'
-        elif min_items == 0 and max_items is None:
-            return f'{item_rule}*'
-        else:
-            return f'{item_rule}{{{min_items},{max_items if max_items is not None else ""}}}'
-
-    result = item_rule + ' ' + _build_repetition(f'({separator_rule} {item_rule})', min_items - 1 if min_items > 0 else 0, max_items - 1 if max_items is not None else None)
-    return f'({result})?' if min_items == 0 else result
-
-def _generate_min_max_int(min_value: Optional[int], max_value: Optional[int], out: list, decimals_left: int = 16, top_level: bool = True):
-    def digit_range(from_char: str, to_char: str):
-        out.append("[")
-        if from_char == to_char:
-            out.append(from_char)
-        else:
-            out.append(from_char)
-            out.append("-")
-            out.append(to_char)
-        out.append("]")
-
-    def more_digits(min_digits: int, max_digits: int):
-        out.append("[0-9]")
-        if min_digits == max_digits and min_digits == 1:
-            return
-        out.append("{")
-        out.append(str(min_digits))
-        if max_digits != min_digits:
-            out.append(",")
-            if max_digits != sys.maxsize:
-                out.append(str(max_digits))
-        out.append("}")
-
-    def uniform_range(from_str: str, to_str: str):
-        i = 0
-        while i < len(from_str) and from_str[i] == to_str[i]:
-            i += 1
-        if i > 0:
-            out.append("\"")
-            out.append(from_str[:i])
-            out.append("\"")
-        if i < len(from_str):
-            if i > 0:
-                out.append(" ")
-            sub_len = len(from_str) - i - 1
-            if sub_len > 0:
-                from_sub = from_str[i+1:]
-                to_sub = to_str[i+1:]
-                sub_zeros = "0" * sub_len
-                sub_nines = "9" * sub_len
-
-                to_reached = False
-                out.append("(")
-                if from_sub == sub_zeros:
-                    digit_range(from_str[i], chr(ord(to_str[i]) - 1))
-                    out.append(" ")
-                    more_digits(sub_len, sub_len)
-                else:
-                    out.append("[")
-                    out.append(from_str[i])
-                    out.append("] ")
-                    out.append("(")
-                    uniform_range(from_sub, sub_nines)
-                    out.append(")")
-                    if ord(from_str[i]) < ord(to_str[i]) - 1:
-                        out.append(" | ")
-                        if to_sub == sub_nines:
-                            digit_range(chr(ord(from_str[i]) + 1), to_str[i])
-                            to_reached = True
-                        else:
-                            digit_range(chr(ord(from_str[i]) + 1), chr(ord(to_str[i]) - 1))
-                        out.append(" ")
-                        more_digits(sub_len, sub_len)
-                if not to_reached:
-                    out.append(" | ")
-                    digit_range(to_str[i], to_str[i])
-                    out.append(" ")
-                    uniform_range(sub_zeros, to_sub)
-                out.append(")")
-            else:
-                out.append("[")
-                out.append(from_str[i])
-                out.append("-")
-                out.append(to_str[i])
-                out.append("]")
-
-    if min_value is not None and max_value is not None:
-        if min_value < 0 and max_value < 0:
-            out.append("\"-\" (")
-            _generate_min_max_int(-max_value, -min_value, out, decimals_left, top_level=True)
-            out.append(")")
-            return
-
-        if min_value < 0:
-            out.append("\"-\" (")
-            _generate_min_max_int(0, -min_value, out, decimals_left, top_level=True)
-            out.append(") | ")
-            min_value = 0
-
-        min_s = str(min_value)
-        max_s = str(max_value)
-        min_digits = len(min_s)
-        max_digits = len(max_s)
-
-        for digits in range(min_digits, max_digits):
-            uniform_range(min_s, "9" * digits)
-            min_s = "1" + "0" * digits
-            out.append(" | ")
-        uniform_range(min_s, max_s)
-        return
-
-    less_decimals = max(decimals_left - 1, 1)
-
-    if min_value is not None:
-        if min_value < 0:
-            out.append("\"-\" (")
-            _generate_min_max_int(None, -min_value, out, decimals_left, top_level=False)
-            out.append(") | [0] | [1-9] ")
-            more_digits(0, decimals_left - 1)
-        elif min_value == 0:
-            if top_level:
-                out.append("[0] | [1-9] ")
-                more_digits(0, less_decimals)
-            else:
-                more_digits(1, decimals_left)
-        elif min_value <= 9:
-            c = str(min_value)
-            range_start = '1' if top_level else '0'
-            if c > range_start:
-                digit_range(range_start, chr(ord(c) - 1))
-                out.append(" ")
-                more_digits(1, less_decimals)
-                out.append(" | ")
-            digit_range(c, "9")
-            out.append(" ")
-            more_digits(0, less_decimals)
-        else:
-            min_s = str(min_value)
-            length = len(min_s)
-            c = min_s[0]
-
-            if c > "1":
-                digit_range("1" if top_level else "0", chr(ord(c) - 1))
-                out.append(" ")
-                more_digits(length, less_decimals)
-                out.append(" | ")
-            digit_range(c, c)
-            out.append(" (")
-            _generate_min_max_int(int(min_s[1:]), None, out, less_decimals, top_level=False)
-            out.append(")")
-            if c < "9":
-                out.append(" | ")
-                digit_range(chr(ord(c) + 1), "9")
-                out.append(" ")
-                more_digits(length - 1, less_decimals)
-        return
-
-    if max_value is not None:
-        if max_value >= 0:
-            if top_level:
-                out.append("\"-\" [1-9] ")
-                more_digits(0, less_decimals)
-                out.append(" | ")
-            _generate_min_max_int(0, max_value, out, decimals_left, top_level=True)
-        else:
-            out.append("\"-\" (")
-            _generate_min_max_int(-max_value, None, out, decimals_left, top_level=False)
-            out.append(")")
-        return
-
-    raise RuntimeError("At least one of min_value or max_value must be set")
-
-class BuiltinRule:
-    def __init__(self, content: str, deps: list | None = None):
-        self.content = content
-        self.deps = deps or []
-
-# Constraining spaces to prevent model "running away".
-SPACE_RULE = '| " " | "\\n"{1,2} [ \\t]{0,20}'
-
-PRIMITIVE_RULES = {
-    'boolean'      : BuiltinRule('("true" | "false")', []),
-    'decimal-part' : BuiltinRule('[0-9]{1,16}', []),
-    'integral-part': BuiltinRule('[0] | [1-9] [0-9]{0,15}', []),
-    'number'       : BuiltinRule('("-"? integral-part) ("." decimal-part)? ([eE] [-+]? integral-part)?', ['integral-part', 'decimal-part']),
-    'integer'      : BuiltinRule('("-"? integral-part)', ['integral-part']),
-    'value'        : BuiltinRule('object | array | string | number | boolean | null', ['object', 'array', 'string', 'number', 'boolean', 'null']),
-    'object'       : BuiltinRule('"{" space ( string ":" space value ("," space string ":" space value)* )? space "}"', ['string', 'value']),
-    'array'        : BuiltinRule('"[" space ( value ("," space value)* )? space "]"', ['value']),
-    'uuid'         : BuiltinRule(r'"\"" [0-9a-fA-F]{8} "-" [0-9a-fA-F]{4} "-" [0-9a-fA-F]{4} "-" [0-9a-fA-F]{4} "-" [0-9a-fA-F]{12} "\""', []),
-    'char'         : BuiltinRule(r'[^"\\\x7F\x00-\x1F] | [\\] (["\\bfnrt] | "u" [0-9a-fA-F]{4})', []),
-    'string'       : BuiltinRule(r'"\"" char* "\""', ['char']),
-    'null'         : BuiltinRule('"null"', []),
-}
-
-# TODO: support "uri", "email" string formats
-STRING_FORMAT_RULES = {
-    'date'            : BuiltinRule('[0-9]{4} "-" ( "0" [1-9] | "1" [0-2] ) "-" ( \"0\" [1-9] | [1-2] [0-9] | "3" [0-1] )', []),
-    'time'            : BuiltinRule('([01] [0-9] | "2" [0-3]) ":" [0-5] [0-9] ":" [0-5] [0-9] ( "." [0-9]{3} )? ( "Z" | ( "+" | "-" ) ( [01] [0-9] | "2" [0-3] ) ":" [0-5] [0-9] )', []),
-    'date-time'       : BuiltinRule('date "T" time', ['date', 'time']),
-    'date-string'     : BuiltinRule('"\\"" date "\\""', ['date']),
-    'time-string'     : BuiltinRule('"\\"" time "\\""', ['time']),
-    'date-time-string': BuiltinRule('"\\"" date-time "\\""', ['date-time']),
-}
-
-DOTALL = '[\\U00000000-\\U0010FFFF]'
-DOT = '[^\\x0A\\x0D]'
-
-RESERVED_NAMES = set(["root", "dot", *PRIMITIVE_RULES.keys(), *STRING_FORMAT_RULES.keys()])
-
-INVALID_RULE_CHARS_RE = re.compile(r'[^a-zA-Z0-9-]+')
-GRAMMAR_LITERAL_ESCAPE_RE = re.compile(r'[\r\n"\\]')
-GRAMMAR_RANGE_LITERAL_ESCAPE_RE = re.compile(r'[\r\n"\]\-\\]')
-GRAMMAR_LITERAL_ESCAPES = {'\r': '\\r', '\n': '\\n', '"': '\\"', '-': '\\-', ']': '\\]', '\\': '\\\\'}
-
-NON_LITERAL_SET = set('|.()[]{}*+?')
-ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS = set('^$.[]()|{}*+?')
-
-
-class SchemaConverter:
-    def __init__(self, *, prop_order, allow_fetch, dotall, raw_pattern):
-        self._prop_order = prop_order
-        self._allow_fetch = allow_fetch
-        self._dotall = dotall
-        self._raw_pattern = raw_pattern
-        self._rules = {
-            'space': SPACE_RULE,
-        }
-        self._refs = {}
-        self._refs_being_resolved = set()
-
-    def _format_literal(self, literal):
-        escaped = GRAMMAR_LITERAL_ESCAPE_RE.sub(
-            lambda m: GRAMMAR_LITERAL_ESCAPES.get(m.group(0)) or m.group(0), literal
-        )
-        return f'"{escaped}"'
-
-    def not_literal(self, literal: str, dotall: bool = True, maybe_escaped_underscores = False) -> str:
-        '''
-            not_literal('a') -> '[^a]'
-            not_literal('abc') -> '([^a] | "a" ([^b] | "b" ([^c])?)?)?'
-        '''
-        assert len(literal) > 0, 'Empty literal not supported'
-        def recurse(i: int):
-            c = literal[i]
-            if maybe_escaped_underscores and c == '_':
-                yield f'[^{c}\\\\]'
-                yield ' | '
-                yield f'"\\\\"? "{c}"'
-            else:
-                yield f'[^{c}]'
-            if i < len(literal) - 1:
-                yield ' | '
-                yield self._format_literal(c)
-                yield ' ('
-                yield from recurse(i + 1)
-                yield ')?'
-
-        return ''.join(('(', *recurse(0), ')'))
-
-    def _not_strings(self, strings):
-        class TrieNode:
-            def __init__(self):
-                self.children = {}
-                self.is_end_of_string = False
-
-            def insert(self, string):
-                node = self
-                for c in string:
-                    node = node.children.setdefault(c, TrieNode())
-                node.is_end_of_string = True
-
-        trie = TrieNode()
-        for s in strings:
-            trie.insert(s)
-
-        char_rule = self._add_primitive('char', PRIMITIVE_RULES['char'])
-        out = ['["] ( ']
-
-        def visit(node):
-            rejects = []
-            first = True
-            for c in sorted(node.children.keys()):
-                child = node.children[c]
-                rejects.append(c)
-                if first:
-                    first = False
-                else:
-                    out.append(' | ')
-                out.append(f'[{c}]')
-                if child.children:
-                    out.append(f' (')
-                    visit(child)
-                    out.append(')')
-                elif child.is_end_of_string:
-                    out.append(f' {char_rule}+')
-            if node.children:
-                if not first:
-                    out.append(' | ')
-                out.append(f'[^"{"".join(rejects)}] {char_rule}*')
-        visit(trie)
-
-        out.append(f' ){"" if trie.is_end_of_string else "?"} ["]')
-        return ''.join(out)
-
-    def _add_rule(self, name, rule):
-        esc_name = INVALID_RULE_CHARS_RE.sub('-', name)
-        if esc_name not in self._rules or self._rules[esc_name] == rule:
-            key = esc_name
-        else:
-            i = 0
-            while f'{esc_name}{i}' in self._rules and self._rules[f'{esc_name}{i}'] != rule:
-                i += 1
-            key = f'{esc_name}{i}'
-        self._rules[key] = rule
-        return key
-
-    def resolve_refs(self, schema: dict, url: str):
-        '''
-            Resolves all $ref fields in the given schema, fetching any remote schemas,
-            replacing $ref with absolute reference URL and populating self._refs with the
-            respective referenced (sub)schema dictionaries.
-        '''
-        def visit(n: dict):
-            if isinstance(n, list):
-                return [visit(x) for x in n]
-            elif isinstance(n, dict):
-                ref = n.get('$ref')
-                if ref is not None and ref not in self._refs:
-                    if ref.startswith('https://'):
-                        assert self._allow_fetch, 'Fetching remote schemas is not allowed (use --allow-fetch for force)'
-                        import requests
-
-                        frag_split = ref.split('#')
-                        base_url = frag_split[0]
-
-                        target = self._refs.get(base_url)
-                        if target is None:
-                            target = self.resolve_refs(requests.get(ref).json(), base_url)
-                            self._refs[base_url] = target
-
-                        if len(frag_split) == 1 or frag_split[-1] == '':
-                            return target
-                    elif ref.startswith('#/'):
-                        target = schema
-                        ref = f'{url}{ref}'
-                        n['$ref'] = ref
-                    else:
-                        raise ValueError(f'Unsupported ref {ref}')
-
-                    for sel in ref.split('#')[-1].split('/')[1:]:
-                        assert target is not None, f'Error resolving ref {ref}: {sel} not in {target}'
-                        if isinstance(target, list):
-                            try:
-                                sel_index = int(sel)
-                            except ValueError:
-                                raise ValueError(f'Error resolving ref {ref}: {sel} not in {target}')
-                            assert 0 <= sel_index < len(target), f'Error resolving ref {ref}: {sel} not in {target}'
-                            target = target[sel_index]
-                        else:
-                            assert sel in target, f'Error resolving ref {ref}: {sel} not in {target}'
-                            target = target[sel]
-
-                    self._refs[ref] = target
-                else:
-                    for v in n.values():
-                        visit(v)
-
-            return n
-        return visit(schema)
-
-    def _generate_union_rule(self, name, alt_schemas):
-        return ' | '.join((
-            self.visit(alt_schema, f'{name}{"-" if name else "alternative-"}{i}')
-            for i, alt_schema in enumerate(alt_schemas)
-        ))
-
-    def _visit_pattern(self, pattern, name):
-        '''
-            Transforms a regular expression pattern into a GBNF rule.
-
-            Input: https://json-schema.org/understanding-json-schema/reference/regular_expressions
-            Output: https://github.com/ggml-org/llama.cpp/blob/master/grammars/README.md
-
-            Unsupported features: negative/positive lookaheads, greedy/non-greedy modifiers.
-
-            Mostly a 1:1 translation, except for {x} / {x,} / {x,y} quantifiers for which
-            we define sub-rules to keep the output lean.
-        '''
-
-        assert pattern.startswith('^') and pattern.endswith('$'), 'Pattern must start with "^" and end with "$"'
-        pattern = pattern[1:-1]
-        sub_rule_ids = {}
-
-        i = 0
-        length = len(pattern)
-
-        def to_rule(s: tuple[str, bool]) -> str:
-            (txt, is_literal) = s
-            return "\"" + txt + "\"" if is_literal else txt
-
-        def transform() -> tuple[str, bool]:
-            '''
-                Parse a unit at index i (advancing it), and return its string representation + whether it's a literal.
-            '''
-            nonlocal i
-            nonlocal pattern
-            nonlocal sub_rule_ids
-
-            start = i
-            # For each component of this sequence, store its string representation and whether it's a literal.
-            # We only need a flat structure here to apply repetition operators to the last item, and
-            # to merge literals at the and (we're parsing grouped ( sequences ) recursively and don't treat '|' specially
-            # (GBNF's syntax is luckily very close to regular expressions!)
-            seq: list[tuple[str, bool]] = []
-
-            def get_dot():
-                if self._dotall:
-                    rule = DOTALL
-                else:
-                    # Accept any character... except \n and \r line break chars (\x0A and \xOD)
-                    rule = DOT
-                return self._add_rule(f'dot', rule)
-
-            def join_seq():
-                nonlocal seq
-                ret = []
-                for is_literal, g in itertools.groupby(seq, lambda x: x[1]):
-                    if is_literal:
-                        ret.append((''.join(x[0] for x in g), True))
-                    else:
-                        ret.extend(g)
-                if len(ret) == 1:
-                    return ret[0]
-                return (' '.join(to_rule(x) for x in seq), False)
-
-            while i < length:
-                c = pattern[i]
-                if c == '.':
-                    seq.append((get_dot(), False))
-                    i += 1
-                elif c == '(':
-                    i += 1
-                    if i < length:
-                        assert pattern[i] != '?', f'Unsupported pattern syntax "{pattern[i]}" at index {i} of /{pattern}/'
-                    seq.append((f'({to_rule(transform())})', False))
-                elif c == ')':
-                    i += 1
-                    assert start > 0 and pattern[start-1] == '(', f'Unbalanced parentheses; start = {start}, i = {i}, pattern = {pattern}'
-                    return join_seq()
-                elif c == '[':
-                    square_brackets = c
-                    i += 1
-                    while i < length and pattern[i] != ']':
-                        if pattern[i] == '\\':
-                            square_brackets += pattern[i:i+2]
-                            i += 2
-                        else:
-                            square_brackets += pattern[i]
-                            i += 1
-                    assert i < length, f'Unbalanced square brackets; start = {start}, i = {i}, pattern = {pattern}'
-                    square_brackets += ']'
-                    i += 1
-                    seq.append((square_brackets, False))
-                elif c == '|':
-                    seq.append(('|', False))
-                    i += 1
-                elif c in ('*', '+', '?'):
-                    seq[-1] = (to_rule(seq[-1]) + c, False)
-                    i += 1
-                elif c == '{':
-                    curly_brackets = c
-                    i += 1
-                    while i < length and pattern[i] != '}':
-                        curly_brackets += pattern[i]
-                        i += 1
-                    assert i < length, f'Unbalanced curly brackets; start = {start}, i = {i}, pattern = {pattern}'
-                    curly_brackets += '}'
-                    i += 1
-                    nums = [s.strip() for s in curly_brackets[1:-1].split(',')]
-                    min_times = 0
-                    max_times = None
-                    try:
-                        if len(nums) == 1:
-                            min_times = int(nums[0])
-                            max_times = min_times
-                        else:
-                            assert len(nums) == 2
-                            min_times = int(nums[0]) if nums[0] else 0
-                            max_times = int(nums[1]) if nums[1] else None
-                    except ValueError:
-                        raise ValueError(f'Invalid quantifier {curly_brackets} in /{pattern}/')
-
-                    (sub, sub_is_literal) = seq[-1]
-
-                    if not sub_is_literal:
-                        id = sub_rule_ids.get(sub)
-                        if id is None:
-                            id = self._add_rule(f'{name}-{len(sub_rule_ids) + 1}', sub)
-                            sub_rule_ids[sub] = id
-                        sub = id
-
-                    seq[-1] = (_build_repetition(f'"{sub}"' if sub_is_literal else sub, min_times, max_times), False)
-                else:
-                    literal = ''
-                    while i < length:
-                        if pattern[i] == '\\' and i < length - 1:
-                            next = pattern[i + 1]
-                            if next in ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS:
-                                i += 1
-                                literal += pattern[i]
-                                i += 1
-                            else:
-                                literal += pattern[i:i+2]
-                                i += 2
-                        elif pattern[i] == '"' and not self._raw_pattern:
-                            literal += '\\"'
-                            i += 1
-                        elif pattern[i] not in NON_LITERAL_SET and \
-                                (i == length - 1 or literal == '' or pattern[i+1] == '.' or pattern[i+1] not in NON_LITERAL_SET):
-                            literal += pattern[i]
-                            i += 1
-                        else:
-                            break
-                    if literal:
-                        seq.append((literal, True))
-
-            return join_seq()
-
-        return self._add_rule(
-            name,
-            to_rule(transform()) if self._raw_pattern \
-                else "\"\\\"\" (" + to_rule(transform()) + ") \"\\\"\"")
-
-
-    def _resolve_ref(self, ref):
-        ref_fragment = ref.split('#')[-1]
-        ref_name = 'ref' + re.sub(r'[^a-zA-Z0-9-]+', '-', ref_fragment)
-        if ref_name not in self._rules and ref not in self._refs_being_resolved:
-            self._refs_being_resolved.add(ref)
-            resolved = self._refs[ref]
-            ref_name = self.visit(resolved, ref_name)
-            self._refs_being_resolved.remove(ref)
-        return ref_name
-
-    def _generate_constant_rule(self, value):
-        return self._format_literal(json.dumps(value))
-
-    def visit(self, schema, name):
-        schema_type = schema.get('type')
-        schema_format = schema.get('format')
-        rule_name = name + '-' if name in RESERVED_NAMES else name or 'root'
-
-        if (ref := schema.get('$ref')) is not None:
-            return self._add_rule(rule_name, self._resolve_ref(ref))
-
-        elif 'oneOf' in schema or 'anyOf' in schema:
-            return self._add_rule(rule_name, self._generate_union_rule(name, schema.get('oneOf') or schema['anyOf']))
-
-        elif isinstance(schema_type, list):
-            return self._add_rule(rule_name, self._generate_union_rule(name, [{**schema, 'type': t} for t in schema_type]))
-
-        elif 'const' in schema:
-            return self._add_rule(rule_name, self._generate_constant_rule(schema['const']))
-
-        elif 'enum' in schema:
-            rule = '(' + ' | '.join((self._generate_constant_rule(v) for v in schema['enum'])) + ')'
-            return self._add_rule(rule_name, rule)
-
-        elif schema_type in (None, 'object') and \
-             ('properties' in schema or \
-              ('additionalProperties' in schema and schema['additionalProperties'] is not True)):
-            required = set(schema.get('required', []))
-            properties = list(schema.get('properties', {}).items())
-            return self._add_rule(rule_name, self._build_object_rule(properties, required, name, schema.get('additionalProperties')))
-
-        elif schema_type in (None, 'object', 'string') and 'allOf' in schema:
-            required = set()
-            properties = []
-            enum_sets = []
-            hybrid_name = name
-            def add_component(comp_schema, is_required):
-                if (ref := comp_schema.get('$ref')) is not None:
-                    comp_schema = self._refs[ref]
-
-                if 'properties' in comp_schema:
-                    for prop_name, prop_schema in comp_schema['properties'].items():
-                        properties.append((prop_name, prop_schema))
-                        if is_required:
-                            required.add(prop_name)
-
-                if 'enum' in comp_schema:
-                    enum_sets.append(set(comp_schema['enum']))
-
-            for t in schema['allOf']:
-                if 'anyOf' in t:
-                    for tt in t['anyOf']:
-                        add_component(tt, is_required=False)
-                else:
-                    add_component(t, is_required=True)
-
-            if enum_sets:
-                enum_intersection = enum_sets[0]
-                for s in enum_sets[1:]:
-                    enum_intersection &= s
-
-                if enum_intersection:
-                    rule = '(' + ' | '.join((self._generate_constant_rule(v) for v in sorted(enum_intersection))) + ')'
-                    return self._add_rule(rule_name, rule)
-
-            return self._add_rule(rule_name, self._build_object_rule(properties, required, hybrid_name, additional_properties=None))
-
-        elif schema_type in (None, 'array') and ('items' in schema or 'prefixItems' in schema):
-            items = schema.get('items', schema.get('prefixItems'))
-            if isinstance(items, list):
-                return self._add_rule(
-                    rule_name,
-                    '"[" space ' +
-                    ' "," space '.join(
-                        self.visit(item, f'{name}{"-" if name else ""}tuple-{i}')
-                        for i, item in enumerate(items)) +
-                    ' space "]"')
-            else:
-                item_rule_name = self.visit(items, f'{name}{"-" if name else ""}item')
-                min_items = schema.get("minItems", 0)
-                max_items = schema.get("maxItems")
-                return self._add_rule(rule_name, '"[" space ' + _build_repetition(item_rule_name, min_items, max_items, separator_rule='"," space') + ' space "]"')
-
-        elif schema_type in (None, 'string') and 'pattern' in schema:
-            return self._visit_pattern(schema['pattern'], rule_name)
-
-        elif schema_type in (None, 'string') and re.match(r'^uuid[1-5]?$', schema_format or ''):
-            return self._add_primitive(
-                'root' if rule_name == 'root' else schema_format,
-                PRIMITIVE_RULES['uuid']
-            )
-
-        elif schema_type in (None, 'string') and f'{schema_format}-string' in STRING_FORMAT_RULES:
-            prim_name = f'{schema_format}-string'
-            return self._add_rule(rule_name, self._add_primitive(prim_name, STRING_FORMAT_RULES[prim_name]))
-
-        elif schema_type == 'string' and ('minLength' in schema or 'maxLength' in schema):
-            char_rule = self._add_primitive('char', PRIMITIVE_RULES['char'])
-            min_len = schema.get('minLength', 0)
-            max_len = schema.get('maxLength')
-
-            return self._add_rule(rule_name, r'"\"" ' + _build_repetition(char_rule, min_len, max_len) + r' "\""')
-
-        elif schema_type in (None, 'integer') and \
-                ('minimum' in schema or 'exclusiveMinimum' in schema or 'maximum' in schema or 'exclusiveMaximum' in schema):
-            min_value = None
-            max_value = None
-            if 'minimum' in schema:
-                min_value = schema['minimum']
-            elif 'exclusiveMinimum' in schema:
-                min_value = schema['exclusiveMinimum'] + 1
-            if 'maximum' in schema:
-                max_value = schema['maximum']
-            elif 'exclusiveMaximum' in schema:
-                max_value = schema['exclusiveMaximum'] - 1
-
-            out = ["("]
-            _generate_min_max_int(min_value, max_value, out)
-            out.append(")")
-            return self._add_rule(rule_name, ''.join(out))
-
-        elif (schema_type == 'object') or (len(schema) == 0):
-            return self._add_rule(rule_name, self._add_primitive('object', PRIMITIVE_RULES['object']))
-
-        elif schema_type is None and isinstance(schema, dict):
-            # No type constraint and no recognized structural keywords (e.g. {"description": "..."}).
-            # Per JSON Schema semantics this is equivalent to {} and accepts any value.
-            return self._add_rule(rule_name, self._add_primitive('value', PRIMITIVE_RULES['value']))
-
-        else:
-            assert schema_type in PRIMITIVE_RULES, f'Unrecognized schema: {schema}'
-            # TODO: support minimum, maximum, exclusiveMinimum, exclusiveMaximum at least for zero
-            return self._add_primitive('root' if rule_name == 'root' else schema_type, PRIMITIVE_RULES[schema_type])
-
-    def _add_primitive(self, name: str, rule: BuiltinRule):
-        n = self._add_rule(name, rule.content)
-
-        for dep in rule.deps:
-            dep_rule = PRIMITIVE_RULES.get(dep) or STRING_FORMAT_RULES.get(dep)
-            assert dep_rule, f'Rule {dep} not known'
-            if dep not in self._rules:
-                self._add_primitive(dep, dep_rule)
-        return n
-
-    def _build_object_rule(self, properties: List[Tuple[str, Any]], required: Set[str], name: str, additional_properties: Optional[Union[bool, Any]]):
-        prop_order = self._prop_order
-        # sort by position in prop_order (if specified) then by original order
-        sorted_props = [kv[0] for _, kv in sorted(enumerate(properties), key=lambda ikv: (prop_order.get(ikv[1][0], len(prop_order)), ikv[0]))]
-
-        prop_kv_rule_names = {}
-        for prop_name, prop_schema in properties:
-            prop_rule_name = self.visit(prop_schema, f'{name}{"-" if name else ""}{prop_name}')
-            prop_kv_rule_names[prop_name] = self._add_rule(
-                f'{name}{"-" if name else ""}{prop_name}-kv',
-                fr'{self._format_literal(json.dumps(prop_name))} space ":" space {prop_rule_name}'
-            )
-        required_props = [k for k in sorted_props if k in required]
-        optional_props = [k for k in sorted_props if k not in required]
-
-        if additional_properties is not None and additional_properties != False:
-            sub_name = f'{name}{"-" if name else ""}additional'
-            value_rule = self.visit(additional_properties, f'{sub_name}-value') if isinstance(additional_properties, dict) else \
-                self._add_primitive('value', PRIMITIVE_RULES['value'])
-            key_rule = self._add_primitive('string', PRIMITIVE_RULES['string']) if not sorted_props \
-                else self._add_rule(f'{sub_name}-k', self._not_strings(sorted_props))
-
-            prop_kv_rule_names["*"] = self._add_rule(
-                f'{sub_name}-kv',
-                f'{key_rule} ":" space {value_rule}'
-            )
-            optional_props.append("*")
-
-        if not required_props and not optional_props:
-            return '"{" space "}"'
-
-        rule = '"{" space '
-        rule += ' "," space '.join(prop_kv_rule_names[k] for k in required_props)
-
-        if optional_props:
-            rule += ' ('
-            if required_props:
-                rule += ' "," space ( '
-
-            def get_recursive_refs(ks, first_is_optional):
-                [k, *rest] = ks
-                kv_rule_name = prop_kv_rule_names[k]
-                comma_ref = f'( "," space {kv_rule_name} )'
-                if first_is_optional:
-                    res = comma_ref + ('*' if k == '*' else '?')
-                else:
-                    res = kv_rule_name + (' ' + comma_ref + "*" if k == '*' else '')
-                if len(rest) > 0:
-                    res += ' ' + self._add_rule(
-                        f'{name}{"-" if name else ""}{k}-rest',
-                        get_recursive_refs(rest, first_is_optional=True)
-                    )
-                return res
-
-            rule += ' | '.join(
-                get_recursive_refs(optional_props[i:], first_is_optional=False)
-                for i in range(len(optional_props))
-            )
-            if required_props:
-                rule += ' )'
-            rule += ' )?'
-
-        rule += ' space "}"'
-
-        return rule
-
-    def format_grammar(self):
-        return '\n'.join(
-            f'{name} ::= {rule}'
-            for name, rule in sorted(self._rules.items(), key=lambda kv: kv[0])
-        )
-
-
-def main(args_in = None):
-    parser = argparse.ArgumentParser(
-        description='''
-            Generates a grammar (suitable for use in ./llama-cli) that produces JSON conforming to a
-            given JSON schema. Only a subset of JSON schema features are supported; more may be
-            added in the future.
-        ''',
-    )
-    parser.add_argument(
-        '--prop-order',
-        default=[],
-        type=lambda s: s.split(','),
-        help='''
-            comma-separated property names defining the order of precedence for object properties;
-            properties not specified here are given lower precedence than those that are, and
-            are kept in their original order from the schema. Required properties are always
-            given precedence over optional properties.
-        '''
-    )
-    parser.add_argument(
-        '--allow-fetch',
-        action='store_true',
-        default=False,
-        help='Whether to allow fetching referenced schemas over HTTPS')
-    parser.add_argument(
-        '--dotall',
-        action='store_true',
-        default=False,
-        help='Whether to treat dot (".") as matching all chars including line breaks in regular expression patterns')
-    parser.add_argument(
-        '--raw-pattern',
-        action='store_true',
-        default=False,
-        help='Treats string patterns as raw patterns w/o quotes (or quote escapes)')
-
-    parser.add_argument('schema', help='file containing JSON schema ("-" for stdin)')
-    args = parser.parse_args(args_in)
-
-    if args.schema.startswith('https://'):
-        url = args.schema
-        import requests
-        schema = requests.get(url).json()
-    elif args.schema == '-':
-        url = 'stdin'
-        schema = json.load(sys.stdin)
-    else:
-        url = f'file://{args.schema}'
-        with open(args.schema) as f:
-            schema = json.load(f)
-    converter = SchemaConverter(
-        prop_order={name: idx for idx, name in enumerate(args.prop_order)},
-        allow_fetch=args.allow_fetch,
-        dotall=args.dotall,
-        raw_pattern=args.raw_pattern)
-    schema = converter.resolve_refs(schema, url)
-    converter.visit(schema, '')
-    print(converter.format_grammar())
-
-
-if __name__ == '__main__':
-    main()

examples/regex_to_grammar.py

@@ -1,20 +0,0 @@
-import json, subprocess, sys, os
-
-assert len(sys.argv) >= 2
-[_, pattern, *rest] = sys.argv
-
-print(subprocess.check_output(
-    [
-        "python",
-        os.path.join(
-        os.path.dirname(os.path.realpath(__file__)),
-        "json_schema_to_grammar.py"),
-        *rest,
-        "-",
-        "--raw-pattern",
-    ],
-    text=True,
-    input=json.dumps({
-        "type": "string",
-        "pattern": pattern,
-    }, indent=2)))

examples/ts-type-to-grammar.sh

@@ -1,28 +0,0 @@
-#!/usr/bin/env bash
-#
-# ./examples/ts-type-to-grammar.sh "{a:string,b:string,c?:string}"
-# python examples/json_schema_to_grammar.py https://json.schemastore.org/tsconfig.json
-#
-set -euo pipefail
-
-readonly type="$1"
-
-# Create a temporary directory
-TMPDIR=""
-trap 'rm -fR "$TMPDIR"' EXIT
-TMPDIR=$(mktemp -d)
-
-DTS_FILE="$TMPDIR/type.d.ts"
-SCHEMA_FILE="$TMPDIR/schema.json"
-
-echo "export type MyType = $type" > "$DTS_FILE"
-
-# This is a fork of typescript-json-schema, actively maintained as of March 2024:
-# https://github.com/vega/ts-json-schema-generator
-npx ts-json-schema-generator --unstable --no-top-ref --path "$DTS_FILE" --type MyType -e none > "$SCHEMA_FILE"
-
-# Alternative, not actively maintained as of March 2024:
-# https://github.com/YousefED/typescript-json-schema
-# npx typescript-json-schema --defaultProps --required "$DTS_FILE" MyType | tee "$SCHEMA_FILE" >&2
-
-./examples/json_schema_to_grammar.py "$SCHEMA_FILE"

ggml/CMakeLists.txt

@@ -4,7 +4,7 @@ project("ggml" C CXX ASM)
 
 ### GGML Version
 set(GGML_VERSION_MAJOR 0)
-set(GGML_VERSION_MINOR 23)
+set(GGML_VERSION_MINOR 24)
 set(GGML_VERSION_PATCH 0)
 set(GGML_VERSION_BASE "${GGML_VERSION_MAJOR}.${GGML_VERSION_MINOR}.${GGML_VERSION_PATCH}")
 

ggml/src/ggml-backend.cpp

@@ -1705,6 +1705,10 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
                     ggml_tensor * ids_tensor = node->src[2];
                     ggml_backend_t ids_backend = split_backend;
 
+                    if (ggml_nelements(ids_tensor) == 0) {
+                        continue;
+                    }
+
                     // if the ids tensor is also an input of the split, it may not have been copied yet to the split backend
                     // in that case, we use the original ids tensor
                     for (int i = input_id + 1; i < split->n_inputs; i++) {

ggml/src/ggml-cpu/arch/s390/quants.c

@@ -417,6 +417,7 @@ void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo
     sumf = vec_hsum_f32x4(v_acc);
     *s = sumf;
 #else
+    UNUSED(nb);
     UNUSED(x);
     UNUSED(y);
     UNUSED(ib);

ggml/src/ggml-cpu/arch/s390/repack.cpp

@@ -70,6 +70,7 @@ void ggml_quantize_mat_q8_0_4x4(const float * GGML_RESTRICT x, void * GGML_RESTR
 #endif
 }
 
+#if defined(__VXE__) || defined(__VXE2__)
 static inline int16x8_t vxe_dot_acc(const int8x16_t v_x, const int8x16_t v_y, const int16x8_t v_acc) {
     return vec_meadd(v_x, v_y, vec_moadd(v_x, v_y, v_acc));
 }
@@ -84,6 +85,7 @@ static inline int32x4_t vxe_fold(const int16x8_t v_sumi) {
     const int16x8_t v_ones = vec_splats((int16_t)1);
     return vec_add(vec_mule(v_sumi, v_ones), vec_mulo(v_sumi, v_ones));
 }
+#endif
 
 void ggml_gemv_q4_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) {
     const int qk = QK8_0;

ggml/src/ggml-cpu/ops.h

@@ -5,10 +5,10 @@
 //
 // cache line
 //
-
-#if defined(__cpp_lib_hardware_interference_size)
-#define CACHE_LINE_SIZE std::hardware_destructive_interference_size
-#else
+// TODO: rework CACHE_LINE_SIZE so std::hardware_destructive_interference_size
+// can be used consistently between C and C++ TUs; the previous macro form
+// diverged based on include order and undersized the work buffer.
+// ref: https://github.com/ggml-org/llama.cpp/pull/28882
 #if defined(__POWER9_VECTOR__)
 #define CACHE_LINE_SIZE 128
 #elif defined(__VXE__) || defined(__VXE2__)
@@ -16,7 +16,6 @@
 #else
 #define CACHE_LINE_SIZE 64
 #endif
-#endif
 
 static const size_t CACHE_LINE_SIZE_F32 = CACHE_LINE_SIZE/sizeof(float);
 

ggml/src/ggml-cuda/common.cuh

@@ -329,6 +329,12 @@ static bool fp16_mma_hardware_available(const int cc) {
         (GGML_CUDA_CC_IS_MTHREADS(cc) && cc >= GGML_CUDA_CC_QY2);
 }
 
+// To be used for feature selection of external libraries, e.g. cuBLAS.
+static bool fast_bf16_hardware_available(const int cc) {
+        return (GGML_CUDA_CC_IS_AMD(cc) && (cc >= GGML_CUDA_CC_RDNA3 || GGML_CUDA_CC_IS_CDNA(cc)))
+            || (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_AMPERE);
+}
+
 static bool bf16_mma_hardware_available(const int cc) {
     return (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_AMPERE) ||
         GGML_CUDA_CC_IS_CDNA(cc) || cc >= GGML_CUDA_CC_RDNA3 ||

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

@@ -1620,11 +1620,19 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const
 }
 
 static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    const int cc = ggml_cuda_info().devices[ctx.device].cc;
     ggml_type compute_type = src0->type;
     if (ggml_is_quantized(compute_type)) {
-        compute_type = fast_fp16_hardware_available(ggml_cuda_info().devices[ctx.device].cc) ? GGML_TYPE_F16 : GGML_TYPE_F32;
-    } else if (compute_type == GGML_TYPE_F16 && !fast_fp16_hardware_available(ggml_cuda_info().devices[ctx.device].cc)) {
+        compute_type = fast_fp16_hardware_available(cc) ? GGML_TYPE_F16 : GGML_TYPE_F32;
+    } else if (compute_type == GGML_TYPE_F16 && !fast_fp16_hardware_available(cc)) {
         compute_type = GGML_TYPE_F32;
+    } else if (compute_type == GGML_TYPE_BF16 && !fast_bf16_hardware_available(cc)) {
+        if (GGML_CUDA_CC_IS_AMD(cc) && src1->ne[1] > 32) {
+            compute_type = GGML_TYPE_F32;
+        }
+        if (GGML_CUDA_CC_IS_NVIDIA(cc) && src1->ne[1] > (cc >= GGML_CUDA_CC_VOLTA ? 8 : 128)) {
+            compute_type = GGML_TYPE_F32;
+        }
     }
     if (dst->op_params[0] == GGML_PREC_F32) {
         compute_type = GGML_TYPE_F32;

ggml/src/ggml-cuda/mmq-config-gcn.cuh

@@ -0,0 +1,281 @@
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_gcn(ggml_type type, int J, bool fallback) {
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2,  64,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_Q2_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 3,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_MXFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}

ggml/src/ggml-cuda/mmq.cuh

@@ -218,6 +218,7 @@ struct ggml_cuda_mmq_config {
 #include "mmq-config-ampere.cuh"
 #include "mmq-config-blackwell.cuh"
 
+#include "mmq-config-gcn.cuh"
 #include "mmq-config-cdna.cuh"
 #include "mmq-config-rdna2.cuh"
 #include "mmq-config-rdna3.cuh"
@@ -228,6 +229,9 @@ struct ggml_cuda_mmq_config {
 
 static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type type, const int J, const bool fallback, const int cc) {
     if (GGML_CUDA_CC_IS_AMD(cc)) {
+        if (GGML_CUDA_CC_IS_GCN(cc)) {
+            return ggml_cuda_mmq_get_config_gcn(type, J, fallback);
+        }
         if (GGML_CUDA_CC_IS_CDNA(cc)) {
             return ggml_cuda_mmq_get_config_cdna(type, J, fallback);
         }
@@ -256,7 +260,9 @@ static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type ty
 
 static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_type type, int J, bool fallback) {
 #ifdef GGML_USE_HIP
-#ifdef CDNA
+#ifdef GCN
+    return ggml_cuda_mmq_get_config_gcn(type, J, fallback);
+#elif defined(CDNA)
     return ggml_cuda_mmq_get_config_cdna(type, J, fallback);
 #elif defined(RDNA4)
     return ggml_cuda_mmq_get_config_rdna4(type, J, fallback);

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

@@ -344,6 +344,12 @@ struct htp_opformat {
         } else if (htp_op_is_unary(node.opcode)) {
             const auto * kparams = (const struct htp_unary_kernel_params *) node.kernel_params;
             snprintf(str, max_size, "%s vtcm %d", kparams->col_tile ? "wide-row" : "row-block", (int) kparams->vtcm_size);
+        } else if (node.opcode == HTP_OP_MDEV_GROUP && node.node) {
+            snprintf(str, max_size, "idx %d count %d", (int) node.node->op_params[0], (int) node.dst()->ne[1]);
+        } else if ((node.opcode == HTP_OP_FENCE || node.opcode == HTP_OP_CPY_FENCE) && node.node) {
+            snprintf(str, max_size, "seq 0x%x", (uint32_t) node.node->op_params[0]);
+        } else if (node.opcode == HTP_OP_ALLREDUCE && node.node) {
+            snprintf(str, max_size, "seq 0x%x -> 0x%x", (uint32_t) node.node->op_params[0], (uint32_t) node.node->op_params[1]);
         } else {
             snprintf(str, max_size, "----");
         }

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

@@ -3,7 +3,6 @@
 #pragma clang diagnostic ignored "-Wunused-but-set-variable"
 
 #include <HAP_farf.h>
-#include <HAP_perf.h>
 
 #include <math.h>
 #include <string.h>
@@ -15,7 +14,7 @@
 #include "ggml-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
-#include "htp-ops.h"
+#include "hex-common.h"
 #include "htp-tensor.h"
 #include "htp-vtcm.h"
 
@@ -80,6 +79,7 @@ struct htp_act_context {
     uint32_t                 block;
     uint32_t                 src0_nrows;
     uint32_t                 src0_nrows_per_thread;
+    uint32_t                 row_start;
     int                      nc;
 
     uint8_t *                vtcm_src0;
@@ -329,104 +329,104 @@ static void geglu_f32(const float * restrict src0,
     }
 }
 
-#define DEFINE_GLU_PER_THREAD(NAME, OP_STR, CORE_EXPR)                                                                 \
-    static void glu_##NAME##_f32_per_thread(unsigned int nth, unsigned int ith, void * data) {                         \
-        struct htp_act_context * actx = (struct htp_act_context *) data;                                               \
-        htp_act_preamble;                                                                                              \
-                                                                                                                       \
-        struct htp_thread_trace * tr = actx->octx->ctx ? &actx->octx->ctx->trace[ith] : NULL;                          \
-                                                                                                                       \
-        size_t src0_row_size = actx->src0_row_size;                                                                    \
-        size_t src1_row_size = actx->src1_row_size;                                                                    \
-        size_t dst_row_size  = actx->dst_row_size;                                                                     \
-                                                                                                                       \
-        size_t src0_row_stride = actx->src0_row_stride;                                                                \
-        size_t src1_row_stride = actx->src1_row_stride;                                                                \
-                                                                                                                       \
-        const uint32_t src0_nrows            = actx->src0_nrows;                                                       \
-        const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread;                                            \
-                                                                                                                       \
-        const uint32_t src0_start_row = src0_nrows_per_thread * ith;                                                   \
-        const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);                       \
-                                                                                                                       \
-        /* no work for this thread */                                                                                  \
-        if (src0_start_row >= src0_end_row) {                                                                          \
-            return;                                                                                                    \
-        }                                                                                                              \
-                                                                                                                       \
-        const uint8_t * restrict data_src0 = actx->data_src0;                                                          \
-        const uint8_t * restrict data_src1 = actx->data_src1;                                                          \
-        uint8_t * restrict data_dst        = actx->data_dst;                                                           \
-                                                                                                                       \
-        const size_t src0_row_size_aligned = actx->src0_row_size_aligned;                                              \
-        const size_t src1_row_size_aligned = actx->src1_row_size_aligned;                                              \
-        const size_t dst_row_size_aligned  = actx->dst_row_size_aligned;                                               \
-                                                                                                                       \
-        uint8_t * restrict src0_spad_data = actx->vtcm_src0 + (ith * actx->vtcm_src0_size_per_thread);                 \
-        uint8_t * restrict src1_spad_data = actx->vtcm_src1 + (ith * actx->vtcm_src1_size_per_thread);                 \
-        uint8_t * restrict dst_spad_data  = actx->vtcm_dst  + (ith * actx->vtcm_dst_size_per_thread);                  \
-                                                                                                                       \
-        size_t src0_spad_half_size = actx->src0_spad_half_size;                                                        \
-        size_t src1_spad_half_size = actx->src1_spad_half_size;                                                        \
-        size_t dst_spad_half_size  = actx->dst_spad_half_size;                                                         \
-                                                                                                                       \
-        const int BLOCK = actx->block;                                                                                 \
-        if (BLOCK == 0) {                                                                                              \
-            FARF(ERROR,                                                                                                \
-                 OP_STR                                                                                                \
-                 " : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n",      \
-                 actx->vtcm_src0_size_per_thread, src0_row_size_aligned);                                              \
-            return;                                                                                                    \
-        }                                                                                                              \
-                                                                                                                       \
-        dma_queue * dma_queue = actx->octx->ctx->dma[ith];                                                             \
-                                                                                                                       \
-        /* See discussion: https://github.com/ggml-org/llama.cpp/pull/18151#issuecomment-3678235379 */                 \
-        for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) { \
-            const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);                                                 \
-                                                                                                                       \
-            /* Dummy DMA transation for sequencing (interleaving dst,src,dst,...) */                                   \
-            dma_queue_push_vtcm_to_ddr(dma_queue,                                                                      \
-                                       dma_make_ptr(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)),        \
-                                       dst_row_size, dst_row_size_aligned, 0);                                         \
-                                                                                                                       \
-            dma_queue_push(                                                                                            \
-                dma_queue,                                                                                             \
-                dma_make_ptr(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_stride)),   \
-                src0_row_size_aligned, src0_row_stride, src0_row_size, block_size);                                    \
-            dma_queue_push(                                                                                            \
-                dma_queue,                                                                                             \
-                dma_make_ptr(src1_spad_data + (spad_idx * src1_spad_half_size), data_src1 + (ir * src1_row_stride)),   \
-                src1_row_size_aligned, src1_row_stride, src1_row_size, block_size);                                    \
-        }                                                                                                              \
-                                                                                                                       \
-        for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) {                                           \
-            const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);                                                 \
-                                                                                                                       \
-            float * dst_spad  = (float *) dma_queue_pop(dma_queue).src;                                                \
-            float * src0_spad = (float *) dma_queue_pop(dma_queue).dst;                                                \
-            float * src1_spad = (float *) dma_queue_pop(dma_queue).dst;                                                \
-                                                                                                                       \
-            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir);                                                     \
-            CORE_EXPR;                                                                                                 \
-            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir);                                                      \
-                                                                                                                       \
-            dma_queue_push_vtcm_to_ddr(dma_queue, dma_make_ptr(data_dst + (ir * dst_row_size), dst_spad),              \
-                                       dst_row_size, dst_row_size_aligned, block_size);                                \
-                                                                                                                       \
-            /* prefetch N+2 loop iteration if any */                                                                   \
-            const uint32_t pref_block = (ir + BLOCK * 2);                                                              \
-            if (pref_block < src0_end_row) {                                                                           \
-                const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block);                                \
-                dma_queue_push(dma_queue, dma_make_ptr(src0_spad, data_src0 + (pref_block * src0_row_stride)),         \
-                               src0_row_size_aligned, src0_row_stride, src0_row_size, pref_block_size);                \
-                dma_queue_push(dma_queue, dma_make_ptr(src1_spad, data_src1 + (pref_block * src1_row_stride)),         \
-                               src1_row_size_aligned, src1_row_stride, src1_row_size, pref_block_size);                \
-            }                                                                                                          \
-        }                                                                                                              \
-                                                                                                                       \
-        dma_queue_flush(dma_queue);                                                                                    \
-                                                                                                                       \
+#define DEFINE_GLU_PER_THREAD(NAME, OP_STR, CORE_EXPR)                                                                   \
+    static void glu_##NAME##_f32_per_thread(unsigned int nth, unsigned int ith, void * data) {                           \
+        struct htp_act_context * actx = (struct htp_act_context *) data;                                                 \
+        htp_act_preamble;                                                                                                \
+                                                                                                                         \
+        struct htp_thread_trace * tr = actx->octx->ctx ? &actx->octx->ctx->trace[ith] : NULL;                            \
+                                                                                                                         \
+        size_t src0_row_size = actx->src0_row_size;                                                                      \
+        size_t src1_row_size = actx->src1_row_size;                                                                      \
+        size_t dst_row_size  = actx->dst_row_size;                                                                       \
+                                                                                                                         \
+        size_t src0_row_stride = actx->src0_row_stride;                                                                  \
+        size_t src1_row_stride = actx->src1_row_stride;                                                                  \
+                                                                                                                         \
+        const uint32_t src0_nrows            = actx->src0_nrows;                                                         \
+        const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread;                                              \
+                                                                                                                         \
+        const uint32_t src0_start_row = actx->row_start + src0_nrows_per_thread * ith;                                   \
+        const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, actx->row_start + src0_nrows);       \
+                                                                                                                         \
+        /* no work for this thread */                                                                                    \
+        if (src0_start_row >= src0_end_row) {                                                                            \
+            return;                                                                                                      \
+        }                                                                                                                \
+                                                                                                                         \
+        const uint8_t * restrict data_src0 = actx->data_src0;                                                            \
+        const uint8_t * restrict data_src1 = actx->data_src1;                                                            \
+        uint8_t * restrict data_dst        = actx->data_dst;                                                             \
+                                                                                                                         \
+        const size_t src0_row_size_aligned = actx->src0_row_size_aligned;                                                \
+        const size_t src1_row_size_aligned = actx->src1_row_size_aligned;                                                \
+        const size_t dst_row_size_aligned  = actx->dst_row_size_aligned;                                                 \
+                                                                                                                         \
+        uint8_t * restrict src0_spad_data = actx->vtcm_src0 + (ith * actx->vtcm_src0_size_per_thread);                   \
+        uint8_t * restrict src1_spad_data = actx->vtcm_src1 + (ith * actx->vtcm_src1_size_per_thread);                   \
+        uint8_t * restrict dst_spad_data  = actx->vtcm_dst  + (ith * actx->vtcm_dst_size_per_thread);                    \
+                                                                                                                         \
+        size_t src0_spad_half_size = actx->src0_spad_half_size;                                                          \
+        size_t src1_spad_half_size = actx->src1_spad_half_size;                                                          \
+        size_t dst_spad_half_size  = actx->dst_spad_half_size;                                                           \
+                                                                                                                         \
+        const int BLOCK = actx->block;                                                                                   \
+        if (BLOCK == 0) {                                                                                                \
+            FARF(ERROR,                                                                                                  \
+                 OP_STR                                                                                                  \
+                 " : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n",        \
+                 actx->vtcm_src0_size_per_thread, src0_row_size_aligned);                                                \
+            return;                                                                                                      \
+        }                                                                                                                \
+                                                                                                                         \
+        dma_queue * dma_queue = actx->octx->ctx->dma[ith];                                                               \
+                                                                                                                         \
+        /* See discussion: https://github.com/ggml-org/llama.cpp/pull/18151#issuecomment-3678235379 */                   \
+        for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) {   \
+            const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);                                                   \
+                                                                                                                         \
+            /* Dummy DMA transation for sequencing (interleaving dst,src,dst,...) */                                     \
+            dma_queue_push_vtcm_to_ddr(dma_queue,                                                                        \
+                                       dma_make_ptr(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)),          \
+                                       dst_row_size, dst_row_size_aligned, 0);                                           \
+                                                                                                                         \
+            dma_queue_push(                                                                                              \
+                dma_queue,                                                                                               \
+                dma_make_ptr(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_stride)),     \
+                src0_row_size_aligned, src0_row_stride, src0_row_size, block_size);                                      \
+            dma_queue_push(                                                                                              \
+                dma_queue,                                                                                               \
+                dma_make_ptr(src1_spad_data + (spad_idx * src1_spad_half_size), data_src1 + (ir * src1_row_stride)),     \
+                src1_row_size_aligned, src1_row_stride, src1_row_size, block_size);                                      \
+        }                                                                                                                \
+                                                                                                                         \
+        for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) {                                             \
+            const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);                                                   \
+                                                                                                                         \
+            float * dst_spad  = (float *) dma_queue_pop(dma_queue).src;                                                  \
+            float * src0_spad = (float *) dma_queue_pop(dma_queue).dst;                                                  \
+            float * src1_spad = (float *) dma_queue_pop(dma_queue).dst;                                                  \
+                                                                                                                         \
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir);                                                       \
+            CORE_EXPR;                                                                                                   \
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir);                                                        \
+                                                                                                                         \
+            dma_queue_push_vtcm_to_ddr(dma_queue, dma_make_ptr(data_dst + (ir * dst_row_size), dst_spad),                \
+                                       dst_row_size, dst_row_size_aligned, block_size);                                  \
+                                                                                                                         \
+            /* prefetch N+2 loop iteration if any */                                                                     \
+            const uint32_t pref_block = (ir + BLOCK * 2);                                                                \
+            if (pref_block < src0_end_row) {                                                                             \
+                const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block);                                  \
+                dma_queue_push(dma_queue, dma_make_ptr(src0_spad, data_src0 + (pref_block * src0_row_stride)),           \
+                               src0_row_size_aligned, src0_row_stride, src0_row_size, pref_block_size);                  \
+                dma_queue_push(dma_queue, dma_make_ptr(src1_spad, data_src1 + (pref_block * src1_row_stride)),           \
+                               src1_row_size_aligned, src1_row_stride, src1_row_size, pref_block_size);                  \
+            }                                                                                                            \
+        }                                                                                                                \
+                                                                                                                         \
+        dma_queue_flush(dma_queue);                                                                                      \
+                                                                                                                         \
     }
 
 DEFINE_GLU_PER_THREAD(swiglu, "swiglu-f32", swiglu_f32(src0_spad, src1_spad, dst_spad, block_size, actx))
@@ -473,14 +473,30 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
     }
 
     const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
-    const uint32_t n_threads  = MIN(octx->n_threads, src0_nrows);
+    const size_t dst_row_size = dst->ne[0] * SIZEOF_FP32;
+
+    uint32_t row_start = 0;
+    uint32_t nrows     = src0_nrows;
+
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t rows_per_chunk = 0;
+        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_row_size, &rows_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
 
     // row_size   = bytes of useful data per row (what the kernel touches / what DMA copies).
     // row_stride = bytes between successive rows in DDR (may exceed row_size for non-contig src).
-    const size_t nc_bytes    = dst->ne[0] * SIZEOF_FP32;
-    const size_t src0_row_size = nc_bytes;
-    const size_t src1_row_size = nc_bytes;
-    const size_t dst_row_size  = nc_bytes;
+    const size_t nc_bytes        = dst_row_size;
+    const size_t src0_row_size   = nc_bytes;
+    const size_t src1_row_size   = nc_bytes;
     const size_t src0_row_stride = src0->nb[1];
     const size_t src1_row_stride = src1 ? src1->nb[1] : src0->nb[1];
 
@@ -518,7 +534,7 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
     struct htp_act_context actx;
     actx.octx = octx;
 
-    actx.src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads;
+    actx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
 
     actx.src0_row_size = src0_row_size;
     actx.src1_row_size = src1_row_size;
@@ -545,7 +561,8 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
     actx.dst_spad_half_size  = L.dst_bytes_per_thread / 2;
 
     actx.block = actx.src0_spad_half_size / actx.src0_row_size_aligned;
-    actx.src0_nrows = src0_nrows;
+    actx.src0_nrows = nrows;
+    actx.row_start  = row_start;
 
     actx.nc = dst->ne[0];
 
@@ -570,7 +587,7 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
     actx.data_src1 = data_src1;
     actx.data_dst  = (uint8_t *) dst->data;
 
-    worker_pool_run_func(octx->ctx->worker_pool, act_op_func, &actx, n_threads);
+    work_queue_run(octx->ctx->work_queue, act_op_func, &actx, n_threads);
     return HTP_STATUS_OK;
 }
 

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

@@ -17,6 +17,7 @@
 #include "hex-dma.h"
 #include "hex-profile.h"
 #include "allreduce-ops.h"
+#include "htp-fence.h"
 
 struct htp_allreduce_context {
     struct htp_ops_context * octx;
@@ -242,7 +243,42 @@ DEFINE_ALLREDUCE_THREAD_DMA_2D(add_f32,       float,  hvx_add_f32_aaa, 1, 0)
 DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f16, __fp16, hvx_add_f16_aaa, 1, 1)
 DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f32, float,  hvx_add_f32_aaa, 1, 1)
 
+static int validate_allreduce(
+    struct htp_ops_context * octx,
+    const struct htp_allreduce_kernel_params * kparams,
+    uint32_t n_ranks
+) {
+    if (!htp_ops_context_set_n_threads(octx, (uint32_t) kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    if (kparams->vtcm_size_per_thread <= 0 || kparams->vtcm_size <= 0) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD);
+    const size_t n_vtcm_buffers = htp_allreduce_vtcm_buffer_count(
+        n_ranks, octx->n_threads, has_add, kparams->is_row_bcast != 0);
+    const size_t vtcm_size = n_vtcm_buffers * (size_t) kparams->vtcm_size_per_thread;
+    if (vtcm_size != (size_t) kparams->vtcm_size) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+    if (vtcm_size > octx->ctx->vtcm_size) {
+        return HTP_STATUS_VTCM_TOO_SMALL;
+    }
+
+    if (octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_F32) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    return HTP_STATUS_OK;
+}
+
 int op_allreduce(struct htp_ops_context * octx) {
+    if (octx->ctx->mdev.count > 1 && octx->ctx->mdev.idx > 0) {
+        return HTP_STATUS_OK;
+    }
+
     const struct htp_allreduce_kernel_params * kparams = (const struct htp_allreduce_kernel_params *) octx->kernel_params;
     const struct htp_tensor * dst = octx->dst;
 
@@ -253,38 +289,53 @@ int op_allreduce(struct htp_ops_context * octx) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
-    if (dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_F32) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    const uint32_t nelem = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3];
     const uint32_t fence_seq_entry = (uint32_t) octx->op_params[0];
     const uint32_t fence_seq_exit  = (uint32_t) octx->op_params[1];
 
+    const struct htp_tensor * my_sync = octx->src[n_ranks + rank];
+    atomic_uint * my_fence = (atomic_uint *) (uintptr_t) my_sync->data;
+
+    const int status = validate_allreduce(octx, kparams, n_ranks);
+    if (status != HTP_STATUS_OK) {
+        if (status == HTP_STATUS_NO_SUPPORT) {
+            FARF(ERROR, "ggml-hex: allreduce unsupported type %d : rank %u\n", dst->type, rank);
+        }
+        htp_fence_write(my_fence, fence_seq_exit, status);
+        return status;
+    }
+
+    const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD);
+    const uint32_t nelem = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3];
+
     // 1. Entry Barrier: Synchronize all ranks before reading
     struct htp_thread_trace * tr0 = &octx->ctx->trace[0];
     htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
 
-    const struct htp_tensor * my_sync = octx->src[n_ranks + rank];
-    atomic_uint * my_fence = (atomic_uint *) my_sync->data;
-
-    atomic_store(&my_fence[0], fence_seq_entry);
-    asm volatile ("syncht" : : : "memory");
-    Q6_dccleaninva_A((void *) my_fence);
+    htp_fence_write(my_fence, fence_seq_entry, octx->status);
 
     for (uint32_t j = 0; j < n_ranks; j++) {
         if (j == rank) continue;
         const struct htp_tensor * peer_sync = octx->src[n_ranks + j];
-        atomic_uint * peer_fence = (atomic_uint *) peer_sync->data;
+        atomic_uint * peer_fence = (atomic_uint *) (uintptr_t) peer_sync->data;
         uint64_t spins = 0;
         while (1) {
-            Q6_dccleaninva_A((void *) peer_fence);
-            uint32_t val = atomic_load(&peer_fence[0]);
-            if (val == fence_seq_entry || val == fence_seq_exit) {
+            uint32_t peer_seq;
+            uint32_t peer_status;
+            htp_fence_read(peer_fence, &peer_seq, &peer_status);
+            if ((int32_t)(peer_seq - fence_seq_entry) >= 0) {
+                if (peer_status > HTP_STATUS_OK) {
+                    FARF(ERROR, "ggml-hex: allreduce entry peer %u failed with status %u\n", j, peer_status);
+                    htp_fence_write(my_fence, fence_seq_exit, peer_status);
+                    htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
+                    return peer_status;
+                }
                 break;
             }
             if (++spins > HTP_FENCE_TIMEOUT) {
-                FARF(ERROR, "ggml-hex: allreduce entry fence-wait TIMEOUT: rank %u waiting on %u (fence %p seq %u)\n", rank, j, peer_fence, fence_seq_entry);
+                FARF(ERROR, "ggml-hex: allreduce entry fence-wait TIMEOUT : rank %u waiting on %u fence %p seq 0x%x peer-seq 0x%x\n",
+                     rank, j, peer_fence, fence_seq_entry, peer_seq);
+                htp_fence_write(my_fence, fence_seq_exit, HTP_STATUS_INTERNAL_ERR);
+                htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
                 return HTP_STATUS_INTERNAL_ERR;
             }
             hex_pause();
@@ -301,8 +352,6 @@ int op_allreduce(struct htp_ops_context * octx) {
         const uint32_t elems_per_thread     = (uint32_t) kparams->elems_per_thread;
         const uint32_t vtcm_size_per_thread = (uint32_t) kparams->vtcm_size_per_thread;
 
-        const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD);
-
         struct htp_allreduce_context actx;
         actx.octx                 = octx;
         actx.n_ranks              = n_ranks;
@@ -339,6 +388,8 @@ int op_allreduce(struct htp_ops_context * octx) {
                 }
                 break;
             default:
+                FARF(ERROR, "ggml-hex: allreduce unsupported kernel %d : rank %u\n", kparams->kernel_type, rank);
+                htp_fence_write(my_fence, fence_seq_exit, HTP_STATUS_NO_SUPPORT);
                 return HTP_STATUS_NO_SUPPORT;
         }
 
@@ -368,23 +419,31 @@ int op_allreduce(struct htp_ops_context * octx) {
     // 4. Exit Barrier: Synchronize all ranks after writing
     htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
 
-    atomic_store(&my_fence[0], fence_seq_exit);
-    asm volatile ("syncht" : : : "memory");
-    Q6_dccleaninva_A((void *) my_fence);
+    htp_fence_write(my_fence, fence_seq_exit, octx->status);
 
     for (uint32_t j = 0; j < n_ranks; j++) {
         if (j == rank) continue;
         const struct htp_tensor * peer_sync = octx->src[n_ranks + j];
-        atomic_uint * peer_fence = (atomic_uint *) peer_sync->data;
+        atomic_uint * peer_fence = (atomic_uint *) (uintptr_t) peer_sync->data;
         uint64_t spins = 0;
         while (1) {
-            Q6_dccleaninva_A((void *) peer_fence);
-            uint32_t val = atomic_load(&peer_fence[0]);
-            if (val == fence_seq_exit) {
+            uint32_t peer_seq;
+            uint32_t peer_status;
+            htp_fence_read(peer_fence, &peer_seq, &peer_status);
+            if ((int32_t)(peer_seq - fence_seq_exit) >= 0) {
+                if (peer_status > HTP_STATUS_OK) {
+                    FARF(ERROR, "ggml-hex: allreduce exit peer %u failed with status %u\n", j, peer_status);
+                    htp_fence_write(my_fence, fence_seq_exit, peer_status);
+                    htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
+                    return peer_status;
+                }
                 break;
             }
             if (++spins > HTP_FENCE_TIMEOUT) {
-                FARF(ERROR, "ggml-hex: allreduce exit fence-wait TIMEOUT: rank %u waiting on %u (fence %p seq %u)\n", rank, j, peer_fence, fence_seq_exit);
+                FARF(ERROR, "ggml-hex: allreduce exit fence-wait TIMEOUT : rank %u waiting on %u fence %p seq 0x%x peer-seq 0x%x\n",
+                     rank, j, peer_fence, fence_seq_exit, peer_seq);
+                htp_fence_write(my_fence, fence_seq_exit, HTP_STATUS_INTERNAL_ERR);
+                htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
                 return HTP_STATUS_INTERNAL_ERR;
             }
             hex_pause();
@@ -394,5 +453,5 @@ int op_allreduce(struct htp_ops_context * octx) {
 
     htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
 
-    return HTP_STATUS_OK;
+    return octx->status;
 }

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

@@ -2,6 +2,8 @@
 #define ALLREDUCE_OPS_H
 
 #include <stdint.h>
+#include <stddef.h>
+#include <stdbool.h>
 
 #define HTP_ALLREDUCE_MAX_RANKS 4
 
@@ -15,6 +17,15 @@ enum htp_allreduce_kernel_type {
     HTP_ALLREDUCE_KERNEL_DMA_2D,
 };
 
+static inline size_t htp_allreduce_vtcm_buffer_count(
+    uint32_t n_ranks,
+    uint32_t n_threads,
+    bool has_add,
+    bool is_row_bcast
+) {
+    return (size_t) (n_ranks + 1) * n_threads + (has_add ? (is_row_bcast ? 1 : n_threads) : 0);
+}
+
 struct htp_allreduce_kernel_params {
     int32_t rank;
     int32_t n_ranks;

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

@@ -11,9 +11,10 @@
 #include "hvx-utils.h"
 #include "hex-dma.h"
 
+#include "hex-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
-#include "htp-ops.h"
+#include "htp-tensor.h"
 
 #ifndef MIN
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -22,6 +23,9 @@
 struct htp_argsort_context {
     struct htp_ops_context * octx;
     uint32_t                 nrows_per_thread;
+    uint32_t                 total_rows;
+    uint32_t                 row_start;
+    uint32_t                 row_end;
     uint8_t *                vtcm_base;
     size_t                   vtcm_per_thread;
 };
@@ -336,10 +340,9 @@ static void htp_argsort_f32_##ne00##_##order_name(unsigned int n, unsigned int i
     const struct htp_tensor * src0 = octx->src[0];                                                             \
     const struct htp_tensor * dst = octx->dst;                                                                 \
     uint8_t * spad = actx->vtcm_base + actx->vtcm_per_thread * i;                                              \
-    uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3];                                             \
     uint32_t rows_per_thread = actx->nrows_per_thread;                                                         \
-    uint32_t start_row = rows_per_thread * i;                                                                  \
-    uint32_t end_row = MIN(start_row + rows_per_thread, total_rows);                                           \
+    uint32_t start_row = actx->row_start + rows_per_thread * i;                                                \
+    uint32_t end_row = MIN(start_row + rows_per_thread, actx->row_end);                                        \
     size_t values_size = hex_round_up(ne00 * sizeof(float), 128);                                              \
     float * values_buf = (float *) spad;                                                                       \
     int32_t * indices_buf = (int32_t *) (spad + values_size);                                                  \
@@ -386,9 +389,6 @@ static void htp_argsort_f32_fallback(unsigned int n, unsigned int i, void * data
 
     // Dimensions
     uint32_t ne00 = src0->ne[0];
-    uint32_t ne01 = src0->ne[1];
-    uint32_t ne02 = src0->ne[2];
-    uint32_t ne03 = src0->ne[3];
 
     uint32_t nb01 = src0->nb[1];
 
@@ -398,10 +398,9 @@ static void htp_argsort_f32_fallback(unsigned int n, unsigned int i, void * data
     enum ggml_sort_order order = (enum ggml_sort_order) octx->op_params[0];
 
     // Rows to process
-    uint32_t total_rows = ne01 * ne02 * ne03;
     uint32_t rows_per_thread = actx->nrows_per_thread;
-    uint32_t start_row = rows_per_thread * i;
-    uint32_t end_row = MIN(start_row + rows_per_thread, total_rows);
+    uint32_t start_row = actx->row_start + rows_per_thread * i;
+    uint32_t end_row = MIN(start_row + rows_per_thread, actx->row_end);
 
     size_t values_size = hex_round_up(ne00 * sizeof(float), 128);
     uint32_t num_vec_ind_values = hmx_ceil_div(ne00, VLEN/(sizeof(int32_t)));
@@ -451,8 +450,28 @@ int op_argsort(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    const uint32_t total_rows = octx->src[0]->ne[1] * octx->src[0]->ne[2] * octx->src[0]->ne[3];
-    const uint32_t n_threads = MIN(total_rows, octx->n_threads);
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * dst  = octx->dst;
+
+    const uint32_t total_rows  = src0->ne[1] * src0->ne[2] * src0->ne[3];
+    const size_t dst_row_size  = dst->ne[0]  * sizeof(int32_t);
+
+    uint32_t row_start = 0;
+    uint32_t row_end   = total_rows;
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t rows_per_chunk = 0;
+        htp_tensor_mdev_rows_per_chunk(dst, sizeof(int32_t), (uint32_t) dst_row_size, &rows_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        row_end   = range.start + range.count;
+    }
+
+    const uint32_t nrows = row_end - row_start;
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
 
     // Allocate scratchpad
     // We need 1 row of float + 1 row of int32 per thread.
@@ -478,7 +497,10 @@ int op_argsort(struct htp_ops_context * octx) {
 
     struct htp_argsort_context actx;
     actx.octx = octx;
-    actx.nrows_per_thread = (total_rows + n_threads - 1) / n_threads;
+    actx.nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    actx.total_rows       = nrows;
+    actx.row_start        = row_start;
+    actx.row_end          = row_end;
     actx.vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
     actx.vtcm_per_thread = spad_per_thread;
 
@@ -508,7 +530,7 @@ int op_argsort(struct htp_ops_context * octx) {
     }
 
     // Run jobs
-    worker_pool_run_func(octx->ctx->worker_pool, job_func, &actx, n_threads);
+    work_queue_run(octx->ctx->work_queue, job_func, &actx, n_threads);
 
     return HTP_STATUS_OK;
 }

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

@@ -13,9 +13,10 @@
 
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
+#include "hex-common.h"
+#include "hex-profile.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
-#include "htp-ops.h"
 #include "htp-tensor.h"
 
 #ifndef MIN
@@ -36,6 +37,8 @@ struct htp_binary_context {
 
     uint32_t block_max;
     uint32_t nrows_per_thread;
+    uint32_t total_rows;
+    uint32_t row_start;
     size_t   src0_row_size_aligned;
     size_t   src1_row_size_aligned;
     size_t   dst_row_size_aligned;
@@ -48,27 +51,27 @@ struct htp_binary_context {
     const struct htp_tensor * src0 = octx->src[0]; \
     const struct htp_tensor * src1 = octx->src[1]; \
     const struct htp_tensor * dst  = octx->dst;    \
-                                       \
-    const uint32_t ne00 = src0->ne[0]; \
-    const uint32_t ne01 = src0->ne[1]; \
-    const uint32_t ne02 = src0->ne[2]; \
-    const uint32_t ne03 = src0->ne[3]; \
-                                       \
-    const uint32_t ne10 = src1->ne[0]; \
-    const uint32_t ne11 = src1->ne[1]; \
-    const uint32_t ne12 = src1->ne[2]; \
-    const uint32_t ne13 = src1->ne[3]; \
-                                       \
-    const uint32_t nb01 = src0->nb[1]; \
-    const uint32_t nb02 = src0->nb[2]; \
-    const uint32_t nb03 = src0->nb[3]; \
-                                       \
-    const uint32_t nb11 = src1->nb[1]; \
-    const uint32_t nb12 = src1->nb[2]; \
-    const uint32_t nb13 = src1->nb[3]; \
-                                       \
-    const uint32_t nb1 = dst->nb[1];   \
-    const uint32_t nb2 = dst->nb[2];   \
+                                                   \
+    const uint32_t ne00 = src0->ne[0];             \
+    const uint32_t ne01 = src0->ne[1];             \
+    const uint32_t ne02 = src0->ne[2];             \
+    const uint32_t ne03 = src0->ne[3];             \
+                                                   \
+    const uint32_t ne10 = src1->ne[0];             \
+    const uint32_t ne11 = src1->ne[1];             \
+    const uint32_t ne12 = src1->ne[2];             \
+    const uint32_t ne13 = src1->ne[3];             \
+                                                   \
+    const uint32_t nb01 = src0->nb[1];             \
+    const uint32_t nb02 = src0->nb[2];             \
+    const uint32_t nb03 = src0->nb[3];             \
+                                                   \
+    const uint32_t nb11 = src1->nb[1];             \
+    const uint32_t nb12 = src1->nb[2];             \
+    const uint32_t nb13 = src1->nb[3];             \
+                                                   \
+    const uint32_t nb1 = dst->nb[1];               \
+    const uint32_t nb2 = dst->nb[2];               \
     const uint32_t nb3 = dst->nb[3];
 
 static inline uint32_t calc_block_size(struct htp_binary_context * bctx, uint32_t ir, uint32_t end_row, uint32_t ne01, uint32_t ne02) {
@@ -93,87 +96,87 @@ static inline uint32_t calc_block_size(struct htp_binary_context * bctx, uint32_
 }
 
 // Macro for scalar op switch
-#define COMPUTE_SCALAR_OP(DST, SRC, VAL, TYPE, N) \
-    if(TYPE == HTP_TYPE_F32) { \
-        switch (octx->op) { \
-            case HTP_OP_ADD: hvx_add_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break; \
-            case HTP_OP_SUB: hvx_sub_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break; \
-            case HTP_OP_MUL: hvx_mul_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break; \
+#define COMPUTE_SCALAR_OP(DST, SRC, VAL, TYPE, N)                                               \
+    if(TYPE == HTP_TYPE_F32) {                                                                  \
+        switch (octx->op) {                                                                     \
+            case HTP_OP_ADD: hvx_add_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break;          \
+            case HTP_OP_SUB: hvx_sub_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break;          \
+            case HTP_OP_MUL: hvx_mul_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break;          \
             case HTP_OP_DIV: hvx_mul_scalar_f32_aa(DST, SRC, 1.0f / (*(float *)VAL), N); break; \
-            default: break; \
-        } \
-    } \
-    else { \
-        switch (octx->op) { \
-            case HTP_OP_ADD: hvx_add_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break; \
-            case HTP_OP_SUB: hvx_sub_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break; \
-            case HTP_OP_MUL: hvx_mul_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break; \
-            case HTP_OP_DIV: hvx_div_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break; \
-            default: break; \
-        } \
+            default: break;                                                                     \
+        }                                                                                       \
+    }                                                                                           \
+    else {                                                                                      \
+        switch (octx->op) {                                                                     \
+            case HTP_OP_ADD: hvx_add_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break;       \
+            case HTP_OP_SUB: hvx_sub_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break;       \
+            case HTP_OP_MUL: hvx_mul_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break;       \
+            case HTP_OP_DIV: hvx_div_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break;       \
+            default: break;                                                                     \
+        }                                                                                       \
     }
 
 // Macro for vector op switch (All Aligned)
-#define COMPUTE_VECTOR_OP_AAA(DST, SRC0, SRC1, TYPE, N) \
-    if(TYPE == HTP_TYPE_F32) { \
-        switch (octx->op) { \
+#define COMPUTE_VECTOR_OP_AAA(DST, SRC0, SRC1, TYPE, N)                  \
+    if(TYPE == HTP_TYPE_F32) {                                           \
+        switch (octx->op) {                                              \
             case HTP_OP_ADD: hvx_add_f32_aaa(DST, SRC0, SRC1, N); break; \
             case HTP_OP_SUB: hvx_sub_f32_aaa(DST, SRC0, SRC1, N); break; \
             case HTP_OP_MUL: hvx_mul_f32_aaa(DST, SRC0, SRC1, N); break; \
             case HTP_OP_DIV: hvx_div_f32_aaa(DST, SRC0, SRC1, N); break; \
-            default: break; \
-        } \
-    } \
-    else { \
-        switch (octx->op) { \
+            default: break;                                              \
+        }                                                                \
+    }                                                                    \
+    else {                                                               \
+        switch (octx->op) {                                              \
             case HTP_OP_ADD: hvx_add_f16_aaa(DST, SRC0, SRC1, N); break; \
             case HTP_OP_SUB: hvx_sub_f16_aaa(DST, SRC0, SRC1, N); break; \
             case HTP_OP_MUL: hvx_mul_f16_aaa(DST, SRC0, SRC1, N); break; \
             case HTP_OP_DIV: hvx_div_f16_aaa(DST, SRC0, SRC1, N); break; \
-            default: break; \
-        } \
+            default: break;                                              \
+        }                                                                \
     }
 
 // Macro for vector op switch (Dst Aligned, Src0 Aligned, Src1 Unaligned)
-#define COMPUTE_VECTOR_OP_AAU(DST, SRC0, SRC1, TYPE, N) \
-    if(TYPE == HTP_TYPE_F32) { \
-        switch (octx->op) { \
+#define COMPUTE_VECTOR_OP_AAU(DST, SRC0, SRC1, TYPE, N)                  \
+    if(TYPE == HTP_TYPE_F32) {                                           \
+        switch (octx->op) {                                              \
             case HTP_OP_ADD: hvx_add_f32_aau(DST, SRC0, SRC1, N); break; \
             case HTP_OP_SUB: hvx_sub_f32_aau(DST, SRC0, SRC1, N); break; \
             case HTP_OP_MUL: hvx_mul_f32_aau(DST, SRC0, SRC1, N); break; \
             case HTP_OP_DIV: hvx_div_f32_aau(DST, SRC0, SRC1, N); break; \
-            default: break; \
-        } \
-    } \
-    else { \
-        switch (octx->op) { \
+            default: break;                                              \
+        }                                                                \
+    }                                                                    \
+    else {                                                               \
+        switch (octx->op) {                                              \
             case HTP_OP_ADD: hvx_add_f16_aau(DST, SRC0, SRC1, N); break; \
             case HTP_OP_SUB: hvx_sub_f16_aau(DST, SRC0, SRC1, N); break; \
             case HTP_OP_MUL: hvx_mul_f16_aau(DST, SRC0, SRC1, N); break; \
             case HTP_OP_DIV: hvx_div_f16_aau(DST, SRC0, SRC1, N); break; \
-            default: break; \
-        } \
+            default: break;                                              \
+        }                                                                \
     }
 
 // Macro for vector op switch (All Unaligned - generic loop used in element repeat)
-#define COMPUTE_VECTOR_OP_UUU(DST, SRC0, SRC1, TYPE, N) \
-    if(TYPE == HTP_TYPE_F32) { \
-        switch (octx->op) { \
+#define COMPUTE_VECTOR_OP_UUU(DST, SRC0, SRC1, TYPE, N)                  \
+    if(TYPE == HTP_TYPE_F32) {                                           \
+        switch (octx->op) {                                              \
             case HTP_OP_ADD: hvx_add_f32_uuu(DST, SRC0, SRC1, N); break; \
             case HTP_OP_SUB: hvx_sub_f32_uuu(DST, SRC0, SRC1, N); break; \
             case HTP_OP_MUL: hvx_mul_f32_uuu(DST, SRC0, SRC1, N); break; \
             case HTP_OP_DIV: hvx_div_f32_uuu(DST, SRC0, SRC1, N); break; \
-            default: break; \
-        } \
-    } \
-    else { \
-        switch (octx->op) { \
+            default: break;                                              \
+        }                                                                \
+    }                                                                    \
+    else {                                                               \
+        switch (octx->op) {                                              \
             case HTP_OP_ADD: hvx_add_f16_uuu(DST, SRC0, SRC1, N); break; \
             case HTP_OP_SUB: hvx_sub_f16_uuu(DST, SRC0, SRC1, N); break; \
             case HTP_OP_MUL: hvx_mul_f16_uuu(DST, SRC0, SRC1, N); break; \
             case HTP_OP_DIV: hvx_div_f16_uuu(DST, SRC0, SRC1, N); break; \
-            default: break; \
-        } \
+            default: break;                                              \
+        }                                                                \
     }
 
 // 1. Scalar src1 (ne10 == 1)
@@ -184,9 +187,8 @@ static void binary_job_scalar(unsigned int nth, unsigned int ith, void * data) {
 
     const uint32_t src0_type = octx->src[0]->type;
     const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16);
-    const uint32_t total_rows = ne01 * ne02 * ne03;
-    const uint32_t start_row = bctx->nrows_per_thread * ith;
-    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, total_rows);
+    const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith;
+    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
     if (start_row >= end_row) return;
 
     FARF(HIGH, "binary-scalar: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
@@ -222,6 +224,8 @@ static void binary_job_scalar(unsigned int nth, unsigned int ith, void * data) {
     }
 
     // Main loop
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
     for (uint32_t ir = start_row; ir < end_row; ) {
         uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
 
@@ -242,12 +246,14 @@ static void binary_job_scalar(unsigned int nth, unsigned int ith, void * data) {
         uint8_t * src1_ptr = (uint8_t *)src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11;
         uint32_t s1_stride = (ne11 == 1) ? 0 : nb11;
 
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
         for (uint32_t r = 0; r < current_block_size; r++) {
             uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned;
             uint8_t * r_dst  = d_spad + r * bctx->dst_row_size_aligned;
             COMPUTE_SCALAR_OP(r_dst, r_src0, src1_ptr, src0_type, ne00);
             src1_ptr += s1_stride;
         }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
 
         uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
         dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size);
@@ -266,6 +272,7 @@ static void binary_job_scalar(unsigned int nth, unsigned int ith, void * data) {
         }
         ir += current_block_size;
     }
+
     dma_queue_flush(q);
 }
 
@@ -277,9 +284,8 @@ static void binary_job_vector_same_shape(unsigned int nth, unsigned int ith, voi
 
     const uint32_t src0_type = octx->src[0]->type;
     const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16);
-    const uint32_t total_rows = ne01 * ne02 * ne03;
-    const uint32_t start_row = bctx->nrows_per_thread * ith;
-    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, total_rows);
+    const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith;
+    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
     if (start_row >= end_row) return;
 
     FARF(HIGH, "binary-same-shape: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
@@ -323,18 +329,22 @@ static void binary_job_vector_same_shape(unsigned int nth, unsigned int ith, voi
         spad_idx ^= 1;
     }
 
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
     for (uint32_t ir = start_row; ir < end_row; ) {
         uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
         uint8_t * d_spad  = (uint8_t *) dma_queue_pop(q).src;
         uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst;
         uint8_t * s1_spad = (uint8_t *) dma_queue_pop(q).dst;
 
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
         for (uint32_t r = 0; r < current_block_size; r++) {
             uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned;
             uint8_t * r_src1 = s1_spad + r * bctx->src1_row_size_aligned;
             uint8_t * r_dst  = d_spad  + r * bctx->dst_row_size_aligned;
             COMPUTE_VECTOR_OP_AAA(r_dst, r_src0, r_src1, src0_type, ne00);
         }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
 
         uint32_t i03, i02, i01, rem;
         i03 = fastdiv(ir, &bctx->src0_dim12_div);
@@ -366,6 +376,7 @@ static void binary_job_vector_same_shape(unsigned int nth, unsigned int ith, voi
         }
         ir += current_block_size;
     }
+
     dma_queue_flush(q);
 }
 
@@ -377,9 +388,8 @@ static void binary_job_vector_row_broadcast(unsigned int nth, unsigned int ith,
 
     const uint32_t src0_type  = octx->src[0]->type;
     const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16);
-    const uint32_t total_rows = ne01 * ne02 * ne03;
-    const uint32_t start_row  = bctx->nrows_per_thread * ith;
-    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, total_rows);
+    const uint32_t start_row  = bctx->row_start + bctx->nrows_per_thread * ith;
+    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
     if (start_row >= end_row) return;
 
     FARF(HIGH, "binary-row-bcast: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
@@ -416,17 +426,21 @@ static void binary_job_vector_row_broadcast(unsigned int nth, unsigned int ith,
         spad_idx ^= 1;
     }
 
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
     for (uint32_t ir = start_row; ir < end_row; ) {
         uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
         uint8_t * d_spad  = (uint8_t *) dma_queue_pop(q).src;
         uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst;
 
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
         for (uint32_t r = 0; r < current_block_size; r++) {
             uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned;
             uint8_t * r_src1 = (uint8_t *)s1_ptr; // Constant
             uint8_t * r_dst  = d_spad + r * bctx->dst_row_size_aligned;
             COMPUTE_VECTOR_OP_AAA(r_dst, r_src0, r_src1, src0_type, ne00);
         }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
 
         uint32_t i03 = fastdiv(ir, &bctx->src0_dim12_div);
         uint32_t rem = ir - i03 * (ne02 * ne01);
@@ -447,6 +461,7 @@ static void binary_job_vector_row_broadcast(unsigned int nth, unsigned int ith,
         }
         ir += current_block_size;
     }
+
     dma_queue_flush(q);
 }
 
@@ -458,9 +473,8 @@ static void binary_job_vector_complex(unsigned int nth, unsigned int ith, void *
 
     const uint32_t src0_type = octx->src[0]->type;
     const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16);
-    const uint32_t total_rows = ne01 * ne02 * ne03;
-    const uint32_t start_row  = bctx->nrows_per_thread * ith;
-    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, total_rows);
+    const uint32_t start_row  = bctx->row_start + bctx->nrows_per_thread * ith;
+    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
     if (start_row >= end_row) return;
 
     FARF(HIGH, "binary-complex: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
@@ -493,6 +507,8 @@ static void binary_job_vector_complex(unsigned int nth, unsigned int ith, void *
         spad_idx ^= 1;
     }
 
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
     for (uint32_t ir = start_row; ir < end_row; ) {
         uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
         uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src;
@@ -503,6 +519,7 @@ static void binary_job_vector_complex(unsigned int nth, unsigned int ith, void *
         uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div);
         uint32_t i01 = rem - i02 * ne01;
 
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
         for (uint32_t r = 0; r < current_block_size; r++) {
             uint32_t r_i01 = i01 + r;
             uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div);
@@ -516,6 +533,7 @@ static void binary_job_vector_complex(unsigned int nth, unsigned int ith, void *
             // Read src1 from DDR (unaligned)
             COMPUTE_VECTOR_OP_AAU(r_dst, r_src0, r_src1, src0_type, ne00);
         }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
 
         uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
         dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size);
@@ -532,6 +550,7 @@ static void binary_job_vector_complex(unsigned int nth, unsigned int ith, void *
         }
         ir += current_block_size;
     }
+
     dma_queue_flush(q);
 }
 
@@ -544,9 +563,8 @@ static void binary_job_element_repeat(unsigned int nth, unsigned int ith, void *
     const uint32_t src0_type = octx->src[0]->type;
     const uint32_t elem_size_bytes = (src0_type == HTP_TYPE_F32) ? sizeof(float) : sizeof(_Float16);
     const uint32_t row_size_bytes = ne00 * elem_size_bytes;;
-    const uint32_t total_rows = ne01 * ne02 * ne03;
-    const uint32_t start_row  = bctx->nrows_per_thread * ith;
-    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, total_rows);
+    const uint32_t start_row  = bctx->row_start + bctx->nrows_per_thread * ith;
+    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
     if (start_row >= end_row) return;
 
     uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread);
@@ -579,6 +597,8 @@ static void binary_job_element_repeat(unsigned int nth, unsigned int ith, void *
         spad_idx ^= 1;
     }
 
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
     for (uint32_t ir = start_row; ir < end_row; ) {
         uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
         uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src;
@@ -589,6 +609,7 @@ static void binary_job_element_repeat(unsigned int nth, unsigned int ith, void *
         uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div);
         uint32_t i01 = rem - i02 * ne01;
 
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
         for (uint32_t r = 0; r < current_block_size; r++) {
             uint32_t r_i01 = i01 + r;
             uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div);
@@ -606,6 +627,7 @@ static void binary_job_element_repeat(unsigned int nth, unsigned int ith, void *
                 COMPUTE_VECTOR_OP_UUU(r_dst + c * elem_size_bytes, r_src0 + c * elem_size_bytes, r_src1_row, src0_type, len);
             }
         }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
 
         uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
         dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size);
@@ -622,6 +644,7 @@ static void binary_job_element_repeat(unsigned int nth, unsigned int ith, void *
         }
         ir += current_block_size;
     }
+
     dma_queue_flush(q);
 }
 
@@ -650,9 +673,8 @@ static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) {
     const uint32_t nb2 = dst->nb[2];
     const uint32_t nb3 = dst->nb[3];
 
-    const uint32_t total_rows = ne01 * ne02 * ne03;
-    const uint32_t start_row = bctx->nrows_per_thread * ith;
-    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, total_rows);
+    const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith;
+    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
     if (start_row >= end_row) return;
 
     uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread);
@@ -683,6 +705,8 @@ static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) {
         spad_idx ^= 1;
     }
 
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
     for (uint32_t ir = start_row; ir < end_row; ) {
         uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
         uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src;
@@ -693,6 +717,7 @@ static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) {
         uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div);
         uint32_t i01 = rem - i02 * ne01;
 
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
         for (uint32_t r = 0; r < current_block_size; r++) {
             uint32_t r_i01 = i01 + r; // linear within block since we split at ne01
 
@@ -704,6 +729,7 @@ static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) {
 
             hvx_add_f32_aau(r_dst, r_src0, r_src1, ne00);
         }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
 
         uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
         dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, ne00 * sizeof(float), current_block_size);
@@ -720,6 +746,7 @@ static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) {
         }
         ir += current_block_size;
     }
+
     dma_queue_flush(q);
 }
 
@@ -729,15 +756,31 @@ static int execute_op_binary(struct htp_ops_context * octx) {
     const struct htp_tensor * dst  = octx->dst;
 
     const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
-    const uint32_t n_threads  = MIN(octx->n_threads, src0_nrows);
 
-    // Use packed row sizes for VTCM allocation
+    // Use packed row sizes for VTCM allocation and alignment
     const uint32_t src0_type = octx->src[0]->type;
     const size_t elem_size = (src0_type == HTP_TYPE_F32) ? sizeof(float) : sizeof(_Float16);
     const size_t src0_row_size = src0->ne[0] * elem_size;
     const size_t src1_row_size = src1->ne[0] * elem_size;
     const size_t dst_row_size  = dst->ne[0]  * elem_size;
 
+    uint32_t row_start = 0;
+    uint32_t nrows     = src0_nrows;
+
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t rows_per_chunk = 0;
+        htp_tensor_mdev_rows_per_chunk(dst, (uint32_t) elem_size, (uint32_t) dst_row_size, &rows_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+
     size_t src0_row_size_aligned = hex_round_up(src0_row_size, VLEN);
     size_t src1_row_size_aligned = hex_round_up(src1_row_size, VLEN);
     size_t dst_row_size_aligned  = hex_round_up(dst_row_size,  VLEN);
@@ -815,7 +858,9 @@ static int execute_op_binary(struct htp_ops_context * octx) {
 
     struct htp_binary_context bctx;
     bctx.octx                  = octx;
-    bctx.nrows_per_thread      = (src0_nrows + n_threads - 1) / n_threads;
+    bctx.nrows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    bctx.total_rows            = nrows;
+    bctx.row_start             = row_start;
     bctx.block_max             = rows_per_buffer;
     bctx.src0_row_size_aligned = src0_row_size_aligned;
     bctx.src1_row_size_aligned = src1_row_size_aligned;
@@ -850,7 +895,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
         dma_queue_pop(q);
     }
 
-    worker_pool_run_func(octx->ctx->worker_pool, worker_func, &bctx, n_threads);
+    work_queue_run(octx->ctx->work_queue, worker_func, &bctx, n_threads);
 
     return HTP_STATUS_OK;
 }
@@ -870,4 +915,3 @@ int op_binary(struct htp_ops_context * octx) {
 
     return HTP_STATUS_NO_SUPPORT;
 }
-

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

@@ -1,5 +1,8 @@
+#include "hex-common.h"
+#include "hex-profile.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"
 #include "hexagon_types.h"
 #include "hexagon_protos.h"
 #include "hvx_hexagon_protos.h"
@@ -13,6 +16,10 @@ struct htp_concat_context {
     struct htp_ops_context * octx;
     uint32_t dim;
     uint32_t nrows_per_thread;
+    uint32_t row_start;
+    uint32_t nrows;
+    uint32_t elem_start;
+    uint32_t nelems;
     struct fastdiv_values div_ne0;
     struct fastdiv_values div_ne1;
     struct fastdiv_values div_ne2;
@@ -28,10 +35,10 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *
 
     const uint32_t src0_ne0 = src0->ne[0];
     const uint32_t src1_ne0 = src1->ne[0];
-    const uint32_t ne1      = dst->ne[1];
 
-    const uint32_t start_i = ith * cctx->nrows_per_thread;
-    const uint32_t end_i   = (start_i + cctx->nrows_per_thread < ne1) ? (start_i + cctx->nrows_per_thread) : ne1;
+    const uint32_t row_end = cctx->row_start + cctx->nrows;
+    const uint32_t start_i = cctx->row_start + ith * cctx->nrows_per_thread;
+    const uint32_t end_i   = (start_i + cctx->nrows_per_thread < row_end) ? (start_i + cctx->nrows_per_thread) : row_end;
     if (start_i >= end_i) return;
 
     dma_queue * q = octx->ctx->dma[ith];
@@ -51,6 +58,8 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *
     const uint32_t spad0_row_bytes = hex_round_up((src0_ne0 + src1_ne0_padded) * sizeof(float), VLEN);
     uint32_t mu = src1_ne0_padded * spad1_stride;
 
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
     for (uint32_t i = start_i; i < end_i; i += block_i) {
         uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i;
 
@@ -66,6 +75,7 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *
 
         HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride);
 
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
         for (uint32_t j = 0; j < src1_ne0_padded; j += 32) {
             #pragma unroll(4)
             for (uint32_t ii = 0; ii < current_block_i; ii++) {
@@ -75,6 +85,7 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *
                 hvx_vmemu(dst_ptr) = vtcm_tmp[ii];
             }
         }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
 
         dma_queue_pop(q); // src0
 
@@ -95,10 +106,10 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *
 
     const uint32_t src0_ne0 = src0->ne[0];
     const uint32_t src1_ne0 = src1->ne[0];
-    const uint32_t ne1      = dst->ne[1];
 
-    const uint32_t start_i = ith * cctx->nrows_per_thread;
-    const uint32_t end_i   = (start_i + cctx->nrows_per_thread < ne1) ? (start_i + cctx->nrows_per_thread) : ne1;
+    const uint32_t row_end = cctx->row_start + cctx->nrows;
+    const uint32_t start_i = cctx->row_start + ith * cctx->nrows_per_thread;
+    const uint32_t end_i   = (start_i + cctx->nrows_per_thread < row_end) ? (start_i + cctx->nrows_per_thread) : row_end;
     if (start_i >= end_i) return;
 
     dma_queue * q = octx->ctx->dma[ith];
@@ -118,6 +129,8 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *
     const uint32_t spad0_row_bytes = hex_round_up((src0_ne0 + src1_ne0_padded) * sizeof(__fp16), VLEN);
     uint32_t mu = src1_ne0_padded * spad1_stride;
 
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
     for (uint32_t i = start_i; i < end_i; i += block_i) {
         uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i;
 
@@ -133,6 +146,7 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *
 
         HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride);
 
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
         for (uint32_t j = 0; j < src1_ne0_padded; j += 64) {
             #pragma unroll(4)
             for (uint32_t ii = 0; ii < current_block_i; ii++) {
@@ -142,6 +156,7 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *
                 hvx_vmemu(dst_ptr) = vtcm_tmp[ii];
             }
         }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
 
         dma_queue_pop(q); // src0
 
@@ -164,11 +179,14 @@ static void concat_generic(unsigned int nth, unsigned int ith, void * data) {
     const uint32_t type_size = (dst->type == HTP_TYPE_F32 || dst->type == HTP_TYPE_I32) ? 4 : 2;
 
     const uint32_t ne[4] = {dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]};
-    const uint32_t total_elements = ne[0] * ne[1] * ne[2] * ne[3];
-    const uint32_t chunk_size = (total_elements + nth - 1) / nth;
 
-    const uint32_t start_idx = MIN(ith * chunk_size, total_elements);
-    const uint32_t end_idx   = MIN(start_idx + chunk_size, total_elements);
+    // Per-device element range aligned to prevent false sharing
+    const uint32_t elem_start = cctx->elem_start;
+    const uint32_t nelems     = cctx->nelems;
+    const uint32_t chunk_size = (nelems + nth - 1) / nth;
+
+    const uint32_t start_idx = MIN(elem_start + ith * chunk_size, elem_start + nelems);
+    const uint32_t end_idx   = MIN(start_idx + chunk_size, elem_start + nelems);
 
     // Naive scalar element-wise copy
     for (uint32_t idx = start_idx; idx < end_idx; idx++) {
@@ -236,13 +254,28 @@ int op_concat(struct htp_ops_context * octx) {
     void (*worker_func)(unsigned int, unsigned int, void *) = concat_generic;
 
     if (dim == 0 && is_2d && is_src1_transposed && !is_src0_transposed) {
-        n_threads = MIN(dst->ne[1], n_threads);
-        if (n_threads < 1) {
-            n_threads = 1;
+        const uint32_t total_rows = dst->ne[1];
+        const size_t dst_data_row_size = dst->ne[0] * type_size;
+        uint32_t row_start = 0;
+        uint32_t nrows     = total_rows;
+        if (octx->ctx->mdev.count > 1) {
+            uint32_t rows_per_chunk = 0;
+            htp_tensor_mdev_rows_per_chunk(dst, type_size, (uint32_t) dst_data_row_size, &rows_per_chunk);
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            row_start = range.start;
+            nrows     = range.count;
         }
+
+        if (nrows == 0) {
+            return HTP_STATUS_OK;
+        }
+
+        cctx.row_start = row_start;
+        cctx.nrows     = nrows;
+
         uint32_t block_i = (type_size == 4) ? 32 : 64;
 
-        cctx.nrows_per_thread = hmx_ceil_div(dst->ne[1], n_threads);
+        cctx.nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
 
         // Allocate VTCM
         uint32_t spad1_stride = block_i * type_size;
@@ -270,8 +303,26 @@ int op_concat(struct htp_ops_context * octx) {
         } else {
             worker_func = concat_2d_f16_transposed;
         }
+    } else {
+        const uint32_t total_elements = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3];
+        uint32_t elem_start = 0;
+        uint32_t nelems     = total_elements;
+        if (octx->ctx->mdev.count > 1) {
+            const uint32_t elems_per_chunk = HEX_L2_LINE_SIZE / type_size;
+            const bool can_split = htp_tensor_mdev_data_aligned(dst) && htp_tensor_is_contiguous(dst, type_size) && !htp_tensor_is_permuted(dst);
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_elements, can_split ? elems_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            elem_start = range.start;
+            nelems     = range.count;
+        }
+
+        if (nelems == 0) {
+            return HTP_STATUS_OK;
+        }
+
+        cctx.elem_start = elem_start;
+        cctx.nelems     = nelems;
     }
 
-    worker_pool_run_func(octx->ctx->worker_pool, worker_func, &cctx, n_threads);
+    work_queue_run(octx->ctx->work_queue, worker_func, &cctx, n_threads);
     return HTP_STATUS_OK;
 }

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

@@ -16,6 +16,7 @@
 #include "htp-ops.h"
 #include "hvx-utils.h"
 #include "htp-tensor.h"
+#include "htp-fence.h"
 
 struct htp_copy_context {
     struct htp_ops_context * octx;
@@ -29,7 +30,23 @@ struct htp_copy_context {
     uint32_t          src0_blocks_per_row;
     uint32_t          dst_blocks_per_row;
 
+    uint32_t          elem_start;
+    uint32_t          nelem;
+    uint32_t          elem_per_thread;
+
     uint32_t          src0_nrows_per_thread;
+    uint32_t          row_start;
+    uint32_t          nrows;
+
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne02_ne01;
+
+    struct fastdiv_values div_ne0;
+    struct fastdiv_values div_ne1_ne0;
+    struct fastdiv_values div_ne2_ne1_ne0;
+    struct fastdiv_values div_ne00;
+    struct fastdiv_values div_ne01_ne00;
+    struct fastdiv_values div_ne02_ne01_ne00;
 };
 
 #define cpy_preamble                              \
@@ -54,131 +71,113 @@ struct htp_copy_context {
     const uint32_t  nb0 = dst->nb[0];             \
     const uint32_t  nb1 = dst->nb[1];             \
     const uint32_t  nb2 = dst->nb[2];             \
-    const uint32_t  nb3 = dst->nb[3];             \
-                                                  \
-    const uint32_t   nr = ne01;
+    const uint32_t  nb3 = dst->nb[3];
 
-#define DEFINE_CPY_SAMESHAPE(NAME, ELEM_TYPE, ELEM_SIZE)                                                       \
-static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, void * data) {                   \
-    struct htp_copy_context * ct = (struct htp_copy_context *) data;                                           \
-    struct htp_ops_context * octx = ct->octx;                                                                  \
-    cpy_preamble;                                                                                              \
-    const uint32_t dr  = ct->src0_nrows_per_thread;                                                            \
-    const uint32_t ir0 = dr * ith;                                                                             \
-    const uint32_t ir1 = (ir0 + dr) < nr ? (ir0 + dr) : nr;                                                    \
-    if (ir0 >= nr) return;                                                                                     \
-    for (uint32_t i03 = 0; i03 < ne03; i03++) {                                                                \
-        for (uint32_t i02 = 0; i02 < ne02; i02++) {                                                            \
-            _Pragma("unroll(4)")                                                                               \
-            for (uint32_t i01 = ir0; i01 < ir1; i01++) {                                                       \
-                uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;                     \
-                uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;                    \
-                hex_l2fetch(src0_ptr, ne00 * ELEM_SIZE, nb01, 2);                                              \
-                hvx_copy_uu(dst_ptr, src0_ptr, ne00, ELEM_SIZE);                                               \
-            }                                                                                                  \
-        }                                                                                                      \
-    }                                                                                                          \
+#define DEFINE_CPY_SAMESHAPE(NAME, ELEM_TYPE, ELEM_SIZE)                                       \
+static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, void * data) {   \
+    struct htp_copy_context * ct = (struct htp_copy_context *) data;                           \
+    struct htp_ops_context * octx = ct->octx;                                                  \
+    cpy_preamble;                                                                              \
+    const uint32_t dr  = ct->src0_nrows_per_thread;                                            \
+    const uint32_t ir0 = ct->row_start + dr * ith;                                             \
+    const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows);                             \
+    if (ir0 >= ir1) return;                                                                    \
+    const bool contiguous = (nb01 == ne00 * ELEM_SIZE) && (nb1 == nb01) &&                     \
+                            (nb02 == ne01 * nb01)      && (nb2 == nb02) &&                     \
+                            (nb03 == ne02 * nb02)      && (nb3 == nb03);                       \
+    const uint32_t ne02_ne01 = ne02 * ne01;                                                    \
+    uint32_t i03 = fastdiv(ir0, &ct->div_ne02_ne01);                                           \
+    uint32_t rem = ir0 - i03 * ne02_ne01;                                                      \
+    uint32_t i02 = fastdiv(rem, &ct->div_ne01);                                                \
+    uint32_t i01 = rem - i02 * ne01;                                                           \
+    uint8_t * dst_ptr  = (uint8_t *) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;               \
+    uint8_t * src0_ptr = (uint8_t *) src0->data + i01*nb01 + i02*nb02 + i03*nb03;              \
+    if (contiguous) {                                                                          \
+        hvx_copy_uu(dst_ptr, src0_ptr, (ir1 - ir0) * ne00, ELEM_SIZE);                         \
+        return;                                                                                \
+    }                                                                                          \
+    for (uint32_t r = ir0; r < ir1; r++) {                                                     \
+        hex_l2fetch(src0_ptr, ne00 * ELEM_SIZE, nb01, 2);                                      \
+        hvx_copy_uu(dst_ptr, src0_ptr, ne00, ELEM_SIZE);                                       \
+        dst_ptr  += nb1;                                                                       \
+        src0_ptr += nb01;                                                                      \
+        if (++i01 == ne01) {                                                                   \
+            i01 = 0;                                                                           \
+            if (++i02 == ne02) {                                                               \
+                i02 = 0;                                                                       \
+                i03++;                                                                         \
+            }                                                                                  \
+            dst_ptr  = (uint8_t *) dst->data  + i02*nb2  + i03*nb3;                            \
+            src0_ptr = (uint8_t *) src0->data + i02*nb02 + i03*nb03;                           \
+        }                                                                                      \
+    }                                                                                          \
 }
 
 DEFINE_CPY_SAMESHAPE(f32,  float, 4)
 DEFINE_CPY_SAMESHAPE(f16, __fp16, 2)
 
-#define DEFINE_CPY_RESHAPE(NAME, ELEM_TYPE, ELEM_SIZE)                                                         \
-static void cpy_thread_##NAME##_reshape(unsigned int nth, unsigned int ith, void * data) {                     \
-    struct htp_copy_context * ct = (struct htp_copy_context *) data;                                           \
-    struct htp_ops_context * octx = ct->octx;                                                                  \
-    cpy_preamble;                                                                                              \
-    const uint32_t dr  = ct->src0_nrows_per_thread;                                                            \
-    const uint32_t ir0 = dr * ith;                                                                             \
-    const uint32_t ir1 = (ir0 + dr) < nr ? (ir0 + dr) : nr;                                                    \
-    if (ir0 >= nr) return;                                                                                     \
-    const bool src0_contig = (nb00 == ELEM_SIZE)   &&                                                          \
-                             (nb01 == ne00 * nb00) &&                                                          \
-                             (nb02 == ne01 * nb01) &&                                                          \
-                             (nb03 == ne02 * nb02);                                                            \
-    const bool dst_contig  = (nb0  == ELEM_SIZE)   &&                                                          \
-                             (nb1  == ne0  * nb0)  &&                                                          \
-                             (nb2  == ne1  * nb1)  &&                                                          \
-                             (nb3  == ne2  * nb2);                                                             \
-    if (src0_contig && dst_contig) {                                                                           \
-        for (int64_t i03 = 0; i03 < ne03; i03++) {                                                             \
-            for (int64_t i02 = 0; i02 < ne02; i02++) {                                                         \
-                uint8_t * src_ptr = (uint8_t *) src0->data + i03*nb03 + i02*nb02 + ir0*nb01;                   \
-                uint32_t  flat    = ((i03*ne02 + i02)*ne01 + ir0) * ne00;                                      \
-                uint8_t * dst_ptr = (uint8_t *) dst->data  + flat * ELEM_SIZE;                                 \
-                hvx_copy_uu(dst_ptr, src_ptr, (ir1 - ir0) * ne00, ELEM_SIZE);                                  \
-            }                                                                                                  \
-        }                                                                                                      \
-        return;                                                                                                \
-    }                                                                                                          \
-    const bool reshape_flat_fast = (ne03 == 1 && ne2 == 1 && ne3 == 1) &&                                      \
-                                   (ne0 == ne00 * ne01) && (ne1 == ne02) &&                                    \
-                                   (nb00 == ELEM_SIZE) && (nb0 == ELEM_SIZE);                                  \
-    if (reshape_flat_fast) {                                                                                   \
-        for (uint32_t i02 = 0; i02 < ne02; i02++) {                                                            \
-            for (uint32_t i01 = ir0; i01 < ir1; i01++) {                                                       \
-                uint8_t * src0_ptr = (uint8_t *) src0->data + i01 * nb01 + i02 * nb02;                         \
-                uint8_t * dst_ptr  = (uint8_t *) dst->data  + i01 * ne00 * ELEM_SIZE + i02 * nb1;              \
-                hvx_copy_uu(dst_ptr, src0_ptr, ne00, ELEM_SIZE);                                               \
-            }                                                                                                  \
-        }                                                                                                      \
-        return;                                                                                                \
-    }                                                                                                          \
-    int64_t k10 = 0;                                                                                           \
-    int64_t i11 = 0;                                                                                           \
-    int64_t i12 = 0;                                                                                           \
-    int64_t i13 = 0;                                                                                           \
-    const int64_t nk00 = ct->src0_blocks_per_row;                                                              \
-    const int64_t nk0  = ct->dst_blocks_per_row;                                                               \
-    for (int64_t i03 = 0; i03 < ne03; i03++) {                                                                 \
-        for (int64_t i02 = 0; i02 < ne02; i02++) {                                                             \
-            k10 += nk00 * ir0;                                                                                 \
-            while (k10 >= nk0) {                                                                               \
-                k10 -= nk0;                                                                                    \
-                if (++i11 == ne1) {                                                                            \
-                    i11 = 0;                                                                                   \
-                    if (++i12 == ne2) {                                                                        \
-                        i12 = 0;                                                                               \
-                        if (++i13 == ne3) {                                                                    \
-                            i13 = 0;                                                                           \
-                        }                                                                                      \
-                    }                                                                                          \
-                }                                                                                              \
-            }                                                                                                  \
-            for (int64_t i01 = ir0; i01 < ir1; i01++) {                                                        \
-                for (int64_t k00 = 0; k00 < nk00; k00++) {                                                     \
-                    const char * src0_ptr = ((char *) src0->data + k00*nb00 + i01*nb01 + i02*nb02 + i03*nb03); \
-                          char * dst_ptr  = ((char *)  dst->data + k10*nb0  + i11*nb1  + i12*nb2  + i13*nb3);  \
-                    memcpy(dst_ptr, src0_ptr, ELEM_SIZE);                                                      \
-                    if (++k10 == nk0) {                                                                        \
-                        k10 = 0;                                                                               \
-                        if (++i11 == ne1) {                                                                    \
-                            i11 = 0;                                                                           \
-                            if (++i12 == ne2) {                                                                \
-                                i12 = 0;                                                                       \
-                                if (++i13 == ne3) {                                                            \
-                                    i13 = 0;                                                                   \
-                                }                                                                              \
-                            }                                                                                  \
-                        }                                                                                      \
-                    }                                                                                          \
-                }                                                                                              \
-            }                                                                                                  \
-            k10 += nk00 * (ne01 - ir1);                                                                        \
-            while (k10 >= nk0) {                                                                               \
-                k10 -= nk0;                                                                                    \
-                if (++i11 == ne1) {                                                                            \
-                    i11 = 0;                                                                                   \
-                    if (++i12 == ne2) {                                                                        \
-                        i12 = 0;                                                                               \
-                        if (++i13 == ne3) {                                                                    \
-                            i13 = 0;                                                                           \
-                        }                                                                                      \
-                    }                                                                                          \
-                }                                                                                              \
-            }                                                                                                  \
-        }                                                                                                      \
-    }                                                                                                          \
+#define DEFINE_CPY_RESHAPE(NAME, ELEM_TYPE, ELEM_SIZE)                                               \
+static void cpy_thread_##NAME##_reshape(unsigned int nth, unsigned int ith, void * data) {           \
+    struct htp_copy_context * ct = (struct htp_copy_context *) data;                                 \
+    struct htp_ops_context * octx = ct->octx;                                                        \
+    cpy_preamble;                                                                                    \
+    const uint32_t th_nelem = ct->elem_per_thread;                                                   \
+    const uint32_t th_start = ct->elem_start + ith * th_nelem;                                       \
+    const uint32_t th_end   = MIN(th_start + th_nelem, ct->elem_start + ct->nelem);                  \
+    if (th_start >= th_end) return;                                                                  \
+                                                                                                     \
+    const uint32_t ne01_ne00      = ne01 * ne00;                                                     \
+    const uint32_t ne02_ne01_ne00 = ne02 * ne01_ne00;                                                \
+    const uint32_t ne1_ne0        = ne1 * ne0;                                                       \
+    const uint32_t ne2_ne1_ne0    = ne2 * ne1_ne0;                                                   \
+                                                                                                     \
+    uint32_t e = th_start;                                                                           \
+    uint32_t i13 = fastdiv(e, &ct->div_ne2_ne1_ne0);                                                 \
+    uint32_t rem = e - i13 * ne2_ne1_ne0;                                                            \
+    uint32_t i12 = fastdiv(rem, &ct->div_ne1_ne0);                                                   \
+    uint32_t rem2 = rem - i12 * ne1_ne0;                                                             \
+    uint32_t i11 = fastdiv(rem2, &ct->div_ne0);                                                      \
+    uint32_t i10 = rem2 - i11 * ne0;                                                                 \
+                                                                                                     \
+    uint32_t i03 = fastdiv(e, &ct->div_ne02_ne01_ne00);                                              \
+    uint32_t rem_s = e - i03 * ne02_ne01_ne00;                                                       \
+    uint32_t i02 = fastdiv(rem_s, &ct->div_ne01_ne00);                                               \
+    uint32_t rem2_s = rem_s - i02 * ne01_ne00;                                                       \
+    uint32_t i01 = fastdiv(rem2_s, &ct->div_ne00);                                                   \
+    uint32_t i00 = rem2_s - i01 * ne00;                                                              \
+                                                                                                     \
+    char * dst_ptr        = (char *)       dst->data  + i10*nb0  + i11*nb1  + i12*nb2  + i13*nb3;    \
+    const char * src0_ptr = (const char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03;   \
+                                                                                                     \
+    for (; e < th_end; e++) {                                                                        \
+        *((ELEM_TYPE *) dst_ptr) = *((const ELEM_TYPE *) src0_ptr);                                  \
+                                                                                                     \
+        dst_ptr += nb0;                                                                              \
+        if (++i10 == ne0) {                                                                          \
+            i10 = 0;                                                                                 \
+            if (++i11 == ne1) {                                                                      \
+                i11 = 0;                                                                             \
+                if (++i12 == ne2) {                                                                  \
+                    i12 = 0;                                                                         \
+                    i13++;                                                                           \
+                }                                                                                    \
+            }                                                                                        \
+            dst_ptr = (char *) dst->data + i11*nb1 + i12*nb2 + i13*nb3;                              \
+        }                                                                                            \
+                                                                                                     \
+        src0_ptr += nb00;                                                                            \
+        if (++i00 == ne00) {                                                                         \
+            i00 = 0;                                                                                 \
+            if (++i01 == ne01) {                                                                     \
+                i01 = 0;                                                                             \
+                if (++i02 == ne02) {                                                                 \
+                    i02 = 0;                                                                         \
+                    i03++;                                                                           \
+                }                                                                                    \
+            }                                                                                        \
+            src0_ptr = (const char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03;                   \
+        }                                                                                            \
+    }                                                                                                \
 }
 
 DEFINE_CPY_RESHAPE(f32,  float, 4)
@@ -189,22 +188,33 @@ static void cpy_thread_f16_f32_sameshape(unsigned int nth, unsigned int ith, voi
     struct htp_ops_context * octx = ct->octx;
     cpy_preamble;
 
-    // parallelize by src0 rows
     const uint32_t dr  = ct->src0_nrows_per_thread;
-    const uint32_t ir0 = dr * ith;
-    const uint32_t ir1 = (ir0 + dr) < nr ? (ir0 + dr) : nr;
-    if (ir0 >= nr) return;
+    const uint32_t ir0 = ct->row_start + dr * ith;
+    const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows);
+    if (ir0 >= ir1) return;
 
-    // copy by rows
-    for (uint32_t i03 = 0; i03 < ne03; i03++) {
-        for (uint32_t i02 = 0; i02 < ne02; i02++) {
-            #pragma unroll(2)
-            for (uint32_t i01 = ir0; i01 < ir1; i01++) {
-                uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;
-                uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
-                hex_l2fetch(src0_ptr, ne00 * sizeof(float), nb01, 2);
-                hvx_copy_f16_f32_uu(dst_ptr, src0_ptr, ne00);
+    const uint32_t ne02_ne01 = ne02 * ne01;
+    uint32_t i03 = fastdiv(ir0, &ct->div_ne02_ne01);
+    uint32_t rem = ir0 - i03 * ne02_ne01;
+    uint32_t i02 = fastdiv(rem, &ct->div_ne01);
+    uint32_t i01 = rem - i02 * ne01;
+
+    uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;
+    uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
+
+    for (uint32_t r = ir0; r < ir1; r++) {
+        hex_l2fetch(src0_ptr, ne00 * sizeof(float), nb01, 2);
+        hvx_copy_f16_f32_uu(dst_ptr, src0_ptr, ne00);
+        dst_ptr  += nb1;
+        src0_ptr += nb01;
+        if (++i01 == ne01) {
+            i01 = 0;
+            if (++i02 == ne02) {
+                i02 = 0;
+                i03++;
             }
+            dst_ptr  = (uint8_t*) dst->data  + i02*nb2  + i03*nb3;
+            src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03;
         }
     }
 }
@@ -214,22 +224,33 @@ static void cpy_thread_f32_f16_sameshape(unsigned int nth, unsigned int ith, voi
     struct htp_ops_context * octx = ct->octx;
     cpy_preamble;
 
-    // parallelize by src0 rows
     const uint32_t dr  = ct->src0_nrows_per_thread;
-    const uint32_t ir0 = dr * ith;
-    const uint32_t ir1 = (ir0 + dr) < nr ? (ir0 + dr) : nr;
-    if (ir0 >= nr) return;
+    const uint32_t ir0 = ct->row_start + dr * ith;
+    const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows);
+    if (ir0 >= ir1) return;
 
-    // copy by rows
-    for (uint32_t i03 = 0; i03 < ne03; i03++) {
-        for (uint32_t i02 = 0; i02 < ne02; i02++) {
-            #pragma unroll(2)
-            for (uint32_t i01 = ir0; i01 < ir1; i01++) {
-                uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;
-                uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
-                hex_l2fetch(src0_ptr, ne00 * sizeof(__fp16), nb01, 2);
-                hvx_copy_f32_f16_uu(dst_ptr, src0_ptr, ne00);
+    const uint32_t ne02_ne01 = ne02 * ne01;
+    uint32_t i03 = fastdiv(ir0, &ct->div_ne02_ne01);
+    uint32_t rem = ir0 - i03 * ne02_ne01;
+    uint32_t i02 = fastdiv(rem, &ct->div_ne01);
+    uint32_t i01 = rem - i02 * ne01;
+
+    uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;
+    uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
+
+    for (uint32_t r = ir0; r < ir1; r++) {
+        hex_l2fetch(src0_ptr, ne00 * sizeof(__fp16), nb01, 2);
+        hvx_copy_f32_f16_uu(dst_ptr, src0_ptr, ne00);
+        dst_ptr  += nb1;
+        src0_ptr += nb01;
+        if (++i01 == ne01) {
+            i01 = 0;
+            if (++i02 == ne02) {
+                i02 = 0;
+                i03++;
             }
+            dst_ptr  = (uint8_t*) dst->data  + i02*nb2  + i03*nb3;
+            src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03;
         }
     }
 }
@@ -250,15 +271,19 @@ static inline void cpy_dma_sametype_sameshape(
     dma_queue * q = octx->ctx->dma[0];
 
     if (contiguous_outer) {
-        dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data), nb1, nb01, ne00 * elem_size, ne01 * ne02 * ne03);
-        dma_queue_pop(q);
+        if (!dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data), nb1, nb01, ne00 * elem_size, ne01 * ne02 * ne03)) {
+            dma_queue_flush(q);
+            dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data), nb1, nb01, ne00 * elem_size, ne01 * ne02 * ne03);
+        }
+        dma_queue_flush(q);
         return;
     }
 
     for (uint32_t i03 = 0; i03 < ne03; i03++) {
         for (uint32_t i02 = 0; i02 < ne02; i02++) {
-            uint8_t* dst_ptr  = (uint8_t*) dst->data  + i02*nb2  + i03*nb3;
-            uint8_t* src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03;
+            uint8_t * dst_ptr  = (uint8_t *) dst->data  + i02 * nb2  + i03 * nb3;
+            uint8_t * src0_ptr = (uint8_t *) src0->data + i02 * nb02 + i03 * nb03;
+
             if (!dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01)) {
                 dma_queue_flush(q);
                 dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01);
@@ -269,10 +294,9 @@ static inline void cpy_dma_sametype_sameshape(
     dma_queue_flush(q);
 }
 
-int op_cpy(struct htp_ops_context * octx) {
+static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) {
     cpy_preamble;
-
-    const uint32_t n_threads = MIN(nr, octx->n_threads);
+    *use_dma = false;
 
     struct htp_copy_context ct;
     ct.octx = octx;
@@ -296,59 +320,117 @@ int op_cpy(struct htp_ops_context * octx) {
     }
 
     const bool sametype   = (src0->type == dst->type);
-    const bool transposed = (nb00 > nb01) || (nb0 > nb1);
+    const bool transposed = (nb00 > nb01) || (nb0 > nb1) ||
+                            (nb00 != ct.src0_type_size) || (nb0 != ct.dst_type_size) ||
+                            (nb01 < ne00 * ct.src0_type_size) || (nb1 < ne0 * ct.dst_type_size);
     const bool sameshape  = !transposed && (ne00 == ne0 && ne01 == ne1 && ne02 == ne2 && ne03 == ne3);
 
-    ct.src0_nrows_per_thread = (nr + n_threads - 1) / n_threads;
+    const uint32_t n_threads = octx->n_threads;
 
-    worker_callback_t copy_fun = NULL;
-    bool use_dma = false;
+    const bool dst_is_contiguous = htp_tensor_is_contiguous(dst, ct.dst_type_size);
 
-    if (sametype && sameshape) {
-        use_dma = true;
-    } else if (sameshape) {
-        /**/ if (dst->type == HTP_TYPE_F16 && src0->type == HTP_TYPE_F32)
-            copy_fun = cpy_thread_f16_f32_sameshape;
-        else if (dst->type == HTP_TYPE_F32 && src0->type == HTP_TYPE_F16)
-            copy_fun = cpy_thread_f32_f16_sameshape;
-        else
-            return HTP_STATUS_NO_SUPPORT;
-    } else if (sametype) {
-        if (src0->type == HTP_TYPE_F32) {
-            copy_fun = cpy_thread_f32_reshape;
-        } else {
-            copy_fun = cpy_thread_f16_reshape;
+    if (sameshape) {
+        const uint32_t total_rows = ne01 * ne02 * ne03;
+        const uint32_t row_size   = ne00 * ct.dst_type_size;
+
+        ct.div_ne01      = init_fastdiv_values(ne01);
+        ct.div_ne02_ne01 = init_fastdiv_values(ne02 * ne01);
+
+        uint32_t row_start = 0;
+        uint32_t nrows     = total_rows;
+
+        if (octx->ctx->mdev.count > 1) {
+            const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
+            const bool can_split = htp_tensor_mdev_data_aligned(dst) && dst_is_contiguous;
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, can_split ? rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            row_start = range.start;
+            nrows     = range.count;
         }
+
+        if (nrows == 0) {
+            return HTP_STATUS_OK;
+        }
+
+        ct.row_start = row_start;
+        ct.nrows     = nrows;
+        ct.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+
+        if (sametype && octx->ctx->mdev.count <= 1) {
+            *use_dma = true;
+            cpy_dma_sametype_sameshape(octx, dst, src0, ct.src0_type_size, ne00, ne01, ne02, ne03, nb01, nb02, nb03, nb1, nb2, nb3);
+        } else {
+            work_queue_func_t copy_fun = NULL;
+            if (sametype) {
+                copy_fun = (src0->type == HTP_TYPE_F32) ? cpy_thread_f32_sameshape : cpy_thread_f16_sameshape;
+            } else if (dst->type == HTP_TYPE_F16 && src0->type == HTP_TYPE_F32) {
+                copy_fun = cpy_thread_f16_f32_sameshape;
+            } else if (dst->type == HTP_TYPE_F32 && src0->type == HTP_TYPE_F16) {
+                copy_fun = cpy_thread_f32_f16_sameshape;
+            } else {
+                return HTP_STATUS_NO_SUPPORT;
+            }
+            work_queue_run(octx->ctx->work_queue, copy_fun, &ct, n_threads);
+        }
+    } else if (sametype) {
+        const uint32_t total_elems = ne0 * ne1 * ne2 * ne3;
+        const uint32_t elems_per_line = (ct.dst_type_size == 4) ? 32 : 64;
+
+        ct.div_ne0            = init_fastdiv_values(ne0);
+        ct.div_ne1_ne0        = init_fastdiv_values(ne1 * ne0);
+        ct.div_ne2_ne1_ne0    = init_fastdiv_values(ne2 * ne1 * ne0);
+        ct.div_ne00           = init_fastdiv_values(ne00);
+        ct.div_ne01_ne00      = init_fastdiv_values(ne01 * ne00);
+        ct.div_ne02_ne01_ne00 = init_fastdiv_values(ne02 * ne01 * ne00);
+
+        uint32_t elem_start = 0;
+        uint32_t nelem      = total_elems;
+
+        if (octx->ctx->mdev.count > 1) {
+            const bool can_split = htp_tensor_mdev_data_aligned(dst) && dst_is_contiguous;
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_elems, can_split ? elems_per_line : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            elem_start = range.start;
+            nelem      = range.count;
+        }
+
+        if (nelem == 0) {
+            return HTP_STATUS_OK;
+        }
+
+        ct.elem_start      = elem_start;
+        ct.nelem           = nelem;
+        ct.elem_per_thread = fastdiv(nelem + n_threads - 1, &octx->n_threads_div);
+
+        work_queue_func_t copy_fun = (src0->type == HTP_TYPE_F32) ? cpy_thread_f32_reshape : cpy_thread_f16_reshape;
+        work_queue_run(octx->ctx->work_queue, copy_fun, &ct, n_threads);
     } else {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    FARF(HIGH, "cpy-%s-%s: (%ux%ux%ux%u) -> (%ux%ux%ux%u) : use_dma=%d n_threads %u\n",
-         src0->type == HTP_TYPE_F32 ? "f32" : "f16", dst->type == HTP_TYPE_F32 ? "f32" : "f16",
-         ne00, ne01, ne02, ne03, ne0, ne1, ne2, ne3, use_dma, n_threads);
-
-    if (use_dma) {
-        cpy_dma_sametype_sameshape(octx, dst, src0, ct.src0_type_size, ne00, ne01, ne02, ne03, nb01, nb02, nb03, nb1, nb2, nb3);
-    } else {
-        worker_pool_run_func(octx->ctx->worker_pool, copy_fun, &ct, n_threads);
-    }
-
-    const struct htp_tensor *sync = octx->src[1];
-    if (sync && (sync->flags & HTP_TENSOR_FENCE)) {
-        if (!use_dma) {
-            // htp_tensor_flush_all(octx->ctx, octx->dsts, 1);
-            qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
-        }
-
-        atomic_uint * sync_fence = (atomic_uint *) sync->data;
-        const uint32_t seq = (uint32_t) octx->op_params[0];
-
-        atomic_store(&sync_fence[0], seq);
-        asm volatile ("syncht" : : : "memory");
-        Q6_dccleaninva_A((void *) sync_fence);
-
-        FARF(HIGH, "ggml-hex: sync-release : fence %p seq %u\n", sync_fence, seq);
-    }
-
     return HTP_STATUS_OK;
 }
+
+int op_cpy(struct htp_ops_context * octx) {
+    bool use_dma = false;
+    int status = exec_cpy(octx, &use_dma);
+
+    htp_ops_context_set_status(octx, status);
+
+    if (octx->op == HTP_OP_CPY_FENCE) {
+        if (!use_dma) {
+            htp_flush_dirty_ranges(octx->ctx);
+        }
+
+        htp_mdev_group_barrier(octx);
+
+        if (octx->ctx->mdev.idx == 0) {
+            const struct htp_tensor * sync = octx->src[1];
+            const uint32_t seq = (uint32_t) octx->op_params[0];
+            atomic_uint * sync_fence = (atomic_uint *) (uintptr_t) sync->data;
+            htp_fence_write(sync_fence, seq, octx->status);
+
+            FARF(HIGH, "ggml-hex: sync-release : fence %p seq 0x%x status %d\n", sync_fence, seq, octx->status);
+        }
+    }
+
+    return octx->status;
+}

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

@@ -7,6 +7,8 @@
 
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
+#include "hex-common.h"
+#include "hex-profile.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
@@ -17,25 +19,25 @@
 #define htp_cumsum_tensors_preamble                         \
     const struct htp_tensor * restrict src0 = octx->src[0]; \
     const struct htp_tensor * restrict dst  = octx->dst;    \
-                                                     \
-    const uint32_t ne00 = src0->ne[0];               \
-    const uint32_t ne01 = src0->ne[1];               \
-    const uint32_t ne02 = src0->ne[2];               \
-    const uint32_t ne03 = src0->ne[3];               \
-                                                     \
-    const uint32_t ne0 = dst->ne[0];                 \
-    const uint32_t ne1 = dst->ne[1];                 \
-    const uint32_t ne2 = dst->ne[2];                 \
-    const uint32_t ne3 = dst->ne[3];                 \
-                                                     \
-    const uint32_t nb00 = src0->nb[0];               \
-    const uint32_t nb01 = src0->nb[1];               \
-    const uint32_t nb02 = src0->nb[2];               \
-    const uint32_t nb03 = src0->nb[3];               \
-                                                     \
-    const uint32_t nb0 = dst->nb[0];                 \
-    const uint32_t nb1 = dst->nb[1];                 \
-    const uint32_t nb2 = dst->nb[2];                 \
+                                                            \
+    const uint32_t ne00 = src0->ne[0];                      \
+    const uint32_t ne01 = src0->ne[1];                      \
+    const uint32_t ne02 = src0->ne[2];                      \
+    const uint32_t ne03 = src0->ne[3];                      \
+                                                            \
+    const uint32_t ne0 = dst->ne[0];                        \
+    const uint32_t ne1 = dst->ne[1];                        \
+    const uint32_t ne2 = dst->ne[2];                        \
+    const uint32_t ne3 = dst->ne[3];                        \
+                                                            \
+    const uint32_t nb00 = src0->nb[0];                      \
+    const uint32_t nb01 = src0->nb[1];                      \
+    const uint32_t nb02 = src0->nb[2];                      \
+    const uint32_t nb03 = src0->nb[3];                      \
+                                                            \
+    const uint32_t nb0 = dst->nb[0];                        \
+    const uint32_t nb1 = dst->nb[1];                        \
+    const uint32_t nb2 = dst->nb[2];                        \
     const uint32_t nb3 = dst->nb[3];
 
 struct htp_cumsum_context {
@@ -46,6 +48,7 @@ struct htp_cumsum_context {
     size_t          dst_row_size_aligned;
     uint32_t        rows_per_thread;
     uint32_t        total_rows;
+    uint32_t        row_start;
 };
 
 #define htp_cumsum_preamble                                                \
@@ -116,11 +119,8 @@ static inline void hvx_cumsum_row_f32(const float * restrict src, float * restri
 static void cumsum_thread_f32_dma(unsigned int nth, unsigned int ith, void * data) {
     htp_cumsum_preamble;
 
-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
-
-    const uint32_t ir0 = cctx->rows_per_thread * ith;
-    const uint32_t ir1 = MIN(ir0 + cctx->rows_per_thread, cctx->total_rows);
+    const uint32_t ir0 = cctx->row_start + cctx->rows_per_thread * ith;
+    const uint32_t ir1 = MIN(ir0 + cctx->rows_per_thread, cctx->row_start + cctx->total_rows);
 
     if (ir0 >= ir1) {
         return;
@@ -149,11 +149,15 @@ static void cumsum_thread_f32_dma(unsigned int nth, unsigned int ith, void * dat
                                    src_row_size_aligned, src_row_size, 1);
     }
 
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
     for (uint32_t ir = ir0; ir < ir1; ir++) {
         float * dst_spad_row = (float *) dma_queue_pop(dma_queue).src;
         float * src_spad_row = (float *) dma_queue_pop(dma_queue).dst;
 
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
         hvx_cumsum_row_f32(src_spad_row, dst_spad_row, ne00);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
 
         dma_queue_push_vtcm_to_ddr(dma_queue,
                                    dma_make_ptr(dst_data + (ir * dst_row_size), (uint8_t *) dst_spad_row),
@@ -168,12 +172,10 @@ static void cumsum_thread_f32_dma(unsigned int nth, unsigned int ith, void * dat
     }
 
     dma_queue_flush(dma_queue);
-    t2 = HAP_perf_get_qtimer_count();
 
-    FARF(HIGH, "cumsum-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n",
+    FARF(HIGH, "cumsum-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u\n",
          ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ir0, ir1,
-         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
 }
 
 // ---------------------------------------------------------------------------
@@ -183,14 +185,14 @@ static void cumsum_thread_f32_dma(unsigned int nth, unsigned int ith, void * dat
 static void cumsum_thread_f32(unsigned int nth, unsigned int ith, void * data) {
     htp_cumsum_preamble;
 
-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
-
     const uint8_t * src_data = (const uint8_t *) src0->data;
     uint8_t *       dst_data = (uint8_t *) dst->data;
 
-    const uint32_t ir0 = cctx->rows_per_thread * ith;
-    const uint32_t ir1 = MIN(ir0 + cctx->rows_per_thread, cctx->total_rows);
+    const uint32_t ir0 = cctx->row_start + cctx->rows_per_thread * ith;
+    const uint32_t ir1 = MIN(ir0 + cctx->rows_per_thread, cctx->row_start + cctx->total_rows);
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
 
     for (uint32_t ir = ir0; ir < ir1; ir++) {
         const float * restrict src_row = (const float *) (src_data + ir * cctx->src_row_size);
@@ -198,12 +200,11 @@ static void cumsum_thread_f32(unsigned int nth, unsigned int ith, void * data) {
         hvx_cumsum_row_f32(src_row, dst_row, ne00);
     }
 
-    t2 = HAP_perf_get_qtimer_count();
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
 
-    FARF(HIGH, "cumsum-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n",
+    FARF(HIGH, "cumsum-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u\n",
          ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ir0, ir1,
-         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
 }
 
 int op_cumsum_f32(struct htp_ops_context * octx) {
@@ -214,8 +215,25 @@ int op_cumsum_f32(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    const uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3];
-    const uint32_t n_threads  = MIN(octx->n_threads, total_rows);
+    const uint32_t total_rows      = src0->ne[1] * src0->ne[2] * src0->ne[3];
+    const size_t dst_data_row_size = dst->ne[0] * sizeof(float);
+
+    uint32_t row_start = 0;
+    uint32_t nrows     = total_rows;
+
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t rows_per_chunk = 0;
+        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_data_row_size, &rows_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
 
     const size_t src_row_size         = src0->nb[1];
     const size_t dst_row_size         = dst->nb[1];
@@ -240,14 +258,15 @@ int op_cumsum_f32(struct htp_ops_context * octx) {
         .dst_row_size         = dst_row_size,
         .src_row_size_aligned = src_row_size_aligned,
         .dst_row_size_aligned = dst_row_size_aligned,
-        .rows_per_thread      = (total_rows + n_threads - 1) / n_threads,
-        .total_rows           = total_rows,
+        .rows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
+        .total_rows           = nrows,
+        .row_start            = row_start,
     };
 
     if (octx->ctx->vtcm_size < spad_per_thread * n_threads) {
-        worker_pool_run_func(octx->ctx->worker_pool, cumsum_thread_f32, &cctx, n_threads);
+        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32, &cctx, n_threads);
     } else {
-        worker_pool_run_func(octx->ctx->worker_pool, cumsum_thread_f32_dma, &cctx, n_threads);
+        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32_dma, &cctx, n_threads);
     }
 
     return HTP_STATUS_OK;

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

@@ -5,8 +5,11 @@
 
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
+#include "hex-common.h"
+#include "hex-profile.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"
 #include "hvx-types.h"
 #include "hex-utils.h"
 #include "hvx-copy.h"
@@ -15,17 +18,17 @@
 #define htp_diag_tensors_preamble                           \
     const struct htp_tensor * restrict src0 = octx->src[0]; \
     const struct htp_tensor * restrict dst  = octx->dst;    \
-                                                     \
-    const uint32_t ne02 = src0->ne[2];               \
-                                                     \
-    const uint32_t ne0 = dst->ne[0];                 \
-    const uint32_t ne1 = dst->ne[1];                 \
-                                                     \
-    const uint32_t nb02 = src0->nb[2];               \
-    const uint32_t nb03 = src0->nb[3];               \
-                                                     \
-    const uint32_t nb1 = dst->nb[1];                 \
-    const uint32_t nb2 = dst->nb[2];                 \
+                                                            \
+    const uint32_t ne02 = src0->ne[2];                      \
+                                                            \
+    const uint32_t ne0 = dst->ne[0];                        \
+    const uint32_t ne1 = dst->ne[1];                        \
+                                                            \
+    const uint32_t nb02 = src0->nb[2];                      \
+    const uint32_t nb03 = src0->nb[3];                      \
+                                                            \
+    const uint32_t nb1 = dst->nb[1];                        \
+    const uint32_t nb2 = dst->nb[2];                        \
     const uint32_t nb3 = dst->nb[3];
 
 struct htp_diag_context {
@@ -36,6 +39,7 @@ struct htp_diag_context {
     size_t          dst_row_size_aligned;
     uint32_t        batches_per_thread;
     uint32_t        total_batches;
+    uint32_t        batch_start;
 };
 
 #define htp_diag_preamble                                              \
@@ -57,11 +61,8 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
     htp_diag_preamble;
     dma_queue * dma_queue = octx->ctx->dma[ith];
 
-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
-
-    const uint32_t ib0 = dctx->batches_per_thread * ith;
-    const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->total_batches);
+    const uint32_t ib0 = dctx->batch_start + dctx->batches_per_thread * ith;
+    const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->batch_start + dctx->total_batches);
 
     if (ib0 >= ib1) {
         return;
@@ -79,6 +80,8 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
     uint8_t * src_spad = octx->src0_spad.data + (ith * src_batch_size_aligned);
     uint8_t * dst_spad = octx->dst_spad.data  + (ith * dst_row_size_aligned);
 
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
     for (uint32_t ib = ib0; ib < ib1; ib++) {
         const uint32_t i3 = ib / ne02;
         const uint32_t i2 = ib % ne02;
@@ -96,7 +99,9 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
 
         for (uint32_t i1 = 0; i1 < ne1; i1++) {
             // Compute row in VTCM
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (ib * ne1 + i1));
             hvx_diag_row_f32(src_spad_f32, dst_spad_f32, i1, ne0);
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (ib * ne1 + i1));
 
             // Write completed row back to DDR
             uint8_t * dst_row = dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1;
@@ -107,12 +112,9 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
         }
     }
 
-    t2 = HAP_perf_get_qtimer_count();
-
-    FARF(HIGH, "diag-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n",
+    FARF(HIGH, "diag-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u\n",
          ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ib0, ib1,
-         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
 }
 
 // ---------------------------------------------------------------------------
@@ -122,14 +124,14 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
 static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) {
     htp_diag_preamble;
 
-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
-
     const uint8_t * src_data = (const uint8_t *) src0->data;
     uint8_t *       dst_data = (uint8_t *) dst->data;
 
-    const uint32_t ib0 = dctx->batches_per_thread * ith;
-    const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->total_batches);
+    const uint32_t ib0 = dctx->batch_start + dctx->batches_per_thread * ith;
+    const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->batch_start + dctx->total_batches);
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ib0);
 
     for (uint32_t ib = ib0; ib < ib1; ib++) {
         const uint32_t i3 = ib / ne02;
@@ -143,12 +145,11 @@ static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) {
         }
     }
 
-    t2 = HAP_perf_get_qtimer_count();
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ib0);
 
-    FARF(HIGH, "diag-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n",
+    FARF(HIGH, "diag-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u\n",
          ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ib0, ib1,
-         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
 }
 
 int op_diag_f32(struct htp_ops_context * octx) {
@@ -160,7 +161,36 @@ int op_diag_f32(struct htp_ops_context * octx) {
     }
 
     const uint32_t total_batches = src0->ne[2] * src0->ne[3];
-    const uint32_t n_threads     = MIN(octx->n_threads, total_batches);
+    const size_t dst_batch_size  = dst->ne[1] * dst->nb[1];
+
+    uint32_t batch_start = 0;
+    uint32_t nbatches    = total_batches;
+
+    if (octx->ctx->mdev.count > 1) {
+        bool can_split = htp_tensor_mdev_data_aligned(dst) && (dst->ne[0] == 1 || dst->nb[0] == sizeof(float)) && !htp_tensor_is_permuted(dst);
+        uint32_t batches_per_chunk = 1;
+        if (can_split) {
+            if (dst->ne[2] > 1 && (dst->nb[2] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0 &&
+                (dst->ne[3] <= 1 || (dst->nb[3] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0)) {
+                batches_per_chunk = 1;
+            } else if (dst->nb[2] == dst_batch_size &&
+                       (dst->ne[3] <= 1 || dst->nb[3] == dst->nb[2] * dst->ne[2])) {
+                batches_per_chunk = (dst_batch_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(dst_batch_size, HEX_L2_LINE_SIZE)) : 1;
+            } else {
+                can_split = false;
+            }
+        }
+
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_batches, can_split ? batches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        batch_start = range.start;
+        nbatches    = range.count;
+    }
+
+    if (nbatches == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads     = octx->n_threads;
 
     const size_t src_batch_size         = src0->ne[0] * sizeof(float);
     const size_t dst_row_size           = dst->ne[0] * sizeof(float);
@@ -185,14 +215,15 @@ int op_diag_f32(struct htp_ops_context * octx) {
         .dst_row_size           = dst_row_size,
         .src_batch_size_aligned = src_batch_size_aligned,
         .dst_row_size_aligned   = dst_row_size_aligned,
-        .batches_per_thread     = (total_batches + n_threads - 1) / n_threads,
-        .total_batches          = total_batches,
+        .batches_per_thread     = fastdiv(nbatches + n_threads - 1, &octx->n_threads_div),
+        .total_batches          = nbatches,
+        .batch_start            = batch_start,
     };
 
     if (octx->ctx->vtcm_size < spad_per_thread * n_threads) {
-        worker_pool_run_func(octx->ctx->worker_pool, diag_thread_f32, &dctx, n_threads);
+        work_queue_run(octx->ctx->work_queue, diag_thread_f32, &dctx, n_threads);
     } else {
-        worker_pool_run_func(octx->ctx->worker_pool, diag_thread_f32_dma, &dctx, n_threads);
+        work_queue_run(octx->ctx->work_queue, diag_thread_f32_dma, &dctx, n_threads);
     }
 
     return HTP_STATUS_OK;

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

@@ -3,10 +3,11 @@
 #pragma clang diagnostic ignored "-Wunused-but-set-variable"
 
 #include <HAP_farf.h>
-#include <HAP_perf.h>
-
 #include <string.h>
 
+#include "hex-common.h"
+#include "hex-profile.h"
+
 #include "hvx-copy.h"
 #include "hvx-utils.h"
 
@@ -14,28 +15,30 @@
 #include "ggml-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"
 
 // ggml op_params layout for FILL:
 //   op_params[0] (as float) - the scalar fill value
 
-#define fill_preamble \
+#define fill_preamble                          \
     const struct htp_tensor * dst = octx->dst; \
-    \
-    const uint32_t ne0 = dst->ne[0]; \
-    const uint32_t ne1 = dst->ne[1]; \
-    const uint32_t ne2 = dst->ne[2]; \
-    const uint32_t ne3 = dst->ne[3]; \
-    \
-    const uint32_t nb1 = dst->nb[1]; \
-    const uint32_t nb2 = dst->nb[2]; \
-    const uint32_t nb3 = dst->nb[3]; \
-    \
+                                               \
+    const uint32_t ne0 = dst->ne[0];           \
+    const uint32_t ne1 = dst->ne[1];           \
+    const uint32_t ne2 = dst->ne[2];           \
+    const uint32_t ne3 = dst->ne[3];           \
+                                               \
+    const uint32_t nb1 = dst->nb[1];           \
+    const uint32_t nb2 = dst->nb[2];           \
+    const uint32_t nb3 = dst->nb[3];           \
+                                               \
     const uint32_t nr = ne1 * ne2 * ne3;
 
 struct htp_fill_context {
     struct htp_ops_context * octx;
     uint32_t nrows_per_thread;
     uint32_t total_rows;  // ne1 * ne2 * ne3
+    uint32_t row_start;
     bool     opt_path;
     HVX_Vector splat_vec;
     uint32_t   elem_size;
@@ -47,10 +50,15 @@ static void fill_thread(unsigned int nth, unsigned int ith, void * data) {
     fill_preamble;
 
     // Parallelise over the flat row index spanning ne1*ne2*ne3
-    const uint32_t ir0 = fctx->nrows_per_thread * ith;
-    const uint32_t ir1 = MIN(ir0 + fctx->nrows_per_thread, fctx->total_rows);
+    const uint32_t ir0 = fctx->row_start + fctx->nrows_per_thread * ith;
+    const uint32_t ir1 = MIN(ir0 + fctx->nrows_per_thread, fctx->row_start + fctx->total_rows);
 
-    uint64_t t1 = HAP_perf_get_qtimer_count();
+    if (ir0 >= ir1) {
+        return;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
 
     if (fctx->opt_path) {
         // Opt path: tensor is fully contiguous, treat as flat array
@@ -69,9 +77,8 @@ static void fill_thread(unsigned int nth, unsigned int ith, void * data) {
         }
     }
 
-    uint64_t t2 = HAP_perf_get_qtimer_count();
-    FARF(HIGH, "fill %u/%u: rows %u:%u usec %u\n",
-         ith, nth, ir0, ir1, (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir1);
+    FARF(HIGH, "fill %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
 }
 
 int op_fill(struct htp_ops_context * octx) {
@@ -85,8 +92,23 @@ int op_fill(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
+    uint32_t row_start = 0;
+    uint32_t nrows     = nr;
+
+    if (octx->ctx->mdev.count > 1) {
+        const uint32_t row_size = nb1;
+        const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(nr, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
     // nr = ne1*ne2*ne3 (flat row count across all outer dims); parallelise over it.
-    const uint32_t n_threads = MIN(nr, octx->n_threads);
+    const uint32_t n_threads = octx->n_threads;
 
     // Optimize if fully contiguous: skip stride arithmetic, treat as flat array
     const bool opt_path = (nb2 == nb1 * ne1) && (nb3 == nb2 * ne2);
@@ -99,8 +121,9 @@ int op_fill(struct htp_ops_context * octx) {
 
     struct htp_fill_context fctx = {
         .octx             = octx,
-        .nrows_per_thread = (nr + n_threads - 1) / n_threads,
-        .total_rows       = nr,
+        .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
+        .total_rows       = nrows,
+        .row_start        = row_start,
         .opt_path         = opt_path,
     };
 
@@ -117,7 +140,7 @@ int op_fill(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    worker_pool_run_func(octx->ctx->worker_pool, fill_thread, &fctx, n_threads);
+    work_queue_run(octx->ctx->work_queue, fill_thread, &fctx, n_threads);
 
     return HTP_STATUS_OK;
 }

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

@@ -5,7 +5,6 @@
 #include <assert.h>
 #include <HAP_compute_res.h>
 #include <HAP_farf.h>
-#include <HAP_perf.h>
 #include <math.h>
 #include <stdbool.h>
 #include <stdatomic.h>
@@ -75,6 +74,7 @@ struct htp_fa_context {
 
     uint32_t qrows;
     uint32_t qrows_per_thread;
+    uint32_t qrow_start;
 
     bool is_q_fp32;
 
@@ -89,8 +89,6 @@ struct htp_fa_context {
 
     const struct htp_tensor * k;
     const struct htp_tensor * v;
-
-    uint64_t t_start;
 };
 
 struct hmx_fa_context {
@@ -206,10 +204,9 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
     const uint32_t nb3 = dst->nb[3];
 
     // total rows in q
-    const uint32_t nr = factx->qrows;
-    const uint32_t dr = factx->qrows_per_thread;
-    const uint32_t ir0 = dr * ith;
-    const uint32_t ir1 = MIN(ir0 + dr, nr);
+    const uint32_t dr  = factx->qrows_per_thread;
+    const uint32_t ir0 = factx->qrow_start + dr * ith;
+    const uint32_t ir1 = MIN(ir0 + dr, factx->qrow_start + factx->qrows);
 
     if (ir0 >= ir1) return;
 
@@ -1888,6 +1885,24 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     const uint32_t n_threads = factx.n_threads;
     const uint32_t G = factx.G;
 
+    // Multi-device: split Q blocks across devices
+    const uint32_t n_q_blocks = (neq1 + Br - 1) / Br;
+    uint32_t q_start_min = 0;
+    uint32_t q_start_max = neq1;
+
+    if (octx->ctx->mdev.count > 1) {
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(n_q_blocks, htp_tensor_mdev_data_aligned(dst) ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        const uint32_t block_start = range.start;
+        const uint32_t block_end   = range.start + range.count;
+
+        if (block_start >= block_end) {
+            return HTP_STATUS_OK;
+        }
+
+        q_start_min = block_start * Br;
+        q_start_max = MIN(block_end * Br, neq1);
+    }
+
     // ======== VTCM allocation (GQA-aware) ========
     // K/V row sizes drive the DMA descriptors (not the VTCM layout) and are used
     // throughout the KV loop below.
@@ -1977,7 +1992,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     // ======== Main loop ========
     for (uint32_t ib3 = 0; ib3 < neq3; ++ib3) {
         const uint32_t im3 = mask ? fastmodulo(ib3, mask->ne[3], &factx.src3_div3) : 0;
-        for (uint32_t q_start = 0; q_start < neq1; q_start += Br) {
+        for (uint32_t q_start = q_start_min; q_start < q_start_max; q_start += Br) {
             const uint32_t n_rows_q    = hex_smin(Br, neq1 - q_start);
             const size_t   n_rows_g    = n_rows_q * G;
             const size_t   g_br_actual = hex_align_up(n_rows_g, HMX_FP16_TILE_N_ROWS);
@@ -1991,8 +2006,9 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
 
                 // 1. Push Q and KV DMAs for the very first iteration.
                 // Subsequent iterations are enqueued early at the end of the previous iteration.
-                if (ib3 == 0 && q_start == 0 && kv_head == 0) {
-                    const uint8_t * q_ptr = (const uint8_t *) q->data;
+                if (ib3 == 0 && q_start == q_start_min && kv_head == 0) {
+                    const uint8_t * q_ptr = (const uint8_t *) q->data + q_start * q->nb[1] +
+                                            (kv_head * factx.G) * q->nb[2] + ib3 * q->nb[3];
                     const size_t q_row_bytes = q_transposed ? n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
                     const size_t n_rows      = q_transposed ? factx.G : n_rows_q;
                     dma_queue_push(dma, dma_make_ptr(factx.vtcm_q_dma, q_ptr), q_row_bytes, hex_smax(q_src_stride, q_row_bytes), q_row_bytes, n_rows);
@@ -2311,8 +2327,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                 if (next_kv_head >= n_kv_heads) {
                     next_kv_head = 0;
                     next_q_start = q_start + Br;
-                    if (next_q_start >= neq1) {
-                        next_q_start = 0;
+                    if (next_q_start >= q_start_max) {
+                        next_q_start = q_start_min;
                         next_ib3     = ib3 + 1;
                     }
                 }
@@ -2398,6 +2414,10 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     if (kparams->kernel_type == HTP_FA_KERNEL_HMX) {
         return hmx_flash_attn_ext(octx);
     }
@@ -2407,8 +2427,6 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
     factx.k = k;
     factx.v = v;
 
-    factx.t_start = HAP_perf_get_qtimer_count();
-
     factx.src0_div21 = kparams->u.hvx.src0_div21;
     factx.src0_div1  = kparams->u.hvx.src0_div1;
 
@@ -2451,8 +2469,30 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
     }
 
     // total rows in q
-    factx.qrows = kparams->qrows;
-    factx.qrows_per_thread = kparams->qrows_per_thread;
+    const uint32_t neq1 = q->ne[1];
+    const uint32_t neq2 = q->ne[2];
+    const uint32_t neq3 = q->ne[3];
+    const uint32_t total_qrows = neq1 * neq2 * neq3;
+
+    uint32_t qrow_start = 0;
+    uint32_t qrows      = total_qrows;
+
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_mdev_data_aligned(dst) && ((dst->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_qrows, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        qrow_start = range.start;
+        qrows      = range.count;
+    }
+
+    if (qrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+
+    factx.qrows            = qrows;
+    factx.qrow_start       = qrow_start;
+    factx.qrows_per_thread = fastdiv(qrows + n_threads - 1, &octx->n_threads_div);
 
     size_t size_vkq_acc = hex_round_up(v->ne[0] * sizeof(float), 128); // VKQ32
 
@@ -2461,18 +2501,18 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
 
     uint8_t * vtcm_cur = octx->ctx->vtcm_base;
 
-    factx.spad_q = vtcm_seq_alloc(&vtcm_cur, size_q_block * octx->n_threads);
-    factx.spad_k = vtcm_seq_alloc(&vtcm_cur, factx.size_k_block * 2 * octx->n_threads);
-    factx.spad_v = vtcm_seq_alloc(&vtcm_cur, factx.size_v_block * 2 * octx->n_threads);
-    factx.spad_m = vtcm_seq_alloc(&vtcm_cur, (mask ? factx.size_m_block * HVX_FA_DMA_CACHE_SIZE : 0) * octx->n_threads);
-    factx.spad_a = vtcm_seq_alloc(&vtcm_cur, size_vkq_acc * octx->n_threads);
+    factx.spad_q = vtcm_seq_alloc(&vtcm_cur, size_q_block * n_threads);
+    factx.spad_k = vtcm_seq_alloc(&vtcm_cur, factx.size_k_block * 2 * n_threads);
+    factx.spad_v = vtcm_seq_alloc(&vtcm_cur, factx.size_v_block * 2 * n_threads);
+    factx.spad_m = vtcm_seq_alloc(&vtcm_cur, (mask ? factx.size_m_block * HVX_FA_DMA_CACHE_SIZE : 0) * n_threads);
+    factx.spad_a = vtcm_seq_alloc(&vtcm_cur, size_vkq_acc * n_threads);
 
     if ((size_t) (vtcm_cur - octx->ctx->vtcm_base) > octx->ctx->vtcm_size) {
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
     if (!(octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) {
-        work_queue_run(octx->ctx->work_queue, flash_attn_ext_f16_thread, &factx, octx->n_threads);
+        work_queue_run(octx->ctx->work_queue, flash_attn_ext_f16_thread, &factx, n_threads);
     }
 
     return HTP_STATUS_OK;

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

@@ -51,6 +51,7 @@ struct htp_fa_kernel_params {
 
     uint32_t qrows;
     uint32_t qrows_per_thread;
+    uint32_t qrow_start;
     float    m0;
     float    m1;
     uint32_t n_head_log2;

ggml/src/ggml-hexagon/htp/gated-delta-net-ops.c

@@ -4,10 +4,13 @@
 
 #include "hvx-utils.h"
 #include "hex-fastdiv.h"
+#include "hex-common.h"
+#include "hex-profile.h"
 
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
 #include "htp-ctx.h"
+#include "htp-tensor.h"
 
 #ifndef MIN
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -22,6 +25,8 @@ struct htp_gdn_context {
     size_t   state_bytes;
     uint8_t * vtcm_base;
     size_t   vtcm_per_thread;
+    uint32_t row_start;
+    uint32_t nrows;
 };
 
 static inline HVX_Vector gdn_mul_dot_f32(float * restrict dst, const float * restrict mul, const float * restrict dot, uint32_t n) {
@@ -586,8 +591,9 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo
     const uint32_t n_seqs   = v->ne[3];
     const uint32_t K        = octx->op_params[0];
 
-    const uint32_t total_rows = H * n_seqs;
-    if (ith >= total_rows) {
+    const uint32_t row_end = gctx->row_start + gctx->nrows;
+
+    if (ith >= gctx->nrows) {
         return;
     }
 
@@ -621,11 +627,11 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo
     const uint64_t state_seq_stride = state->nb[3] / sizeof(float);
     const uint64_t state_size_per_snap = (uint64_t) S_v * S_v * H * n_seqs;
 
-    uint32_t ir_prefetch = ith;
+    uint32_t ir_prefetch = gctx->row_start + ith;
     int spad_idx = 0;
 
     // Prefetch preamble (up to 2 steps)
-    for (int k = 0; k < 2 && ir_prefetch < total_rows; k++) {
+    for (int k = 0; k < 2 && ir_prefetch < row_end; k++) {
         const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H);
         const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H);
         const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v;
@@ -646,8 +652,11 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo
         spad_idx ^= 1;
     }
 
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (gctx->row_start + ith));
+
     int curr_spad_idx = 0;
-    for (uint32_t ir = ith; ir < total_rows; ir += nth) {
+    for (uint32_t ir = gctx->row_start + ith; ir < row_end; ir += nth) {
         dma_queue_pop(dma);
         dma_queue_pop(dma);
 
@@ -812,7 +821,7 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo
                        S_v * sizeof(float), S_v);
 
         // Prefetch next block (if any)
-        if (ir_prefetch < total_rows) {
+        if (ir_prefetch < row_end) {
             const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H);
             const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H);
             const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v;
@@ -828,6 +837,7 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo
         curr_spad_idx ^= 1;
     }
     dma_queue_flush(dma);
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row_end);
 }
 
 
@@ -847,8 +857,9 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo
     const uint32_t H        = v->ne[1];
     const uint32_t n_seqs   = v->ne[3];
 
-    const uint32_t total_rows = H * n_seqs;
-    if (ith >= total_rows) {
+    const uint32_t row_end = gctx->row_start + gctx->nrows;
+
+    if (ith >= gctx->nrows) {
         return;
     }
 
@@ -881,11 +892,11 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo
 
     const uint64_t state_seq_stride = state->nb[3] / sizeof(float);
 
-    uint32_t ir_prefetch = ith;
+    uint32_t ir_prefetch = gctx->row_start + ith;
     int spad_idx = 0;
 
     // Prefetch preamble (up to 2 steps)
-    for (int k = 0; k < 2 && ir_prefetch < total_rows; k++) {
+    for (int k = 0; k < 2 && ir_prefetch < row_end; k++) {
         const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H);
         const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H);
         const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v;
@@ -906,8 +917,11 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo
         spad_idx ^= 1;
     }
 
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (gctx->row_start + ith));
+
     int curr_spad_idx = 0;
-    for (uint32_t ir = ith; ir < total_rows; ir += nth) {
+    for (uint32_t ir = gctx->row_start + ith; ir < row_end; ir += nth) {
         dma_queue_pop(dma);
         dma_queue_pop(dma);
 
@@ -1057,7 +1071,7 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo
                        S_v * sizeof(float), S_v);
 
         // Prefetch next block (if any)
-        if (ir_prefetch < total_rows) {
+        if (ir_prefetch < row_end) {
             const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H);
             const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H);
             const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v;
@@ -1073,6 +1087,7 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo
         curr_spad_idx ^= 1;
     }
     dma_queue_flush(dma);
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row_end);
 }
 
 
@@ -1085,10 +1100,6 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
     const struct htp_tensor * state = octx->src[5];
     const struct htp_tensor * dst   = octx->dst;
 
-    if (!q || !k || !v || !g || !beta || !state || !dst) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-
     if (q->type != HTP_TYPE_F32 || k->type != HTP_TYPE_F32 || v->type != HTP_TYPE_F32 ||
         g->type != HTP_TYPE_F32 || beta->type != HTP_TYPE_F32 || state->type != HTP_TYPE_F32 ||
         dst->type != HTP_TYPE_F32) {
@@ -1124,16 +1135,37 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
+    const uint32_t total_rows = H * n_seqs;
+
+    uint32_t row_start = 0;
+    uint32_t nrows     = total_rows;
+
+    if (octx->ctx->mdev.count > 1) {
+        const uint32_t head_bytes = S_v * sizeof(float);
+        const uint32_t rows_per_chunk = (head_bytes > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(head_bytes, HEX_L2_LINE_SIZE)) : 1;
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0,
+                                                                             octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+
     struct htp_gdn_context gctx;
     gctx.octx = octx;
-    gctx.rows_per_thread = (H * n_seqs + octx->n_threads - 1) / octx->n_threads;
+    gctx.row_start       = row_start;
+    gctx.nrows           = nrows;
+    gctx.rows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
     gctx.state_bytes = (size_t) S_v * S_v * sizeof(float);
 
     size_t state_aligned = (size_t) S_v * S_v * sizeof(float);
     state_aligned = (state_aligned + 127) & ~(size_t)127;
 
-    assert(octx->ctx->vtcm_base != NULL);
-    assert(octx->ctx->vtcm_size >= 2 * state_aligned * octx->n_threads);
+    assert(octx->ctx->vtcm_size >= 2 * state_aligned * n_threads);
 
     gctx.vtcm_base = octx->ctx->vtcm_base;
     gctx.vtcm_per_thread = 2 * state_aligned;
@@ -1148,9 +1180,9 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
          gctx.vtcm_per_thread * octx->n_threads, octx->n_threads);
 
     if (n_tokens == 1) {
-        worker_pool_run_func(octx->ctx->worker_pool, gated_delta_net_f32_tg_thread, &gctx, octx->n_threads);
+        work_queue_run(octx->ctx->work_queue, gated_delta_net_f32_tg_thread, &gctx, n_threads);
     } else {
-        worker_pool_run_func(octx->ctx->worker_pool, gated_delta_net_f32_pp_thread, &gctx, octx->n_threads);
+        work_queue_run(octx->ctx->work_queue, gated_delta_net_f32_pp_thread, &gctx, n_threads);
     }
 
     return HTP_STATUS_OK;

ggml/src/ggml-hexagon/htp/get-rows-ops.c

@@ -10,6 +10,7 @@
 
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
+#include "hex-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
@@ -23,9 +24,12 @@ struct get_rows_context {
     const struct htp_get_rows_kernel_params * kparams;
     struct htp_get_rows_vtcm_layout vtcm_layout;
     uint8_t * vtcm_base;
+    uint32_t task_start;
+    uint32_t tasks;
+    uint32_t tasks_per_thread;
 };
 
-#define get_rows_preamble \
+#define get_rows_preamble                      \
     const uint32_t ne00 = octx->src[0]->ne[0]; \
     const uint32_t ne01 = octx->src[0]->ne[1]; \
     const uint32_t ne02 = octx->src[0]->ne[2]; \
@@ -61,12 +65,12 @@ static void get_rows_thread_st_##IDX_TYPE(unsigned int nth, unsigned int ith, vo
     struct htp_ops_context * octx = grctx->octx;                                                                       \
     const struct htp_get_rows_kernel_params * kparams = grctx->kparams;                                                \
     get_rows_preamble;                                                                                                 \
-    const uint32_t dr  = kparams->tasks_per_thread;                                                                    \
-    const uint32_t ir0 = dr * ith;                                                                                     \
-    if (ir0 >= kparams->total_tasks) {                                                                                 \
+    const uint32_t dr  = grctx->tasks_per_thread;                                                                      \
+    const uint32_t ir0 = grctx->task_start + dr * ith;                                                                 \
+    if (ir0 >= grctx->task_start + grctx->tasks) {                                                                     \
         return;                                                                                                        \
     }                                                                                                                  \
-    const uint32_t ir1 = MIN(ir0 + dr, kparams->total_tasks);                                                          \
+    const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks);                                              \
     const uint32_t row_size_bytes = htp_tensor_get_row_size(octx->src[0]->type, ne00);                                 \
     dma_queue * dma_queue = octx->ctx->dma[ith];                                                                       \
     for (uint32_t i = ir0; i < ir1; ++i) {                                                                             \
@@ -101,12 +105,12 @@ static void get_rows_thread_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned
     const struct htp_get_rows_kernel_params * kparams = grctx->kparams;                                                \
     get_rows_preamble;                                                                                                 \
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                             \
-    const uint32_t dr  = kparams->tasks_per_thread;                                                                    \
-    const uint32_t ir0 = dr * ith;                                                                                     \
-    if (ir0 >= kparams->total_tasks) {                                                                                 \
+    const uint32_t dr  = grctx->tasks_per_thread;                                                                      \
+    const uint32_t ir0 = grctx->task_start + dr * ith;                                                                 \
+    if (ir0 >= grctx->task_start + grctx->tasks) {                                                                     \
         return;                                                                                                        \
     }                                                                                                                  \
-    const uint32_t ir1 = MIN(ir0 + dr, kparams->total_tasks);                                                          \
+    const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks);                                              \
     const uint32_t chunks_per_row = kparams->chunks_per_row;                                                           \
     const uint32_t chunk_size     = kparams->chunk_size;                                                               \
     dma_queue * dma_queue = octx->ctx->dma[ith];                                                                       \
@@ -225,13 +229,41 @@ int op_get_rows(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
+    const struct htp_tensor * dst = octx->dst;
+    const uint32_t total_tasks    = kparams->total_tasks;
+    const size_t dst_row_size     = htp_tensor_get_row_size(dst->type, dst->ne[0]);
+
+    uint32_t task_start = 0;
+    uint32_t tasks      = total_tasks;
+
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t tasks_per_chunk = 1;
+        htp_tensor_mdev_rows_per_chunk(dst, dst_row_size / dst->ne[0], (uint32_t) dst_row_size, &tasks_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_tasks, tasks_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        task_start = range.start;
+        tasks      = range.count;
+    }
+
+    if (tasks == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    if (!htp_ops_context_set_n_threads(octx, (uint32_t) kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+
     struct get_rows_context grctx;
     grctx.octx = octx;
     grctx.kparams = kparams;
     grctx.vtcm_base = (uint8_t *)octx->ctx->vtcm_base;
+    grctx.task_start = task_start;
+    grctx.tasks = tasks;
+    grctx.tasks_per_thread = fastdiv(tasks + n_threads - 1, &octx->n_threads_div);
 
     const uint32_t ne00 = octx->src[0]->ne[0];
-    htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, octx->src[0]->type, ne00, kparams->n_threads);
+    htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, octx->src[0]->type, ne00, n_threads);
 
     const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32);
 
@@ -247,14 +279,14 @@ int op_get_rows(struct htp_ops_context * octx) {
         }
     }
 
-    FARF(HIGH, "get-rows: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-vtcm-size %zu dst-vtcm-size %zu use_dma=%d n_threads %d\n",
+    FARF(HIGH, "get-rows: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-vtcm-size %zu dst-vtcm-size %zu use-dma %d n-threads %d\n",
          octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3],
          octx->src[1]->ne[0], octx->src[1]->ne[1], octx->src[1]->ne[2], octx->src[1]->ne[3],
          octx->dst->ne[0], octx->dst->ne[1], octx->dst->ne[2], octx->dst->ne[3],
-         grctx.vtcm_layout.src0_bytes_per_thread * kparams->n_threads,
-         grctx.vtcm_layout.dst_bytes_per_thread  * kparams->n_threads,
-         kparams->use_dma, kparams->n_threads);
+         grctx.vtcm_layout.src0_bytes_per_thread * n_threads,
+         grctx.vtcm_layout.dst_bytes_per_thread  * n_threads,
+         kparams->use_dma, n_threads);
 
-    work_queue_run(octx->ctx->work_queue, q_func, &grctx, kparams->n_threads);
+    work_queue_run(octx->ctx->work_queue, q_func, &grctx, n_threads);
     return HTP_STATUS_OK;
 }

ggml/src/ggml-hexagon/htp/hex-common.h

@@ -77,4 +77,13 @@ static inline bool hex_add_overflow(size_t a, size_t b, size_t *out) {
     return false;
 }
 
+static inline uint32_t hex_gcd_u32(uint32_t a, uint32_t b) {
+    while (b != 0) {
+        uint32_t t = b;
+        b = a % b;
+        a = t;
+    }
+    return a;
+}
+
 #endif // HEX_COMMON_H

ggml/src/ggml-hexagon/htp/hex-utils.h

@@ -39,7 +39,6 @@ static inline void hex_l2fetch_block(const void * addr, size_t size) {
 
 #define HEX_L2_LINE_SIZE           128
 #define HEX_L2_BLOCK_SIZE          (HEX_L2_LINE_SIZE * 4) // flush granularity (lines per loop iteration)
-#define HEX_L2_FLUSH_IL_THRESHOLD  1024                   // inline flush threshold
 #define HEX_L2_FLUSH_WQ_THRESHOLD  (4 * 1024)
 #define HEX_L2_FLUSH_ALL_THRESHOLD (4 * 1024 * 1024)
 

ggml/src/ggml-hexagon/htp/hmx-utils.h

@@ -27,7 +27,7 @@ static inline void hmx_init_column_scales(void *out_scales, HVX_Vector v_scale)
 // vscatter offsets for fused dequant+transpose: write K-values directly to [K][N] tile.
 // word[i] = i*128 maps K-row-pair i to byte offset i*128.
 // Column offset (n*4) is added at runtime.  Entries 0..15 cover one tile (region 2047);
-// entries 16..31 cover the next adjacent tile (region 4095) — pick region size at the
+// entries 16..31 cover the next adjacent tile (region 4095) - pick region size at the
 // call site to scatter into one tile (masked) or two contiguous tiles (unmasked).
 static const int32_t hmx_transpose_scatter_offsets[32] __attribute__((aligned(VLEN))) = {
     0 * 128,  1 * 128,  2 * 128,  3 * 128,  4 * 128,  5 * 128,  6 * 128,  7 * 128,  8 * 128,  9 * 128,  10 * 128,
@@ -198,16 +198,16 @@ static inline void hmx_interleave_cols_to_tiles(__fp16 * restrict tiles_out,
 }
 
 // --- HMX inline asm macros for load-store packetization ---
-#define HMX_LOAD_MPY_F16(act, wt, range) \
-    "{\n" \
+#define HMX_LOAD_MPY_F16(act, wt, range)              \
+    "{\n"                                             \
     "    activation.hf = mxmem(" act ", " range ")\n" \
-    "    weight.hf = mxmem(" wt ", " range ")\n" \
+    "    weight.hf = mxmem(" wt ", " range ")\n"      \
     "}\n"
 
-#define HMX_LOAD_MPY_DEEP_F16(act, wt, range) \
-    "{\n" \
+#define HMX_LOAD_MPY_DEEP_F16(act, wt, range)              \
+    "{\n"                                                  \
     "    activation.hf = mxmem(" act ", " range "):deep\n" \
-    "    weight.hf = mxmem(" wt ", " range ")\n" \
+    "    weight.hf = mxmem(" wt ", " range ")\n"           \
     "}\n"
 
 #define HMX_STORE_AFTER_F16(out, scale_reg) \

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

@@ -19,7 +19,7 @@
 #endif
 #define HTP_MAX_MMAPS    16
 
-#define HTP_MAX_DIRTY_RANGES 16
+#define HTP_MAX_DIRTY_RANGES 32
 
 // Memory mapping
 struct htp_mmap {
@@ -29,6 +29,11 @@ struct htp_mmap {
     uint32_t reserved;
 };
 
+struct htp_dirty_range {
+    uint32_t start;
+    uint32_t end;
+};
+
 // Scratchpad state
 struct htp_spad {
     const struct htp_tensor * src;             // original src of the data (for reuse)
@@ -38,6 +43,14 @@ struct htp_spad {
     uint32_t                  size_per_thread; // size per thread
 };
 
+struct htp_mdev_group {
+    uint16_t              idx;
+    uint16_t              count;
+    struct fastdiv_values count_div;
+    uint8_t *             fence_base;
+    uint32_t              fence_seq;
+};
+
 struct htp_context;
 
 // Context while processing an Op
@@ -65,8 +78,10 @@ struct htp_ops_context {
     struct htp_spad src3_spad;
     struct htp_spad dst_spad;
 
-    uint32_t n_threads;
-    uint32_t flags;
+    uint32_t              flags;
+    uint32_t              n_threads;
+    struct fastdiv_values n_threads_div;
+    int                   status;
 };
 
 // Main context for htp DSP backend
@@ -76,6 +91,7 @@ struct htp_context {
     struct htp_mmap        mmap[HTP_MAX_MMAPS];
     dma_queue_t            dma[HTP_MAX_NTHREADS];
     dma_queue_t            dma_cached[HTP_MAX_NTHREADS];
+    struct htp_thread_trace trace[HTP_MAX_NTHREADS + 1];
     work_queue_t           work_queue;
     hmx_queue_t            hmx_queue;
 
@@ -88,7 +104,6 @@ struct htp_context {
     bool                   hmx_enabled;
     bool                   etm;
     uint32_t               profiler;
-    struct htp_thread_trace trace[HTP_MAX_NTHREADS + 1];
 
     uint8_t *              vtcm_base;
     size_t                 vtcm_size;
@@ -97,16 +112,13 @@ struct htp_context {
     atomic_bool            vtcm_needs_release;
 
     uint64_t               max_vmem;
-    struct htp_dirty_range {
-        uint32_t start;
-        uint32_t end;
-        uint32_t bi;
-    } dirty_ranges[HTP_MAX_DIRTY_RANGES];
+    struct htp_dirty_range dirty_ranges[HTP_MAX_DIRTY_RANGES];
 
     // Persistent DDR scratchpad for MUL_MAT_ID mappings
     void *                 ddr_spad_base;
     size_t                 ddr_spad_size;
 
+    struct htp_mdev_group  mdev;
     struct htp_ops_context octx;
 
     qurt_thread_t          main_thread;
@@ -115,6 +127,27 @@ struct htp_context {
     size_t                 footprint;
 };
 
+static inline bool htp_ops_context_set_n_threads(struct htp_ops_context * octx, uint32_t n_threads) {
+    if (n_threads == 0 || n_threads > octx->ctx->n_threads) {
+        return false;
+    }
+
+    if (n_threads != octx->n_threads) {
+        octx->n_threads = n_threads;
+        octx->n_threads_div = n_threads == octx->ctx->n_threads
+            ? octx->ctx->n_threads_div
+            : init_fastdiv_values(n_threads);
+    }
+
+    return true;
+}
+
+static inline void htp_ops_context_set_status(struct htp_ops_context * octx, int status) {
+    if (status > HTP_STATUS_OK && octx->status == HTP_STATUS_OK) {
+        octx->status = status;
+    }
+}
+
 int op_matmul(struct htp_ops_context * octx);
 int op_matmul_id(struct htp_ops_context * octx);
 int op_matmul_nx(struct htp_ops_context * octx);

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

@@ -0,0 +1,89 @@
+#ifndef HTP_FENCE_H
+#define HTP_FENCE_H
+
+#include <stdatomic.h>
+#include <stdint.h>
+
+#include <HAP_farf.h>
+
+#include "hex-utils.h"
+#include "htp-ops.h"
+#include "htp-ctx.h"
+
+static inline atomic_uint * htp_mdev_fence_slot(const void * fence_base, uint32_t idx) {
+    return (atomic_uint *) ((const uint8_t *) fence_base + (size_t) idx * HTP_FENCE_SLOT_SIZE);
+}
+
+static inline void htp_fence_write(void * fence_ptr, uint32_t seq, uint32_t status) {
+    atomic_uint * fence = (atomic_uint *) fence_ptr;
+    atomic_store(&fence[1], status);
+    atomic_store(&fence[0], seq);
+    asm volatile ("syncht" : : : "memory");
+    Q6_dccleaninva_A((void *) fence);
+}
+
+static inline void htp_fence_read(const void * fence_ptr, uint32_t * seq, uint32_t * status) {
+    const atomic_uint * fence = (const atomic_uint *) fence_ptr;
+    Q6_dccleaninva_A((void *) fence);
+    asm volatile ("syncht" : : : "memory");
+    *seq = atomic_load(&fence[0]);
+    *status = atomic_load(&fence[1]);
+}
+
+static inline void htp_mdev_group_barrier(struct htp_ops_context * octx) {
+    struct htp_context * ctx = octx->ctx;
+    if (ctx->mdev.count <= 1) {
+        return;
+    }
+
+    const uint32_t seq = ++ctx->mdev.fence_seq;
+
+    struct htp_thread_trace * tr = &ctx->trace[0];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq);
+
+    const uint32_t mdev_idx   = ctx->mdev.idx;
+    const uint32_t mdev_count = ctx->mdev.count;
+
+    uint8_t * fence_base   = ctx->mdev.fence_base;
+    atomic_uint * my_fence = htp_mdev_fence_slot(fence_base, mdev_idx);
+    htp_fence_write(my_fence, seq, octx->status);
+
+    for (uint32_t d = 0; d < mdev_count; d++) {
+        if (d == mdev_idx) continue;
+        atomic_uint * peer_fence = htp_mdev_fence_slot(fence_base, d);
+        uint64_t spins = 0;
+        while (1) {
+            uint32_t peer_seq;
+            uint32_t peer_status;
+            htp_fence_read(peer_fence, &peer_seq, &peer_status);
+            if ((int32_t)(peer_seq - seq) >= 0) {
+                if (peer_status > HTP_STATUS_OK) {
+                    FARF(ERROR, "ggml-hex: mdev %u peer %u failed with status %u : seq 0x%08x\n",
+                         mdev_idx, d, peer_status, seq);
+                    htp_ops_context_set_status(octx, peer_status);
+                }
+                break;
+            }
+            if (++spins == 10000) {
+                FARF(ALWAYS, "ggml-hex: mdev %u waiting for mdev %u : seq 0x%08x (b %u op %u) my-fence %p peer-fence %p peer-seq 0x%08x (diff %d)\n",
+                     mdev_idx, d, seq, seq >> 12, seq & 0xfff, my_fence, peer_fence, peer_seq, (int32_t)(peer_seq - seq));
+            }
+            if (spins > HTP_FENCE_TIMEOUT) {
+                FARF(ERROR, "ggml-hex: mdev %u timeout waiting for mdev %u : seq 0x%08x (b %u op %u) peer-fence %p peer-seq 0x%08x\n",
+                     mdev_idx, d, seq, seq >> 12, seq & 0xfff, peer_fence, peer_seq);
+                htp_ops_context_set_status(octx, HTP_STATUS_INTERNAL_ERR);
+                break;
+            }
+            hex_pause();
+        }
+    }
+    asm volatile ("syncht" : : : "memory");
+
+    if (octx->status > HTP_STATUS_OK) {
+        htp_fence_write(my_fence, seq, octx->status);
+    }
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq);
+}
+
+#endif // HTP_FENCE_H

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

@@ -77,6 +77,7 @@ enum htp_op_code {
     HTP_OP_GET_ROWS,
     HTP_OP_SCALE,
     HTP_OP_CPY,
+    HTP_OP_CPY_FENCE,
     HTP_OP_ARGSORT,
     HTP_OP_SQR,
     HTP_OP_SQRT,
@@ -100,6 +101,7 @@ enum htp_op_code {
     HTP_OP_ALLREDUCE,
     HTP_OP_ALLREDUCE_ADD,
     HTP_OP_GLU_SWIGLU_CLAMP,
+    HTP_OP_MDEV_GROUP,
 
     HTP_OP_INVALID
 };
@@ -114,6 +116,7 @@ enum htp_op_code {
 #define HTP_OP_MAX_TENSORS 8192 // must stay under 64K (uint16)
 
 #define HTP_FENCE_TIMEOUT  (1000000000ULL)
+#define HTP_FENCE_SLOT_SIZE 128
 
 #define HTP_OP_MAX_VMEM_DEFAULT (3355443200u)
 
@@ -214,30 +217,26 @@ struct htp_prof_desc {
 };
 
 struct htp_opbatch_req {
-    uint32_t id;          // Batch id
+    uint64_t seq;         // Sequence number
     uint32_t n_bufs;      // Number of buffers
     uint32_t n_tensors;   // Number of tensors
     uint32_t n_ops;       // Number of ops
     uint32_t n_traces;    // Number of trace descriptors per thread
-    uint32_t pad;         // unused
-    uint64_t seq;         // Sequence number
     // struct htp_buf_desc  bufs[];    -- dspqueue buf 0
     // struct htp_tensor    tensors[]; -- dspqueue buf 0
     // struct htp_op_desc   ops[];     -- dspqueue buf 0
 };
 
 struct htp_opbatch_rsp {
-    uint32_t id;         // Batch id
-    uint32_t status;     // HTP_STATUS_...
-    uint32_t n_bufs;     // Number of buffers
-    uint32_t n_tensors;  // Number of tensors
-    uint32_t n_ops;      // Number of op profile descriptors
-    uint32_t n_traces[HTP_MAX_NTHREADS + 1];
-    uint32_t usecs;          // Number of usec
-    uint32_t pad;            // align to 8 bytes
+    uint64_t seq;            // Sequence number
     uint64_t cycles_start;   // Start cycle counter
     uint64_t cycles_stop;    // Stop cycle counter
-    uint64_t seq;            // Sequence number
+    uint32_t status;         // HTP_STATUS_...
+    uint32_t n_bufs;         // Number of buffers
+    uint32_t n_tensors;      // Number of tensors
+    uint32_t n_ops;          // Number of op profile descriptors
+    uint32_t usecs;          // Number of usec
+    uint32_t n_traces[HTP_MAX_NTHREADS + 1];
     // struct htp_prof_desc profs[];  -- dspqueue buf 0
 };
 

ggml/src/ggml-hexagon/htp/htp-tensor.c

@@ -20,7 +20,7 @@ struct l2flush_range {
 
 struct l2flush_multi_task {
     struct htp_thread_trace * trace;
-    struct l2flush_range      ranges[HTP_OP_MAX_INPUTS];
+    struct l2flush_range      ranges[HTP_MAX_DIRTY_RANGES];
     uint32_t                  n_ranges;
     uint32_t                  total_blocks;
     uint32_t                  blocks_per_thread;
@@ -73,6 +73,27 @@ static void l2flush_multi_worker(unsigned int n, unsigned int i, void * data) {
     htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, gb_first);
 }
 
+static void merge_dirty_ranges(struct htp_context * ctx) {
+    for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
+        struct htp_dirty_range * r = &ctx->dirty_ranges[i];
+        if (!r->start) continue;
+
+        for (uint32_t j = 0; j < HTP_MAX_DIRTY_RANGES;) {
+            struct htp_dirty_range * s = &ctx->dirty_ranges[j];
+            if (i == j || !s->start || r->end < s->start || s->end < r->start) {
+                j++;
+                continue;
+            }
+
+            r->start = MIN(r->start, s->start);
+            r->end   = MAX(r->end, s->end);
+            s->start = 0;
+            s->end   = 0;
+            j = 0;
+        }
+    }
+}
+
 void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
     const struct htp_tensor * pending[HTP_OP_MAX_OUTPUTS];
     uint32_t n_pending = 0;
@@ -83,11 +104,6 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co
             continue;
         }
 
-        if (t->size <= HEX_L2_FLUSH_IL_THRESHOLD) {
-            hex_l2flush((void *) (uintptr_t) t->data, t->size);
-            continue;
-        }
-
         uint32_t t_start = t->data;
         uint32_t t_end   = t_start + t->size;
 
@@ -110,6 +126,8 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co
         }
     }
 
+    merge_dirty_ranges(ctx);
+
     if (n_pending == 0) {
         return;
     }
@@ -132,8 +150,8 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co
             struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
             r->start = pending[i]->data;
             r->end   = pending[i]->data + pending[i]->size;
-            r->bi    = pending[i]->bi;
         }
+        merge_dirty_ranges(ctx);
         return;
     }
 
@@ -151,12 +169,12 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co
             struct htp_dirty_range * r = &ctx->dirty_ranges[i];
             r->start = pending[i]->data;
             r->end   = pending[i]->data + pending[i]->size;
-            r->bi    = pending[i]->bi;
         }
+        merge_dirty_ranges(ctx);
         return;
     }
 
-    if (total_evict_size > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1 && n_evict <= HTP_OP_MAX_INPUTS) {
+    if (total_evict_size > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1 && n_evict <= HTP_MAX_DIRTY_RANGES) {
         struct l2flush_multi_task task;
         task.trace    = ctx->trace;
         task.n_ranges = n_evict;
@@ -195,7 +213,6 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co
         struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
         r->start = pending[i]->data;
         r->end   = pending[i]->data + pending[i]->size;
-        r->bi    = pending[i]->bi;
     }
 
     for (uint32_t i = 0; i < n_empty; i++) {
@@ -203,8 +220,9 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co
         struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
         r->start = pending[n_evict + i]->data;
         r->end   = pending[n_evict + i]->data + pending[n_evict + i]->size;
-        r->bi    = pending[n_evict + i]->bi;
     }
+
+    merge_dirty_ranges(ctx);
 }
 
 static void make_tensor_clean(struct htp_context * ctx, const struct htp_tensor * t) {
@@ -242,15 +260,77 @@ static inline bool is_tensor_dirty(struct htp_context * ctx, const struct htp_te
     return false;
 }
 
+static void flush_dirty_ranges(struct htp_context * ctx, const struct htp_dirty_range * ranges, uint32_t n_ranges, uint64_t total_dirty) {
+    if (total_dirty >= HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
+        struct l2flush_multi_task task;
+        task.trace    = ctx->trace;
+        task.n_ranges = n_ranges;
+
+        uint32_t block_acc = 0;
+        for (uint32_t i = 0; i < n_ranges; i++) {
+            const struct htp_dirty_range * r = &ranges[i];
+            struct l2flush_range * rg = &task.ranges[i];
+            rg->start = hex_align_down((size_t) r->start, HEX_L2_LINE_SIZE);
+            rg->end   = hex_align_up((size_t) r->end, HEX_L2_LINE_SIZE);
+            rg->block_first = block_acc;
+            rg->n_blocks = (rg->end - rg->start + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
+            block_acc += rg->n_blocks;
+        }
+
+        task.total_blocks      = block_acc;
+        task.blocks_per_thread = fastdiv(block_acc + ctx->n_threads - 1, &ctx->n_threads_div);
+
+        work_queue_run(ctx->work_queue, l2flush_multi_worker, &task, ctx->n_threads);
+    } else {
+        struct htp_thread_trace * tr = &ctx->trace[0];
+        htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
+        for (uint32_t i = 0; i < n_ranges; i++) {
+            const struct htp_dirty_range * r = &ranges[i];
+            hex_l2flush((void *) (uintptr_t) r->start, r->end - r->start);
+        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
+    }
+}
+
+void htp_flush_dirty_ranges(struct htp_context * ctx) {
+    struct htp_dirty_range ranges[HTP_MAX_DIRTY_RANGES];
+    uint32_t n_ranges = 0;
+    uint64_t total_dirty = 0;
+
+    for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
+        const struct htp_dirty_range * r = &ctx->dirty_ranges[i];
+        if (!r->start) {
+            continue;
+        }
+        ranges[n_ranges++] = *r;
+        total_dirty += r->end - r->start;
+    }
+
+    if (total_dirty == 0) {
+        return;
+    }
+
+    if (total_dirty > HEX_L2_FLUSH_ALL_THRESHOLD) {
+        flush_all_dcache(ctx);
+        return;
+    }
+
+    flush_dirty_ranges(ctx, ranges, n_ranges, total_dirty);
+    memset(ctx->dirty_ranges, 0, sizeof(ctx->dirty_ranges));
+}
+
 void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
     const struct htp_tensor * dirty_tensors[HTP_OP_MAX_INPUTS];
+    struct htp_dirty_range ranges[HTP_OP_MAX_INPUTS];
     uint32_t n_dirty = 0;
     uint64_t total_dirty = 0;
 
     for (uint32_t i = 0; i < n; i++) {
         const struct htp_tensor * t = tensors[i];
-        if (t && !(t->flags & (HTP_TENSOR_WEIGHT | HTP_TENSOR_FENCE)) && is_tensor_dirty(ctx, t)) {
+        if (t && is_tensor_dirty(ctx, t)) {
             dirty_tensors[n_dirty++] = t;
+            ranges[n_dirty - 1].start = t->data;
+            ranges[n_dirty - 1].end   = t->data + t->size;
             total_dirty += t->size;
         }
     }
@@ -264,37 +344,8 @@ void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * co
         return;
     }
 
-    if (total_dirty >= HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
-        struct l2flush_multi_task task;
-        task.trace    = ctx->trace;
-        task.n_ranges = 0;
-
-        uint32_t block_acc = 0;
-        for (uint32_t i = 0; i < n_dirty; i++) {
-            const struct htp_tensor * t = dirty_tensors[i];
-            make_tensor_clean(ctx, t);
-
-            struct l2flush_range * rg = &task.ranges[task.n_ranges++];
-            rg->start = hex_align_down((size_t) t->data, HEX_L2_LINE_SIZE);
-            rg->end   = hex_align_up((size_t) t->data + t->size, HEX_L2_LINE_SIZE);
-            rg->block_first = block_acc;
-            rg->n_blocks = (rg->end - rg->start + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
-            block_acc += rg->n_blocks;
-        }
-
-        task.total_blocks      = block_acc;
-        task.blocks_per_thread = fastdiv(block_acc + ctx->n_threads - 1, &ctx->n_threads_div);
-
-        work_queue_run(ctx->work_queue, l2flush_multi_worker, &task, ctx->n_threads);
-        return;
-    }
-
-    struct htp_thread_trace * tr = &ctx->trace[0];
+    flush_dirty_ranges(ctx, ranges, n_dirty, total_dirty);
     for (uint32_t i = 0; i < n_dirty; i++) {
-        const struct htp_tensor * t = dirty_tensors[i];
-        htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
-        hex_l2flush((void *) (uintptr_t) t->data, t->size);
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
-        make_tensor_clean(ctx, t);
+        make_tensor_clean(ctx, dirty_tensors[i]);
     }
 }

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

@@ -2,8 +2,20 @@
 #define HTP_TENSOR_H
 
 #include <stdint.h>
+#include <stdbool.h>
 #include "htp-ops.h"
 #include "hex-bitmap.h"
+#include "hex-common.h"
+#include "hex-fastdiv.h"
+
+enum {
+    HTP_TENSOR_MDEV_LINE_SIZE = 128,
+};
+
+struct htp_tensor_mdev_range {
+    uint32_t start;
+    uint32_t count;
+};
 
 static inline void * htp_tensor_data(const struct htp_tensor * t) {
     return (void *) (uintptr_t) t->data;
@@ -13,6 +25,102 @@ static inline uint32_t * htp_tensor_flags(const struct htp_tensor * t) {
     return (uint32_t *) &t->flags;
 }
 
+static inline bool htp_tensor_is_contiguous(const struct htp_tensor * t, uint32_t type_size) {
+    uint32_t next_nb = type_size;
+    if (t->ne[0] != 1 && t->nb[0] != next_nb) {
+        return false;
+    }
+    next_nb *= t->ne[0];
+    for (int i = 1; i < HTP_OP_MAX_DIMS; i++) {
+        if (t->ne[i] != 1 && t->nb[i] != next_nb) {
+            return false;
+        }
+        next_nb *= t->ne[i];
+    }
+    return true;
+}
+
+static inline bool htp_tensor_is_permuted(const struct htp_tensor * t) {
+    return t->nb[0] > t->nb[1] || t->nb[1] > t->nb[2] || t->nb[2] > t->nb[3];
+}
+
+static inline bool htp_tensor_mdev_data_aligned(const struct htp_tensor * t) {
+    return ((uintptr_t) t->data & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0;
+}
+
+static inline bool htp_tensor_can_row_partition(const struct htp_tensor * t, uint32_t elem_size) {
+    if (!htp_tensor_mdev_data_aligned(t)) {
+        return false;
+    }
+    if (t->ne[0] != 1 && t->nb[0] != elem_size) {
+        return false;
+    }
+    if (htp_tensor_is_permuted(t)) {
+        return false;
+    }
+    if (t->ne[1] > 1 && (t->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) != 0) return false;
+    if (t->ne[2] > 1 && (t->nb[2] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) != 0) return false;
+    if (t->ne[3] > 1 && (t->nb[3] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) != 0) return false;
+    return true;
+}
+
+static inline bool htp_tensor_mdev_rows_per_chunk(const struct htp_tensor * t, uint32_t elem_size, uint32_t row_size, uint32_t * rows_per_chunk) {
+    *rows_per_chunk = 0;
+
+    if (!htp_tensor_mdev_data_aligned(t)) {
+        return false;
+    }
+    if (t->ne[0] != 1 && t->nb[0] != elem_size) {
+        return false;
+    }
+    if (htp_tensor_is_permuted(t)) {
+        return false;
+    }
+    if (t->ne[1] > 1 && (t->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0 &&
+        (t->ne[2] <= 1 || (t->nb[2] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0) &&
+        (t->ne[3] <= 1 || (t->nb[3] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0)) {
+        *rows_per_chunk = 1;
+        return true;
+    }
+    if (t->nb[1] == row_size &&
+        (t->ne[2] <= 1 || t->nb[2] == t->nb[1] * t->ne[1]) &&
+        (t->ne[3] <= 1 || t->nb[3] == t->nb[2] * t->ne[2])) {
+        *rows_per_chunk = (row_size > 0) ? (HTP_TENSOR_MDEV_LINE_SIZE / hex_gcd_u32(row_size, HTP_TENSOR_MDEV_LINE_SIZE)) : 1;
+        return true;
+    }
+    return false;
+}
+
+static inline struct htp_tensor_mdev_range htp_tensor_mdev_partition(uint32_t total_units, uint32_t units_per_chunk, uint32_t mdev_idx, uint32_t mdev_count, const struct fastdiv_values * mdev_count_div) {
+    struct htp_tensor_mdev_range range = { 0, total_units };
+
+    if (mdev_count <= 1) {
+        return range;
+    }
+
+    if (units_per_chunk == 0) {
+        range.start = (mdev_idx == 0) ? 0 : total_units;
+        range.count = (mdev_idx == 0) ? total_units : 0;
+        return range;
+    }
+
+    const uint32_t total_chunks = total_units / units_per_chunk;
+    if (total_chunks < mdev_count) {
+        range.start = (mdev_idx == 0) ? 0 : total_units;
+        range.count = (mdev_idx == 0) ? total_units : 0;
+        return range;
+    }
+
+    const uint32_t chunks_per_mdev = fastdiv(total_chunks + mdev_count - 1, mdev_count_div);
+    range.start = MIN(mdev_idx * chunks_per_mdev * units_per_chunk, total_units);
+    if (mdev_idx == mdev_count - 1) {
+        range.count = total_units - range.start;
+    } else {
+        range.count = MIN(chunks_per_mdev * units_per_chunk, total_units - range.start);
+    }
+    return range;
+}
+
 static inline uint32_t htp_tensor_get_row_size(int type, uint32_t ne00) {
     switch (type) {
         case HTP_TYPE_F32:  return ne00 * 4;
@@ -23,6 +131,7 @@ static inline uint32_t htp_tensor_get_row_size(int type, uint32_t ne00) {
 }
 
 struct htp_context;
+void htp_flush_dirty_ranges(struct htp_context * ctx);
 void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
 void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
 

ggml/src/ggml-hexagon/htp/hvx-arith.h

@@ -16,25 +16,25 @@
 #define UNUSED(x) (void)(x)
 
 #define hvx_arith_loop_body(dst_type, src0_type, src1_type, elem_size, vec_store, vec_op) \
-    do {                                                                       \
-        dst_type * vdst  = (dst_type *) dst;                                   \
-        src0_type * vsrc0 = (src0_type *) src0;                                \
-        src1_type * vsrc1 = (src1_type *) src1;                                \
-                                                                               \
-        const uint32_t epv  = 128 / (elem_size);                               \
-        const uint32_t nvec = n / epv;                                         \
-        const uint32_t nloe = n % epv;                                         \
-                                                                               \
-        uint32_t i = 0;                                                        \
-                                                                               \
-        _Pragma("unroll(4)")                                                   \
-        for (; i < nvec; i++) {                                                \
-            vdst[i] = vec_op(vsrc0[i], vsrc1[i]);                              \
-        }                                                                      \
-        if (nloe) {                                                            \
-            HVX_Vector v = vec_op(vsrc0[i], vsrc1[i]);                         \
-            vec_store((void *) &vdst[i], nloe * (elem_size), v);               \
-        }                                                                      \
+    do {                                                                                  \
+        dst_type * vdst  = (dst_type *) dst;                                              \
+        src0_type * vsrc0 = (src0_type *) src0;                                           \
+        src1_type * vsrc1 = (src1_type *) src1;                                           \
+                                                                                          \
+        const uint32_t epv  = 128 / (elem_size);                                          \
+        const uint32_t nvec = n / epv;                                                    \
+        const uint32_t nloe = n % epv;                                                    \
+                                                                                          \
+        uint32_t i = 0;                                                                   \
+                                                                                          \
+        _Pragma("unroll(4)")                                                              \
+        for (; i < nvec; i++) {                                                           \
+            vdst[i] = vec_op(vsrc0[i], vsrc1[i]);                                         \
+        }                                                                                 \
+        if (nloe) {                                                                       \
+            HVX_Vector v = vec_op(vsrc0[i], vsrc1[i]);                                    \
+            vec_store((void *) &vdst[i], nloe * (elem_size), v);                          \
+        }                                                                                 \
     } while(0)
 
 #if __HVX_ARCH__ < 79
@@ -56,43 +56,43 @@
 #define HVX_OP_MUL_F16(a, b) hvx_vec_mul_f16_f16(a, b)
 
 // Generic macro to define alignment permutations for an op
-#define DEFINE_HVX_BINARY_OP_VARIANTS(OP_NAME, OP_MACRO, ELEM_TYPE) \
-static inline void OP_NAME##_aaa(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
-    assert((uintptr_t) dst % 128 == 0); \
-    assert((uintptr_t) src0 % 128 == 0); \
-    assert((uintptr_t) src1 % 128 == 0); \
-    hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \
-} \
-static inline void OP_NAME##_aau(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
-    assert((uintptr_t) dst % 128 == 0); \
-    assert((uintptr_t) src0 % 128 == 0); \
-    hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \
-} \
-static inline void OP_NAME##_aua(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
-    assert((uintptr_t) dst % 128 == 0); \
-    assert((uintptr_t) src1 % 128 == 0); \
-    hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \
-} \
-static inline void OP_NAME##_auu(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
-    assert((uintptr_t) dst % 128 == 0); \
-    hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \
-} \
-static inline void OP_NAME##_uaa(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
-    assert((uintptr_t) src0 % 128 == 0); \
-    assert((uintptr_t) src1 % 128 == 0); \
-    hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \
-} \
-static inline void OP_NAME##_uau(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
-    assert((uintptr_t) src0 % 128 == 0); \
-    hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \
-} \
-static inline void OP_NAME##_uua(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
-    assert((uintptr_t) src1 % 128 == 0); \
-    hvx_arith_loop_body(HVX_UVector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \
-} \
-static inline void OP_NAME##_uuu(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
+#define DEFINE_HVX_BINARY_OP_VARIANTS(OP_NAME, OP_MACRO, ELEM_TYPE)                                           \
+static inline void OP_NAME##_aaa(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
+    assert((uintptr_t) dst % 128 == 0);                                                                       \
+    assert((uintptr_t) src0 % 128 == 0);                                                                      \
+    assert((uintptr_t) src1 % 128 == 0);                                                                      \
+    hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO);    \
+}                                                                                                             \
+static inline void OP_NAME##_aau(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
+    assert((uintptr_t) dst % 128 == 0);                                                                       \
+    assert((uintptr_t) src0 % 128 == 0);                                                                      \
+    hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO);   \
+}                                                                                                             \
+static inline void OP_NAME##_aua(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
+    assert((uintptr_t) dst % 128 == 0);                                                                       \
+    assert((uintptr_t) src1 % 128 == 0);                                                                      \
+    hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO);   \
+}                                                                                                             \
+static inline void OP_NAME##_auu(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
+    assert((uintptr_t) dst % 128 == 0);                                                                       \
+    hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO);  \
+}                                                                                                             \
+static inline void OP_NAME##_uaa(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
+    assert((uintptr_t) src0 % 128 == 0);                                                                      \
+    assert((uintptr_t) src1 % 128 == 0);                                                                      \
+    hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO);   \
+}                                                                                                             \
+static inline void OP_NAME##_uau(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
+    assert((uintptr_t) src0 % 128 == 0);                                                                      \
+    hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO);  \
+}                                                                                                             \
+static inline void OP_NAME##_uua(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
+    assert((uintptr_t) src1 % 128 == 0);                                                                      \
+    hvx_arith_loop_body(HVX_UVector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO);  \
+}                                                                                                             \
+static inline void OP_NAME##_uuu(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
     hvx_arith_loop_body(HVX_UVector, HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \
-} \
+}                                                                                                             \
 
 DEFINE_HVX_BINARY_OP_VARIANTS(hvx_add_f32, HVX_OP_ADD_F32, float)
 DEFINE_HVX_BINARY_OP_VARIANTS(hvx_sub_f32, HVX_OP_SUB_F32, float)
@@ -103,25 +103,25 @@ DEFINE_HVX_BINARY_OP_VARIANTS(hvx_sub_f16, HVX_OP_SUB_F16, _Float16)
 DEFINE_HVX_BINARY_OP_VARIANTS(hvx_mul_f16, HVX_OP_MUL_F16, _Float16)
 
 // Dispatcher logic
-#define HVX_BINARY_DISPATCHER(OP_NAME) \
+#define HVX_BINARY_DISPATCHER(OP_NAME)                                                                                                       \
 static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, const uint32_t num_elems) { \
-    if (hex_is_aligned((void *) dst, 128)) { \
-        if (hex_is_aligned((void *) src0, 128)) { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aaa(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_aau(dst, src0, src1, num_elems); \
-        } else { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aua(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_auu(dst, src0, src1, num_elems); \
-        } \
-    } else { \
-        if (hex_is_aligned((void *) src0, 128)) { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uaa(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_uau(dst, src0, src1, num_elems); \
-        } else { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uua(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_uuu(dst, src0, src1, num_elems); \
-        } \
-    } \
+    if (hex_is_aligned((void *) dst, 128)) {                                                                                                 \
+        if (hex_is_aligned((void *) src0, 128)) {                                                                                            \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aaa(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_aau(dst, src0, src1, num_elems);                                               \
+        } else {                                                                                                                             \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aua(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_auu(dst, src0, src1, num_elems);                                               \
+        }                                                                                                                                    \
+    } else {                                                                                                                                 \
+        if (hex_is_aligned((void *) src0, 128)) {                                                                                            \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uaa(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_uau(dst, src0, src1, num_elems);                                               \
+        } else {                                                                                                                             \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uua(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_uuu(dst, src0, src1, num_elems);                                               \
+        }                                                                                                                                    \
+    }                                                                                                                                        \
 }
 
 HVX_BINARY_DISPATCHER(hvx_add_f32)
@@ -166,44 +166,44 @@ static inline void hvx_mul_mul_f32_aa(uint8_t * restrict dst, const uint8_t * re
 
 // Scalar Operations
 
-#define hvx_scalar_loop_body(dst_type, src_type, elem_size, vec_store, scalar_op_macro)   \
-    do {                                                                       \
-        dst_type * restrict vdst = (dst_type *) dst;                           \
-        src_type * restrict vsrc = (src_type *) src;                           \
-                                                                               \
-        const uint32_t epv  = 128 / (elem_size);                               \
-        const uint32_t nvec = n / epv;                                         \
-        const uint32_t nloe = n % epv;                                         \
-                                                                               \
-        uint32_t i = 0;                                                        \
-                                                                               \
-        _Pragma("unroll(4)")                                                   \
-        for (; i < nvec; i++) {                                                \
-            HVX_Vector v = vsrc[i];                                            \
-            vdst[i] = scalar_op_macro(v);                                      \
-        }                                                                      \
-        if (nloe) {                                                            \
-            HVX_Vector v = vsrc[i];                                            \
-            v = scalar_op_macro(v);                                            \
-            vec_store((void *) &vdst[i], nloe * (elem_size), v);               \
-        }                                                                      \
+#define hvx_scalar_loop_body(dst_type, src_type, elem_size, vec_store, scalar_op_macro) \
+    do {                                                                                \
+        dst_type * restrict vdst = (dst_type *) dst;                                    \
+        src_type * restrict vsrc = (src_type *) src;                                    \
+                                                                                        \
+        const uint32_t epv  = 128 / (elem_size);                                        \
+        const uint32_t nvec = n / epv;                                                  \
+        const uint32_t nloe = n % epv;                                                  \
+                                                                                        \
+        uint32_t i = 0;                                                                 \
+                                                                                        \
+        _Pragma("unroll(4)")                                                            \
+        for (; i < nvec; i++) {                                                         \
+            HVX_Vector v = vsrc[i];                                                     \
+            vdst[i] = scalar_op_macro(v);                                               \
+        }                                                                               \
+        if (nloe) {                                                                     \
+            HVX_Vector v = vsrc[i];                                                     \
+            v = scalar_op_macro(v);                                                     \
+            vec_store((void *) &vdst[i], nloe * (elem_size), v);                        \
+        }                                                                               \
     } while(0)
 
-#define HVX_OP_ADD_SCALAR_F32(v) \
-    ({ \
+#define HVX_OP_ADD_SCALAR_F32(v)                                   \
+    ({                                                             \
         const HVX_VectorPred pred_inf = Q6_Q_vcmp_eq_VwVw(inf, v); \
-        HVX_Vector out = HVX_OP_ADD_F32(v, val_vec); \
-        Q6_V_vmux_QVV(pred_inf, inf, out); \
+        HVX_Vector out = HVX_OP_ADD_F32(v, val_vec);               \
+        Q6_V_vmux_QVV(pred_inf, inf, out);                         \
     })
 
 #define HVX_OP_MUL_SCALAR_F32(v) HVX_OP_MUL_F32(v, val_vec)
 #define HVX_OP_SUB_SCALAR_F32(v) HVX_OP_SUB_F32(v, val_vec)
 
-#define HVX_OP_ADD_SCALAR_F16(v) \
-    ({ \
+#define HVX_OP_ADD_SCALAR_F16(v)                                   \
+    ({                                                             \
         const HVX_VectorPred pred_inf = Q6_Q_vcmp_eq_VhVh(inf, v); \
-        HVX_Vector out = HVX_OP_ADD_F16(v, val_vec); \
-        Q6_V_vmux_QVV(pred_inf, inf, out); \
+        HVX_Vector out = HVX_OP_ADD_F16(v, val_vec);               \
+        Q6_V_vmux_QVV(pred_inf, inf, out);                         \
     })
 
 #define HVX_OP_MUL_SCALAR_F16(v) HVX_OP_MUL_F16(v, val_vec)
@@ -212,31 +212,31 @@ static inline void hvx_mul_mul_f32_aa(uint8_t * restrict dst, const uint8_t * re
 // Scalar Variants
 
 // Generic macro to define alignment permutations for an op
-#define DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(OP_NAME, OP_MACRO, SPLAT_MACRO, ELEM_TYPE) \
+#define DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(OP_NAME, OP_MACRO, SPLAT_MACRO, ELEM_TYPE)                                  \
 static inline void OP_NAME##_aa(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \
-    const HVX_Vector val_vec = SPLAT_MACRO(val); \
-    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \
-    assert((uintptr_t) dst % 128 == 0); \
-    assert((uintptr_t) src % 128 == 0); \
-    hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \
-} \
+    const HVX_Vector val_vec = SPLAT_MACRO(val);                                                                         \
+    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf);                                                \
+    assert((uintptr_t) dst % 128 == 0);                                                                                  \
+    assert((uintptr_t) src % 128 == 0);                                                                                  \
+    hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO);                          \
+}                                                                                                                        \
 static inline void OP_NAME##_au(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \
-    const HVX_Vector val_vec = SPLAT_MACRO(val); \
-    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \
-    assert((uintptr_t) dst % 128 == 0); \
-    hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \
-} \
+    const HVX_Vector val_vec = SPLAT_MACRO(val);                                                                         \
+    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf);                                                \
+    assert((uintptr_t) dst % 128 == 0);                                                                                  \
+    hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO);                         \
+}                                                                                                                        \
 static inline void OP_NAME##_ua(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \
-    const HVX_Vector val_vec = SPLAT_MACRO(val); \
-    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \
-    assert((uintptr_t) src % 128 == 0); \
-    hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \
-} \
+    const HVX_Vector val_vec = SPLAT_MACRO(val);                                                                         \
+    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf);                                                \
+    assert((uintptr_t) src % 128 == 0);                                                                                  \
+    hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO);                         \
+}                                                                                                                        \
 static inline void OP_NAME##_uu(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \
-    const HVX_Vector val_vec = SPLAT_MACRO(val); \
-    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \
-    hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \
-} \
+    const HVX_Vector val_vec = SPLAT_MACRO(val);                                                                         \
+    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf);                                                \
+    hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO);                        \
+}                                                                                                                        \
 
 DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_add_scalar_f32, HVX_OP_ADD_SCALAR_F32, hvx_vec_splat_f32, float)
 DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_sub_scalar_f32, HVX_OP_SUB_SCALAR_F32, hvx_vec_splat_f32, float)
@@ -247,17 +247,17 @@ DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_sub_scalar_f16, HVX_OP_SUB_SCALAR_F16,
 DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_mul_scalar_f16, HVX_OP_MUL_SCALAR_F16, hvx_vec_splat_f16, _Float16)
 
 // Dispatcher logic
-#define HVX_BINARY_SCALAR_DISPATCHER(OP_NAME, ELEM_TYPE) \
+#define HVX_BINARY_SCALAR_DISPATCHER(OP_NAME, ELEM_TYPE)                                                                          \
 static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, const uint32_t num_elems) { \
-    if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) { \
-        OP_NAME##_aa(dst, src, val, num_elems); \
-    } else if (hex_is_aligned((void *) dst, 128)) { \
-        OP_NAME##_au(dst, src, val, num_elems); \
-    } else if (hex_is_aligned((void *) src, 128)) { \
-        OP_NAME##_ua(dst, src, val, num_elems); \
-    } else { \
-        OP_NAME##_uu(dst, src, val, num_elems); \
-    } \
+    if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) {                                                 \
+        OP_NAME##_aa(dst, src, val, num_elems);                                                                                   \
+    } else if (hex_is_aligned((void *) dst, 128)) {                                                                               \
+        OP_NAME##_au(dst, src, val, num_elems);                                                                                   \
+    } else if (hex_is_aligned((void *) src, 128)) {                                                                               \
+        OP_NAME##_ua(dst, src, val, num_elems);                                                                                   \
+    } else {                                                                                                                      \
+        OP_NAME##_uu(dst, src, val, num_elems);                                                                                   \
+    }                                                                                                                             \
 }
 
 HVX_BINARY_SCALAR_DISPATCHER(hvx_add_scalar_f32, float)
@@ -350,12 +350,12 @@ static inline void hvx_max_scalar_f32(uint8_t * restrict dst, const uint8_t * re
 
 // CLAMP Scalar variants
 
-#define HVX_OP_CLAMP_SCALAR(v) \
-    ({ \
+#define HVX_OP_CLAMP_SCALAR(v)                                           \
+    ({                                                                   \
         HVX_VectorPred pred_cap_right = Q6_Q_vcmp_gt_VsfVsf(v, max_vec); \
         HVX_VectorPred pred_cap_left  = Q6_Q_vcmp_gt_VsfVsf(min_vec, v); \
-        HVX_Vector tmp = Q6_V_vmux_QVV(pred_cap_right, max_vec, v); \
-        Q6_V_vmux_QVV(pred_cap_left, min_vec, tmp); \
+        HVX_Vector tmp = Q6_V_vmux_QVV(pred_cap_right, max_vec, v);      \
+        Q6_V_vmux_QVV(pred_cap_left, min_vec, tmp);                      \
     })
 
 static inline void hvx_clamp_scalar_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, uint32_t n) {

ggml/src/ggml-hexagon/htp/hvx-div.h

@@ -219,64 +219,64 @@ static inline HVX_Vector hvx_vec_hybrid_div_f16(HVX_Vector vec1, HVX_Vector vec2
     } while(0)
 
 // Generic macro to define alignment permutations for an op
-#define DEFINE_HVX_DIV_OP_VARIANTS(OP_NAME, OP_LOOP_BODY) \
+#define DEFINE_HVX_DIV_OP_VARIANTS(OP_NAME, OP_LOOP_BODY)                                                                            \
 static inline void OP_NAME##_aaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
-    assert((uintptr_t) dst % 128 == 0); \
-    assert((uintptr_t) src0 % 128 == 0); \
-    assert((uintptr_t) src1 % 128 == 0); \
-    OP_LOOP_BODY(HVX_Vector, HVX_Vector, HVX_Vector, hvx_vec_store_a); \
-} \
+    assert((uintptr_t) dst % 128 == 0);                                                                                              \
+    assert((uintptr_t) src0 % 128 == 0);                                                                                             \
+    assert((uintptr_t) src1 % 128 == 0);                                                                                             \
+    OP_LOOP_BODY(HVX_Vector, HVX_Vector, HVX_Vector, hvx_vec_store_a);                                                               \
+}                                                                                                                                    \
 static inline void OP_NAME##_aau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
-    assert((uintptr_t) dst % 128 == 0); \
-    assert((uintptr_t) src0 % 128 == 0); \
-    OP_LOOP_BODY(HVX_Vector, HVX_Vector, HVX_UVector, hvx_vec_store_a); \
-} \
+    assert((uintptr_t) dst % 128 == 0);                                                                                              \
+    assert((uintptr_t) src0 % 128 == 0);                                                                                             \
+    OP_LOOP_BODY(HVX_Vector, HVX_Vector, HVX_UVector, hvx_vec_store_a);                                                              \
+}                                                                                                                                    \
 static inline void OP_NAME##_aua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
-    assert((uintptr_t) dst % 128 == 0); \
-    assert((uintptr_t) src1 % 128 == 0); \
-    OP_LOOP_BODY(HVX_Vector, HVX_UVector, HVX_Vector, hvx_vec_store_a); \
-} \
+    assert((uintptr_t) dst % 128 == 0);                                                                                              \
+    assert((uintptr_t) src1 % 128 == 0);                                                                                             \
+    OP_LOOP_BODY(HVX_Vector, HVX_UVector, HVX_Vector, hvx_vec_store_a);                                                              \
+}                                                                                                                                    \
 static inline void OP_NAME##_auu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
-    assert((uintptr_t) dst % 128 == 0); \
-    OP_LOOP_BODY(HVX_Vector, HVX_UVector, HVX_UVector, hvx_vec_store_a); \
-} \
+    assert((uintptr_t) dst % 128 == 0);                                                                                              \
+    OP_LOOP_BODY(HVX_Vector, HVX_UVector, HVX_UVector, hvx_vec_store_a);                                                             \
+}                                                                                                                                    \
 static inline void OP_NAME##_uaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
-    assert((uintptr_t) src0 % 128 == 0); \
-    assert((uintptr_t) src1 % 128 == 0); \
-    OP_LOOP_BODY(HVX_UVector, HVX_Vector, HVX_Vector, hvx_vec_store_u); \
-} \
+    assert((uintptr_t) src0 % 128 == 0);                                                                                             \
+    assert((uintptr_t) src1 % 128 == 0);                                                                                             \
+    OP_LOOP_BODY(HVX_UVector, HVX_Vector, HVX_Vector, hvx_vec_store_u);                                                              \
+}                                                                                                                                    \
 static inline void OP_NAME##_uau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
-    assert((uintptr_t) src0 % 128 == 0); \
-    OP_LOOP_BODY(HVX_UVector, HVX_Vector, HVX_UVector, hvx_vec_store_u); \
-} \
+    assert((uintptr_t) src0 % 128 == 0);                                                                                             \
+    OP_LOOP_BODY(HVX_UVector, HVX_Vector, HVX_UVector, hvx_vec_store_u);                                                             \
+}                                                                                                                                    \
 static inline void OP_NAME##_uua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
-    assert((uintptr_t) src1 % 128 == 0); \
-    OP_LOOP_BODY(HVX_UVector, HVX_UVector, HVX_Vector, hvx_vec_store_u); \
-} \
+    assert((uintptr_t) src1 % 128 == 0);                                                                                             \
+    OP_LOOP_BODY(HVX_UVector, HVX_UVector, HVX_Vector, hvx_vec_store_u);                                                             \
+}                                                                                                                                    \
 static inline void OP_NAME##_uuu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
-    OP_LOOP_BODY(HVX_UVector, HVX_UVector, HVX_UVector, hvx_vec_store_u); \
-} \
+    OP_LOOP_BODY(HVX_UVector, HVX_UVector, HVX_UVector, hvx_vec_store_u);                                                            \
+}                                                                                                                                    \
 
 // Dispatcher logic
-#define HVX_DIV_DISPATCHER(OP_NAME) \
+#define HVX_DIV_DISPATCHER(OP_NAME)                                                                                                          \
 static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, const uint32_t num_elems) { \
-    if (hex_is_aligned((void *) dst, 128)) { \
-        if (hex_is_aligned((void *) src0, 128)) { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aaa(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_aau(dst, src0, src1, num_elems); \
-        } else { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aua(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_auu(dst, src0, src1, num_elems); \
-        } \
-    } else { \
-        if (hex_is_aligned((void *) src0, 128)) { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uaa(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_uau(dst, src0, src1, num_elems); \
-        } else { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uua(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_uuu(dst, src0, src1, num_elems); \
-        } \
-    } \
+    if (hex_is_aligned((void *) dst, 128)) {                                                                                                 \
+        if (hex_is_aligned((void *) src0, 128)) {                                                                                            \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aaa(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_aau(dst, src0, src1, num_elems);                                               \
+        } else {                                                                                                                             \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aua(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_auu(dst, src0, src1, num_elems);                                               \
+        }                                                                                                                                    \
+    } else {                                                                                                                                 \
+        if (hex_is_aligned((void *) src0, 128)) {                                                                                            \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uaa(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_uau(dst, src0, src1, num_elems);                                               \
+        } else {                                                                                                                             \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uua(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_uuu(dst, src0, src1, num_elems);                                               \
+        }                                                                                                                                    \
+    }                                                                                                                                        \
 }
 
 DEFINE_HVX_DIV_OP_VARIANTS(hvx_div_f32, hvx_div_f32_loop_body)

ggml/src/ggml-hexagon/htp/hvx-inverse.h

@@ -169,36 +169,36 @@ static inline HVX_Vector hvx_vec_inverse_f16_guard(HVX_Vector v_sf, HVX_Vector n
     } while(0)
 
 // Generic macro to define alignment permutations for an op
-#define DEFINE_HVX_INV_OP_VARIANTS(OP_NAME, OP_LOOP_BODY) \
+#define DEFINE_HVX_INV_OP_VARIANTS(OP_NAME, OP_LOOP_BODY)                                           \
 static inline void OP_NAME##_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \
-    assert((uintptr_t) dst % 128 == 0); \
-    assert((uintptr_t) src % 128 == 0); \
-    OP_LOOP_BODY(HVX_Vector, HVX_Vector, hvx_vec_store_a); \
-} \
+    assert((uintptr_t) dst % 128 == 0);                                                             \
+    assert((uintptr_t) src % 128 == 0);                                                             \
+    OP_LOOP_BODY(HVX_Vector, HVX_Vector, hvx_vec_store_a);                                          \
+}                                                                                                   \
 static inline void OP_NAME##_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \
-    assert((uintptr_t) dst % 128 == 0); \
-    OP_LOOP_BODY(HVX_Vector, HVX_UVector, hvx_vec_store_a); \
-} \
+    assert((uintptr_t) dst % 128 == 0);                                                             \
+    OP_LOOP_BODY(HVX_Vector, HVX_UVector, hvx_vec_store_a);                                         \
+}                                                                                                   \
 static inline void OP_NAME##_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \
-    assert((uintptr_t) src % 128 == 0); \
-    OP_LOOP_BODY(HVX_UVector, HVX_Vector, hvx_vec_store_u); \
-} \
+    assert((uintptr_t) src % 128 == 0);                                                             \
+    OP_LOOP_BODY(HVX_UVector, HVX_Vector, hvx_vec_store_u);                                         \
+}                                                                                                   \
 static inline void OP_NAME##_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \
-    OP_LOOP_BODY(HVX_UVector, HVX_UVector, hvx_vec_store_u); \
-} \
+    OP_LOOP_BODY(HVX_UVector, HVX_UVector, hvx_vec_store_u);                                        \
+}                                                                                                   \
 
 // Dispatcher logic
-#define HVX_INV_DISPATCHER(OP_NAME) \
+#define HVX_INV_DISPATCHER(OP_NAME)                                                                          \
 static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t num_elems) { \
-    if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) { \
-        OP_NAME##_aa(dst, src, num_elems); \
-    } else if (hex_is_aligned((void *) dst, 128)) { \
-        OP_NAME##_au(dst, src, num_elems); \
-    } else if (hex_is_aligned((void *) src, 128)) { \
-        OP_NAME##_ua(dst, src, num_elems); \
-    } else { \
-        OP_NAME##_uu(dst, src, num_elems); \
-    } \
+    if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) {                            \
+        OP_NAME##_aa(dst, src, num_elems);                                                                   \
+    } else if (hex_is_aligned((void *) dst, 128)) {                                                          \
+        OP_NAME##_au(dst, src, num_elems);                                                                   \
+    } else if (hex_is_aligned((void *) src, 128)) {                                                          \
+        OP_NAME##_ua(dst, src, num_elems);                                                                   \
+    } else {                                                                                                 \
+        OP_NAME##_uu(dst, src, num_elems);                                                                   \
+    }                                                                                                        \
 }
 
 DEFINE_HVX_INV_OP_VARIANTS(hvx_inverse_f32, hvx_inverse_f32_loop_body)

ggml/src/ggml-hexagon/htp/hvx-scale.h

@@ -68,30 +68,30 @@ static inline void hvx_scale_f32(uint8_t * restrict dst, const uint8_t * restric
     }
 }
 
-#define hvx_scale_offset_f32_loop_body(dst_type, src_type, vec_store)                \
-    do {                                                                             \
-        dst_type * restrict vdst = (dst_type *) dst;                                 \
-        src_type * restrict vsrc = (src_type *) src;                                 \
-                                                                                     \
-        HVX_Vector vs = hvx_vec_splat_f32(scale);                                    \
-        HVX_Vector vo = hvx_vec_splat_f32(offset);                                   \
-                                                                                     \
-        const uint32_t elem_size = sizeof(float);                                    \
-        const uint32_t epv = 128 / elem_size;                                        \
-        const uint32_t nvec = n / epv;                                               \
-        const uint32_t nloe = n % epv;                                               \
-                                                                                     \
-        uint32_t i = 0;                                                              \
-                                                                                     \
-        _Pragma("unroll(4)")                                                         \
-        for (; i < nvec; ++i) {                                                      \
+#define hvx_scale_offset_f32_loop_body(dst_type, src_type, vec_store)                     \
+    do {                                                                                  \
+        dst_type * restrict vdst = (dst_type *) dst;                                      \
+        src_type * restrict vsrc = (src_type *) src;                                      \
+                                                                                          \
+        HVX_Vector vs = hvx_vec_splat_f32(scale);                                         \
+        HVX_Vector vo = hvx_vec_splat_f32(offset);                                        \
+                                                                                          \
+        const uint32_t elem_size = sizeof(float);                                         \
+        const uint32_t epv = 128 / elem_size;                                             \
+        const uint32_t nvec = n / epv;                                                    \
+        const uint32_t nloe = n % epv;                                                    \
+                                                                                          \
+        uint32_t i = 0;                                                                   \
+                                                                                          \
+        _Pragma("unroll(4)")                                                              \
+        for (; i < nvec; ++i) {                                                           \
             HVX_Vector v = Q6_Vqf32_vadd_Vqf32Vsf(Q6_Vqf32_vmpy_VsfVsf(vsrc[i], vs), vo); \
-            vdst[i] = Q6_Vsf_equals_Vqf32(v);                                        \
-        }                                                                            \
-        if (nloe) {                                                                  \
+            vdst[i] = Q6_Vsf_equals_Vqf32(v);                                             \
+        }                                                                                 \
+        if (nloe) {                                                                       \
             HVX_Vector v = Q6_Vqf32_vadd_Vqf32Vsf(Q6_Vqf32_vmpy_VsfVsf(vsrc[i], vs), vo); \
-            vec_store((void *) &vdst[i], nloe * elem_size, Q6_Vsf_equals_Vqf32(v));  \
-        }                                                                            \
+            vec_store((void *) &vdst[i], nloe * elem_size, Q6_Vsf_equals_Vqf32(v));       \
+        }                                                                                 \
     } while(0)
 
 static inline void hvx_scale_offset_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) {

ggml/src/ggml-hexagon/htp/hvx-sigmoid.h

@@ -68,50 +68,50 @@ static inline HVX_Vector hvx_vec_tanh_f32(HVX_Vector x) {
     return Q6_Vsf_equals_Vqf32(res);
 }
 
-#define hvx_sigmoid_loop_body(dst_type, src_type, vec_store)    \
-    do {                                                        \
-        dst_type * restrict vdst = (dst_type *) dst;            \
-        src_type * restrict vsrc = (src_type *) src;            \
-                                                                \
-        const HVX_Vector one     = hvx_vec_splat_f32(1.f);      \
-        const HVX_Vector max_exp = hvx_vec_splat_f32(87.f);     \
-        const HVX_Vector min_exp = hvx_vec_splat_f32(-87.f);    \
-                                                                \
-        const uint32_t epv  = 128 / sizeof(float);              \
-        const uint32_t nvec = n / epv;                          \
-        const uint32_t nloe = n % epv;                          \
-                                                                \
-        uint32_t i = 0;                                         \
-                                                                \
-        _Pragma("unroll(4)")                                    \
-        for (; i < nvec; i++) {                                 \
-             vdst[i] = hvx_vec_fast_sigmoid_f32_guard(vsrc[i], one, max_exp, min_exp); \
-        }                                                       \
-        if (nloe) {                                             \
+#define hvx_sigmoid_loop_body(dst_type, src_type, vec_store)                                  \
+    do {                                                                                      \
+        dst_type * restrict vdst = (dst_type *) dst;                                          \
+        src_type * restrict vsrc = (src_type *) src;                                          \
+                                                                                              \
+        const HVX_Vector one     = hvx_vec_splat_f32(1.f);                                    \
+        const HVX_Vector max_exp = hvx_vec_splat_f32(87.f);                                   \
+        const HVX_Vector min_exp = hvx_vec_splat_f32(-87.f);                                  \
+                                                                                              \
+        const uint32_t epv  = 128 / sizeof(float);                                            \
+        const uint32_t nvec = n / epv;                                                        \
+        const uint32_t nloe = n % epv;                                                        \
+                                                                                              \
+        uint32_t i = 0;                                                                       \
+                                                                                              \
+        _Pragma("unroll(4)")                                                                  \
+        for (; i < nvec; i++) {                                                               \
+             vdst[i] = hvx_vec_fast_sigmoid_f32_guard(vsrc[i], one, max_exp, min_exp);        \
+        }                                                                                     \
+        if (nloe) {                                                                           \
              HVX_Vector tmp = hvx_vec_fast_sigmoid_f32_guard(vsrc[i], one, max_exp, min_exp); \
-             vec_store((void *) &vdst[i], nloe * sizeof(float), tmp); \
-        }                                                       \
+             vec_store((void *) &vdst[i], nloe * sizeof(float), tmp);                         \
+        }                                                                                     \
     } while(0)
 
-#define hvx_tanh_loop_body(dst_type, src_type, vec_store)       \
-    do {                                                        \
-        dst_type * restrict vdst = (dst_type *) dst;            \
-        src_type * restrict vsrc = (src_type *) src;            \
-                                                                \
-        const uint32_t epv  = 128 / sizeof(float);              \
-        const uint32_t nvec = n / epv;                          \
-        const uint32_t nloe = n % epv;                          \
-                                                                \
-        uint32_t i = 0;                                         \
-                                                                \
-        _Pragma("unroll(4)")                                    \
-        for (; i < nvec; i++) {                                 \
-             vdst[i] = hvx_vec_tanh_f32(vsrc[i]);               \
-        }                                                       \
-        if (nloe) {                                             \
-             HVX_Vector tmp = hvx_vec_tanh_f32(vsrc[i]);        \
+#define hvx_tanh_loop_body(dst_type, src_type, vec_store)             \
+    do {                                                              \
+        dst_type * restrict vdst = (dst_type *) dst;                  \
+        src_type * restrict vsrc = (src_type *) src;                  \
+                                                                      \
+        const uint32_t epv  = 128 / sizeof(float);                    \
+        const uint32_t nvec = n / epv;                                \
+        const uint32_t nloe = n % epv;                                \
+                                                                      \
+        uint32_t i = 0;                                               \
+                                                                      \
+        _Pragma("unroll(4)")                                          \
+        for (; i < nvec; i++) {                                       \
+             vdst[i] = hvx_vec_tanh_f32(vsrc[i]);                     \
+        }                                                             \
+        if (nloe) {                                                   \
+             HVX_Vector tmp = hvx_vec_tanh_f32(vsrc[i]);              \
              vec_store((void *) &vdst[i], nloe * sizeof(float), tmp); \
-        }                                                       \
+        }                                                             \
     } while(0)
 
 static inline void hvx_sigmoid_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {

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

@@ -3,11 +3,12 @@
 #pragma clang diagnostic ignored "-Wunused-but-set-variable"
 
 #include <HAP_farf.h>
-#include <HAP_perf.h>
 #include <hexagon_protos.h>
 #include <hexagon_types.h>
 #include <string.h>
 
+#include "hex-common.h"
+
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
 #include "htp-ctx.h"
@@ -16,14 +17,19 @@
 #include "hex-dma.h"
 #include "hex-profile.h"
 #include "htp-vtcm.h"
+#include "htp-tensor.h"
 
 struct htp_im2col_context {
     struct htp_ops_context * octx;
+    uint32_t                 patch_base;           // first patch index assigned to this dev
+    uint32_t                 npatches;             // number of patches assigned to this dev
     uint32_t                 npatches_per_thread;  // patches = N*OH*OW (pure-DDR kernel)
 
-    uint32_t pe_rows_per_thread;                   // N*OH rows per worker
-    uint32_t pe_src_row_bytes;                     // one output row's source: IC*KH*IW*4, rounded 256
-    uint32_t pe_dst_row_bytes;                     // one output row's dst: OW*patch_stride*2, rounded 256
+    uint32_t pe_row_base;              // first N*OH row index assigned to this dev (DMA path)
+    uint32_t pe_nrows;                 // number of N*OH rows assigned to this dev (DMA path)
+    uint32_t pe_rows_per_thread;       // N*OH rows per worker
+    uint32_t pe_src_row_bytes;         // one output row's source: IC*KH*IW*4, rounded 256
+    uint32_t pe_dst_row_bytes;         // one output row's dst: OW*patch_stride*2, rounded 256
 
     // Patch-embed DMA path VTCM ping-pong.
     uint8_t * pe_vtcm_src;                         // base of the 2x src buffers region
@@ -58,33 +64,27 @@ static inline void htp_im2col_vtcm_layout_build(struct htp_im2col_vtcm_layout *
         struct htp_im2col_context * ictx        = (struct htp_im2col_context *) data;                     \
         struct htp_ops_context *    octx        = ictx->octx;                                             \
         struct htp_thread_trace * restrict tr   = &octx->ctx->trace[ith];                                 \
+        const struct htp_tensor * restrict src0 = octx->src[0];                                           \
         const struct htp_tensor * restrict src1 = octx->src[1];                                           \
         const struct htp_tensor * restrict dst  = octx->dst;                                              \
-        const int32_t  s0                       = octx->op_params[0];                                     \
-        const int32_t  s1                       = octx->op_params[1];                                     \
-        const int32_t  p0                       = octx->op_params[2];                                     \
-        const int32_t  p1                       = octx->op_params[3];                                     \
-        const int32_t  d0                       = octx->op_params[4];                                     \
-        const int32_t  d1                       = octx->op_params[5];                                     \
-        const uint32_t N                        = src1->ne[3];                                            \
-        const uint32_t IC                       = src1->ne[2];                                            \
-        const uint32_t IH                       = src1->ne[1];                                            \
-        const uint32_t IW                       = src1->ne[0];                                            \
-        const uint32_t KH                       = octx->src[0]->ne[1];                                    \
-        const uint32_t KW                       = octx->src[0]->ne[0];                                    \
+        const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1];                                   \
+        const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3];                                   \
+        const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5];                                   \
+        const uint32_t N = src1->ne[3], IC = src1->ne[2], IH = src1->ne[1], IW = src1->ne[0];             \
+        const uint32_t KH                       = src0->ne[1], KW = src0->ne[0];                          \
         const uint32_t OH                       = dst->ne[2];                                             \
         const uint32_t OW                       = dst->ne[1];                                             \
         const uint32_t patch_stride             = IC * KH * KW;                                           \
         const float * restrict src_data         = (const float *) src1->data;                             \
         DST_CTYPE * restrict dst_data           = (DST_CTYPE *) dst->data;                                \
-        const uint32_t npatches                 = N * OH * OW;                                            \
-        const uint32_t patch_start              = ictx->npatches_per_thread * ith;                        \
-        const uint32_t patch_end                = MIN(patch_start + ictx->npatches_per_thread, npatches); \
-        if (patch_start >= patch_end) {                                                                   \
+        const uint32_t patch_end                = ictx->patch_base + ictx->npatches;                      \
+        const uint32_t patch_start              = ictx->patch_base + ictx->npatches_per_thread * ith;     \
+        const uint32_t patch_stop               = MIN(patch_start + ictx->npatches_per_thread, patch_end);\
+        if (patch_start >= patch_stop) {                                                                  \
             return;                                                                                       \
         }                                                                                                 \
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, patch_start);                                   \
-        for (uint32_t p = patch_start; p < patch_end; p++) {                                              \
+        for (uint32_t p = patch_start; p < patch_stop; p++) {                                             \
             const uint32_t iow             = p % OW;                                                      \
             const uint32_t ioh             = (p / OW) % OH;                                               \
             const uint32_t in              = p / (OW * OH);                                               \
@@ -154,10 +154,10 @@ IM2COL_PATCHEMBED_BODY(im2col_patchembed_f32_thread, float, hvx_copy_f32_uu, hvx
         uint8_t *      dst_base         = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread;                    \
         float *        srcb             = (float *) src_base;                                                        \
         DST_CTYPE *    dstb             = (DST_CTYPE *) dst_base;                                                    \
-        const uint32_t nrows            = N * OH;                                                                    \
+        const uint32_t row_end_max      = ictx->pe_row_base + ictx->pe_nrows;                                        \
         const uint32_t per_thread       = ictx->pe_rows_per_thread;                                                  \
-        const uint32_t row_start        = per_thread * ith;                                                          \
-        const uint32_t row_end          = MIN(row_start + per_thread, nrows);                                        \
+        const uint32_t row_start        = ictx->pe_row_base + per_thread * ith;                                      \
+        const uint32_t row_end          = MIN(row_start + per_thread, row_end_max);                                  \
         if (row_start >= row_end)                                                                                    \
             return;                                                                                                  \
         for (uint32_t r = row_start; r < row_end; r++) {                                                             \
@@ -266,26 +266,55 @@ int op_im2col(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    const uint32_t N         = src1->ne[3];
-    const uint32_t OH        = dst->ne[2];
-    const uint32_t OW        = dst->ne[1];
-    const uint32_t npatches  = N * OH * OW;
-    const uint32_t n_threads = MIN(octx->n_threads, npatches);
-
-    if ((octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) || n_threads == 0) {
+    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
         return HTP_STATUS_OK;
     }
 
+    const uint32_t N             = src1->ne[3];
+    const uint32_t OH            = dst->ne[2];
+    const uint32_t OW            = dst->ne[1];
+    const uint32_t total_patches = N * OH * OW;
+    const uint32_t total_rows    = N * OH;
+
+    uint32_t patch_base = 0;
+    uint32_t npatches   = total_patches;
+    if (octx->ctx->mdev.count > 1) {
+        const uint32_t patch_size = dst->nb[1];
+        const uint32_t patches_per_chunk = (patch_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(patch_size, HEX_L2_LINE_SIZE)) : 1;
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_patches, htp_tensor_mdev_data_aligned(dst) ? patches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        patch_base = range.start;
+        npatches   = range.count;
+    }
+
+    uint32_t row_base = 0;
+    uint32_t nrows    = total_rows;
+    if (octx->ctx->mdev.count > 1) {
+        const uint32_t row_size = dst->nb[2];
+        const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_base = range.start;
+        nrows    = range.count;
+    }
+
+    if (npatches == 0 && nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = MIN(octx->n_threads, MAX(npatches, 1));
+
     struct htp_im2col_context ictx = { 0 };
-    ictx.octx                      = octx;
-    ictx.npatches_per_thread       = (npatches + n_threads - 1) / n_threads;
+    ictx.octx                = octx;
+    ictx.patch_base          = patch_base;
+    ictx.npatches            = npatches;
+    ictx.npatches_per_thread = (npatches + n_threads - 1) / n_threads;
 
     // Clean non-overlapping patch-embed -> DMA kernel (if it fits VTCM);
     // everything else (padding/dilation/stride edges) -> pure-DDR kernel.
-    if (im2col_use_patchembed_dma(octx)) {
-        const uint32_t nrows = N * OH;
-        const uint32_t pth   = MIN(octx->n_threads, nrows);
+    if (im2col_use_patchembed_dma(octx) && nrows > 0) {
+        const uint32_t pth = MIN(octx->n_threads, nrows);
         if (pth > 0 && im2col_patchembed_dma_fits(octx, &ictx, pth)) {
+            ictx.pe_row_base        = row_base;
+            ictx.pe_nrows           = nrows;
             ictx.pe_rows_per_thread = (nrows + pth - 1) / pth;
             if (dst->type == HTP_TYPE_F16) {
                 work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_thread, &ictx, pth);
@@ -297,6 +326,10 @@ int op_im2col(struct htp_ops_context * octx) {
         // else: doesn't fit -> fall through to the pure-DDR kernel below.
     }
 
+    if (npatches == 0) {
+        return HTP_STATUS_OK;
+    }
+
     if (dst->type == HTP_TYPE_F16) {
         work_queue_run(octx->ctx->work_queue, im2col_patchembed_thread, &ictx, n_threads);
     } else {

ggml/src/ggml-hexagon/htp/main.c

@@ -34,6 +34,7 @@
 #include "work-queue.h"
 #include "hex-profile.h"
 #include "allreduce-ops.h"
+#include "htp-fence.h"
 
 #define HMX_QUEUE_CAPACITY     16
 #define HMX_QUEUE_STACK_SIZE   16384
@@ -710,22 +711,43 @@ static inline void profile_stop(uint32_t mode, struct profile_data * d) {
 static int op_fence(struct htp_ops_context * octx) {
     struct htp_context *ctx = octx->ctx;
     struct htp_thread_trace * tr = &ctx->trace[0];
-    const uint32_t seq = (uint32_t) octx->op_params[0];
+    const uint32_t seq  = (uint32_t) octx->op_params[0];
+    const uint32_t mode = (uint32_t) octx->op_params[1];
 
     htp_trace_event_start(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq);
 
     const struct htp_tensor * sync = octx->src[0];
-    atomic_uint * sync_fence = (atomic_uint *) sync->data;
+    atomic_uint * sync_fence = (atomic_uint *) (uintptr_t) sync->data;
+
+    if (mode == 1) {
+        htp_flush_dirty_ranges(ctx);
+
+        htp_mdev_group_barrier(octx);
+
+        if (ctx->mdev.idx == 0) {
+            htp_fence_write(sync_fence, seq, octx->status);
+        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq);
+        FARF(HIGH, "ggml-hex: sync-signal : fence %p seq 0x%x status %d\n", sync_fence, seq, octx->status);
+        return octx->status;
+    }
+
+    int status = HTP_STATUS_OK;
     uint64_t spins = 0;
     while (1) {
-        Q6_dccleaninva_A((void *) sync_fence);
-        asm volatile ("syncht" : : : "memory");
-        uint32_t val = atomic_load(&sync_fence[0]);
-        if ((int32_t)(val - seq) >= 0) {
+        uint32_t sync_seq;
+        uint32_t sync_status;
+        htp_fence_read(sync_fence, &sync_seq, &sync_status);
+        if ((int32_t)(sync_seq - seq) >= 0) {
+            if (sync_status > HTP_STATUS_OK) {
+                FARF(ERROR, "ggml-hex: sync-wait peer failed with status %u : fence %p seq 0x%x\n", sync_status, sync_fence, seq);
+                status = sync_status;
+            }
             break;
         }
         if (++spins > HTP_FENCE_TIMEOUT) {
-            FARF(ERROR, "ggml-hex: sync-wait TIMEOUT : fence %p spins %llu seq %u\n", sync_fence, spins, seq);
+            FARF(ERROR, "ggml-hex: sync-wait TIMEOUT : fence %p spins %llu seq 0x%x\n", sync_fence, spins, seq);
+            status = HTP_STATUS_INTERNAL_ERR;
             break;
         }
         hex_pause();
@@ -733,12 +755,27 @@ static int op_fence(struct htp_ops_context * octx) {
 
     htp_trace_event_stop(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq);
 
-    FARF(HIGH, "ggml-hex: sync-done : fence %p spins %llu seq %u\n", sync_fence, spins, seq);
+    FARF(HIGH, "ggml-hex: sync-done : fence %p spins %llu seq 0x%x\n", sync_fence, spins, seq);
+    return status;
+}
+
+static int op_mdev_group(struct htp_ops_context * octx) {
+    struct htp_context * ctx = octx->ctx;
+    const struct htp_tensor * sync = octx->src[0];
+    ctx->mdev.idx   = (uint16_t) octx->op_params[0];
+    ctx->mdev.count = (uint16_t) sync->ne[1];
+    if (ctx->mdev.count > 1) {
+        ctx->mdev.count_div = init_fastdiv_values(ctx->mdev.count);
+        ctx->mdev.fence_base = (uint8_t *) sync->data;
+    }
     return HTP_STATUS_OK;
 }
 
 static int execute_op(struct htp_ops_context * octx) {
     switch (octx->op) {
+        case HTP_OP_MDEV_GROUP:
+            return op_mdev_group(octx);
+
         case HTP_OP_FENCE:
             return op_fence(octx);
 
@@ -812,6 +849,7 @@ static int execute_op(struct htp_ops_context * octx) {
             return op_sum_rows(octx);
 
         case HTP_OP_CPY:
+        case HTP_OP_CPY_FENCE:
             return op_cpy(octx);
 
         case HTP_OP_REPEAT:
@@ -855,7 +893,7 @@ static int execute_op(struct htp_ops_context * octx) {
     }
 
     FARF(ERROR, "Unknown Op %u", octx->op);
-    return -1;
+    return HTP_STATUS_NO_SUPPORT;
 }
 
 static inline bool reuse_buf(struct htp_context *ctx, uint32_t *m_reuse, struct htp_buf_desc *b) {
@@ -984,11 +1022,19 @@ static void prep_tensors(struct htp_context *ctx, struct htp_buf_desc *bufs, str
     }
 }
 
-static int proc_op_req(struct htp_ops_context * octx, struct htp_tensor *tens, uint32_t idx, struct htp_op_desc * op) {
-    memcpy(octx->op_params, op->params, sizeof(octx->op_params));
+static void mdev_group_init(struct htp_context * ctx, const struct htp_opbatch_req * req) {
+    memset(&ctx->mdev, 0, sizeof(ctx->mdev));
+    ctx->mdev.fence_seq = (uint32_t)((req->seq & 0xfffff) << 12);
+}
+
+static int proc_op_req(struct htp_ops_context * octx, struct htp_buf_desc * bufs, uint32_t n_bufs,
+                       struct htp_tensor * tens, uint32_t idx, struct htp_op_desc * op) {
+    memcpy(octx->op_params,     op->params, sizeof(octx->op_params));
     memcpy(octx->kernel_params, op->kernel_params, sizeof(octx->kernel_params));
-    octx->flags = op->flags;
-    octx->op    = op->opcode;
+    octx->flags         = op->flags;
+    octx->op            = op->opcode;
+    octx->n_threads     = octx->ctx->n_threads;
+    octx->n_threads_div = octx->ctx->n_threads_div;
 
     FARF(HIGH, "proc-op #%u: opcode %u flags 0x%x", idx, octx->op, octx->flags);
 
@@ -1027,9 +1073,13 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_tensor *tens, u
             dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
     }
 
+    htp_tensor_dirty_all(octx->ctx, octx->dsts, HTP_OP_MAX_OUTPUTS);
+
+    htp_mdev_group_barrier(octx);
+
     int status = execute_op(octx);
 
-    htp_tensor_dirty_all(octx->ctx, octx->dsts, HTP_OP_MAX_OUTPUTS);
+    htp_ops_context_set_status(octx, status);
 
     octx->src0_spad.src = NULL;
     octx->src1_spad.src = NULL;
@@ -1037,7 +1087,7 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_tensor *tens, u
     octx->src3_spad.src = NULL;
     octx->dst_spad.src  = NULL;
 
-    return status;
+    return octx->status;
 }
 
 static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_req * req, const struct dspqueue_buffer * dbuf) {
@@ -1059,7 +1109,7 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
         return;
     }
 
-    FARF(HIGH, "processing opbatch #%u: n-bufs %u n-tensors %u n-ops %u n-traces %u : m-size %u b-size %u t-size %u o-size %u", req->id,
+    FARF(HIGH, "processing opbatch #%llu: n-bufs %u n-tensors %u n-ops %u n-traces %u : m-size %u b-size %u t-size %u o-size %u", (unsigned long long) req->seq,
             n_bufs, n_tens, n_ops, req->n_traces, dbuf->size, b_size, t_size, o_size);
 
     // Setup descriptor pointers
@@ -1096,8 +1146,11 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
 
     struct htp_ops_context *octx = &ctx->octx;
     memset(octx, 0, sizeof(*octx));
-    octx->n_threads = ctx->n_threads;
-    octx->ctx       = ctx;
+    octx->n_threads     = ctx->n_threads;
+    octx->n_threads_div = ctx->n_threads_div;
+    octx->ctx           = ctx;
+
+    mdev_group_init(ctx, req);
 
     work_queue_wakeup(ctx->work_queue);
     if (ctx->hmx_queue) {
@@ -1105,15 +1158,18 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
     }
 
     int op_status = HTP_STATUS_OK;
-    for (uint32_t i = 0; i < n_ops && op_status == HTP_STATUS_OK; i++) {
+    octx->status  = HTP_STATUS_OK;
+    for (uint32_t i = 0; i < n_ops; i++) {
         struct profile_data prof;
 
         profile_start(ctx->profiler, &prof);
 
-        op_status = proc_op_req(octx, tens, i, &ops[i]);
+        op_status = proc_op_req(octx, bufs, n_bufs, tens, i, &ops[i]);
 
         profile_stop(ctx->profiler, &prof);
 
+        htp_ops_context_set_status(octx, op_status);
+
         if (ctx->profiler) {
             pds[i].opcode = ops[i].opcode;
             pds[i].usecs  = prof.usecs;
@@ -1136,19 +1192,20 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
     qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
     htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
 
+    htp_mdev_group_barrier(octx);
+
     profile_stop(HTP_PROF_BASIC, &batch_prof);
 
     struct htp_opbatch_rsp rsp;
     memset(&rsp, 0, sizeof(rsp));
-    rsp.id           = req->id;
-    rsp.status       = op_status;
+    rsp.seq          = req->seq;
+    rsp.status       = octx->status;
     rsp.n_bufs       = n_bufs;
     rsp.n_tensors    = n_tens;
     rsp.n_ops        = n_ops;
     rsp.usecs        = batch_prof.usecs;
     rsp.cycles_start = batch_prof.cycles_start;
     rsp.cycles_stop  = batch_prof.cycles_stop;
-    rsp.seq          = req->seq;
 
     if (ctx->profiler == HTP_PROF_TRACE) {
         for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {

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

@@ -21,6 +21,7 @@
 #include "ggml-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"
 #include "matmul-ops.h"
 #include "htp-vtcm.h"
 
@@ -89,6 +90,8 @@ struct htp_mm_context {
 
     // Precomputed values
     uint32_t src0_nrows_per_thread;
+    uint32_t src0_row_start;
+    uint32_t src0_row_end;
     uint32_t src0_row_size_padded;
     uint32_t src1_nrows;
 
@@ -135,6 +138,23 @@ struct htp_mm_context {
     uint32_t vtcm_dst_size_per_thread;
 };
 
+static int htp_mm_init_context(
+    struct htp_ops_context * octx,
+    const struct htp_mm_kernel_params * kparams
+) {
+    if (!htp_ops_context_set_n_threads(octx, (uint32_t) kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    if (kparams->n_hmx) {
+        if (kparams->n_act_threads <= 0 || kparams->n_act_threads > (int32_t) octx->n_threads) {
+            return HTP_STATUS_INVAL_PARAMS;
+        }
+    }
+
+    return HTP_STATUS_OK;
+}
+
 // vdelta control to expand first 32 e8m0 values into 32 uint32 elements
 static const uint8_t __attribute__((aligned(128))) expand_x32_e8m0[128] = {
     0x00, 0x00, 0x00, 0x00, 0x01, 0x04, 0x00, 0x00, 0x02, 0x00, 0x08, 0x08, 0x01, 0x02, 0x00, 0x04, 0x04, 0x00, 0x00,
@@ -238,22 +258,24 @@ static void hvx_mm_4d(unsigned int nth, unsigned int ith, void * data) {
     // This is the size of the rest of the dimensions of the result
     const uint32_t nr1 = ne1 * ne2 * ne3;
 
+    const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start;
+
     // distribute the thread work across the inner or outer loop based on which one is larger
     uint32_t dr0, dr1, ith0, ith1;
     if (nr0 > nr1) {
-        dr0  = fastdiv(nr0 + nth - 1, &octx->ctx->n_threads_div);
+        dr0  = fastdiv(src0_nrows + nth - 1, &octx->n_threads_div);
         dr1  = nr1;
         ith0 = ith;
         ith1 = 0;
     } else {
-        dr0  = nr0;
-        dr1  = fastdiv(nr1 + nth - 1, &octx->ctx->n_threads_div);
+        dr0  = src0_nrows;
+        dr1  = fastdiv(nr1 + nth - 1, &octx->n_threads_div);
         ith0 = 0;
         ith1 = ith;
     }
 
-    const uint32_t ir0_start = dr0 * ith0;
-    const uint32_t ir0_end   = MIN(ir0_start + dr0, nr0);
+    const uint32_t ir0_start = mmctx->src0_row_start + dr0 * ith0;
+    const uint32_t ir0_end   = MIN(ir0_start + dr0, mmctx->src0_row_end);
 
     const uint32_t ir1_start = dr1 * ith1;
     const uint32_t ir1_end   = MIN(ir1_start + dr1, nr1);
@@ -312,11 +334,11 @@ static void hvx_mm_4d(unsigned int nth, unsigned int ith, void * data) {
 static void hvx_mm_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                                        \
     htp_matmul_preamble;                                                                                                          \
                                                                                                                                   \
-    const uint32_t src0_nrows = ne01 * ne02 * ne03;                                                                               \
+    const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start;                                                      \
     const uint32_t src1_nrows = ne11 * ne12 * ne13;                                                                               \
                                                                                                                                   \
-    const uint32_t src0_start_row  = src0_nrows_per_thread * ith;                                                                 \
-    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);                                     \
+    const uint32_t src0_start_row  = mmctx->src0_row_start + src0_nrows_per_thread * ith;                                         \
+    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end);                            \
                                                                                                                                   \
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                                        \
                                                                                                                                   \
@@ -414,10 +436,10 @@ static void hvx_mm_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void
 static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                                        \
     htp_matmul_preamble;                                                                                                          \
                                                                                                                                   \
-    const uint32_t src0_nrows = ne01;                                                                                             \
+    const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start;                                                      \
                                                                                                                                   \
-    const uint32_t src0_start_row  = src0_nrows_per_thread * ith;                                                                 \
-    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);                                     \
+    const uint32_t src0_start_row  = mmctx->src0_row_start + src0_nrows_per_thread * ith;                                         \
+    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end);                            \
                                                                                                                                   \
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                                        \
                                                                                                                                   \
@@ -549,12 +571,22 @@ static void hvx_mm_nx_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, v
         uint32_t n_k_tiles_w = ne00 / 32;                                                                                         \
         uint32_t tile_row_stride = n_k_tiles_w * tile_size;                                                                       \
                                                                                                                                   \
-        const uint32_t src0_nrows = ne01 * src_w->ne[2] * src_w->ne[3];                                                           \
-        uint32_t src0_nrows_per_thread = fastdiv(src0_nrows + nth - 1, &octx->ctx->n_threads_div);                                \
+        uint32_t src0_start_row = 0;                                                                                              \
+        uint32_t src0_end_row   = ne01;                                                                                           \
+        if (octx->ctx->mdev.count > 1) {                                                                                          \
+            const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float));                                              \
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(ne01, can_split ? 32 : 0,                        \
+                                                         octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); \
+            src0_start_row = range.start;                                                                                         \
+            src0_end_row   = range.start + range.count;                                                                           \
+        }                                                                                                                         \
+                                                                                                                                  \
+        const uint32_t nrows = src0_end_row - src0_start_row;                                                                     \
+        uint32_t src0_nrows_per_thread = fastdiv(nrows + nth - 1, &octx->n_threads_div);                                          \
         src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32);                                                          \
                                                                                                                                   \
-        const uint32_t start_row = src0_nrows_per_thread * ith;                                                                   \
-        const uint32_t end_row   = MIN(start_row + src0_nrows_per_thread, src0_nrows);                                            \
+        const uint32_t start_row = src0_start_row + src0_nrows_per_thread * ith;                                                  \
+        const uint32_t end_row   = MIN(start_row + src0_nrows_per_thread, src0_end_row);                                          \
         if (start_row >= end_row) continue;                                                                                       \
                                                                                                                                   \
         uint32_t ct_start = start_row / 32;                                                                                       \
@@ -735,11 +767,11 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) {
     assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0);
     const uint32_t prefetch_mask = n_prefetch - 1;
 
-    const uint32_t src0_nrows = ne01 * ne02 * ne03;  // src0 rows
-    const uint32_t src1_nrows = ne11 * ne12 * ne13;  // src1 rows
+    const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start;  // src0 rows
+    const uint32_t src1_nrows = ne11 * ne12 * ne13;                          // src1 rows
 
-    const uint32_t src0_start_row  = src0_nrows_per_thread * ith;
-    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+    const uint32_t src0_start_row  = mmctx->src0_row_start + src0_nrows_per_thread * ith;
+    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end);
     const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U);
 
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
@@ -781,7 +813,7 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) {
         const uint8_t * ss0 = dma_queue_pop(dma_queue).dst;
 
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0);
-        // Process src1 columns in pairs (2×2 tiling)
+        // Process src1 columns in pairs (2x2 tiling)
         uint32_t ir1 = 0;
         for (; ir1 + 1 < src1_nrows; ir1 += 2) {
             const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride);
@@ -791,7 +823,7 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) {
             mmctx->vec_dot_2x2(ne00, &dst_row0[ir0], &dst_row1[ir0], ss0, ss0 + src0_stride, src1_col0, src1_col1);
         }
 
-        // Handle remaining src1 rows (fallback to 2×1)
+        // Handle remaining src1 rows (fallback to 2x1)
         for (; ir1 < src1_nrows; ++ir1) {
             const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride);
             float * restrict dst_row          = (float *) (dst->data + (ir1 * dst_row_size));
@@ -833,10 +865,10 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) {
 static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) {
     htp_matmul_preamble;
 
-    const uint32_t src0_nrows = ne01;
+    const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start;
 
-    const uint32_t src0_start_row  = src0_nrows_per_thread * ith;
-    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+    const uint32_t src0_start_row  = mmctx->src0_row_start + src0_nrows_per_thread * ith;
+    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end);
 
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
@@ -943,13 +975,10 @@ static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) {
 
     const struct htp_tensor * restrict ids = octx->src[2];
 
-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
-
-    const uint32_t src0_nrows      = ne01;  // src0 rows per expert
+    const uint32_t src0_nrows      = mmctx->src0_row_end - mmctx->src0_row_start;  // src0 rows per expert
     const uint32_t src1_nrows      = ne11;
-    const uint32_t src0_start_row  = src0_nrows_per_thread * ith;
-    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+    const uint32_t src0_start_row  = mmctx->src0_row_start + src0_nrows_per_thread * ith;
+    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end);
 
     hvx_mm_run_quant_task(mmctx, ith);
 
@@ -1036,9 +1065,9 @@ static void hvx_mv_id(unsigned int nth, unsigned int ith, void * data) {
 
     const struct htp_tensor * restrict ids = octx->src[2];
 
-    const uint32_t src0_nrows      = ne01;  // src0 rows per expert
-    const uint32_t src0_start_row  = src0_nrows_per_thread * ith;
-    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+    const uint32_t src0_nrows      = mmctx->src0_row_end - mmctx->src0_row_start;  // src0 rows per expert
+    const uint32_t src0_start_row  = mmctx->src0_row_start + src0_nrows_per_thread * ith;
+    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end);
 
     hvx_mm_run_quant_task(mmctx, ith);
 
@@ -1143,12 +1172,22 @@ static void hvx_mv_id_nx(unsigned int nth, unsigned int ith, void * data) {
             const struct htp_tensor * restrict dst   = octx->dsts[p];
             if (!src_w || !dst) continue;
 
-            const uint32_t src0_nrows = src_w->ne[1];
-            uint32_t src0_nrows_per_thread = fastdiv(src0_nrows + nth - 1, &octx->ctx->n_threads_div);
+            const uint32_t ne01 = src_w->ne[1];
+            uint32_t start_row = 0;
+            uint32_t end_row   = ne01;
+            if (octx->ctx->mdev.count > 1) {
+                const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float));
+                const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(ne01, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+                start_row = range.start;
+                end_row   = range.start + range.count;
+            }
+
+            const uint32_t nrows = end_row - start_row;
+            uint32_t src0_nrows_per_thread = fastdiv(nrows + nth - 1, &octx->n_threads_div);
             src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32);
 
-            const uint32_t src0_start_row = src0_nrows_per_thread * ith;
-            const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+            const uint32_t src0_start_row = start_row + src0_nrows_per_thread * ith;
+            const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, end_row);
             if (src0_start_row >= src0_end_row) continue;
 
             const uint8_t * restrict src0_row = (const uint8_t *) src_w->data + eid * src_w->nb[2];
@@ -1227,12 +1266,22 @@ static void hvx_mm_id_nx(unsigned int nth, unsigned int ith, void * data) {
             const struct htp_tensor * restrict dst   = octx->dsts[p];
             if (!src_w || !dst) continue;
 
-            const uint32_t src0_nrows = src_w->ne[1];
-            uint32_t src0_nrows_per_thread = fastdiv(src0_nrows + nth - 1, &octx->ctx->n_threads_div);
+            const uint32_t ne01 = src_w->ne[1];
+            uint32_t start_row = 0;
+            uint32_t end_row   = ne01;
+            if (octx->ctx->mdev.count > 1) {
+                const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float));
+                const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(ne01, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+                start_row = range.start;
+                end_row   = range.start + range.count;
+            }
+
+            const uint32_t nrows = end_row - start_row;
+            uint32_t src0_nrows_per_thread = fastdiv(nrows + nth - 1, &octx->n_threads_div);
             src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32);
 
-            const uint32_t src0_start_row = src0_nrows_per_thread * ith;
-            const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+            const uint32_t src0_start_row = start_row + src0_nrows_per_thread * ith;
+            const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, end_row);
             if (src0_start_row >= src0_end_row) continue;
 
             const uint8_t * src0_row = (const uint8_t *) src_w->data + cur_a * src_w->nb[2];
@@ -1323,15 +1372,33 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
 
     const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
 
-    const uint32_t src0_nrows = ne01 * ne02 * ne03;
+    const uint32_t src0_nrows = ne01;
     const uint32_t src1_nrows = ne11 * ne12 * ne13;
 
+    uint32_t src0_row_start = 0;
+    uint32_t src0_row_end   = src0_nrows;
+
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float));
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        src0_row_start = range.start;
+        src0_row_end   = range.start + range.count;
+    }
+
+    if (src0_row_start >= src0_row_end) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t nrows = src0_row_end - src0_row_start;
+    mmctx->src0_row_start = src0_row_start;
+    mmctx->src0_row_end   = src0_row_end;
+
     bool is_repacked = (src0->type == HTP_TYPE_Q4_0 || src0->type == HTP_TYPE_Q4_1 ||
                         src0->type == HTP_TYPE_Q8_0 || src0->type == HTP_TYPE_IQ4_NL ||
                         src0->type == HTP_TYPE_MXFP4);
 
     // Compute src0_nrows_per_thread
-    mmctx->src0_nrows_per_thread  = fastdiv(src0_nrows + octx->n_threads - 1, &octx->ctx->n_threads_div);
+    mmctx->src0_nrows_per_thread  = fastdiv(nrows + octx->n_threads - 1, &octx->n_threads_div);
     if (is_repacked) {
         mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
     } else {
@@ -1503,13 +1570,13 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
         kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
         mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
     } else {
-        mmctx->vtcm_src1_size_per_thread = fastdiv(L.src1_bytes, &octx->ctx->n_threads_div);
+        mmctx->vtcm_src1_size_per_thread = fastdiv(L.src1_bytes, &octx->n_threads_div);
     }
 
-    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->ctx->n_threads_div);
-    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->ctx->n_threads_div);
+    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div);
+    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->n_threads_div);
 
-    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
+    const size_t vtcm_size = L.total_bytes;
 
     FARF(HIGH, "matmul-%s : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu (%zu)\n", mmctx->type,
          L.src0_bytes, L.src1_bytes, L.dst_bytes, vtcm_size);
@@ -1583,13 +1650,21 @@ static void hvx_mm_nx_2d(unsigned int nth, unsigned int ith, void * data) {
 
         const uint32_t ne00 = src_w->ne[0];
         const uint32_t ne01 = src_w->ne[1];
-        const uint32_t src0_nrows = ne01 * src_w->ne[2] * src_w->ne[3];
+        uint32_t start_row = 0;
+        uint32_t end_row   = ne01;
+        if (octx->ctx->mdev.count > 1) {
+            const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float));
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(ne01, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            start_row = range.start;
+            end_row   = range.start + range.count;
+        }
 
-        uint32_t src0_nrows_per_thread = fastdiv(src0_nrows + nth - 1, &octx->ctx->n_threads_div);
+        const uint32_t nrows = end_row - start_row;
+        uint32_t src0_nrows_per_thread = fastdiv(nrows + nth - 1, &octx->n_threads_div);
         src0_nrows_per_thread += (src0_nrows_per_thread & 1);
 
-        const uint32_t src0_start_row  = src0_nrows_per_thread * ith;
-        const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+        const uint32_t src0_start_row  = start_row + src0_nrows_per_thread * ith;
+        const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, end_row);
         const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U);
         if (src0_start_row >= src0_end_row) continue;
 
@@ -2638,10 +2713,6 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k
     const struct htp_tensor * restrict src0 = octx->src[0];
     const struct htp_tensor * restrict act  = octx->src[n_weights];
 
-    if (!src0 || !act) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-
     const int weight_type = (int) src0->type;
     const int k           = (int) act->ne[0];
     const int k_valid     = (int) act->ne[0];
@@ -2714,16 +2785,31 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k
 
     hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00));  // scale: 1.0, bias: 0.0 in FP16
 
-    FARF(HIGH, "hmx-mm-nx-2d: n_weights %u m %d k %d wtype %d mc %d nc %d vtcm %zu/%zu",
-         n_weights, m, k, weight_type, m_chunk_n_rows, n_chunk_n_cols, L.total_bytes, vtcm_budget);
+    int m_start = 0;
+    int m_rows  = m;
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_can_row_partition(octx->dsts[0], sizeof(float));
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition((uint32_t) m, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        m_start = (int) range.start;
+        m_rows  = (int) range.count;
+    }
+
+    if (m_rows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    FARF(HIGH, "hmx-mm-nx-2d: n_weights %u m %d (%d..%d) k %d wtype %d mc %d nc %d vtcm %zu/%zu",
+         n_weights, m, m_start, m_start + m_rows, k, weight_type, m_chunk_n_rows, n_chunk_n_cols, L.total_bytes, vtcm_budget);
 
     htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
 
+    const size_t mr_end = (size_t)(m_start + m_rows);
+
     if (pipeline) {
         hmx_matmul_job_t job_slots[2];
 
-        for (size_t mr = 0; mr < (size_t) m; mr += m_chunk_n_rows) {
-            const size_t n_rows = hex_smin(m - mr, m_chunk_n_rows);
+        for (size_t mr = (size_t) m_start; mr < mr_end; mr += m_chunk_n_rows) {
+            const size_t n_rows = hex_smin(mr_end - mr, m_chunk_n_rows);
 
             void *vtcm_weight_bufs[2] = { vtcm_scratch0, vtcm_scratch1 };
             void *vtcm_output_bufs[2] = { vtcm_output,   vtcm_scratch2 };
@@ -2822,8 +2908,8 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k
         }
     } else {
         hmx_matmul_job_t job;
-        for (size_t mr = 0; mr < (size_t) m; mr += m_chunk_n_rows) {
-            const size_t n_rows = hex_smin(m - mr, m_chunk_n_rows);
+        for (size_t mr = (size_t) m_start; mr < mr_end; mr += m_chunk_n_rows) {
+            const size_t n_rows = hex_smin(mr_end - mr, m_chunk_n_rows);
 
             struct activation_transfer_params act_params = {
                 .ctx = ctx,
@@ -3095,7 +3181,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_
                             int chunk_dst_cols = params->n - (int)nc;
                             if (chunk_dst_cols > 0) {
                                 transfer_output_chunk_threaded(ctx, output, src2_chunk, vtcm_output, (int) n_rows, (int) n_cols,
-                                                               params->dst_stride, params->src2_stride, chunk_dst_cols, ctx->n_threads);
+                                                               params->dst_stride, params->src2_stride, chunk_dst_cols, n_threads);
                             }
                         }
                     }
@@ -3216,7 +3302,10 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx,
                                          int weight_type,
                                          const struct mmid_row_mapping *matrix_rows,
                                          int cur_a,
-                                         int mapping_stride) {
+                                         int mapping_stride,
+                                         int m_start,
+                                         int m_end,
+                                         int n_threads) {
     struct htp_thread_trace * tr = &ctx->trace[0];
     htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
 
@@ -3247,7 +3336,6 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx,
     const int n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS;
     const struct fastdiv_values n_k_tiles_div = init_fastdiv_values(n_k_tiles);
 
-    const int n_threads = ctx->n_threads;
     const bool is_quant   = (weight_type != HTP_TYPE_F16 && weight_type != HTP_TYPE_F32);
 
     const size_t vec_dot_size = k * sizeof(__fp16);
@@ -3303,8 +3391,8 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx,
 
     hmx_matmul_job_t job;
 
-    for (size_t mr = 0; mr < (size_t) m_padded; mr += m_chunk_n_rows) {
-        const size_t n_rows = hex_smin(m_padded - mr, m_chunk_n_rows);
+    for (size_t mr = (size_t) m_start; mr < (size_t) m_end; mr += m_chunk_n_rows) {
+        const size_t n_rows = hex_smin((size_t) m_end - mr, m_chunk_n_rows);
         const size_t n_row_tiles = hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS);
 
         transfer_activation_chunk_gathered_threaded(
@@ -3368,31 +3456,48 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k
     const int act_stride = (int)(src1->nb[1] / sizeof(float));
     const int wgt_stride = (int)(src0->nb[1] / sizeof(__fp16));
 
+    int m_start = 0;
+    int m_rows  = m_total;
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float));
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition((uint32_t) m_total, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        m_start = (int) range.start;
+        m_rows  = (int) range.count;
+    }
+
+    if (m_rows == 0) {
+        return HTP_STATUS_OK;
+    }
+
     const float * src2_ptr = NULL;
     uint32_t src2_stride = 0;
     size_t src2_nb2 = 0;
     size_t src2_nb3 = 0;
     if (src2) {
-        src2_ptr = (const float *) src2->data;
         src2_stride = (src2->ne[1] == 1) ? 0 : (uint32_t) (src2->nb[1] / sizeof(float));
+        src2_ptr = (const float *) src2->data + m_start * src2_stride;
         src2_nb2 = (src2->ne[2] == 1) ? 0 : src2->nb[2];
         src2_nb3 = (src2->ne[3] == 1) ? 0 : src2->nb[3];
     }
 
+    const int dst_stride = (int)(dst->nb[1] / sizeof(float));
+    float       * dst_ptr = (float *)       dst->data  + m_start * dst_stride;
+    const float * act_ptr = (const float *) src1->data + m_start * act_stride;
+
     int ret = -1;
-    const int n_threads = MIN(kparams->n_threads, (int) octx->n_threads);
+    const int n_threads = kparams->n_threads;
     if (kparams->kernel_type == HTP_MM_KERNEL_HMX_F16_BATCHED) {
         hmx_mm_f16_f32_batched_params_t batch_params = {
-            .dst             = (float *) dst->data,
+            .dst             = dst_ptr,
             .src2            = src2_ptr,
-            .activation      = (float *) src1->data,
+            .activation      = act_ptr,
             .weight          = (const __fp16 *) src0->data,
-            .m               = m_total,
+            .m               = m_rows,
             .k               = k,
             .n               = n,
             .act_stride      = act_stride,
             .weight_stride   = wgt_stride,
-            .dst_stride      = (int) (dst->nb[1] / sizeof(float)),
+            .dst_stride      = dst_stride,
             .src2_stride     = src2_stride,
             .ne02            = ne02,
             .ne03            = ne03,
@@ -3420,9 +3525,9 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k
                                      kparams->vtcm_size);
     } else {
         ret = hmx_mm_2d_f32(
-            octx->ctx, (float*) dst->data, src2_ptr, (float*) src1->data, (const uint8_t *) src0->data,
-            m_total, k, n, act_stride, (int) src0->nb[1], (int) src0->type, (int) src1->ne[0],
-            (int)(dst->nb[1] / sizeof(float)), src2_stride, (int)dst->ne[0],
+            octx->ctx, dst_ptr, src2_ptr, act_ptr, (const uint8_t *) src0->data,
+            m_rows, k, n, act_stride, (int) src0->nb[1], (int) src0->type, (int) src1->ne[0],
+            dst_stride, src2_stride, (int)dst->ne[0],
             kparams->m_chunk, kparams->n_chunk, kparams->pipeline, n_threads,
             kparams->n_act_threads,
             &kparams->div_n_act_threads,
@@ -3441,6 +3546,11 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k
 int op_matmul(struct htp_ops_context * octx) {
     const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
 
+    const int status = htp_mm_init_context(octx, kparams);
+    if (status != HTP_STATUS_OK) {
+        return status;
+    }
+
     if (kparams->n_hmx) {
         return hmx_mm_op_matmul(octx, kparams);
     }
@@ -3463,6 +3573,16 @@ static int hmx_mm_op_matmul_id(
         const int32_t cne1 = matrix_row_counts[cur_a];
         if (cne1 == 0) continue;
 
+        const int m_padded = hex_align_up(cne1, 32);
+        int m_start = 0, m_end = m_padded;
+        if (octx->ctx->mdev.count > 1) {
+            const bool can_split = htp_tensor_mdev_data_aligned(dst) && (uint32_t) cne1 >= octx->ctx->mdev.count;
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition((uint32_t) m_padded, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            m_start = (int) range.start;
+            m_end   = (int) (range.start + range.count);
+        }
+        if (m_start >= m_end) continue;
+
         int ret = hmx_mm_id_2d_f32(octx->ctx, (float*) dst->data, (float*) src1->data,
                                    (const uint8_t *) src0->data + cur_a * nb02,
                                    cne1, ne00, ne01,
@@ -3471,7 +3591,8 @@ static int hmx_mm_op_matmul_id(
                                    nb11, nb12,
                                    nb1, nb2,
                                    (int) src0->nb[1], (int) src0->type,
-                                   matrix_rows, cur_a, mmctx->mapping_stride);
+                                   matrix_rows, cur_a, mmctx->mapping_stride,
+                                   m_start, m_end, (int) octx->n_threads);
         if (ret != 0) {
             FARF(ERROR, "HMX matmul failed for expert %u, error %d\n", cur_a, ret);
             return HTP_STATUS_NO_SUPPORT;
@@ -3524,7 +3645,7 @@ static int hvx_mm_matmul_id(
     htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
                                  0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false);
 
-    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
+    const size_t vtcm_size = L.total_bytes;
 
     FARF(HIGH, "matmul-id-%s : src0-spad-size %zu src1-spad-size %zu src2-spad-size 0 dst-spad-size %zu (%zu)\n", mmctx->type,
          L.src0_bytes, L.src1_bytes, L.dst_bytes, vtcm_size);
@@ -3554,10 +3675,10 @@ static int hvx_mm_matmul_id(
     mmctx->vtcm_src0_stride = src0_row_size_padded;
     mmctx->vtcm_src1_stride = src1_row_size;
 
-    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->ctx->n_threads_div);
+    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div);
     mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
     mmctx->vtcm_src2_size_per_thread = 0;
-    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->ctx->n_threads_div);
+    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->n_threads_div);
 
     mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
     mmctx->quant_task_func = quant_task_func;
@@ -3587,6 +3708,20 @@ static int hmx_mm_op_matmul_id_nx(
         const int32_t cne1 = matrix_row_counts[cur_a];
         if (cne1 == 0) continue;
 
+        const int m_padded = hex_align_up(cne1, 32);
+        int m_start = 0, m_end = m_padded;
+        if (octx->ctx->mdev.count > 1) {
+            bool can_split = (uint32_t) cne1 >= octx->ctx->mdev.count;
+            for (uint32_t p = 0; p < n_weights && can_split; ++p) {
+                const struct htp_tensor * restrict dst = octx->dsts[p];
+                can_split = !dst || htp_tensor_mdev_data_aligned(dst);
+            }
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition((uint32_t) m_padded, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            m_start = (int) range.start;
+            m_end   = (int) (range.start + range.count);
+        }
+        if (m_start >= m_end) continue;
+
         for (uint32_t p = 0; p < n_weights; ++p) {
             const struct htp_tensor * restrict src_w = octx->src[p];
             const struct htp_tensor * restrict dst   = octx->dsts[p];
@@ -3600,7 +3735,8 @@ static int hmx_mm_op_matmul_id_nx(
                                        act->nb[1], act->nb[2],
                                        dst->nb[1], dst->nb[2],
                                        (int) src_w->nb[1], (int) src_w->type,
-                                       matrix_rows, cur_a, mmctx->mapping_stride);
+                                       matrix_rows, cur_a, mmctx->mapping_stride,
+                                       m_start, m_end, (int) octx->n_threads);
             if (ret != 0) {
                 FARF(ERROR, "HMX matmul ID NX failed for expert %u weight %u, error %d\n", cur_a, p, ret);
                 return HTP_STATUS_NO_SUPPORT;
@@ -3656,7 +3792,7 @@ static int hvx_mm_matmul_id_nx(
     htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, act->ne[0], src1_nrows, octx->n_threads,
                                  0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false);
 
-    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
+    const size_t vtcm_size = L.total_bytes;
 
     if (octx->ctx->vtcm_size < vtcm_size) {
         FARF(ERROR, "matmul-id-nx: current VTCM reservation %zu is too small, needed %zu\n",
@@ -3678,9 +3814,9 @@ static int hvx_mm_matmul_id_nx(
     mmctx->vtcm_src0_stride = 0;
     mmctx->vtcm_src1_stride = src1_row_size;
 
-    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->ctx->n_threads_div);
+    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div);
     mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
-    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->ctx->n_threads_div);
+    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->n_threads_div);
 
     mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
     mmctx->quant_task_func         = quant_task_func;
@@ -3769,16 +3905,21 @@ static inline void scan_expert_ids(
 int op_matmul_id(struct htp_ops_context * octx) {
     htp_matmul_tensors_preamble;
 
+    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
+    struct htp_mm_context mmctx_struct = {0};
+    struct htp_mm_context * mmctx = &mmctx_struct;
+
+    const int status = htp_mm_init_context(octx, kparams);
+    if (status != HTP_STATUS_OK) {
+        return status;
+    }
+
     struct htp_thread_trace * tr = &octx->ctx->trace[0];
     htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
 
-    struct htp_mm_context mmctx_struct = {0};
-    struct htp_mm_context * mmctx = &mmctx_struct;
     mmctx->octx = octx;
     mmctx->act = src1;
 
-    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
-
     const struct htp_tensor * restrict ids = octx->src[2];
 
     const size_t src0_row_size = nb01;
@@ -3789,9 +3930,6 @@ int op_matmul_id(struct htp_ops_context * octx) {
     const uint32_t src0_nrows = ne01;  // per expert
     const uint32_t src1_nrows = ne11 * ne12 * ne13;
 
-    mmctx->src0_nrows_per_thread = fastdiv(src0_nrows + octx->n_threads - 1, &octx->ctx->n_threads_div);
-    mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
-
     // row groups
     const int n_ids = ids->ne[0];  // n_expert_used
     const int n_as  = ne02;        // n_expert
@@ -3843,6 +3981,29 @@ int op_matmul_id(struct htp_ops_context * octx) {
     if (kparams->n_hmx) {
         s = hmx_mm_op_matmul_id(octx, mmctx);
     } else {
+        uint32_t src0_row_start = 0;
+        uint32_t src0_row_end   = src0_nrows;
+        if (octx->ctx->mdev.count > 1) {
+            const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float));
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            src0_row_start = range.start;
+            src0_row_end   = range.start + range.count;
+        }
+
+        if (src0_row_start >= src0_row_end) {
+            if (mapping_buf != octx->ctx->ddr_spad_base) {
+                free(mapping_buf);
+            }
+            return HTP_STATUS_OK;
+        }
+
+        const uint32_t nrows = src0_row_end - src0_row_start;
+        mmctx->src0_row_start = src0_row_start;
+        mmctx->src0_row_end   = src0_row_end;
+
+        mmctx->src0_nrows_per_thread = fastdiv(nrows + octx->n_threads - 1, &octx->n_threads_div);
+        mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
+
         if (hvx_mm_init_vec_dot(mmctx, src0->type) == 0) {
             s = hvx_mm_matmul_id(octx, mmctx, src1_nrows > 1 ? hvx_mm_id : hvx_mv_id);
         } else {
@@ -3858,29 +4019,31 @@ int op_matmul_id(struct htp_ops_context * octx) {
 }
 
 int op_matmul_id_nx(struct htp_ops_context * octx) {
+    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
+    struct htp_mm_context mmctx_struct = {0};
+    struct htp_mm_context * mmctx = &mmctx_struct;
+
+    const int status = htp_mm_init_context(octx, kparams);
+    if (status != HTP_STATUS_OK) {
+        return status;
+    }
+
     struct htp_thread_trace * tr = &octx->ctx->trace[0];
     htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
 
-    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
+    mmctx->octx = octx;
     const uint32_t n_weights = kparams->n_weights;
     const struct htp_tensor * restrict src0 = octx->src[0];
     const struct htp_tensor * restrict act  = octx->src[n_weights];
     const struct htp_tensor * restrict ids  = octx->src[n_weights + 1];
 
-    struct htp_mm_context mmctx_struct = {0};
-    struct htp_mm_context * mmctx = &mmctx_struct;
-    mmctx->octx = octx;
     mmctx->act = act;
 
     const size_t src0_row_size = src0->nb[1];
     const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128);
 
-    const uint32_t src0_nrows = src0->ne[1];
     const uint32_t src1_nrows = act->ne[1] * act->ne[2] * act->ne[3];
 
-    mmctx->src0_nrows_per_thread = fastdiv(src0_nrows + octx->n_threads - 1, &octx->ctx->n_threads_div);
-    mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
-
     const int n_ids = ids->ne[0];
     const int n_as  = src0->ne[2];
 
@@ -3946,6 +4109,12 @@ int op_matmul_id_nx(struct htp_ops_context * octx) {
 }
 int op_matmul_nx(struct htp_ops_context * octx) {
     const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
+
+    const int status = htp_mm_init_context(octx, kparams);
+    if (status != HTP_STATUS_OK) {
+        return status;
+    }
+
     if (kparams->n_hmx) {
         return hmx_mm_nx_2d_f32(octx, kparams);
     }
@@ -4012,7 +4181,7 @@ int op_matmul_nx(struct htp_ops_context * octx) {
     htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, act->ne[0], src1_nrows, octx->n_threads,
                                  0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, true);
 
-    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
+    const size_t vtcm_size = L.total_bytes;
 
     if (octx->ctx->vtcm_size < vtcm_size) {
         FARF(ERROR, "matmul-nx: current VTCM reservation %zu is too small, needed %zu\n",
@@ -4034,9 +4203,9 @@ int op_matmul_nx(struct htp_ops_context * octx) {
     mmctx->vtcm_src0_stride = is_repacked ? 0 : src0_row_size_padded;
     mmctx->vtcm_src1_stride = src1_row_size;
 
-    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->ctx->n_threads_div);
+    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div);
     mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
-    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->ctx->n_threads_div);
+    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->n_threads_div);
 
     mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
     mmctx->quant_task_func = quant_task_func;