tests : remove stale comment (#29140)

* ggml-metal : derive non-empty fusion ops from ops_all

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* ggml-metal : drop _all suffix from fusion op pattern vectors

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

.devops/openvino.Dockerfile

@@ -1,5 +1,5 @@
-ARG OPENVINO_VERSION_MAJOR=2026.3.1
-ARG OPENVINO_VERSION_FULL=2026.3.1.22476.56d9685302d
+ARG OPENVINO_VERSION_MAJOR=2026.4
+ARG OPENVINO_VERSION_FULL=2026.4.0.22959.99c81491cc3
 ARG UBUNTU_VERSION=24.04
 
 # Intel GPU driver versions. https://github.com/intel/compute-runtime/releases
@@ -10,9 +10,9 @@ ARG COMPUTE_RUNTIME_VERSION_FULL=26.31.39395.13-0
 ARG IGDGMM_VERSION=22.10.0
 
 # Intel NPU driver versions. https://github.com/intel/linux-npu-driver/releases
-ARG NPU_DRIVER_VERSION=v1.35.0
-ARG NPU_DRIVER_FULL=v1.35.0.20260722-29947505341
-ARG LIBZE1_VERSION=1.28.2-1~24.04~ppa1
+ARG NPU_DRIVER_VERSION=v1.38.0
+ARG NPU_DRIVER_FULL=v1.38.0.20260910-34487311128
+ARG LIBZE1_VERSION=1.32.0-1~24.04~ppa1
 
 # Optional proxy build arguments
 ARG http_proxy=
@@ -173,7 +173,7 @@ RUN --mount=type=cache,target=/var/cache/intel-npu,sharing=locked \
     fi; \
     DEB=/var/cache/intel-npu/libze1_${LIBZE1_VERSION}_amd64.deb; \
     if [ ! -f "$DEB" ]; then \
-        wget -q -O "$DEB" https://snapshot.ppa.launchpadcontent.net/kobuk-team/intel-graphics/ubuntu/20260606T100000Z/pool/main/l/level-zero-loader/libze1_${LIBZE1_VERSION}_amd64.deb; \
+        wget -q -O "$DEB" https://snapshot.ppa.launchpadcontent.net/kobuk-team/intel-graphics/ubuntu/20260830T100000Z/pool/main/l/level-zero-loader/libze1_${LIBZE1_VERSION}_amd64.deb; \
     fi; \
     mkdir /tmp/npu/ && cd /tmp/npu/ && tar -xf "$TGZ" && cp "$DEB" .; \
     apt-get update; \

.github/workflows/build-and-test-snapdragon.yml

@@ -29,7 +29,7 @@ concurrency:
 
 jobs:
   android-ndk-snapdragon:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04  # previously ubuntu-latest
     container:
       image: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.7'
     defaults:
@@ -59,7 +59,7 @@ jobs:
           path: pkg-snapdragon/llama.cpp
 
   linux-iot-snapdragon:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04  # previously ubuntu-latest
     container:
       image: 'ghcr.io/snapdragon-toolchain/arm64-linux:v0.7'
     defaults:

.github/workflows/build-android.yml

@@ -33,7 +33,7 @@ env:
 
 jobs:
   default:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04  # previously ubuntu-latest
 
     steps:
       - name: Clone
@@ -49,7 +49,7 @@ jobs:
           distribution: zulu
 
       - name: Setup Android SDK
-        uses: android-actions/setup-android@40fd30fb8d7440372e1316f5d1809ec01dcd3699 # v4.0.1
+        uses: android-actions/setup-android@be39fa834029ff78f1a44aa3bb0819b8fc2bd8fd # v4.0.4
         with:
           log-accepted-android-sdk-licenses: false
 
@@ -59,7 +59,7 @@ jobs:
           ./gradlew build --no-daemon
 
   ndk:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04  # previously ubuntu-latest
     container:
       image: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.3'
     defaults:
@@ -93,7 +93,7 @@ jobs:
           path: pkg-adb/llama.cpp
 
   arm64:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04  # previously ubuntu-latest
 
     env:
       NDK_VERSION: "29.0.14206865"
@@ -123,7 +123,7 @@ jobs:
           distribution: temurin
 
       - name: Setup Android SDK
-        uses: android-actions/setup-android@40fd30fb8d7440372e1316f5d1809ec01dcd3699 # v4.0.1
+        uses: android-actions/setup-android@be39fa834029ff78f1a44aa3bb0819b8fc2bd8fd # v4.0.4
         with:
           log-accepted-android-sdk-licenses: false
 

.github/workflows/build-cache.yml

@@ -41,8 +41,8 @@ jobs:
 
     env:
       # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile
-      OPENVINO_VERSION_MAJOR: "2026.3.1"
-      OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d"
+      OPENVINO_VERSION_MAJOR: "2026.4"
+      OPENVINO_VERSION_FULL: "2026.4.0.22959.99c81491cc3"
 
     steps:
       - name: Clone
@@ -69,8 +69,8 @@ jobs:
 
     env:
       # Sync versions in build.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile
-      OPENVINO_VERSION_MAJOR: "2026.3.1"
-      OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d"
+      OPENVINO_VERSION_MAJOR: "2026.4"
+      OPENVINO_VERSION_FULL: "2026.4.0.22959.99c81491cc3"
 
     steps:
       - name: Clone

.github/workflows/build-openvino.yml

@@ -41,8 +41,8 @@ jobs:
 
     env:
       # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile
-      OPENVINO_VERSION_MAJOR: "2026.3.1"
-      OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d"
+      OPENVINO_VERSION_MAJOR: "2026.4"
+      OPENVINO_VERSION_FULL: "2026.4.0.22959.99c81491cc3"
 
     steps:
       - name: Clone
@@ -96,8 +96,8 @@ jobs:
 
     env:
       # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile
-      OPENVINO_VERSION_MAJOR: "2026.3.1"
-      OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d"
+      OPENVINO_VERSION_MAJOR: "2026.4"
+      OPENVINO_VERSION_FULL: "2026.4.0.22959.99c81491cc3"
 
     steps:
       - name: Clone

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

@@ -412,8 +412,8 @@ jobs:
 
     env:
       # Sync versions in build.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile
-      OPENVINO_VERSION_MAJOR: "2026.3.1"
-      OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d"
+      OPENVINO_VERSION_MAJOR: "2026.4"
+      OPENVINO_VERSION_FULL: "2026.4.0.22959.99c81491cc3"
 
     steps:
       - name: Clone

.github/workflows/check-vendor.yml

@@ -19,7 +19,7 @@ on:
 
 jobs:
   check-vendor:
-    runs-on: [self-hosted, fast]
+    runs-on: ubuntu-slim
 
     steps:
       - name: Checkout

.github/workflows/copilot-setup-steps.yml

@@ -11,10 +11,12 @@ on:
     paths:
       - .github/workflows/copilot-setup-steps.yml
 
+cache-mode: none
+
 jobs:
   # The job MUST be called `copilot-setup-steps` or it will not be picked up by Copilot.
   copilot-setup-steps:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
 
     # Set the permissions to the lowest permissions possible needed for your steps.
     # Copilot will be given its own token for its operations.
@@ -31,8 +33,8 @@ jobs:
       - name: ccache
         uses: ggml-org/ccache-action@v1.2.24
         with:
-          key: copilot-setup-steps
-          evict-old-files: 1d
+          restore: false
+          save: false
 
       - name: Dependencies
         id: depends

.github/workflows/gguf-publish.yml

@@ -21,7 +21,7 @@ on:
 jobs:
   deploy:
 
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04  # previously ubuntu-latest
 
     steps:
     - uses: actions/checkout@v6

.github/workflows/make-release.yml

@@ -13,6 +13,16 @@ on:
         required: true
         type: boolean
         default: true
+      skip_apiabi_check:
+        description: 'Skip API/ABI compatibility check'
+        required: false
+        type: boolean
+        default: false
+      apiabi_compare_tag:
+        description: 'Tag to compare against for API/ABI check (default: latest release)'
+        required: false
+        type: string
+        default: ''
 
 env:
   GH_TOKEN: ${{ github.token }}
@@ -23,7 +33,7 @@ permissions:
 
 jobs:
   make-release:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04  # previously ubuntu-latest
 
     steps:
       - name: Checkout
@@ -33,12 +43,18 @@ jobs:
           ref: ${{ inputs.commit != '' && inputs.commit || github.ref_name }}
           fetch-depth: 0
 
+      - name: Install API/ABI check tools
+        if: ${{ github.event.inputs.skip_apiabi_check != 'true' }}
+        run: sudo apt-get install -y abi-compliance-checker abigail-tools
+
       - name: Run release checks
         id: checks
         run: bash scripts/make-release-checks.sh ${{ github.event.inputs.dry_run == 'true' && '--dry-run' || '' }}
         env:
           GITHUB_REPOSITORY: ${{ github.repository }}
           RELEASE_BRANCH: ${{ github.ref_name }}
+          SKIP_APIABI_CHECK: ${{ github.event.inputs.skip_apiabi_check }}
+          APIABI_COMPARE_TAG: ${{ github.event.inputs.apiabi_compare_tag }}
 
       - name: Create release tag
         if: ${{ github.event.inputs.dry_run == 'false' }}

.github/workflows/pre-tokenizer-hashes.yml

@@ -12,7 +12,7 @@ on:
 
 jobs:
     pre-tokenizer-hashes:
-        runs-on: [self-hosted, fast]
+        runs-on: ubuntu-slim
 
         steps:
         - name: Checkout repository

.github/workflows/python-check-requirements.yml

@@ -20,7 +20,7 @@ concurrency:
 
 jobs:
   python-check-requirements:
-    runs-on: [self-hosted, CPU, fast]
+    runs-on: ${{ 'ubuntu-24.04-arm' || 'ubuntu-24.04' }}
     name: check-requirements
     steps:
       - name: Check out source repository

.github/workflows/python-lint.yml

@@ -21,7 +21,7 @@ concurrency:
 
 jobs:
   flake8-lint:
-    runs-on: [self-hosted, fast]
+    runs-on: ubuntu-slim
     name: Lint
     steps:
       - name: Check out source repository

.github/workflows/python-type-check.yml

@@ -22,7 +22,7 @@ concurrency:
 
 jobs:
   python-type-check:
-    runs-on: [self-hosted, fast]
+    runs-on: ubuntu-slim
     name: python type-check
     steps:
       - name: Check out source repository

.github/workflows/release.yml

@@ -106,6 +106,7 @@ jobs:
         uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-${{ matrix.os }}-${{ matrix.arch }}
+          evict-old-files: 1d
 
       - name: Build
         id: cmake_build
@@ -190,6 +191,7 @@ jobs:
         uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-${{ matrix.os }}-cpu
+          evict-old-files: 1d
 
       - name: Build
         id: cmake_build
@@ -275,6 +277,7 @@ jobs:
         uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-${{ matrix.os }}-vulkan
+          evict-old-files: 1d
 
       - name: Build
         id: cmake_build
@@ -327,13 +330,13 @@ jobs:
             defines: '-DGGML_CUDA_CUB_3DOT2=ON'
           - build: 'x64'
             os: ubuntu-24.04
-            cuda: '13.3.1'
-            label: '13.3'
+            cuda: '13.4.1'
+            label: '13.4'
             defines: ''
           - build: 'arm64'
             os: ubuntu-24.04-arm
-            cuda: '13.3.1'
-            label: '13.3'
+            cuda: '13.4.1'
+            label: '13.4'
             defines: ''
 
     runs-on: ${{ matrix.os }}
@@ -453,7 +456,7 @@ jobs:
     needs: [check-release, ui-build]
     if: ${{ needs.check-release.outputs.should_release == 'true' }}
 
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04  # previously ubuntu-latest
 
     #permissions:
     #  actions: write
@@ -481,10 +484,9 @@ jobs:
           distribution: temurin
 
       - name: Setup Android SDK
-        uses: android-actions/setup-android@40fd30fb8d7440372e1316f5d1809ec01dcd3699 # v4.0.1
+        uses: android-actions/setup-android@be39fa834029ff78f1a44aa3bb0819b8fc2bd8fd # v4.0.4
         with:
           log-accepted-android-sdk-licenses: false
-          packages: 'platform-tools'
 
       - name: Install NDK
         run: |
@@ -501,6 +503,7 @@ jobs:
       #  uses: ggml-org/ccache-action@v1.2.24
       #  with:
       #    key: release-android-arm64
+      #    evict-old-files: 1d
 
       - name: Build
         id: cmake_build
@@ -555,8 +558,8 @@ jobs:
 
     env:
       # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile
-      OPENVINO_VERSION_MAJOR: "2026.3.1"
-      OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d"
+      OPENVINO_VERSION_MAJOR: "2026.4"
+      OPENVINO_VERSION_FULL: "2026.4.0.22959.99c81491cc3"
 
     steps:
       - name: Set OpenVINO version output
@@ -579,6 +582,7 @@ jobs:
         uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-ubuntu-24.04-openvino-release-no-preset-v1
+          evict-old-files: 1d
 
       - name: Dependencies
         run: |
@@ -669,8 +673,8 @@ jobs:
 
     env:
       # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile
-      OPENVINO_VERSION_MAJOR: "2026.3.1"
-      OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d"
+      OPENVINO_VERSION_MAJOR: "2026.4"
+      OPENVINO_VERSION_FULL: "2026.4.0.22959.99c81491cc3"
 
     steps:
       - name: Set OpenVINO version output
@@ -822,6 +826,7 @@ jobs:
         uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-windows-2025-vs2026-${{ matrix.arch }}-cpu
+          evict-old-files: 1d
 
       - name: Build
         shell: cmd
@@ -1066,6 +1071,7 @@ jobs:
       #  uses: ggml-org/ccache-action@v1.2.24
       #  with:
       #    key: release-windows-2025-${{ matrix.arch }}-${{ matrix.backend }}
+      #    evict-old-files: 1d
 
       - name: Install OpenCL Headers and Libs
         id: install_opencl
@@ -1154,6 +1160,7 @@ jobs:
         uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-windows-2022-${{ matrix.arch }}-cuda-${{ matrix.cuda }}
+          evict-old-files: 1d
 
       - name: Build
         id: cmake_build
@@ -1250,6 +1257,7 @@ jobs:
         uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-windows-2022-x64-sycl
+          evict-old-files: 1d
 
       - name: Build
         id: cmake_build
@@ -1368,6 +1376,7 @@ jobs:
         uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-ubuntu-24.04-sycl-${{ matrix.build }}
+          evict-old-files: 1d
 
       - name: Build
         id: cmake_build
@@ -1845,8 +1854,8 @@ jobs:
             - [Ubuntu x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-x64.tar.gz)
             - [Ubuntu arm64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-arm64.tar.gz)
             - [Ubuntu x64 (CUDA 12)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-12.8-x64.tar.gz) - [CUDA 12.8 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-12.8-x64.tar.gz)
-            - [Ubuntu x64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.3-x64.tar.gz) - [CUDA 13.3 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.3-x64.tar.gz)
-            - [Ubuntu arm64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.3-arm64.tar.gz) - [CUDA 13.3 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.3-arm64.tar.gz)
+            - [Ubuntu x64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.4-x64.tar.gz) - [CUDA 13.4 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.4-x64.tar.gz)
+            - [Ubuntu arm64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.4-arm64.tar.gz) - [CUDA 13.4 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.4-arm64.tar.gz)
             - [Ubuntu x64 (ROCm 10.0)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-10.0-x64.tar.gz)
             - [Ubuntu x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-openvino-${{ needs.ubuntu-24-openvino.outputs.openvino_version }}-x64.tar.gz)
             - [Ubuntu x64 (SYCL FP32)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp32-x64.tar.gz)

.github/workflows/update-ops-docs.yml

@@ -16,7 +16,7 @@ on:
 
 jobs:
     update-ops-docs:
-        runs-on: [self-hosted, fast, ARM64]
+        runs-on: ubuntu-slim
 
         steps:
         - name: Checkout repository

.github/workflows/winget.yml

@@ -8,7 +8,7 @@ on:
 jobs:
   update:
     name: Update Winget Package
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04  # previously ubuntu-latest
     if: github.repository_owner == 'ggml-org'
 
     steps:

CODEOWNERS

@@ -97,7 +97,6 @@
 /src/models/                            @CISC
 /tests/                                 @ggerganov
 /tests/test-chat.*                      @pwilkin
-/tests/test-llama-archs.cpp             @JohannesGaessler
 /tools/batched-bench/                   @ggerganov
 /tools/cli/                             @ngxson
 /tools/completion/                      @ggerganov

common/arg.h

@@ -122,6 +122,8 @@ struct common_params_context {
 
 // parse input arguments from CLI
 // if one argument has invalid value, it will automatically display usage of the specific argument (and not the full usage message)
+// TODO: this function can load ggml backend (by calling llama_support_rpc)
+//       this is a side-effect that should be avoided
 bool common_params_parse(int argc, char ** argv, common_params & params, llama_example ex, void(*print_usage)(int, char **) = nullptr);
 
 // load all backends and print the list of available (non-CPU) devices to stdout

common/chat-peg-parser.cpp

@@ -318,13 +318,13 @@ void common_chat_peg_mapper::map(const common_peg_ast_node & node) {
     bool is_content   = node.tag == common_chat_peg_builder::CONTENT;
 
     if (is_reasoning) { // GPT OSS can have more than 1 reasoning block, so concatenate here
-        result.reasoning_content += std::string(node.text);
+        result.reasoning_content += node.sanitized_text();
     }
 
     if (is_content) {
         // Concatenate content from multiple content nodes (e.g., when reasoning markers
         // are preserved before content markers in reasoning_format=NONE mode)
-        result.content += std::string(node.text);
+        result.content += node.sanitized_text();
     }
 
     // Handle tool-related tags (supporting both JSON and tagged formats)
@@ -1058,12 +1058,12 @@ void common_chat_peg_gemma4_mapper::visit(const common_peg_ast_arena & arena, co
     const auto & node = arena.get(id);
 
     if (node.tag == "reasoning") {
-        result.reasoning_content += std::string(node.text);
+        result.reasoning_content += node.sanitized_text();
         return;
     }
 
     if (node.tag == "content") {
-        result.content += std::string(node.text);
+        result.content += node.sanitized_text();
         return;
     }
 
@@ -1206,12 +1206,12 @@ void common_chat_peg_minimax_m3_mapper::visit(const common_peg_ast_arena & arena
     const auto & node = arena.get(id);
 
     if (node.tag == common_chat_peg_builder::REASONING) {
-        result.reasoning_content += std::string(node.text);
+        result.reasoning_content += node.sanitized_text();
         return;
     }
 
     if (node.tag == common_chat_peg_builder::CONTENT) {
-        result.content += std::string(node.text);
+        result.content += node.sanitized_text();
         return;
     }
 

common/chat.cpp

@@ -1133,6 +1133,14 @@ std::optional<common_chat_params> common_chat_try_specialized_template(
         return common_chat_params_init_kimi_k3(tmpl, params);
     }
 
+    // Ling 3.0 / Bailing V3 - <role>X</role> sections with <arg_key>/<arg_value> tagged
+    // tool calls. <role> sections are unique to this family among the tagged-arg templates.
+    if (src.find("<role>ASSISTANT</role>") != std::string::npos &&
+        src.find("<arg_key>") != std::string::npos) {
+        LOG_DBG("Using specialized template: Ling 3.0 (Bailing V3)\n");
+        return common_chat_params_init_ling3(tmpl, params);
+    }
+
     // Cohere2 MoE / North Code - marker-wrapped format with <|START_TEXT|> content and
     // <|START_ACTION|> JSON tool calls. <|START_TEXT|> is unique to this template (the older
     // Command-R templates use <|START_RESPONSE|>).

common/hf-cache.cpp

@@ -30,8 +30,8 @@ namespace hf_cache {
 
 namespace fs = std::filesystem;
 
-static fs::path get_cache_directory() {
-    static const fs::path cache = []() {
+std::string get_cache_path() {
+    static const std::string cache = []() {
         struct {
             const char * var;
             fs::path path;
@@ -46,14 +46,14 @@ static fs::path get_cache_directory() {
         for (const auto & entry : entries) {
             if (auto * p = std::getenv(entry.var); p && *p) {
                 fs::path base(p);
-                return entry.path.empty() ? base : base / entry.path;
+                return (entry.path.empty() ? base : base / entry.path).string();
             }
         }
 #ifndef _WIN32
         const struct passwd * pw = getpwuid(getuid());
 
         if (pw && pw->pw_dir && *pw->pw_dir) {
-            return fs::path(pw->pw_dir) / ".cache" / "huggingface" / "hub";
+            return (fs::path(pw->pw_dir) / ".cache" / "huggingface" / "hub").string();
         }
 #endif
         throw std::runtime_error("Failed to determine HF cache directory");
@@ -80,7 +80,7 @@ static std::string repo_to_folder_name(const std::string & repo_id) {
 }
 
 static fs::path get_repo_path(const std::string & repo_id) {
-    return get_cache_directory() / repo_to_folder_name(repo_id);
+    return fs::path(get_cache_path()) / repo_to_folder_name(repo_id);
 }
 
 static bool is_hex_char(const char c) {
@@ -393,8 +393,8 @@ static std::string get_cached_ref(const fs::path & repo_path) {
 }
 
 hf_files get_cached_files(const std::string & repo_id) {
-    fs::path cache_dir = get_cache_directory();
-    if (!fs::exists(cache_dir)) {
+    const fs::path cache_path = get_cache_path();
+    if (!fs::exists(cache_path)) {
         return {};
     }
 
@@ -405,7 +405,7 @@ hf_files get_cached_files(const std::string & repo_id) {
 
     hf_files files;
 
-    for (const auto & repo : fs::directory_iterator(cache_dir)) {
+    for (const auto & repo : fs::directory_iterator(cache_path)) {
         if (!repo.is_directory()) {
             continue;
         }

common/hf-cache.h

@@ -32,4 +32,7 @@ std::string finalize_file(const hf_file & file);
 // Remove the entire cached directory for a repo, returns true if removed
 bool remove_cached_repo(const std::string & repo_id);
 
+// Returns the HuggingFace hub cache path
+std::string get_cache_path();
+
 } // namespace hf_cache

common/json-schema-to-grammar.cpp

@@ -321,7 +321,8 @@ static size_t gbnf_escape_length(const std::string & pattern, size_t pos) {
         case 'x': n_hex = 2; break;
         case 'u': n_hex = 4; break;
         case 'U': n_hex = 8; break;
-        case 't': case 'r': case 'n': case '\\': case '"': case '[': case ']':
+        // keep in sync with parse_char() in src/llama-grammar.cpp
+        case 't': case 'r': case 'n': case '\\': case '"': case '[': case ']': case '-':
             return 2;
         default:
             return 0;

common/json.h

@@ -82,6 +82,9 @@ struct common_json_value {
     // note: a nested pair {"a", "b"} does not build, use common_json::array({"a", "b"}) for an array
     common_json_value(std::initializer_list<common_json_item> items);
 
+    template <typename T, typename std::enable_if<std::is_enum<T>::value, int>::type = 0>
+    common_json_value(T val) : common_json_value((typename std::underlying_type<T>::type) val) {}
+
     template <typename T, typename std::enable_if<std::is_integral<T>::value && !std::is_same<T, bool>::value, int>::type = 0>
     common_json_value(T val) : type(std::is_signed<T>::value ? VAL_INT : VAL_UINT) {
         if (std::is_signed<T>::value) {
@@ -111,6 +114,7 @@ struct common_json_item {
 // the types common_json_value holds on its own
 // anything else reaches its common_json ctor and recurses forever
 template <typename T> struct common_json_is_value : std::integral_constant<bool,
+    std::is_enum<T>::value ||
     std::is_arithmetic<T>::value ||
     std::is_same<T, std::nullptr_t>::value ||
     std::is_same<T, std::string>::value ||

common/parsers/deepseek.cpp

@@ -104,6 +104,12 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
     const std::string GEN_PROMPT   = "<｜Assistant｜>";
     const std::string TC_SEPARATOR = "\n\n";
 
+    // lets the server find user turns in the prompt and place context checkpoints there
+    data.message_delimiters = {
+        { COMMON_CHAT_ROLE_ASSISTANT, GEN_PROMPT },
+        { COMMON_CHAT_ROLE_USER, "<｜User｜>" },
+    };
+
     data.prompt = common_chat_template_direct_apply_impl(
         tmpl, inputs, adjusted_messages, std::nullopt, additional_context);
     data.generation_prompt = common_chat_template_generation_prompt_impl(

common/parsers/gemma4.cpp

@@ -272,6 +272,10 @@ common_chat_params common_chat_params_init_gemma4(const common_chat_template &
                 /* max = */ inputs.parallel_tool_calls ? -1 : 1
             ));
 
+            if (inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED) {
+                return start + thought + tool_call;
+            }
+
             auto scan_to_toolcall = p.rule("scan-to-toolcall", p.until("<|tool_call>"));
             auto content = p.rule("content", p.content(p.until_one_of({"<|channel>", "<channel|>", "<|tool_call>"})));
             auto message = p.rule("message", thought + content);

common/parsers/ling3.cpp

@@ -0,0 +1,194 @@
+#include "parsers.h"
+
+// Ling 3.0 / Bailing V3 - <role>X</role> sections with tagged tool calls:
+//   assistant := [<think> ... </think>] [content] {<tool_call>name
+//                  <arg_key>k</arg_key>\n<arg_value>v</arg_value> ...</tool_call>}
+// The generation prompt ends with "<role>ASSISTANT</role>\n<think>", so the model
+// never emits the opening think tag, and a tool call can arrive before any
+// </think>. Reasoning therefore terminates at the think close tag or at a tool
+// call start, like the Qwen3-Coder and Kimi K3 parsers. With thinking off the
+// template pre-closes the think block instead, and the model emits bare content.
+common_chat_params common_chat_params_init_ling3(const common_chat_template &          tmpl,
+                                                        const autoparser::generation_params & inputs) {
+    common_chat_params data;
+
+    data.prompt            = common_chat_template_direct_apply_impl(tmpl, inputs);
+    data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs);
+    data.format            = COMMON_CHAT_FORMAT_PEG_NATIVE;
+    data.supports_thinking = true;
+
+    const std::string ROLE        = "<role>ASSISTANT</role>";
+    const std::string THINK_START = "<think>";
+    const std::string THINK_END   = "</think>";
+    const std::string CALL_START  = "<tool_call>";
+    const std::string CALL_END    = "</tool_call>";
+    const std::string ARG_KEY     = "<arg_key>";
+    const std::string ARG_KEY_END = "</arg_key>";
+    const std::string ARG_VAL     = "<arg_value>";
+    const std::string ROLE_END    = "<|role_end|>";
+    const std::string ARG_VAL_END = "</arg_value>";
+
+    data.preserved_tokens = {
+        THINK_START, THINK_END, CALL_START, CALL_END,
+        ARG_KEY, ARG_KEY_END, ARG_VAL, ARG_VAL_END, ROLE_END,
+    };
+
+    data.thinking_start_tag = THINK_START;
+    // Support both </think> and <tool_call> as reasoning end sequences: a call
+    // can be emitted before the think block is closed.
+    data.thinking_end_tags  = { THINK_END, CALL_START };
+
+    data.message_delimiters = {
+        { COMMON_CHAT_ROLE_ASSISTANT, "<role>ASSISTANT</role>" },
+        { COMMON_CHAT_ROLE_USER,      "<role>HUMAN</role>"     },
+        { COMMON_CHAT_ROLE_TOOL,      "<role>OBSERVATION</role>" },
+        { COMMON_CHAT_ROLE_SYSTEM,    "<role>SYSTEM</role>"    },
+    };
+
+    // the model may spell the end-of-turn control token out as text tokens,
+    // which does not stop generation; a literal stop string catches it either
+    // way (as the Laguna patch does for its </assistant> token)
+    data.additional_stops = { ROLE_END };
+
+    if (inputs.has_continuation()) {
+        const auto & msg = inputs.continue_msg;
+
+        data.generation_prompt = ROLE + "\n" + THINK_START + msg.reasoning_content;
+        if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) {
+            data.generation_prompt += THINK_END + msg.render_content();
+        }
+
+        data.prompt += data.generation_prompt;
+    }
+
+    // The generation prompt pre-opens the think block when thinking is on, so
+    // the opening tag is optional here and reasoning runs until </think> or a
+    // tool call start; with thinking off the template pre-closes the block and
+    // everything the model emits is content.
+    bool think_open = false;
+    if (inputs.has_continuation()) {
+        think_open = inputs.continue_final_message != COMMON_CHAT_CONTINUATION_CONTENT;
+    } else {
+        auto last_open  = data.generation_prompt.rfind(THINK_START);
+        auto last_close = data.generation_prompt.rfind(THINK_END);
+        think_open      = last_open != std::string::npos &&
+                     (last_close == std::string::npos || last_open > last_close);
+    }
+
+    auto has_tools         = inputs.tools.is_array() && !inputs.tools.empty();
+    auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
+    auto include_grammar   = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE;
+
+    auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
+        auto end = p.end();
+
+        // the effective parse input is generation_prompt + model output, so the
+        // assistant opener is optionally consumed here
+        auto opener = p.optional(p.literal(ROLE) + p.optional(p.space()));
+
+        // the generation prompt pre-opens the think block, so the opening tag
+        // is optional; a missing close tag does not swallow a tool call
+        auto body_end   = think_open ? p.until_one_of({ THINK_END, CALL_START }) : p.until_one_of({ THINK_END });
+        auto think_body = extract_reasoning ? p.reasoning(body_end) : p.content(body_end);
+
+        auto reasoning = p.optional(p.optional(p.literal(THINK_START)) + think_body +
+                                    p.optional(p.literal(THINK_END)));
+
+        // content between the think block and the first tool call, plus any
+        // trailing text after the last tool call, are plain content
+        auto content = p.optional(p.content(p.until_one_of({ CALL_START })));
+
+        // a trailing end-of-turn token is consumed instead of leaking into content
+        auto tail = p.optional(p.content(p.until(ROLE_END))) + p.optional(p.literal(ROLE_END));
+
+        if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) {
+            return opener + reasoning + tail + end;
+        }
+
+        auto tool_choices = p.choice();
+        auto arg_close    = p.tool_arg_close(p.literal(ARG_VAL_END));
+        auto arg_string   = p.rule("ling3-arg-string",
+                                   p.tool_arg_string_value(p.until(ARG_VAL_END)) + arg_close);
+
+        foreach_function(inputs.tools, [&](const json & tool) {
+            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;
+
+            // each argument may be preceded by whitespace: the model emits
+            // newlines between arguments, the template history does not
+            foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
+                auto rule_name = "ling3-arg-" + name + "-" + param.name;
+
+                auto types = param.schema->value_types();
+
+                // string arguments are raw text up to the closing tag, other
+                // types parse as JSON per their schema; each alternative
+                // consumes the closing tag itself so a JSON prefix can not
+                // commit the choice before the tag matches
+                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 parser tries the JSON alternative first to type the value
+                    arg_value = p.gbnf(p.atomic(p.tool_arg_json_value(p.schema(p.json(), rule_name + "-schema", doc, *param.schema)) + arg_close) | arg_string,
+                                       "ling3-arg-string");
+                }
+
+                auto arg = p.rule(rule_name,
+                    p.optional(p.space()) +
+                    p.tool_arg(p.tool_arg_open(p.literal(ARG_KEY) + p.tool_arg_name(p.literal(param.name)) +
+                                               p.literal(ARG_KEY_END)) +
+                               p.optional(p.space()) + p.literal(ARG_VAL) +
+                               arg_value));
+
+                (param.required ? required_args : optional_args).push_back(arg);
+            });
+
+            // required arguments in any order (as Qwen3-Coder does), then
+            // optional ones in any order and number
+            auto args = p.permute("ling3-" + name + "-args", required_args);
+            if (!optional_args.empty()) {
+                args = args + p.zero_or_more(p.choice(optional_args));
+            }
+
+            auto call = p.tool(p.tool_open(p.literal(CALL_START) + p.tool_name(p.literal(name)) +
+                                           p.optional(p.space())) +
+                               p.tool_args(args) +
+                               p.tool_close(p.optional(p.space()) + p.literal(CALL_END)));
+
+            tool_choices |= p.rule("ling3-tool-" + name, call);
+        });
+
+        auto calls = inputs.parallel_tool_calls ?
+                     tool_choices + p.zero_or_more(p.space() + tool_choices) :
+                     tool_choices;
+
+        auto tools_section = p.trigger_rule("ling3-tool-call", calls + p.space() +
+                                        p.optional(p.content(p.until(ROLE_END))) + p.optional(p.literal(ROLE_END)));
+
+        auto tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? tools_section :
+                                                                             p.optional(tools_section);
+
+        return opener + reasoning + content + tools + tail + end;
+    });
+
+    data.parser = parser.save();
+
+    if (include_grammar) {
+        data.grammar_lazy = inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_REQUIRED;
+        data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
+            parser.build_grammar(builder, data.grammar_lazy);
+        });
+
+        data.grammar_triggers = {
+            { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, CALL_START },
+        };
+    }
+
+    return data;
+}

common/parsers/parsers.h

@@ -63,6 +63,8 @@ common_chat_params common_chat_params_init_kimi_k2(const common_chat_template &
 
 common_chat_params common_chat_params_init_kimi_k3(const common_chat_template & tmpl, const autoparser::generation_params & inputs);
 
+common_chat_params common_chat_params_init_ling3(const common_chat_template & tmpl, const autoparser::generation_params & inputs);
+
 // tool_list_tokens preserves the LFM2 system tool-list markers; LFM2.5 renders without them
 common_chat_params common_chat_params_init_lfm2(const common_chat_template & tmpl, const autoparser::generation_params & inputs, bool tool_list_tokens);
 

common/parsers/sources.cmake

@@ -11,6 +11,7 @@ set(LLAMA_CHAT_PARSERS_SOURCES
     ${CMAKE_CURRENT_LIST_DIR}/gpt-oss.cpp
     ${CMAKE_CURRENT_LIST_DIR}/kimi-k2.cpp
     ${CMAKE_CURRENT_LIST_DIR}/kimi-k3.cpp
+    ${CMAKE_CURRENT_LIST_DIR}/ling3.cpp
     ${CMAKE_CURRENT_LIST_DIR}/lfm2.cpp
     ${CMAKE_CURRENT_LIST_DIR}/minicpm5.cpp
     ${CMAKE_CURRENT_LIST_DIR}/minimax-m3.cpp

common/peg-parser.cpp

@@ -166,6 +166,25 @@ common_peg_ast_id common_peg_ast_arena::find_by_rule(const common_peg_ast_node &
     return COMMON_PEG_INVALID_AST_ID;
 }
 
+std::string common_peg_ast_node::sanitized_text() const {
+    if (invalid_utf8.empty()) {
+        return std::string(text);
+    }
+
+    std::string out;
+    out.reserve(text.size() + 2 * invalid_utf8.size());
+
+    size_t seg_start = start;
+    for (const auto & invalid : invalid_utf8) {
+        out.append(text.data() + (seg_start - start), invalid.pos - seg_start);
+        out.append("\xEF\xBF\xBD");
+        seg_start = invalid.pos + invalid.len;
+    }
+    out.append(text.data() + (seg_start - start), end - seg_start);
+
+    return out;
+}
+
 void common_peg_ast_arena::visit(common_peg_ast_id id, const common_peg_ast_visitor & visitor) const {
     if (id == COMMON_PEG_INVALID_AST_ID) {
         return;
@@ -282,6 +301,7 @@ struct parser_executor {
 
         auto pos = start_pos;
         std::vector<common_peg_ast_id> nodes;
+        std::vector<common_peg_invalid_utf8> invalid_utf8;
 
         for (size_t i = 0; i < p.children.size(); i++) {
             const auto & child_id = p.children[i];
@@ -306,13 +326,14 @@ struct parser_executor {
             if (!result.nodes.empty()) {
                 nodes.insert(nodes.end(), result.nodes.begin(), result.nodes.end());
             }
+            invalid_utf8.insert(invalid_utf8.end(), result.invalid_utf8.begin(), result.invalid_utf8.end());
 
             if (result.need_more_input()) {
                 ctx.parse_depth--;
                 if (ctx.is_debug()) {
                     fprintf(stderr, "%sSEQ -> NEED_MORE\n", debug_indent().c_str());
                 }
-                return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, result.end, std::move(nodes));
+                return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, result.end, std::move(nodes), std::move(invalid_utf8));
             }
 
             pos = result.end;
@@ -322,7 +343,7 @@ struct parser_executor {
         if (ctx.is_debug()) {
             fprintf(stderr, "%sSEQ -> SUCCESS at %zu->%zu\n", debug_indent().c_str(), start_pos, pos);
         }
-        return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos, std::move(nodes));
+        return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos, std::move(nodes), std::move(invalid_utf8));
     }
 
     common_peg_parse_result operator()(const common_peg_choice_parser & p) {
@@ -370,6 +391,7 @@ struct parser_executor {
         auto pos = start_pos;
         int match_count = 0;
         std::vector<common_peg_ast_id> nodes;
+        std::vector<common_peg_invalid_utf8> invalid_utf8;
 
         // Try to match up to max_count times (or unlimited if max_count is -1)
         while (p.max_count == -1 || match_count < p.max_count) {
@@ -400,6 +422,7 @@ struct parser_executor {
                 if (!result.nodes.empty()) {
                     nodes.insert(nodes.end(), result.nodes.begin(), result.nodes.end());
                 }
+                invalid_utf8.insert(invalid_utf8.end(), result.invalid_utf8.begin(), result.invalid_utf8.end());
 
                 pos = result.end;
                 match_count++;
@@ -410,13 +433,14 @@ struct parser_executor {
                 if (!result.nodes.empty()) {
                     nodes.insert(nodes.end(), result.nodes.begin(), result.nodes.end());
                 }
+                invalid_utf8.insert(invalid_utf8.end(), result.invalid_utf8.begin(), result.invalid_utf8.end());
 
                 ctx.parse_depth--;
                 if (ctx.is_debug()) {
                     fprintf(stderr, "%sREPEAT -> NEED_MORE (count=%d, nodes=%zu)\n", debug_indent().c_str(),
                             match_count, nodes.size());
                 }
-                return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, result.end, std::move(nodes));
+                return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, result.end, std::move(nodes), std::move(invalid_utf8));
             }
 
             // Child failed - stop trying
@@ -434,7 +458,7 @@ struct parser_executor {
                     fprintf(stderr, "%sREPEAT -> NEED_MORE (not enough matches: %d < %d)\n", debug_indent().c_str(),
                             match_count, p.min_count);
                 }
-                return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, pos, std::move(nodes));
+                return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, pos, std::move(nodes), std::move(invalid_utf8));
             }
             if (ctx.is_debug()) {
                 fprintf(stderr, "%sREPEAT -> FAIL (not enough matches: %d < %d)\n", debug_indent().c_str(), match_count,
@@ -448,7 +472,7 @@ struct parser_executor {
             fprintf(stderr, "%sREPEAT -> SUCCESS (count=%d, nodes=%zu)\n", debug_indent().c_str(), match_count,
                     nodes.size());
         }
-        return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos, std::move(nodes));
+        return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos, std::move(nodes), std::move(invalid_utf8));
     }
 
     common_peg_parse_result operator()(const common_peg_and_parser & p) {
@@ -664,23 +688,23 @@ struct parser_executor {
         // Scan input and check for delimiters
         size_t pos = start_pos;
         size_t last_valid_pos = start_pos;
+        std::vector<common_peg_invalid_utf8> invalid_utf8;
 
         while (pos < ctx.input.size()) {
             auto utf8_result = common_parse_utf8_codepoint(ctx.input, pos);
 
-            if (utf8_result.status == utf8_parse_result::INCOMPLETE) {
-                // Incomplete UTF-8 sequence
-                if (!ctx.is_lenient()) {
-                    // Input is complete but UTF-8 is incomplete = malformed
-                    return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos);
-                }
-                // Return what we have so far (before incomplete sequence)
-                return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, last_valid_pos);
+            if (utf8_result.status == utf8_parse_result::INCOMPLETE && ctx.is_lenient()) {
+                // The rest of the sequence may still arrive, return what we have so far
+                return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, last_valid_pos, {}, std::move(invalid_utf8));
             }
 
-            if (utf8_result.status == utf8_parse_result::INVALID) {
-                // Malformed UTF-8
-                return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos);
+            if (utf8_result.status != utf8_parse_result::SUCCESS) {
+                // Malformed UTF-8, or a sequence truncated by the end of a complete input.
+                // A delimiter cannot start inside bytes that fail to decode, so consume them and move on
+                invalid_utf8.push_back({pos, utf8_result.bytes_consumed});
+                pos += utf8_result.bytes_consumed;
+                last_valid_pos = pos;
+                continue;
             }
 
             // Check if a delimiter starts at this position
@@ -688,12 +712,12 @@ struct parser_executor {
 
             if (match == common_trie::COMPLETE_MATCH) {
                 // Found a complete delimiter, return everything before it
-                return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos);
+                return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos, {}, std::move(invalid_utf8));
             }
 
             if (match == common_trie::PARTIAL_MATCH) {
                 // Found a partial match extending to end of input, return everything before it
-                return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos);
+                return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos, {}, std::move(invalid_utf8));
             }
 
             pos += utf8_result.bytes_consumed;
@@ -702,9 +726,9 @@ struct parser_executor {
 
         if (last_valid_pos == ctx.input.size() && ctx.is_lenient()) {
             // Reached the end of a partial stream, there might still be more input that we need to consume.
-            return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, last_valid_pos);
+            return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, last_valid_pos, {}, std::move(invalid_utf8));
         }
-        return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, last_valid_pos);
+        return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, last_valid_pos, {}, std::move(invalid_utf8));
     }
 
     common_peg_parse_result operator()(const common_peg_schema_parser & p) {
@@ -728,10 +752,11 @@ struct parser_executor {
                 result.end,
                 text,
                 std::move(result.nodes),
-                result.need_more_input()
+                result.need_more_input(),
+                result.invalid_utf8
             );
 
-            return common_peg_parse_result(result.type, result.start, result.end, { node_id });
+            return common_peg_parse_result(result.type, result.start, result.end, { node_id }, std::move(result.invalid_utf8));
         }
 
         return result;
@@ -757,10 +782,11 @@ struct parser_executor {
                 result.end,
                 text,
                 std::move(result.nodes),
-                result.need_more_input()
+                result.need_more_input(),
+                result.invalid_utf8
             );
 
-            return common_peg_parse_result(result.type, result.start, result.end, { node_id });
+            return common_peg_parse_result(result.type, result.start, result.end, { node_id }, std::move(result.invalid_utf8));
         }
 
         return result;

common/peg-parser.h

@@ -72,6 +72,12 @@ enum common_peg_parse_result_type {
 
 const char * common_peg_parse_result_type_name(common_peg_parse_result_type type);
 
+// A run of input bytes that does not decode as UTF-8
+struct common_peg_invalid_utf8 {
+    size_t pos;
+    size_t len;
+};
+
 struct common_peg_ast_node {
     common_peg_ast_id id;
     std::string rule;
@@ -82,6 +88,12 @@ struct common_peg_ast_node {
     std::vector<common_peg_ast_id> children;
 
     bool is_partial = false;
+
+    // Invalid UTF-8 inside the node, in ascending order
+    std::vector<common_peg_invalid_utf8> invalid_utf8;
+
+    // Returns the text with every invalid run replaced by U+FFFD
+    std::string sanitized_text() const;
 };
 
 struct common_peg_parse_result;
@@ -98,10 +110,11 @@ class common_peg_ast_arena {
         size_t end,
         std::string_view text,
         std::vector<common_peg_ast_id> children,
-        bool is_partial = false
+        bool is_partial = false,
+        std::vector<common_peg_invalid_utf8> invalid_utf8 = {}
     ) {
         common_peg_ast_id id = nodes_.size();
-        nodes_.push_back({id, rule, tag, start, end, text, std::move(children), is_partial});
+        nodes_.push_back({id, rule, tag, start, end, text, std::move(children), is_partial, std::move(invalid_utf8)});
         return id;
     }
 
@@ -127,6 +140,9 @@ struct common_peg_parse_result {
 
     std::vector<common_peg_ast_id> nodes;
 
+    // Invalid UTF-8 consumed by this result, carried up to the enclosing AST nodes
+    std::vector<common_peg_invalid_utf8> invalid_utf8;
+
     common_peg_parse_result() = default;
 
     common_peg_parse_result(common_peg_parse_result_type type, size_t start)
@@ -135,8 +151,8 @@ struct common_peg_parse_result {
     common_peg_parse_result(common_peg_parse_result_type type, size_t start, size_t end)
         : type(type), start(start), end(end) {}
 
-    common_peg_parse_result(common_peg_parse_result_type type, size_t start, size_t end, std::vector<common_peg_ast_id> nodes)
-        : type(type), start(start), end(end), nodes(std::move(nodes)) {}
+    common_peg_parse_result(common_peg_parse_result_type type, size_t start, size_t end, std::vector<common_peg_ast_id> nodes, std::vector<common_peg_invalid_utf8> invalid_utf8 = {})
+        : type(type), start(start), end(end), nodes(std::move(nodes)), invalid_utf8(std::move(invalid_utf8)) {}
 
     bool fail() const { return type == COMMON_PEG_PARSE_RESULT_FAIL; }
     bool need_more_input() const { return type == COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT; }
@@ -430,6 +446,7 @@ class common_peg_parser_builder {
     common_peg_parser space() { return add(common_peg_space_parser{}); }
 
     // Matches all characters until a delimiter is found (delimiter not consumed).
+    // Invalid UTF-8 is consumed and recorded on the AST nodes.
     //   S -> (!delim .)*
     common_peg_parser until(const std::string & delimiter) { return add(common_peg_until_parser{{delimiter}}); }
 

common/unicode.cpp

@@ -26,16 +26,16 @@ utf8_parse_result common_parse_utf8_codepoint(std::string_view input, size_t off
 
     // Invalid: continuation byte as first byte
     if (!(input[offset] & 0x40)) {
-        return utf8_parse_result(utf8_parse_result::INVALID);
+        return utf8_parse_result(utf8_parse_result::INVALID, 0, 1);
     }
 
     // 2-byte sequence
     if (!(input[offset] & 0x20)) {
         if (offset + 1 >= input.size()) {
-            return utf8_parse_result(utf8_parse_result::INCOMPLETE);
+            return utf8_parse_result(utf8_parse_result::INCOMPLETE, 0, 1);
         }
         if ((input[offset + 1] & 0xc0) != 0x80) {
-            return utf8_parse_result(utf8_parse_result::INVALID);
+            return utf8_parse_result(utf8_parse_result::INVALID, 0, 1);
         }
         auto result = ((input[offset] & 0x1f) << 6) | (input[offset + 1] & 0x3f);
         return utf8_parse_result(utf8_parse_result::SUCCESS, result, 2);
@@ -43,11 +43,14 @@ utf8_parse_result common_parse_utf8_codepoint(std::string_view input, size_t off
 
     // 3-byte sequence
     if (!(input[offset] & 0x10)) {
-        if (offset + 2 >= input.size()) {
-            return utf8_parse_result(utf8_parse_result::INCOMPLETE);
-        }
-        if ((input[offset + 1] & 0xc0) != 0x80 || (input[offset + 2] & 0xc0) != 0x80) {
-            return utf8_parse_result(utf8_parse_result::INVALID);
+        // Check one byte at a time so a bad byte is reported before a short input
+        for (size_t i = 1; i < 3; i++) {
+            if (offset + i >= input.size()) {
+                return utf8_parse_result(utf8_parse_result::INCOMPLETE, 0, i);
+            }
+            if ((input[offset + i] & 0xc0) != 0x80) {
+                return utf8_parse_result(utf8_parse_result::INVALID, 0, i);
+            }
         }
         auto result = ((input[offset] & 0x0f) << 12) | ((input[offset + 1] & 0x3f) << 6) | (input[offset + 2] & 0x3f);
         return utf8_parse_result(utf8_parse_result::SUCCESS, result, 3);
@@ -55,18 +58,20 @@ utf8_parse_result common_parse_utf8_codepoint(std::string_view input, size_t off
 
     // 4-byte sequence
     if (!(input[offset] & 0x08)) {
-        if (offset + 3 >= input.size()) {
-            return utf8_parse_result(utf8_parse_result::INCOMPLETE);
-        }
-        if ((input[offset + 1] & 0xc0) != 0x80 || (input[offset + 2] & 0xc0) != 0x80 || (input[offset + 3] & 0xc0) != 0x80) {
-            return utf8_parse_result(utf8_parse_result::INVALID);
+        for (size_t i = 1; i < 4; i++) {
+            if (offset + i >= input.size()) {
+                return utf8_parse_result(utf8_parse_result::INCOMPLETE, 0, i);
+            }
+            if ((input[offset + i] & 0xc0) != 0x80) {
+                return utf8_parse_result(utf8_parse_result::INVALID, 0, i);
+            }
         }
         auto result = ((input[offset] & 0x07) << 18) | ((input[offset + 1] & 0x3f) << 12) | ((input[offset + 2] & 0x3f) << 6) | (input[offset + 3] & 0x3f);
         return utf8_parse_result(utf8_parse_result::SUCCESS, result, 4);
     }
 
     // Invalid first byte
-    return utf8_parse_result(utf8_parse_result::INVALID);
+    return utf8_parse_result(utf8_parse_result::INVALID, 0, 1);
 }
 
 bool common_utf8_is_complete(const std::string & s) {

common/unicode.h

@@ -9,7 +9,7 @@
 
 struct utf8_parse_result {
     uint32_t codepoint;      // Decoded codepoint (only valid if status == SUCCESS)
-    size_t bytes_consumed;   // How many bytes this codepoint uses (1-4)
+    size_t bytes_consumed;   // How many bytes this codepoint uses (1-4), or the length of the valid prefix if status != SUCCESS
     enum status { SUCCESS, INCOMPLETE, INVALID } status;
 
     utf8_parse_result(enum status s, uint32_t cp = 0, size_t bytes = 0)

conversion/base.py

@@ -1861,6 +1861,9 @@ class TextModel(ModelBase):
         if chkhsh == "972da7b59cec44d1f0a490a86c96df53859e486e481563e5dddac155013d87ac":
             # ref: https://huggingface.co/poolside/Laguna-XS.2
             res = "laguna"
+        if chkhsh == "653660222fb704f61cbf2b618a8ae6502b7f8b20c980f9a5de07ed78e13319cd":
+            # ref: https://huggingface.co/ufakai/ufakzeka-1
+            res = "ufakzeka"
 
         if res is None:
             logger.warning("\n")

convert_hf_to_gguf_update.py

@@ -163,6 +163,7 @@ models = [
     {"name": "granite-embed-multi-311m", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/ibm-granite/granite-embedding-311m-multilingual-r2", },
     {"name": "mellum2",          "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/JetBrains/Mellum2-12B-A2.5B-Base"},
     {"name": "laguna",           "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/poolside/Laguna-XS.2", },
+    {"name": "ufakzeka",         "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/ufakai/ufakzeka-1", },
 ]
 
 # some models are known to be broken upstream, so we will skip them as exceptions

docs/backend/OPENVINO.md

@@ -12,6 +12,8 @@ The OpenVINO backend is implemented in `ggml/src/ggml-openvino` and provides a t
 - Compiles and caches the model for the target device.
 - Binds GGML tensor memory to OpenVINO inference tensors and runs inference.
 
+For guidance on contributing to the OpenVINO backend, see the [OpenVINO Backend Contributing Guide](https://github.com/ravi9/llamacpp-ov-dev-guide/blob/main/contributing-llamacpp-ov.md).
+
 ## Contents
 
 - [Supported Devices](#supported-devices)
@@ -96,7 +98,7 @@ Although, the validated models below were tested with `llama-cli` using the `Q4_
   - **SL** = Stateless (`GGML_OPENVINO_STATEFUL_EXECUTION=0`)
   - **SF** = Stateful (`GGML_OPENVINO_STATEFUL_EXECUTION=1`)
   - Note: The NPU operates in stateless mode only.
-- **Validation system:** Intel® Core™ Ultra 5 238V (Lunar Lake) | 32 GB RAM | Ubuntu 24.04 | Intel OpenCL GPU Driver 26.31.39395.13-0 | Intel NPU Driver 1.35.0.
+- **Validation system:** Intel® Core™ Ultra 5 238V (Lunar Lake) | 32 GB RAM | Ubuntu 24.04 | Intel Graphics Compiler 2.41.5 | Intel OpenCL GPU Driver 26.31.39395.13-0 | Intel NPU Driver 1.38.0.
 - See [Known Limitations](#known-limitations) for context on observed failures.
 
 | Model | CPU (SL / SF) | GPU (SL / SF) | NPU (SL) |
@@ -117,9 +119,9 @@ Although, the validated models below were tested with `llama-cli` using the `Q4_
 | [lmstudio-community/Qwen3.5-9B-Q4_K_M](https://huggingface.co/lmstudio-community/Qwen3.5-9B-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ |
 |  |  |  |  |
 | [unsloth/gemma-3-4b-it-Q4_K_M](https://huggingface.co/unsloth/gemma-3-4b-it-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
-| [bartowski/google_gemma-4-E2B-it-Q4_K_M](https://huggingface.co/bartowski/google_gemma-4-E2B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ |
-| [bartowski/google_gemma-4-E4B-it-Q4_K_M](https://huggingface.co/bartowski/google_gemma-4-E4B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✓ |
-| [bartowski/gemma-4-12B-it-Q4_K_M](https://huggingface.co/bartowski/gemma-4-12B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✓ |
+| [bartowski/google_gemma-4-E2B-it-Q4_K_M](https://huggingface.co/bartowski/google_gemma-4-E2B-it-GGUF) | ✓ / ✓ | ✓ / ✓ | ✗ |
+| [bartowski/google_gemma-4-E4B-it-Q4_K_M](https://huggingface.co/bartowski/google_gemma-4-E4B-it-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
+| [bartowski/gemma-4-12B-it-Q4_K_M](https://huggingface.co/bartowski/gemma-4-12B-it-GGUF) | ✓ / ✓ | ✓ / ✓ | ✗ |
 |  |  |  |  |
 | [bartowski/Phi-3-mini-4k-instruct-Q4_K_M](https://huggingface.co/bartowski/Phi-3-mini-4k-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
 | [bartowski/Phi-3.5-mini-instruct-Q4_K_M](https://huggingface.co/bartowski/Phi-3.5-mini-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
@@ -132,9 +134,9 @@ Although, the validated models below were tested with `llama-cli` using the `Q4_
 | [bartowski/DeepSeek-R1-Distill-Llama-8B-Q4_K_M](https://huggingface.co/bartowski/DeepSeek-R1-Distill-Llama-8B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
 | [bartowski/DeepSeek-R1-Distill-Qwen-7B-Q4_K_M](https://huggingface.co/bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
 |  |  |  |  |
-| [ibm-granite/granite-4.0-350m-Q4_K_M](https://huggingface.co/ibm-granite/granite-4.0-350m-GGUF) | ✓ / ✓ | ✗ / ✗ | ✓ |
+| [ibm-granite/granite-4.0-350m-Q4_K_M](https://huggingface.co/ibm-granite/granite-4.0-350m-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
 | [ibm-granite/granite-4.0-micro-Q4_K_M](https://huggingface.co/ibm-granite/granite-4.0-micro-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
-| [ibm-granite/granite-4.0-1b-Q4_K_M](https://huggingface.co/ibm-granite/granite-4.0-1b-GGUF) | ✓ / ✓ | ✗ / ✗ | ✗ |
+| [ibm-granite/granite-4.0-1b-Q4_K_M](https://huggingface.co/ibm-granite/granite-4.0-1b-GGUF) | ✓ / ✓ | ✓ / ✓ | ✗ |
 | [ibm-research/granite-3.2-8b-instruct-Q4_K_M](https://huggingface.co/ibm-research/granite-3.2-8b-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
 |  |  |  |  |
 | [HuggingFaceTB/smollm2-1.7b-instruct-q4_k_m](https://huggingface.co/HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
@@ -242,8 +244,8 @@ chmod +x build-llamacpp-ov.sh
 # ============================================
 set -euo pipefail
 
-OPENVINO_VERSION_MAJOR="2026.3.1"
-OPENVINO_VERSION_FULL="2026.3.1.22476.56d9685302d"
+OPENVINO_VERSION_MAJOR="2026.4"
+OPENVINO_VERSION_FULL="2026.4.0.22959.99c81491cc3"
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 OPENVINO_INSTALL_DIR="/opt/intel/openvino_${OPENVINO_VERSION_MAJOR}"
@@ -340,7 +342,7 @@ echo "  ./build/ReleaseOV/bin/llama-cli -m model.gguf"
 ```
 
 > [!NOTE]
-> The script pins OpenVINO `2026.3.1` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release.
+> The script pins OpenVINO `2026.4` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release.
 
 </details>
 
@@ -370,8 +372,8 @@ REM ============================================
 REM llama.cpp OpenVINO Build Script (Ninja)
 REM ============================================
 
-set "OPENVINO_VERSION_MAJOR=2026.3.1"
-set "OPENVINO_VERSION_FULL=2026.3.1.22476.56d9685302d"
+set "OPENVINO_VERSION_MAJOR=2026.4"
+set "OPENVINO_VERSION_FULL=2026.4.0.22959.99c81491cc3"
 
 set "SCRIPT_DIR=%~dp0"
 set "VCPKG_DIR=C:\vcpkg"
@@ -550,7 +552,7 @@ endlocal
 ```
 
 > [!NOTE]
-> The script pins OpenVINO `2026.3.1` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release. From any new shell, source the matching `setupvars` script via the junction — `call "C:\Intel\openvino\setupvars.bat"` from `cmd`, or `& "C:\Intel\openvino\setupvars.ps1"` from PowerShell. If `winget` cannot register Visual Studio Build Tools on first run, install them once manually and re-run the script from an elevated **Developer Command Prompt for VS 2022**.
+> The script pins OpenVINO `2026.4` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release. From any new shell, source the matching `setupvars` script via the junction — `call "C:\Intel\openvino\setupvars.bat"` from `cmd`, or `& "C:\Intel\openvino\setupvars.ps1"` from PowerShell. If `winget` cannot register Visual Studio Build Tools on first run, install them once manually and re-run the script from an elevated **Developer Command Prompt for VS 2022**.
 
 </details>
 

docs/backend/snapdragon/README.md

@@ -327,6 +327,9 @@ on 4 physical NPUs, or `--devices 'HTP0[0-1:0],HTP1[0-1:1]'` on 2 physical NPUs
 - `GGML_HEXAGON_HOSTBUF=1` (default: 0, disabled)
   Enables allocating host buffers for debugging. By default, host buffers are disabled.
 
+- `GGML_HEXAGON_DMA64=0` (default: enabled on v81+)
+  Disables 64-bit DMA for model weights. Set to `1` to enable it explicitly on a supported architecture.
+
 - `GGML_HEXAGON_VERBOSE=1`
   Enables verbose logging of Ops from the backend. Example output:
 

docs/backend/snapdragon/developer.md

@@ -146,6 +146,28 @@ Writing high-performance operators for Hexagon requires following specific guide
   python3 scripts/snapdragon/ggml-hexagon-align-macros.py --fix ggml/src/ggml-hexagon/htp/
   ```
 
+### Binary Inspection and Spill Analysis
+
+Use [`scripts/snapdragon/ggml-hexagon-inspect.py`](../../../scripts/snapdragon/ggml-hexagon-inspect.py) to audit Hexagon binaries for register
+spills, unexpected float promotions, or disassembly:
+
+- Always verify that compute kernels have zero in-loop vector spills (`--spills --strict`) and no float promotions (`--promotions`).
+- Avoid excessive loop unrolling (`#pragma unroll`), which increases register pressure and causes spills.
+
+```bash
+# Check for vector and scalar register spills
+python3 scripts/snapdragon/ggml-hexagon-inspect.py --spills --strict --func "^compute_"
+
+# Check for float promotions
+python3 scripts/snapdragon/ggml-hexagon-inspect.py --promotions --func "^compute_"
+
+# Disassemble with annotated loops and spill markers
+python3 scripts/snapdragon/ggml-hexagon-inspect.py --disasm compute_same_shape_div_f32
+
+# Resolve crash addresses to function symbols and lines
+python3 scripts/snapdragon/ggml-hexagon-inspect.py --addr2line 0x51a30 0x5ba54
+```
+
 ## Multi-Device Partitioning (mdev)
 
 Multi-device (mdev) mode enables row-level tensor parallel execution across multiple physical NPU cores or virtual NPU

docs/ops.md

@@ -121,7 +121,7 @@ Legend:
 |                       SWIGLU_OAI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             TANH | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |               TIMESTEP_EMBEDDING | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            TOP_K | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | 🟡 | 🟡 | ✅ | ❌ | ❌ |
+|                            TOP_K | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ✅ | ❌ | 🟡 | 🟡 | ✅ | ❌ | ❌ |
 |                              TRI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                            TRUNC | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          UPSCALE | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |

docs/ops/Hexagon.csv

@@ -13936,247 +13936,247 @@
 "HTP0","ARGSORT","type=f32,ne=[2049,2,1,3],order=1","support","1","yes","HTP"
 "HTP0","ARGSORT","type=f32,ne=[2,8,8192,1],order=1","support","1","yes","HTP"
 "HTP0","ARGSORT","type=f32,ne=[2048,512,1,1],order=1","support","1","yes","HTP"
-"HTP0","TOP_K","type=f32,ne=[1,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[12,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[13,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[13,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[15,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[15,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4,1,1,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[15,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[19,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[19,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8,1,1,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[19,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8,1,1,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[19,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[27,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[27,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16,1,1,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[27,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16,1,1,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[27,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16,1,1,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[27,1,2,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[43,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[43,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32,1,1,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[43,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32,1,1,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[43,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32,1,1,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[43,1,2,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[64,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[75,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[64,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[75,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[64,1,1,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[75,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[64,1,1,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[75,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[64,1,1,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[75,1,2,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[128,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[139,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[128,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[139,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[128,1,1,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[139,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[128,1,1,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[139,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[128,1,1,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[139,1,2,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[128,1,1,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[139,1,2,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[256,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[267,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[256,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[267,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[256,1,1,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[267,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[256,1,1,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[267,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[256,1,1,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[267,1,2,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[256,1,1,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[267,1,2,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[512,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[523,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[512,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[523,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[512,1,1,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[523,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[512,1,1,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[523,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[512,1,1,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[523,1,2,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[512,1,1,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[523,1,2,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[512,1,1,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[523,1,2,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=1023,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=1023,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=1023,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=1023,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=1023,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=1023,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=1023,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=1023,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=1023,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=1023,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=9999,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=9999,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=1023,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=1023,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=9999,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=9999,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=1,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[1,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[12,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[13,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[13,1,2,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[15,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[15,1,2,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4,1,1,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[15,1,2,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[19,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[19,1,2,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8,1,1,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[19,1,2,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8,1,1,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[19,1,2,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[27,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[27,1,2,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16,1,1,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[27,1,2,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16,1,1,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[27,1,2,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16,1,1,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[27,1,2,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[43,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[43,1,2,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32,1,1,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[43,1,2,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32,1,1,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[43,1,2,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32,1,1,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[43,1,2,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[64,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[75,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[64,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[75,1,2,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[64,1,1,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[75,1,2,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[64,1,1,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[75,1,2,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[64,1,1,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[75,1,2,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[128,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[139,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[128,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[139,1,2,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[128,1,1,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[139,1,2,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[128,1,1,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[139,1,2,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[128,1,1,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[139,1,2,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[128,1,1,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[139,1,2,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[256,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[267,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[256,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[267,1,2,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[256,1,1,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[267,1,2,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[256,1,1,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[267,1,2,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[256,1,1,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[267,1,2,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[256,1,1,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[267,1,2,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[512,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[523,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[512,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[523,1,2,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[512,1,1,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[523,1,2,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[512,1,1,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[523,1,2,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[512,1,1,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[523,1,2,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[512,1,1,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[523,1,2,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[512,1,1,1],k=500,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[523,1,2,1],k=500,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=500,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=500,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=1023,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1035,1,2,1],k=1023,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=500,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=500,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2048,1,1,1],k=1023,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2059,1,2,1],k=1023,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=500,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=500,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=1023,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4107,1,2,1],k=1023,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=500,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=500,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=1023,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8203,1,2,1],k=1023,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=500,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=500,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=1023,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=1023,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=9999,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16395,1,2,1],k=9999,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=100,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=500,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=500,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=1023,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=1023,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32768,1,1,1],k=9999,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[32779,1,2,1],k=9999,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=1,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[65547,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=2,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=2,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[65547,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=3,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=3,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[65547,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=7,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=7,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[65547,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=15,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=15,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[65547,1,2,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=100,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=100,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[65547,1,2,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=500,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=500,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[65547,1,2,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=1023,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=1023,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[65547,1,2,1],k=1023,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=9999,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[65536,1,1,1],k=9999,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[65547,1,2,1],k=9999,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=1,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=1,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[131083,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=2,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=2,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[131083,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=3,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=3,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[131083,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=7,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=7,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[131083,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=15,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=15,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[131083,1,2,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=100,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=100,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[131083,1,2,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=500,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=500,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[131083,1,2,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=1023,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=1023,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[131083,1,2,1],k=1023,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=9999,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[131072,1,1,1],k=9999,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[131083,1,2,1],k=9999,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=1,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=1,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[262155,1,2,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=2,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=2,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[262155,1,2,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=3,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=3,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[262155,1,2,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=7,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=7,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[262155,1,2,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=15,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=15,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[262155,1,2,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=100,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=100,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[262155,1,2,1],k=100,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=500,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=500,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[262155,1,2,1],k=500,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=1023,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=1023,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[262155,1,2,1],k=1023,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=9999,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[262144,1,1,1],k=9999,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[262155,1,2,1],k=9999,ties=0","support","0","no","HTP"
 "HTP0","TOP_K","type=f32,ne=[524288,1,1,1],k=1,ties=0","support","0","no","HTP"
 "HTP0","TOP_K","type=f32,ne=[524299,1,2,1],k=1,ties=0","support","0","no","HTP"
@@ -14196,99 +14196,105 @@
 "HTP0","TOP_K","type=f32,ne=[524299,1,2,1],k=1023,ties=0","support","0","no","HTP"
 "HTP0","TOP_K","type=f32,ne=[524288,1,1,1],k=9999,ties=0","support","0","no","HTP"
 "HTP0","TOP_K","type=f32,ne=[524299,1,2,1],k=9999,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[202048,1,1,1],k=4,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[151936,1,1,1],k=4,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=4,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8193,1,1,1],k=4,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[202048,1,1,1],k=4,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[151936,1,1,1],k=4,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=4,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8193,1,1,1],k=4,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[202048,8,1,1],k=4,ties=0","support","0","no","HTP"
 "HTP0","TOP_K","type=f32,ne=[151936,8,1,1],k=4,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,8,1,1],k=4,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8193,8,1,1],k=4,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,8,1,1],k=4,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8193,8,1,1],k=4,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[202048,16,1,1],k=4,ties=0","support","0","no","HTP"
 "HTP0","TOP_K","type=f32,ne=[151936,16,1,1],k=4,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,16,1,1],k=4,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8193,16,1,1],k=4,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[202048,1,1,1],k=8,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[151936,1,1,1],k=8,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=8,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8193,1,1,1],k=8,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,16,1,1],k=4,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8193,16,1,1],k=4,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[202048,1,1,1],k=8,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[151936,1,1,1],k=8,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=8,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8193,1,1,1],k=8,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[202048,8,1,1],k=8,ties=0","support","0","no","HTP"
 "HTP0","TOP_K","type=f32,ne=[151936,8,1,1],k=8,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,8,1,1],k=8,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8193,8,1,1],k=8,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,8,1,1],k=8,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8193,8,1,1],k=8,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[202048,16,1,1],k=8,ties=0","support","0","no","HTP"
 "HTP0","TOP_K","type=f32,ne=[151936,16,1,1],k=8,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,16,1,1],k=8,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8193,16,1,1],k=8,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[202048,1,1,1],k=16,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[151936,1,1,1],k=16,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=16,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8193,1,1,1],k=16,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,16,1,1],k=8,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8193,16,1,1],k=8,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[202048,1,1,1],k=16,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[151936,1,1,1],k=16,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=16,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8193,1,1,1],k=16,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[202048,8,1,1],k=16,ties=0","support","0","no","HTP"
 "HTP0","TOP_K","type=f32,ne=[151936,8,1,1],k=16,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,8,1,1],k=16,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8193,8,1,1],k=16,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,8,1,1],k=16,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8193,8,1,1],k=16,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[202048,16,1,1],k=16,ties=0","support","0","no","HTP"
 "HTP0","TOP_K","type=f32,ne=[151936,16,1,1],k=16,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,16,1,1],k=16,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8193,16,1,1],k=16,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[202048,1,1,1],k=32,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[151936,1,1,1],k=32,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=32,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8193,1,1,1],k=32,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,16,1,1],k=16,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8193,16,1,1],k=16,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[202048,1,1,1],k=32,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[151936,1,1,1],k=32,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,1,1,1],k=32,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8193,1,1,1],k=32,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[202048,8,1,1],k=32,ties=0","support","0","no","HTP"
 "HTP0","TOP_K","type=f32,ne=[151936,8,1,1],k=32,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,8,1,1],k=32,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8193,8,1,1],k=32,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,8,1,1],k=32,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8193,8,1,1],k=32,ties=0","support","1","yes","HTP"
 "HTP0","TOP_K","type=f32,ne=[202048,16,1,1],k=32,ties=0","support","0","no","HTP"
 "HTP0","TOP_K","type=f32,ne=[151936,16,1,1],k=32,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8192,16,1,1],k=32,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[8193,16,1,1],k=32,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16,10,10,10],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[60,10,10,10],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1023,2,1,3],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1024,2,1,3],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1025,2,1,3],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2047,2,1,3],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2048,2,1,3],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2049,2,1,3],k=1,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16,10,10,10],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[60,10,10,10],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1023,2,1,3],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1024,2,1,3],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1025,2,1,3],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2047,2,1,3],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2048,2,1,3],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2049,2,1,3],k=2,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16,10,10,10],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[60,10,10,10],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1023,2,1,3],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1024,2,1,3],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1025,2,1,3],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2047,2,1,3],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2048,2,1,3],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2049,2,1,3],k=3,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16,10,10,10],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[60,10,10,10],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1023,2,1,3],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1024,2,1,3],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1025,2,1,3],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2047,2,1,3],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2048,2,1,3],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2049,2,1,3],k=7,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16,10,10,10],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[60,10,10,10],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1023,2,1,3],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1024,2,1,3],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[1025,2,1,3],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2047,2,1,3],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2048,2,1,3],k=15,ties=0","support","0","no","HTP"
-"HTP0","TOP_K","type=f32,ne=[2049,2,1,3],k=15,ties=0","support","0","no","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,16,1,1],k=32,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8193,16,1,1],k=32,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16,10,10,10],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[60,10,10,10],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1023,2,1,3],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1024,2,1,3],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1025,2,1,3],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2047,2,1,3],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2048,2,1,3],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2049,2,1,3],k=1,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16,10,10,10],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[60,10,10,10],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1023,2,1,3],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1024,2,1,3],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1025,2,1,3],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2047,2,1,3],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2048,2,1,3],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2049,2,1,3],k=2,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16,10,10,10],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[60,10,10,10],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1023,2,1,3],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1024,2,1,3],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1025,2,1,3],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2047,2,1,3],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2048,2,1,3],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2049,2,1,3],k=3,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16,10,10,10],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[60,10,10,10],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1023,2,1,3],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1024,2,1,3],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1025,2,1,3],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2047,2,1,3],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2048,2,1,3],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2049,2,1,3],k=7,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16,10,10,10],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[60,10,10,10],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1023,2,1,3],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1024,2,1,3],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1025,2,1,3],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[16384,1,1,1],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2047,2,1,3],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2048,2,1,3],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2049,2,1,3],k=15,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[1024,1,1,1],k=1024,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[2048,2,1,1],k=1024,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[4096,1,1,1],k=2048,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[8192,2,1,1],k=2051,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[33024,1,1,1],k=2051,ties=0","support","1","yes","HTP"
+"HTP0","TOP_K","type=f32,ne=[33024,4,1,1],k=2051,ties=0","support","1","yes","HTP"
 "HTP0","UPSCALE","type=f32,ne=[512,512,3,2],scale_factor=2,mode=nearest,transpose=0","support","0","no","HTP"
 "HTP0","UPSCALE","type=f32,ne=[512,512,3,2],scale_factor=2,mode=nearest,transpose=1","support","0","no","HTP"
 "HTP0","UPSCALE","type=f32,ne=[2,5,7,11],ne_tgt=[5,7,11,13],mode=nearest","support","0","no","HTP"

examples/convert-llama2c-to-ggml/CMakeLists.txt

@@ -1,5 +1,7 @@
-set(TARGET llama-convert-llama2c-to-ggml)
-add_executable(${TARGET} convert-llama2c-to-ggml.cpp)
-install(TARGETS ${TARGET} RUNTIME)
-target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT})
-target_compile_features(${TARGET} PRIVATE cxx_std_17)
+if (GGML_CPU)
+    set(TARGET llama-convert-llama2c-to-ggml)
+    add_executable(${TARGET} convert-llama2c-to-ggml.cpp)
+    install(TARGETS ${TARGET} RUNTIME)
+    target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT})
+    target_compile_features(${TARGET} PRIVATE cxx_std_17)
+endif()

examples/parallel/parallel.cpp

@@ -11,6 +11,7 @@
 #include <clocale>
 #include <cmath>
 #include <cstdio>
+#include <random>
 #include <string>
 #include <vector>
 #include <ctime>
@@ -156,7 +157,7 @@ static std::vector<std::string> split_string(const std::string& input, char deli
 int main(int argc, char ** argv) {
     std::setlocale(LC_NUMERIC, "C");
 
-    srand(1234);
+    std::mt19937 rng(1234);
 
     common_params params;
 
@@ -321,7 +322,7 @@ int main(int argc, char ** argv) {
                     client.t_start_prompt = ggml_time_us();
                     client.t_start_gen    = 0;
 
-                    client.input    = k_prompts[rand() % k_prompts.size()];
+                    client.input    = k_prompts[rng() % k_prompts.size()];
                     client.response = "";
 
                     // construct the prompt:
@@ -334,10 +335,10 @@ int main(int argc, char ** argv) {
                         client.prompt += k_system;
                     }
 
-                    const int n_junk_cur = rand() % n_junk;
+                    const int n_junk_cur = rng() % n_junk;
 
                     for (int i = 0; i < n_junk_cur; ++i) {
-                        const int r = rand() % k_questions.size();
+                        const int r = rng() % k_questions.size();
                         client.prompt += "User:\n" + k_questions[r] + "\nAssistant:\n " + k_answers[r] + "\n";
                     }
                     client.prompt += "User:\n" + client.input + "\nAssistant:\n";

ggml/include/ggml-sycl.h

@@ -25,7 +25,7 @@ GGML_BACKEND_API bool ggml_backend_is_sycl(ggml_backend_t backend);
 GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_sycl_buffer_type(int device);
 
 // split tensor buffer that splits matrices by rows across multiple devices
-GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_sycl_split_buffer_type(const float * tensor_split);
+GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_sycl_split_buffer_type([[maybe_unused]] int main_device, const float * tensor_split);
 
 // Tensor parallelism (--split-mode tensor): comm_init/free/allreduce_tensor
 // trio queried by the meta-backend via ggml_backend_reg_get_proc_address.

ggml/src/ggml-backend.cpp

@@ -1630,7 +1630,10 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) {
             ggml_backend_synchronize(sched->backends[i]);
         }
 
-        ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids);
+        if (!ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids)) {
+            GGML_LOG_ERROR("%s: failed to reserve graph buffers\n", __func__);
+            return false;
+        }
         if (!ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) {
             GGML_LOG_ERROR("%s: failed to allocate graph\n", __func__);
             return false;

ggml/src/ggml-cpu/arch-fallback.h

@@ -39,6 +39,8 @@
 #define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
 #define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8
+#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0
+#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0
 #define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0
 #define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
 #define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0
@@ -55,6 +57,8 @@
 #define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0
 #define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
 #define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
+#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0
+#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
 #define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0
 #define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
 #define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0
@@ -87,6 +91,8 @@
 // repack.cpp
 #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
+#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0
+#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0
 #define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0
 #define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
 #define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K
@@ -98,6 +104,8 @@
 #define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0
 #define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
 #define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
+#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0
+#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
 #define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0
 #define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
 #define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K
@@ -124,6 +132,8 @@
 #define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
 #define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8
+#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0
+#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0
 #define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0
 #define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
 #define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0
@@ -140,6 +150,8 @@
 #define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0
 #define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
 #define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
+#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0
+#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
 #define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0
 #define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
 #define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0
@@ -171,6 +183,8 @@
 #define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
 #define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8
+#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0
+#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0
 #define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0
 #define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
 #define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0
@@ -187,6 +201,8 @@
 #define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0
 #define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
 #define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
+#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0
+#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
 #define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0
 #define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
 #define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0
@@ -213,6 +229,8 @@
 #define ggml_quantize_mat_q8_K_4x1_generic ggml_quantize_mat_q8_K_4x1
 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
 #define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8
+#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0
+#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0
 #define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0
 #define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
 #define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K
@@ -228,6 +246,8 @@
 #define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0
 #define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
 #define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
+#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0
+#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
 #define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0
 #define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
 #define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
@@ -262,6 +282,8 @@
 #define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
 #define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8
+#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0
+#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0
 #define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
 #define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0
 #define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K
@@ -277,6 +299,8 @@
 #define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0
 #define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
 #define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
+#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0
+#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
 #define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
 #define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0
 #define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
@@ -314,6 +338,8 @@
 #define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
 #define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8
+#define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0
+#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0
 #define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0
 #define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
 #define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0
@@ -330,6 +356,8 @@
 #define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0
 #define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
 #define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
+#define ggml_gemm_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0
+#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
 #define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0
 #define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
 #define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0

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

@@ -48,6 +48,24 @@ static inline void decode_q_Kx8_6bit_scales(const uint8_t * scales_in, int16x8_t
 }
 #endif
 
+#if defined(__aarch64__) && defined(__ARM_NEON) && (defined(__ARM_FEATURE_DOTPROD) || defined(__ARM_FEATURE_MATMUL_INT8))
+#define B1(c,s,n)  0x ## n ## c ,  0x ## n ## s
+#define B2(c,s,n) B1(c,s,n ## c), B1(c,s,n ## s)
+#define B3(c,s,n) B2(c,s,n ## c), B2(c,s,n ## s)
+#define B4(c,s,n) B3(c,s,n ## c), B3(c,s,n ## s)
+#define B5(c,s,n) B4(c,s,n ## c), B4(c,s,n ## s)
+#define B6(c,s,n) B5(c,s,n ## c), B5(c,s,n ## s)
+#define B7(c,s,n) B6(c,s,n ## c), B6(c,s,n ## s)
+#define B8(c,s  ) B7(c,s,     c), B7(c,s,     s)
+
+static const uint64_t table_q1_signs[256] = { B8(ff, 01) };
+
+static inline int8x16_t ggml_q1_0_unpack_pair(uint8_t bits0, uint8_t bits1) {
+    return vreinterpretq_s8_u8(vcombine_u8(vcreate_u8(table_q1_signs[bits0]),
+                                           vcreate_u8(table_q1_signs[bits1])));
+}
+#endif
+
 void ggml_quantize_mat_q8_0_4x4(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) {
     assert(QK8_0 == 32);
     assert(k % QK8_0 == 0);
@@ -1823,6 +1841,132 @@ void ggml_gemv_q8_0_4x8_q8_0(int                        n,
     ggml_gemv_q8_0_4x8_q8_0_generic(n, s, bs, vx, vy, nr, nc);
 }
 
+void ggml_gemv_q1_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                = QK1_0;
+    const int nb                = n / qk;
+    const int ncols_interleaved = 4;
+
+    assert(n % qk == 0);
+    assert(nc % ncols_interleaved == 0);
+
+    UNUSED(nb);
+    UNUSED(ncols_interleaved);
+
+#if defined(__aarch64__) && defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD)
+    for (int c = 0; c < nc; c += ncols_interleaved) {
+        const block_q1_0x4 * b_ptr = (const block_q1_0x4 *) vx + (c / ncols_interleaved) * nb;
+        const block_q8_0 * a_ptr = (const block_q8_0 *) vy;
+        float32x4_t acc = vdupq_n_f32(0);
+
+        for (int l = 0; l < nb; l++) {
+            const float32x4_t b_d = vcvt_f32_f16(vld1_f16((const float16_t *) b_ptr[l].d));
+            float32x4_t accb = vdupq_n_f32(0);
+
+            for (int k = 0; k < 4; k++) {
+                const block_q8_0 * GGML_RESTRICT a_blk = a_ptr + l * 4 + k;
+                const float ad = GGML_CPU_FP16_TO_FP32(a_blk->d);
+                int32x4_t ret = vdupq_n_s32(0);
+
+                for (int tile = 0; tile < 8; tile += 4) {
+                    const int8x16_t signs0 = ggml_q1_0_unpack_pair(b_ptr[l].qs[k * 16 + 2 * (tile + 0) + 0],
+                                                                    b_ptr[l].qs[k * 16 + 2 * (tile + 0) + 1]);
+                    const int8x16_t signs1 = ggml_q1_0_unpack_pair(b_ptr[l].qs[k * 16 + 2 * (tile + 1) + 0],
+                                                                    b_ptr[l].qs[k * 16 + 2 * (tile + 1) + 1]);
+                    const int8x16_t signs2 = ggml_q1_0_unpack_pair(b_ptr[l].qs[k * 16 + 2 * (tile + 2) + 0],
+                                                                    b_ptr[l].qs[k * 16 + 2 * (tile + 2) + 1]);
+                    const int8x16_t signs3 = ggml_q1_0_unpack_pair(b_ptr[l].qs[k * 16 + 2 * (tile + 3) + 0],
+                                                                    b_ptr[l].qs[k * 16 + 2 * (tile + 3) + 1]);
+                    const int8x16_t q_tiles = vld1q_s8(a_blk->qs + tile * 4);
+
+                    ret = vdotq_laneq_s32(ret, signs0, q_tiles, 0);
+                    ret = vdotq_laneq_s32(ret, signs1, q_tiles, 1);
+                    ret = vdotq_laneq_s32(ret, signs2, q_tiles, 2);
+                    ret = vdotq_laneq_s32(ret, signs3, q_tiles, 3);
+                }
+
+                accb = vfmaq_n_f32(accb, vcvtq_f32_s32(ret), ad);
+            }
+            acc = vfmaq_f32(acc, accb, b_d);
+        }
+        vst1q_f32(s, acc);
+        s += ncols_interleaved;
+    }
+    return;
+#endif
+    ggml_gemv_q1_0_4x4_q8_0_generic(n, s, bs, vx, vy, nr, nc);
+}
+
+void ggml_gemv_q1_0_4x8_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                = QK1_0;
+    const int nb                = n / qk;
+    const int ncols_interleaved = 4;
+
+    assert(n % qk == 0);
+    assert(nc % ncols_interleaved == 0);
+
+    UNUSED(nb);
+    UNUSED(ncols_interleaved);
+
+#if defined(__aarch64__) && defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD)
+    for (int c = 0; c < nc; c += ncols_interleaved) {
+        const block_q1_0x4 * b_ptr = (const block_q1_0x4 *) vx + (c / ncols_interleaved) * nb;
+        const block_q8_0 * a_ptr = (const block_q8_0 *) vy;
+        float32x4_t acc = vdupq_n_f32(0);
+
+        for (int l = 0; l < nb; l++) {
+            const float32x4_t b_d = vcvt_f32_f16(vld1_f16((const float16_t *) b_ptr[l].d));
+            float32x4_t accb = vdupq_n_f32(0);
+
+            for (int k = 0; k < 4; ++k) {
+                const block_q8_0 * GGML_RESTRICT a_blk = a_ptr + l * 4 + k;
+                const uint8_t * GGML_RESTRICT b_qs = (const uint8_t *) b_ptr[l].qs + k * 16;
+                const float ad = GGML_CPU_FP16_TO_FP32(a_blk->d);
+
+                int8x8x4_t  a_chunks = vld1_s8_x4(a_blk->qs);
+                int8x16_t   a0       = vcombine_s8(a_chunks.val[0], a_chunks.val[0]);
+                int8x16_t   a1       = vcombine_s8(a_chunks.val[1], a_chunks.val[1]);
+                int8x16_t   a2       = vcombine_s8(a_chunks.val[2], a_chunks.val[2]);
+                int8x16_t   a3       = vcombine_s8(a_chunks.val[3], a_chunks.val[3]);
+
+                int32x4_t ret0 = vdupq_n_s32(0);
+                int32x4_t ret1 = vdupq_n_s32(0);
+
+                ret0 = vdotq_s32(ret0, ggml_q1_0_unpack_pair(b_qs[0],  b_qs[1]),  a0);
+                ret1 = vdotq_s32(ret1, ggml_q1_0_unpack_pair(b_qs[2],  b_qs[3]),  a0);
+                ret0 = vdotq_s32(ret0, ggml_q1_0_unpack_pair(b_qs[4],  b_qs[5]),  a1);
+                ret1 = vdotq_s32(ret1, ggml_q1_0_unpack_pair(b_qs[6],  b_qs[7]),  a1);
+                ret0 = vdotq_s32(ret0, ggml_q1_0_unpack_pair(b_qs[8],  b_qs[9]),  a2);
+                ret1 = vdotq_s32(ret1, ggml_q1_0_unpack_pair(b_qs[10], b_qs[11]), a2);
+                ret0 = vdotq_s32(ret0, ggml_q1_0_unpack_pair(b_qs[12], b_qs[13]), a3);
+                ret1 = vdotq_s32(ret1, ggml_q1_0_unpack_pair(b_qs[14], b_qs[15]), a3);
+
+                accb = vfmaq_n_f32(accb, vcvtq_f32_s32(vpaddq_s32(ret0, ret1)), ad);
+            }
+
+            acc = vfmaq_f32(acc, accb, b_d);
+        }
+
+        vst1q_f32(s, acc);
+        s += ncols_interleaved;
+    }
+    return;
+#endif
+
+    ggml_gemv_q1_0_4x8_q8_0_generic(n, s, bs, vx, vy, nr, nc);
+}
+
 void ggml_gemm_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;
     const int nb = n / qk;
@@ -5154,3 +5298,168 @@ void ggml_gemm_q8_0_4x8_q8_0(int                        n,
 #endif  // defined(__aarch64__) && defined(__ARM_NEON) && defined(__ARM_FEATURE_MATMUL_INT8)
     ggml_gemm_q8_0_4x8_q8_0_generic(n, s, bs, vx, vy, nr, nc);
 }
+
+void ggml_gemm_q1_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                = QK1_0;
+    const int nb                = n / qk;
+    const int ncols_interleaved = 4;
+
+    assert(n % qk == 0);
+    assert(nr % 4 == 0);
+    assert(nc % ncols_interleaved == 0);
+
+    UNUSED(nb);
+    UNUSED(ncols_interleaved);
+
+#if defined(__aarch64__) && defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD)
+    for (int y = 0; y < nr / 4; y++) {
+        const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (4 * y * nb);
+        for (int x = 0; x < nc / ncols_interleaved; x++) {
+            const block_q1_0x4 * b_ptr = (const block_q1_0x4 *) vx + (x * nb);
+
+            float32x4_t sumf[4];
+            for (int m = 0; m < 4; m++) {
+                sumf[m] = vdupq_n_f32(0);
+            }
+
+            for (int l = 0; l < nb; l++) {
+                float32x4_t b_d = vcvt_f32_f16(vld1_f16((const float16_t *) b_ptr[l].d));
+                float32x4_t blockf_0 = vdupq_n_f32(0);
+                float32x4_t blockf_1 = vdupq_n_f32(0);
+                float32x4_t blockf_2 = vdupq_n_f32(0);
+                float32x4_t blockf_3 = vdupq_n_f32(0);
+
+                for (int k = 0; k < 4; ++k) {
+                    const block_q8_0x4 * GGML_RESTRICT a_blk = a_ptr + 4 * l + k;
+                    float32x4_t a_d = vcvt_f32_f16(vld1_f16((const float16_t *) a_blk->d));
+
+                    int32x4_t sumi_0 = vdupq_n_s32(0);
+                    int32x4_t sumi_1 = vdupq_n_s32(0);
+                    int32x4_t sumi_2 = vdupq_n_s32(0);
+                    int32x4_t sumi_3 = vdupq_n_s32(0);
+
+                    for (int tile = 0; tile < 8; ++tile) {
+                        const int8x16_t signs = ggml_q1_0_unpack_pair(b_ptr[l].qs[k * 16 + 2 * tile + 0],
+                                                                       b_ptr[l].qs[k * 16 + 2 * tile + 1]);
+                        const int8x16_t a_tile = vld1q_s8(a_blk->qs + tile * 16);
+
+                        sumi_0 = vdotq_laneq_s32(sumi_0, signs, a_tile, 0);
+                        sumi_1 = vdotq_laneq_s32(sumi_1, signs, a_tile, 1);
+                        sumi_2 = vdotq_laneq_s32(sumi_2, signs, a_tile, 2);
+                        sumi_3 = vdotq_laneq_s32(sumi_3, signs, a_tile, 3);
+                    }
+
+                    blockf_0 = vfmaq_laneq_f32(blockf_0, vcvtq_f32_s32(sumi_0), a_d, 0);
+                    blockf_1 = vfmaq_laneq_f32(blockf_1, vcvtq_f32_s32(sumi_1), a_d, 1);
+                    blockf_2 = vfmaq_laneq_f32(blockf_2, vcvtq_f32_s32(sumi_2), a_d, 2);
+                    blockf_3 = vfmaq_laneq_f32(blockf_3, vcvtq_f32_s32(sumi_3), a_d, 3);
+                }
+
+                sumf[0] = vfmaq_f32(sumf[0], blockf_0, b_d);
+                sumf[1] = vfmaq_f32(sumf[1], blockf_1, b_d);
+                sumf[2] = vfmaq_f32(sumf[2], blockf_2, b_d);
+                sumf[3] = vfmaq_f32(sumf[3], blockf_3, b_d);
+            }
+
+            for (int m = 0; m < 4; m++) {
+                vst1q_f32(s + (y * 4 + m) * bs + x * 4, sumf[m]);
+            }
+        }
+    }
+    return;
+#endif
+    ggml_gemm_q1_0_4x4_q8_0_generic(n, s, bs, vx, vy, nr, nc);
+}
+
+void ggml_gemm_q1_0_4x8_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                = QK1_0;
+    const int nb                = n / qk;
+    const int ncols_interleaved = 4;
+
+    assert(n % qk == 0);
+    assert(nr % 4 == 0);
+    assert(nc % ncols_interleaved == 0);
+
+    UNUSED(nb);
+    UNUSED(ncols_interleaved);
+
+#if defined(__aarch64__) && defined(__ARM_NEON) && defined(__ARM_FEATURE_MATMUL_INT8)
+    for (int y = 0; y < nr / 4; y++) {
+        const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (4 * y * nb);
+
+        for (int x = 0; x < nc / ncols_interleaved; x++) {
+            const block_q1_0x4 * b_ptr = (const block_q1_0x4 *) vx + (x * nb);
+
+            float32x4_t sumf[4];
+            for (int m = 0; m < 4; ++m) {
+                sumf[m] = vdupq_n_f32(0);
+            }
+
+            for (int l = 0; l < nb; l++) {
+                const float32x4_t b_d = vcvt_f32_f16(vld1_f16((const float16_t *) b_ptr[l].d));
+                float32x4_t blockf[4];
+                for (int m = 0; m < 4; ++m) {
+                    blockf[m] = vdupq_n_f32(0);
+                }
+
+                for (int k = 0; k < 4; ++k) {
+                    const block_q8_0x4 * GGML_RESTRICT a_blk = a_ptr + 4 * l + k;
+                    const uint8_t * GGML_RESTRICT b_qs = (const uint8_t *) b_ptr[l].qs + k * 16;
+
+                    int32x4_t acc[4];
+                    for (int i = 0; i < 4; ++i) {
+                        acc[i] = vdupq_n_s32(0);
+                    }
+
+                    for (int chunk = 0; chunk < 4; ++chunk) {
+                        const int8x16_t a01 = vld1q_s8(a_blk->qs + chunk * 32);
+                        const int8x16_t a23 = vld1q_s8(a_blk->qs + chunk * 32 + 16);
+                        const int8x16_t b01 = ggml_q1_0_unpack_pair(b_qs[chunk * 4 + 0], b_qs[chunk * 4 + 1]);
+                        const int8x16_t b23 = ggml_q1_0_unpack_pair(b_qs[chunk * 4 + 2], b_qs[chunk * 4 + 3]);
+
+                        acc[0] = vmmlaq_s32(acc[0], a01, b01);
+                        acc[1] = vmmlaq_s32(acc[1], a01, b23);
+                        acc[2] = vmmlaq_s32(acc[2], a23, b01);
+                        acc[3] = vmmlaq_s32(acc[3], a23, b23);
+                    }
+
+                    const int32x4_t row0 = vcombine_s32(vget_low_s32(acc[0]), vget_low_s32(acc[1]));
+                    const int32x4_t row1 = vcombine_s32(vget_high_s32(acc[0]), vget_high_s32(acc[1]));
+                    const int32x4_t row2 = vcombine_s32(vget_low_s32(acc[2]), vget_low_s32(acc[3]));
+                    const int32x4_t row3 = vcombine_s32(vget_high_s32(acc[2]), vget_high_s32(acc[3]));
+                    const float32x4_t a_d = vcvt_f32_f16(vld1_f16((const float16_t *) a_blk->d));
+
+                    blockf[0] = vfmaq_laneq_f32(blockf[0], vcvtq_f32_s32(row0), a_d, 0);
+                    blockf[1] = vfmaq_laneq_f32(blockf[1], vcvtq_f32_s32(row1), a_d, 1);
+                    blockf[2] = vfmaq_laneq_f32(blockf[2], vcvtq_f32_s32(row2), a_d, 2);
+                    blockf[3] = vfmaq_laneq_f32(blockf[3], vcvtq_f32_s32(row3), a_d, 3);
+                }
+
+                sumf[0] = vfmaq_f32(sumf[0], blockf[0], b_d);
+                sumf[1] = vfmaq_f32(sumf[1], blockf[1], b_d);
+                sumf[2] = vfmaq_f32(sumf[2], blockf[2], b_d);
+                sumf[3] = vfmaq_f32(sumf[3], blockf[3], b_d);
+            }
+
+            for (int m = 0; m < 4; ++m) {
+                vst1q_f32(s + (y * 4 + m) * bs + x * 4, sumf[m]);
+            }
+        }
+    }
+    return;
+#endif
+
+    ggml_gemm_q1_0_4x8_q8_0_generic(n, s, bs, vx, vy, nr, nc);
+}

ggml/src/ggml-cpu/ggml-cpu.c

@@ -1329,7 +1329,9 @@ UseGgmlGemm1:;
         const size_t nbw3 = nbw2*ne12;
 
         assert(params->wsize >= ne13*nbw3);
-        GGML_ASSERT(src1->type == GGML_TYPE_F32);
+        GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);
+        // the F16 path below writes plain floats into wdata, so it needs an F32 vec_dot_type
+        GGML_ASSERT(src1->type == GGML_TYPE_F32 || vec_dot_type == GGML_TYPE_F32);
 
     #if 0
         for (int64_t i13 = 0; i13 < ne13; ++i13) {
@@ -1348,9 +1350,15 @@ UseGgmlGemm1:;
                     size_t bs = ggml_blck_size(vec_dot_type);
                     int64_t ne10_block_start = (ith * ne10/bs) / nth;
                     int64_t ne10_block_end   = ((ith + 1) * ne10/bs) / nth;
-                    from_float((float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11 + ne10_block_start*bs*nb10),
-                               (void *)               (wdata + i13*nbw3 + i12*nbw2 + i11*nbw1 + ne10_block_start*nbw0),
-                               (ne10_block_end - ne10_block_start) * bs);
+                    const char * src1_block = (const char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11 + ne10_block_start*bs*nb10;
+                    char * dst_block = wdata + i13*nbw3 + i12*nbw2 + i11*nbw1 + ne10_block_start*nbw0;
+                    const int64_t n_block = (ne10_block_end - ne10_block_start) * bs;
+
+                    if (src1->type == GGML_TYPE_F32) {
+                        from_float((const float *) src1_block, dst_block, n_block);
+                    } else {
+                        ggml_cpu_fp16_to_fp32((const ggml_fp16_t *) src1_block, (float *) dst_block, n_block);
+                    }
                 }
             }
         }

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

@@ -451,6 +451,10 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
                 op->type != GGML_TYPE_IQ1_S   &&
                 op->type != GGML_TYPE_IQ1_M; // missing type_traits.from_float
         case GGML_OP_MUL_MAT:
+            if (ggml_get_op_params_i32(op, 1) == GGML_HINT_SRC0_IS_HADAMARD &&
+                src0->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32) {
+                return src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16;
+            }
             return src1->type == GGML_TYPE_F32 || src1->type == ggml_get_type_traits_cpu(src0->type)->vec_dot_type;
         case GGML_OP_SOFT_MAX_BACK: {
             if (op->src[0]->type != GGML_TYPE_F32 || op->src[1]->type != GGML_TYPE_F32) {

ggml/src/ggml-cpu/ops.cpp

@@ -12015,11 +12015,20 @@ void ggml_compute_forward_opt_step_sgd(const ggml_compute_params * params, ggml_
     }
 }
 
-static void ggml_compute_forward_fwht_f32(const ggml_compute_params * params, ggml_tensor * dst) {
+static inline float ggml_fwht_load(const float value) {
+    return value;
+}
+
+static inline float ggml_fwht_load(const ggml_fp16_t value) {
+    return ggml_fp16_to_fp32(value);
+}
+
+template<typename src_t>
+static void ggml_compute_forward_fwht_impl(const ggml_compute_params * params, ggml_tensor * dst) {
     const ggml_tensor * src0 = dst->src[0];
     const ggml_tensor * src1 = dst->src[1];
 
-    GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    GGML_ASSERT(src1->type == (std::is_same_v<src_t, float> ? GGML_TYPE_F32 : GGML_TYPE_F16));
     GGML_ASSERT(dst->type == GGML_TYPE_F32);
 
     GGML_TENSOR_BINARY_OP_LOCALS
@@ -12046,11 +12055,11 @@ static void ggml_compute_forward_fwht_f32(const ggml_compute_params * params, gg
         const int64_t i12 = (r - i13 * ne11 * ne12) / ne11;
         const int64_t i11 = r - i13 * ne11 * ne12 - i12 * ne11;
 
-        const float * src_row = (const float *) ((const char *) src1->data + i11 * nb11 + i12 * nb12 + i13 * nb13);
+        const src_t * src_row = (const src_t *) ((const char *) src1->data + i11 * nb11 + i12 * nb12 + i13 * nb13);
         float * dst_row = (float *) ((char *) dst->data + i11 * nb1 + i12 * nb2 + i13 * nb3);
 
         for (int64_t j = 0; j < n; j++) {
-            dst_row[j] = src_row[j] * scale;
+            dst_row[j] = ggml_fwht_load(src_row[j]) * scale;
         }
 
         // Scalar passes
@@ -12097,12 +12106,17 @@ void ggml_compute_forward_fwht(const ggml_compute_params * params, ggml_tensor *
     switch (src1->type) {
         case GGML_TYPE_F32:
             {
-                ggml_compute_forward_fwht_f32(params, dst);
+                ggml_compute_forward_fwht_impl<float>(params, dst);
+            }
+            break;
+        case GGML_TYPE_F16:
+            {
+                ggml_compute_forward_fwht_impl<ggml_fp16_t>(params, dst);
             }
             break;
         default:
             {
-                GGML_ABORT("fatal error - fwht is F32 only");
+                GGML_ABORT("fatal error - fwht supports F32 and F16 input");
             }
     }
 }

ggml/src/ggml-cpu/repack.cpp

@@ -1365,6 +1365,133 @@ void ggml_gemv_q8_0_4x8_q8_0_generic(int                        n,
     }
 }
 
+void ggml_gemv_q1_0_4x4_q8_0_generic(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                = QK1_0;
+    const int nb                = n / qk;
+    const int ncols_interleaved = 4;
+
+    assert(nr == 1);
+    assert(n % qk == 0);
+    assert(nc % ncols_interleaved == 0);
+
+    UNUSED(bs);
+    UNUSED(nr);
+
+    float sumf[4];
+
+    const block_q8_0 * a_ptr = (const block_q8_0 *) vy;
+    for (int x = 0; x < nc / ncols_interleaved; x++) {
+        const block_q1_0x4 * b_ptr = (const block_q1_0x4 *) vx + (x * nb);
+
+        for (int j = 0; j < ncols_interleaved; j++) {
+            sumf[j] = 0.0;
+        }
+
+        for (int l = 0; l < nb; l++) {
+            const float d0[4] = {
+                GGML_CPU_FP16_TO_FP32(b_ptr[l].d[0]),
+                GGML_CPU_FP16_TO_FP32(b_ptr[l].d[1]),
+                GGML_CPU_FP16_TO_FP32(b_ptr[l].d[2]),
+                GGML_CPU_FP16_TO_FP32(b_ptr[l].d[3]),
+            };
+
+            for (int k = 0; k < QK1_0 / QK8_0; ++k) {
+                const block_q8_0 * GGML_RESTRICT a_blk = a_ptr + l * (QK1_0 / QK8_0) + k;
+                const float d1 = GGML_CPU_FP16_TO_FP32(a_blk->d);
+                const float scale[4] = { d0[0] * d1, d0[1] * d1, d0[2] * d1, d0[3] * d1 };
+
+                for (int tile = 0; tile < QK8_0 / 4; ++tile) {
+                    const uint8_t bits_lo = b_ptr[l].qs[k * 16 + 2 * tile + 0];
+                    const uint8_t bits_hi = b_ptr[l].qs[k * 16 + 2 * tile + 1];
+
+                    for (int p = 0; p < 4; ++p) {
+                        const float q = (float) a_blk->qs[tile * 4 + p];
+
+                        sumf[0] += ((bits_lo & (1u << p))       ? scale[0] : -scale[0]) * q;
+                        sumf[1] += ((bits_lo & (1u << (4 + p))) ? scale[1] : -scale[1]) * q;
+                        sumf[2] += ((bits_hi & (1u << p))       ? scale[2] : -scale[2]) * q;
+                        sumf[3] += ((bits_hi & (1u << (4 + p))) ? scale[3] : -scale[3]) * q;
+                    }
+                }
+            }
+        }
+
+        for (int j = 0; j < ncols_interleaved; j++) {
+            s[x * ncols_interleaved + j] = sumf[j];
+        }
+    }
+}
+
+void ggml_gemv_q1_0_4x8_q8_0_generic(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                = QK1_0;
+    const int nb                = n / qk;
+    const int ncols_interleaved = 4;
+    const int blocklen          = 8;
+
+    assert(nr == 1);
+    assert(n % qk == 0);
+    assert(nc % ncols_interleaved == 0);
+
+    UNUSED(bs);
+    UNUSED(nr);
+
+    float sumf[4];
+
+    const block_q8_0 * a_ptr = (const block_q8_0 *) vy;
+    for (int x = 0; x < nc / ncols_interleaved; x++) {
+        const block_q1_0x4 * b_ptr = (const block_q1_0x4 *) vx + (x * nb);
+
+        for (int j = 0; j < ncols_interleaved; j++) {
+            sumf[j] = 0.0f;
+        }
+
+        for (int l = 0; l < nb; l++) {
+            const float d0[4] = {
+                GGML_CPU_FP16_TO_FP32(b_ptr[l].d[0]),
+                GGML_CPU_FP16_TO_FP32(b_ptr[l].d[1]),
+                GGML_CPU_FP16_TO_FP32(b_ptr[l].d[2]),
+                GGML_CPU_FP16_TO_FP32(b_ptr[l].d[3]),
+            };
+
+            for (int k = 0; k < qk / blocklen; ++k) {
+                const block_q8_0 * GGML_RESTRICT a_blk = a_ptr + l * (qk / QK8_0) + k / (QK8_0 / blocklen);
+                const float d1 = GGML_CPU_FP16_TO_FP32(a_blk->d);
+                const float scale[4] = { d0[0] * d1, d0[1] * d1, d0[2] * d1, d0[3] * d1 };
+                const uint8_t bits0 = b_ptr[l].qs[k * ncols_interleaved + 0];
+                const uint8_t bits1 = b_ptr[l].qs[k * ncols_interleaved + 1];
+                const uint8_t bits2 = b_ptr[l].qs[k * ncols_interleaved + 2];
+                const uint8_t bits3 = b_ptr[l].qs[k * ncols_interleaved + 3];
+                const int q_offset = (k % (QK8_0 / blocklen)) * blocklen;
+
+                for (int p = 0; p < blocklen; ++p) {
+                    const float q = (float) a_blk->qs[q_offset + p];
+
+                    sumf[0] += ((bits0 & (1u << p)) ? scale[0] : -scale[0]) * q;
+                    sumf[1] += ((bits1 & (1u << p)) ? scale[1] : -scale[1]) * q;
+                    sumf[2] += ((bits2 & (1u << p)) ? scale[2] : -scale[2]) * q;
+                    sumf[3] += ((bits3 & (1u << p)) ? scale[3] : -scale[3]) * q;
+                }
+            }
+        }
+
+        for (int j = 0; j < ncols_interleaved; j++) {
+            s[x * ncols_interleaved + j] = sumf[j];
+        }
+    }
+}
+
 // Only enable these for RISC-V.
 #if defined __riscv_zvfh
 void ggml_gemv_q4_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) {
@@ -2383,6 +2510,176 @@ void ggml_gemm_q8_0_4x8_q8_0_generic(int                        n,
     }
 }
 
+void ggml_gemm_q1_0_4x4_q8_0_generic(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                = QK1_0;
+    const int nb                = n / qk;
+    const int ncols_interleaved = 4;
+
+    assert(n % qk == 0);
+    assert(nr % 4 == 0);
+    assert(nc % ncols_interleaved == 0);
+
+    float sumf[4][4];
+
+    for (int y = 0; y < nr / 4; y++) {
+        const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (4 * y * nb);
+        for (int x = 0; x < nc / ncols_interleaved; x++) {
+            const block_q1_0x4 * b_ptr = (const block_q1_0x4 *) vx + (x * nb);
+
+            for (int m = 0; m < 4; m++) {
+                for (int j = 0; j < ncols_interleaved; j++) {
+                    sumf[m][j] = 0.0;
+                }
+            }
+
+            for (int l = 0; l < nb; l++) {
+                const float d0[4] = {
+                    GGML_CPU_FP16_TO_FP32(b_ptr[l].d[0]),
+                    GGML_CPU_FP16_TO_FP32(b_ptr[l].d[1]),
+                    GGML_CPU_FP16_TO_FP32(b_ptr[l].d[2]),
+                    GGML_CPU_FP16_TO_FP32(b_ptr[l].d[3]),
+                };
+
+                for (int k = 0; k < QK1_0 / QK8_0; ++k) {
+                    const block_q8_0x4 * GGML_RESTRICT a_blk = a_ptr + 4 * l + k;
+                    const float a_d[4] = {
+                        GGML_CPU_FP16_TO_FP32(a_blk->d[0]),
+                        GGML_CPU_FP16_TO_FP32(a_blk->d[1]),
+                        GGML_CPU_FP16_TO_FP32(a_blk->d[2]),
+                        GGML_CPU_FP16_TO_FP32(a_blk->d[3]),
+                    };
+
+                    for (int tile = 0; tile < QK8_0 / 4; ++tile) {
+                        const uint8_t bits_lo = b_ptr[l].qs[k * 16 + 2 * tile + 0];
+                        const uint8_t bits_hi = b_ptr[l].qs[k * 16 + 2 * tile + 1];
+                        const int tile_offset = tile * 16;
+
+                        for (int p = 0; p < 4; ++p) {
+                            const int8_t q_row[4] = {
+                                a_blk->qs[tile_offset + 0 * 4 + p],
+                                a_blk->qs[tile_offset + 1 * 4 + p],
+                                a_blk->qs[tile_offset + 2 * 4 + p],
+                                a_blk->qs[tile_offset + 3 * 4 + p],
+                            };
+                            const int sign[4] = {
+                                (bits_lo & (1u << p))       ? 1 : -1,
+                                (bits_lo & (1u << (4 + p))) ? 1 : -1,
+                                (bits_hi & (1u << p))       ? 1 : -1,
+                                (bits_hi & (1u << (4 + p))) ? 1 : -1,
+                            };
+
+                            for (int m = 0; m < 4; ++m) {
+                                const float row_scale = a_d[m];
+                                sumf[m][0] += sign[0] * q_row[m] * d0[0] * row_scale;
+                                sumf[m][1] += sign[1] * q_row[m] * d0[1] * row_scale;
+                                sumf[m][2] += sign[2] * q_row[m] * d0[2] * row_scale;
+                                sumf[m][3] += sign[3] * q_row[m] * d0[3] * row_scale;
+                            }
+                        }
+                    }
+                }
+            }
+
+            for (int m = 0; m < 4; m++) {
+                for (int j = 0; j < ncols_interleaved; j++) {
+                    s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j];
+                }
+            }
+        }
+    }
+}
+
+void ggml_gemm_q1_0_4x8_q8_0_generic(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                = QK1_0;
+    const int nb                = n / qk;
+    const int ncols_interleaved = 4;
+    const int blocklen          = 8;
+
+    assert(n % qk == 0);
+    assert(nr % 4 == 0);
+    assert(nc % ncols_interleaved == 0);
+
+    float sumf[4][4];
+
+    for (int y = 0; y < nr / 4; y++) {
+        const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (4 * y * nb);
+        for (int x = 0; x < nc / ncols_interleaved; x++) {
+            const block_q1_0x4 * b_ptr = (const block_q1_0x4 *) vx + (x * nb);
+
+            for (int m = 0; m < 4; m++) {
+                for (int j = 0; j < ncols_interleaved; j++) {
+                    sumf[m][j] = 0.0f;
+                }
+            }
+
+            for (int l = 0; l < nb; l++) {
+                const float d0[4] = {
+                    GGML_CPU_FP16_TO_FP32(b_ptr[l].d[0]),
+                    GGML_CPU_FP16_TO_FP32(b_ptr[l].d[1]),
+                    GGML_CPU_FP16_TO_FP32(b_ptr[l].d[2]),
+                    GGML_CPU_FP16_TO_FP32(b_ptr[l].d[3]),
+                };
+
+                for (int k = 0; k < qk / blocklen; ++k) {
+                    const block_q8_0x4 * GGML_RESTRICT a_blk = a_ptr + 4 * l + k / (QK8_0 / blocklen);
+                    const float a_d[4] = {
+                        GGML_CPU_FP16_TO_FP32(a_blk->d[0]),
+                        GGML_CPU_FP16_TO_FP32(a_blk->d[1]),
+                        GGML_CPU_FP16_TO_FP32(a_blk->d[2]),
+                        GGML_CPU_FP16_TO_FP32(a_blk->d[3]),
+                    };
+                    const uint8_t bits0 = b_ptr[l].qs[k * ncols_interleaved + 0];
+                    const uint8_t bits1 = b_ptr[l].qs[k * ncols_interleaved + 1];
+                    const uint8_t bits2 = b_ptr[l].qs[k * ncols_interleaved + 2];
+                    const uint8_t bits3 = b_ptr[l].qs[k * ncols_interleaved + 3];
+                    const int q_offset = (k % (QK8_0 / blocklen)) * 4 * blocklen;
+
+                    for (int p = 0; p < blocklen; ++p) {
+                        const int8_t q_row[4] = {
+                            a_blk->qs[q_offset + 0 * blocklen + p],
+                            a_blk->qs[q_offset + 1 * blocklen + p],
+                            a_blk->qs[q_offset + 2 * blocklen + p],
+                            a_blk->qs[q_offset + 3 * blocklen + p],
+                        };
+                        const int sign[4] = {
+                            (bits0 & (1u << p)) ? 1 : -1,
+                            (bits1 & (1u << p)) ? 1 : -1,
+                            (bits2 & (1u << p)) ? 1 : -1,
+                            (bits3 & (1u << p)) ? 1 : -1,
+                        };
+
+                        for (int m = 0; m < 4; ++m) {
+                            const float row_scale = a_d[m];
+                            sumf[m][0] += sign[0] * q_row[m] * d0[0] * row_scale;
+                            sumf[m][1] += sign[1] * q_row[m] * d0[1] * row_scale;
+                            sumf[m][2] += sign[2] * q_row[m] * d0[2] * row_scale;
+                            sumf[m][3] += sign[3] * q_row[m] * d0[3] * row_scale;
+                        }
+                    }
+                }
+            }
+
+            for (int m = 0; m < 4; m++) {
+                for (int j = 0; j < ncols_interleaved; j++) {
+                    s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j];
+                }
+            }
+        }
+    }
+}
+
 // Only enable these for RISC-V.
 #if defined __riscv_zvfh
 void ggml_gemm_q4_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) {
@@ -2739,6 +3036,50 @@ static block_q8_0x4 make_block_q8_0x4(block_q8_0 * in, unsigned int blck_size_in
     return out;
 }
 
+static block_q1_0x4 make_block_q1_0x4(block_q1_0 * in, unsigned int blck_size_interleave) {
+    block_q1_0x4 out;
+
+    for (int i = 0; i < 4; i++) {
+        out.d[i] = in[i].d;
+    }
+
+    GGML_ASSERT(blck_size_interleave == 4 || blck_size_interleave == 8);
+
+    if (blck_size_interleave == 4) {
+        for (int k = 0; k < QK1_0 / QK8_0; ++k) {
+            for (int tile = 0; tile < QK8_0 / 4; ++tile) {
+                uint8_t packed_lo = 0;
+                uint8_t packed_hi = 0;
+
+                const int weight_base = k * QK8_0 + tile * 4;
+                for (int pos = 0; pos < 4; ++pos) {
+                    const int weight_idx = weight_base + pos;
+                    const int byte_idx   = weight_idx / 8;
+                    const int bit_idx    = weight_idx % 8;
+
+                    packed_lo |= ((in[0].qs[byte_idx] >> bit_idx) & 1u) << pos;
+                    packed_lo |= ((in[1].qs[byte_idx] >> bit_idx) & 1u) << (4 + pos);
+                    packed_hi |= ((in[2].qs[byte_idx] >> bit_idx) & 1u) << pos;
+                    packed_hi |= ((in[3].qs[byte_idx] >> bit_idx) & 1u) << (4 + pos);
+                }
+
+                out.qs[k * 16 + 2 * tile + 0] = packed_lo;
+                out.qs[k * 16 + 2 * tile + 1] = packed_hi;
+            }
+        }
+        return out;
+    }
+
+    for (int byte_idx = 0; byte_idx < QK1_0 / 8; ++byte_idx) {
+        out.qs[byte_idx * 4 + 0] = in[0].qs[byte_idx];
+        out.qs[byte_idx * 4 + 1] = in[1].qs[byte_idx];
+        out.qs[byte_idx * 4 + 2] = in[2].qs[byte_idx];
+        out.qs[byte_idx * 4 + 3] = in[3].qs[byte_idx];
+    }
+
+    return out;
+}
+
 static block_q4_0x4 make_block_q4_0x4(block_q4_0 * in, int blck_size_interleave) {
     block_q4_0x4 out;
 
@@ -3509,6 +3850,38 @@ static int repack_q8_0_to_q8_0_4_bl(struct ggml_tensor *       t,
     return 0;
 }
 
+static int repack_q1_0_to_q1_0_4_bl(struct ggml_tensor *       t,
+                                    int                        interleave_block,
+                                    const void * GGML_RESTRICT data,
+                                    size_t                     data_size) {
+    GGML_ASSERT(t->type == GGML_TYPE_Q1_0);
+    GGML_ASSERT(interleave_block == 4 || interleave_block == 8);
+    constexpr int nrows_interleaved = 4;
+
+    block_q1_0x4 *     dst = (block_q1_0x4 *) t->data;
+    const block_q1_0 * src = (const block_q1_0 *) data;
+    block_q1_0         dst_tmp[4];
+    int                nrow    = ggml_nrows(t);
+    int                nblocks = t->ne[0] / QK1_0;
+
+    GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q1_0));
+
+    if (t->ne[1] % nrows_interleaved != 0) {
+        return -1;
+    }
+
+    for (int b = 0; b < nrow; b += nrows_interleaved) {
+        for (int64_t x = 0; x < nblocks; x++) {
+            for (int i = 0; i < nrows_interleaved; i++) {
+                dst_tmp[i] = src[x + i * nblocks];
+            }
+            *dst++ = make_block_q1_0x4(dst_tmp, interleave_block);
+        }
+        src += nrows_interleaved * nblocks;
+    }
+    return 0;
+}
+
 static block_q8_0x16 make_block_q8_0x16(block_q8_0 * in, unsigned int blck_size_interleave) {
     block_q8_0x16 out;
 
@@ -3865,6 +4238,14 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS>
 int repack(struct ggml_tensor *, const void *, size_t);
 
 // TODO: generalise.
+template <> int repack<block_q1_0, 4, 4>(struct ggml_tensor * t, const void * data, size_t data_size) {
+    return repack_q1_0_to_q1_0_4_bl(t, 4, data, data_size);
+}
+
+template <> int repack<block_q1_0, 8, 4>(struct ggml_tensor * t, const void * data, size_t data_size) {
+    return repack_q1_0_to_q1_0_4_bl(t, 8, data, data_size);
+}
+
 template <> int repack<block_q4_0, 4, 4>(struct ggml_tensor * t, const void * data, size_t data_size) {
     return repack_q4_0_to_q4_0_4_bl(t, 4, data, data_size);
 }
@@ -3960,6 +4341,14 @@ template <> int repack<block_q2_K, 1, 16>(struct ggml_tensor * t, const void * d
 template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PARAM_TYPE>
 void gemv(int, float *, size_t, const void *, const void *, int, int);
 
+template <> void gemv<block_q1_0, 4, 4, GGML_TYPE_Q8_0>(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) {
+    ggml_gemv_q1_0_4x4_q8_0(n, s, bs, vx, vy, nr, nc);
+}
+
+template <> void gemv<block_q1_0, 8, 4, GGML_TYPE_Q8_0>(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) {
+    ggml_gemv_q1_0_4x8_q8_0(n, s, bs, vx, vy, nr, nc);
+}
+
 template <> void gemv<block_q4_0, 4, 4, GGML_TYPE_Q8_0>(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) {
     ggml_gemv_q4_0_4x4_q8_0(n, s, bs, vx, vy, nr, nc);
 }
@@ -4057,6 +4446,14 @@ template <> void gemv<block_q2_K, 1, 16, GGML_TYPE_Q8_K>(int n, float * s, size_
 template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PARAM_TYPE>
 void gemm(int, float *, size_t, const void *, const void *, int, int);
 
+template <> void gemm<block_q1_0, 4, 4, GGML_TYPE_Q8_0>(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) {
+    ggml_gemm_q1_0_4x4_q8_0(n, s, bs, vx, vy, nr, nc);
+}
+
+template <> void gemm<block_q1_0, 8, 4, GGML_TYPE_Q8_0>(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) {
+    ggml_gemm_q1_0_4x8_q8_0(n, s, bs, vx, vy, nr, nc);
+}
+
 template <> void gemm<block_q4_0, 4, 4, GGML_TYPE_Q8_0>(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) {
     ggml_gemm_q4_0_4x4_q8_0(n, s, bs, vx, vy, nr, nc);
 }
@@ -4526,6 +4923,10 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PAR
 }  // namespace ggml::cpu::repack
 
 static const ggml::cpu::tensor_traits * ggml_repack_get_optimal_repack_type(const struct ggml_tensor * cur) {
+    // instance for Q1_0
+    static const ggml::cpu::repack::tensor_traits<block_q1_0, 4, 4, GGML_TYPE_Q8_0> q1_0_4x4_q8_0;
+    static const ggml::cpu::repack::tensor_traits<block_q1_0, 8, 4, GGML_TYPE_Q8_0> q1_0_4x8_q8_0;
+
     // instance for Q4
     static const ggml::cpu::repack::tensor_traits<block_q4_0, 4, 4, GGML_TYPE_Q8_0> q4_0_4x4_q8_0;
     static const ggml::cpu::repack::tensor_traits<block_q4_0, 8, 4, GGML_TYPE_Q8_0> q4_0_4x8_q8_0;
@@ -4723,6 +5124,17 @@ static const ggml::cpu::tensor_traits * ggml_repack_get_optimal_repack_type(cons
             }
             #endif
         }
+    } else if (cur->type == GGML_TYPE_Q1_0) {
+        if (ggml_cpu_has_neon() && ggml_cpu_has_matmul_int8()) {
+            if (cur->ne[1] % 4 == 0) {
+                return &q1_0_4x8_q8_0;
+            }
+        }
+        if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) {
+            if (cur->ne[1] % 4 == 0) {
+                return &q1_0_4x4_q8_0;
+            }
+        }
     }
 
     return nullptr;

ggml/src/ggml-cpu/repack.h

@@ -11,6 +11,9 @@
 ggml_backend_buffer_type_t ggml_backend_cpu_repack_buffer_type(void);
 
 template <int K> constexpr int QK_0() {
+    if constexpr (K == 1) {
+        return QK1_0;
+    }
     if constexpr (K == 4) {
         return QK4_0;
     }
@@ -26,6 +29,7 @@ template <int K, int N> struct block {
 };
 
 // control size
+static_assert(sizeof(block<1, 4>) == 4 * sizeof(ggml_half) + QK1_0 / 2, "wrong block<1,4> size/padding");
 static_assert(sizeof(block<4, 4>) == 4 * sizeof(ggml_half) + QK8_0 * 2, "wrong block<4,4> size/padding");
 static_assert(sizeof(block<4, 8>) == 8 * sizeof(ggml_half) + QK8_0 * 4, "wrong block<4,8> size/padding");
 static_assert(sizeof(block<4, 16>) == 16 * sizeof(ggml_half) + QK8_0 * 8, "wrong block<4,16> size/padding");
@@ -33,6 +37,7 @@ static_assert(sizeof(block<8, 4>) == 4 * sizeof(ggml_half) + QK8_0 * 4, "wrong b
 static_assert(sizeof(block<8, 8>) == 8 * sizeof(ggml_half) + QK8_0 * 8, "wrong block<8,8> size/padding");
 static_assert(sizeof(block<8, 16>) == 16 * sizeof(ggml_half) + QK8_0 * 16, "wrong block<8,16> size/padding");
 
+using block_q1_0x4 = block<1, 4>;
 using block_q4_0x4 = block<4, 4>;
 using block_q4_0x8 = block<4, 8>;
 using block_q4_0x16 = block<4, 16>;
@@ -141,6 +146,8 @@ void ggml_quantize_mat_q8_0_4x4(const float * GGML_RESTRICT x, void * GGML_RESTR
 void ggml_quantize_mat_q8_0_4x8(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k);
 void ggml_quantize_mat_q8_K_4x4(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k);
 void ggml_quantize_mat_q8_K_4x8(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k);
+void ggml_gemv_q1_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);
+void ggml_gemv_q1_0_4x8_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);
 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);
 void ggml_gemv_q4_0_4x8_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);
 void ggml_gemv_q4_0_8x8_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);
@@ -157,6 +164,8 @@ void ggml_gemv_mxfp4_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const v
 void ggml_gemv_mxfp4_8x8_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);
 void ggml_gemv_q8_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);
 void ggml_gemv_q8_0_4x8_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);
+void ggml_gemm_q1_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);
+void ggml_gemm_q1_0_4x8_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);
 void ggml_gemm_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);
 void ggml_gemm_q4_0_4x8_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);
 void ggml_gemm_q4_0_8x8_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);
@@ -193,6 +202,8 @@ void ggml_quantize_mat_q8_0_4x4_generic(const float * GGML_RESTRICT x, void * GG
 void ggml_quantize_mat_q8_0_4x8_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k);
 void ggml_quantize_mat_q8_K_4x4_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k);
 void ggml_quantize_mat_q8_K_4x8_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k);
+void ggml_gemv_q1_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc);
+void ggml_gemv_q1_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc);
 void ggml_gemv_q4_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc);
 void ggml_gemv_q4_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc);
 void ggml_gemv_q4_0_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc);
@@ -209,6 +220,8 @@ void ggml_gemv_mxfp4_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs,
 void ggml_gemv_mxfp4_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc);
 void ggml_gemv_q8_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc);
 void ggml_gemv_q8_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc);
+void ggml_gemm_q1_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc);
+void ggml_gemm_q1_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc);
 void ggml_gemm_q4_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc);
 void ggml_gemm_q4_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc);
 void ggml_gemm_q4_0_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc);

ggml/src/ggml-cuda/argsort.cu

@@ -51,9 +51,12 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool,
                               cudaStream_t     stream) {
     ggml_cuda_pool_alloc<int>   temp_indices_alloc(pool, ncols * nrows);
     ggml_cuda_pool_alloc<float> temp_keys_alloc(pool, ncols * nrows);
+    // Device*Sort algorithms currently do not allow for in-place sorting/aliasing of input/outputs
+    ggml_cuda_pool_alloc<float> temp_keys_out_alloc(pool, ncols * nrows);
 
     int *   temp_indices = temp_indices_alloc.get();
     float * temp_keys    = temp_keys_alloc.get();
+    float * temp_keys_out = temp_keys_out_alloc.get();
 
     static const int block_size = 256;
     const dim3 grid_size((ncols + block_size - 1) / block_size, nrows);
@@ -85,18 +88,18 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool,
 
     if (order == GGML_SORT_ORDER_ASC) {
         if (nrows == 1) {
-            CUDA_CHECK(DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, temp_keys, temp_keys,  // keys (in-place)
+            CUDA_CHECK(DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, temp_keys, temp_keys_out,  // keys in, keys out
                                                   temp_indices, dst,  // values (indices)
                                                   ncols, 0, sizeof(float) * 8, stream));
         } else if (is_capturing) {
             CUDA_CHECK(DeviceSegmentedRadixSort::SortPairs(
-                nullptr, temp_storage_bytes, temp_keys, temp_keys,  // keys (in-place)
+                nullptr, temp_storage_bytes, temp_keys, temp_keys_out,  // keys in, keys out
                 temp_indices, dst,                                  // values (indices)
                 ncols * nrows, nrows,                               // num items, num segments
                 offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream));
         } else {
             CUDA_CHECK(DeviceSegmentedSort::SortPairs(nullptr, temp_storage_bytes, temp_keys,
-                                                      temp_keys,             // keys (in-place)
+                                                      temp_keys_out, // keys out
                                                       temp_indices, dst,     // values (indices)
                                                       ncols * nrows, nrows,  // num items, num segments
                                                       offset_iterator, offset_iterator + 1, stream));
@@ -104,15 +107,15 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool,
     } else {
         if (nrows == 1) {
             CUDA_CHECK(DeviceRadixSort::SortPairsDescending(nullptr, temp_storage_bytes, temp_keys,
-                                                            temp_keys,          // keys (in-place)
+                                                            temp_keys_out, // keys out
                                                             temp_indices, dst,  // values (indices)
                                                             ncols, 0, sizeof(float) * 8, stream));
         } else if (is_capturing) {
             CUDA_CHECK(DeviceSegmentedRadixSort::SortPairsDescending(
-                nullptr, temp_storage_bytes, temp_keys, temp_keys, temp_indices, dst, ncols * nrows, nrows,
+                nullptr, temp_storage_bytes, temp_keys, temp_keys_out, temp_indices, dst, ncols * nrows, nrows,
                 offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream));
         } else {
-            CUDA_CHECK(DeviceSegmentedSort::SortPairsDescending(nullptr, temp_storage_bytes, temp_keys, temp_keys,
+            CUDA_CHECK(DeviceSegmentedSort::SortPairsDescending(nullptr, temp_storage_bytes, temp_keys, temp_keys_out,
                                                                 temp_indices, dst, ncols * nrows, nrows,
                                                                 offset_iterator, offset_iterator + 1, stream));
         }
@@ -124,31 +127,31 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool,
     if (order == GGML_SORT_ORDER_ASC) {
         if (nrows == 1) {
             CUDA_CHECK(DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys,
-                                                  temp_keys,          // keys (in-place)
+                                                  temp_keys_out, // keys out
                                                   temp_indices, dst,  // values (indices)
                                                   ncols, 0, sizeof(float) * 8, stream));
         } else if (is_capturing) {
-            CUDA_CHECK(DeviceSegmentedRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, temp_keys,
+            CUDA_CHECK(DeviceSegmentedRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, temp_keys_out,
                                                            temp_indices, dst, ncols * nrows, nrows, offset_iterator,
                                                            offset_iterator + 1, 0, sizeof(float) * 8, stream));
         } else {
-            CUDA_CHECK(DeviceSegmentedSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, temp_keys,
+            CUDA_CHECK(DeviceSegmentedSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, temp_keys_out,
                                                       temp_indices, dst, ncols * nrows, nrows, offset_iterator,
                                                       offset_iterator + 1, stream));
         }
     } else {
         if (nrows == 1) {
             CUDA_CHECK(DeviceRadixSort::SortPairsDescending(d_temp_storage, temp_storage_bytes, temp_keys,
-                                                            temp_keys,          // keys (in-place)
+                                                            temp_keys_out, // keys out
                                                             temp_indices, dst,  // values (indices)
                                                             ncols, 0, sizeof(float) * 8, stream));
         } else if (is_capturing) {
             CUDA_CHECK(DeviceSegmentedRadixSort::SortPairsDescending(
-                d_temp_storage, temp_storage_bytes, temp_keys, temp_keys, temp_indices, dst, ncols * nrows, nrows,
+                d_temp_storage, temp_storage_bytes, temp_keys, temp_keys_out, temp_indices, dst, ncols * nrows, nrows,
                 offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream));
         } else {
             CUDA_CHECK(DeviceSegmentedSort::SortPairsDescending(d_temp_storage, temp_storage_bytes, temp_keys,
-                                                                temp_keys, temp_indices, dst, ncols * nrows, nrows,
+                                                                temp_keys_out, temp_indices, dst, ncols * nrows, nrows,
                                                                 offset_iterator, offset_iterator + 1, stream));
         }
     }

ggml/src/ggml-cuda/fattn-common.cuh

@@ -719,7 +719,7 @@ static __global__ void flash_attn_mask_to_KV_max(
 }
 
 void ggml_cuda_flash_attn_ext_compact_mask(
-        const ggml_tensor * mask, int32_t * indices, int32_t n_kv_max, cudaStream_t stream);
+        const ggml_tensor * mask, int32_t * indices, int32_t * counts, int32_t n_queries, int32_t ncols1, int32_t n_kv_max, cudaStream_t stream);
 
 template<int D, int ncols1, int ncols2> // D == head size
 __launch_bounds__(D, 1)
@@ -1092,14 +1092,18 @@ void launch_fattn(
     const int ntiles_z_gqa = ((gqa_ratio + ncols2 - 1) / ncols2);
     const int ntiles_dst   = ntiles_x * ntiles_z_gqa * K->ne[2] * Q->ne[3];
 
-    const int32_t n_kv_max = use_sparse ? ggml_get_op_params_i32(KQV, 4) : 0;
+    // sparse: a query tile of ncols1 queries shares one index list, the union of the queries' visible columns
+    int32_t n_kv_max = 0;
     if (use_sparse) {
         GGML_ASSERT(mask != nullptr);
-        GGML_ASSERT(n_kv_max > 0);
-        const size_t mask_rows = size_t(mask->ne[1]) * mask->ne[3];
+        const int32_t n_kv_max_query = ggml_get_op_params_i32(KQV, 4);
+        GGML_ASSERT(n_kv_max_query > 0);
+        n_kv_max = std::min<int64_t>(K->ne[1], int64_t(ncols1)*n_kv_max_query);
 
-        KV_max.alloc(size_t(n_kv_max) * mask_rows);
-        ggml_cuda_flash_attn_ext_compact_mask(mask, KV_max.ptr, n_kv_max, main_stream);
+        const size_t n_lists = size_t(ntiles_x) * mask->ne[3];
+
+        KV_max.alloc(size_t(n_kv_max)*n_lists + n_lists);
+        ggml_cuda_flash_attn_ext_compact_mask(mask, KV_max.ptr, KV_max.ptr + size_t(n_kv_max)*n_lists, Q->ne[1], ncols1, n_kv_max, main_stream);
     }
 
     // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped.
@@ -1231,8 +1235,8 @@ void launch_fattn(
 
     GGML_ASSERT(block_dim.x % warp_size == 0);
 
-        ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks_num, block_dim, nbytes_shared, main_stream);
-        ggml_cuda_kernel_launch(fattn_kernel, launch_params,
+    ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks_num, block_dim, nbytes_shared, main_stream);
+    ggml_cuda_kernel_launch(fattn_kernel, launch_params,
         (const char *) Q->data,
         K_data,
         V_data,

ggml/src/ggml-cuda/fattn-mma-f16.cuh

@@ -68,16 +68,16 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 64, 128, 2,  32,  96,  64,  64, 2, true);
 
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256,  8, 128, 2,  64, 128, 128, 128, 2, true);
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16,  64, 4,  32, 128, 128, 128, 2, true);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16, 256, 1,  64, 128, 128, 128, 2, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 128, 2,  32, 128, 128, 128, 2, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 128, 2,  32, 128, 128, 128, 2, true);
 
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 32, 128, 2,  32, 128, 128, 128, 1, false);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 64, 256, 1,  32, 128, 128, 128, 1, false);
 
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512,  8,  64, 4,  32, 256, 256, 128, 1, false);
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 16,  64, 4,  32, 256, 256, 128, 1, false);
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 32, 128, 2,  32, 128, 128, 128, 1, false);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512,  8, 128, 2,  64, 128, 128, 128, 1, false);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 16, 256, 1,  64, 128, 128, 128, 1, false);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 32, 256, 1,  32, 128, 128, 128, 1, false);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 64, 256, 1,  32, 128, 128, 128, 1, false);
 
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512,  8,  64, 4,  32, 288, 256, 128, 1, false);
@@ -1066,7 +1066,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
 }
 
 #if defined(TURING_MMA_AVAILABLE)
-template<int DV, int ncols> struct mma_tile_sizes {
+template<int DKQ, int ncols> struct mma_tile_sizes {
     using T_A_KQ  = tile<16,  8, half2>; // row-major
     using T_B_KQ  = tile<16,  8, half2>; // column-major
     using T_C_KQ  = tile<16, 16, float>; // column-major
@@ -1074,7 +1074,33 @@ template<int DV, int ncols> struct mma_tile_sizes {
     using T_B_VKQ = tile<16,  8, half2>; // column-major
     using T_C_VKQ = tile<16,  8, half2>; // column-major
 };
-template<int DV> struct mma_tile_sizes<DV, 8> {
+// If there are only 8 columns, use thinner B tiles to avoid wasting compute:
+template<int DKQ> struct mma_tile_sizes<DKQ, 8> {
+    using T_A_KQ  = tile<16,  8, half2>; // row-major
+    using T_B_KQ  = tile< 8,  8, half2>; // column-major
+    using T_C_KQ  = tile<16,  8, float>; // row-major
+    using T_A_VKQ = tile<16,  8, half2>; // row-major
+    using T_B_VKQ = tile< 8,  8, half2>; // column-major
+    using T_C_VKQ = tile<16,  4, half2>; // row-major
+};
+// For very large head sizes, use thinner B tiles to reduce register pressure:
+template<> struct mma_tile_sizes<256, 16> {
+    using T_A_KQ  = tile<16,  8, half2>; // row-major
+    using T_B_KQ  = tile< 8,  8, half2>; // column-major
+    using T_C_KQ  = tile<16,  8, float>; // row-major
+    using T_A_VKQ = tile<16,  8, half2>; // row-major
+    using T_B_VKQ = tile< 8,  8, half2>; // column-major
+    using T_C_VKQ = tile<16,  4, half2>; // row-major
+};
+template<> struct mma_tile_sizes<512, 16> {
+    using T_A_KQ  = tile<16,  8, half2>; // row-major
+    using T_B_KQ  = tile< 8,  8, half2>; // column-major
+    using T_C_KQ  = tile<16,  8, float>; // row-major
+    using T_A_VKQ = tile<16,  8, half2>; // row-major
+    using T_B_VKQ = tile< 8,  8, half2>; // column-major
+    using T_C_VKQ = tile<16,  4, half2>; // row-major
+};
+template<> struct mma_tile_sizes<512, 32> {
     using T_A_KQ  = tile<16,  8, half2>; // row-major
     using T_B_KQ  = tile< 8,  8, half2>; // column-major
     using T_C_KQ  = tile<16,  8, float>; // row-major
@@ -1084,7 +1110,7 @@ template<int DV> struct mma_tile_sizes<DV, 8> {
 };
 #elif defined(AMD_WMMA_AVAILABLE)
 #ifdef RDNA3
-template<int DV, int ncols> struct mma_tile_sizes {
+template<int DKQ, int ncols> struct mma_tile_sizes {
     using T_A_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
     using T_B_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
     using T_C_KQ  = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
@@ -1109,7 +1135,7 @@ template<int ncols> struct mma_tile_sizes<112, ncols> {
     using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
 };
 #else
-template<int DV, int ncols> struct mma_tile_sizes {
+template<int DKQ, int ncols> struct mma_tile_sizes {
     using T_A_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR>;           // row-major
     using T_B_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR>;           // column-major
     using T_C_KQ  = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;           // column-major
@@ -1135,7 +1161,7 @@ template<int ncols> struct mma_tile_sizes<112, ncols> {
 };
 #endif // RDNA3
 #elif defined(AMD_MFMA_AVAILABLE)
-template<int DV, int ncols> struct mma_tile_sizes {
+template<int DKQ, int ncols> struct mma_tile_sizes {
     using T_A_KQ  = tile<16,  8, half2>; // row-major
     using T_B_KQ  = tile<16,  8, half2>; // column-major
     using T_C_KQ  = tile<16, 16, float>; // column-major
@@ -1144,7 +1170,7 @@ template<int DV, int ncols> struct mma_tile_sizes {
     using T_C_VKQ = tile<16, 16, float>; // column-major
 };
 #else // Volta
-template<int DV, int ncols> struct mma_tile_sizes {
+template<int DKQ, int ncols> struct mma_tile_sizes {
     using T_A_KQ  = tile< 8,  4, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
     using T_B_KQ  = tile<32,  4, half2, DATA_LAYOUT_I_MAJOR>;          // column-major
     using T_C_KQ  = tile<32,  8, float, DATA_LAYOUT_I_MAJOR>;          // column-major
@@ -1185,12 +1211,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
 
     constexpr int warp_size = ggml_cuda_get_physical_warp_size();
     constexpr int ncols = ncols1 * ncols2;
-    using     T_A_KQ    = typename mma_tile_sizes<DV, ncols>::T_A_KQ;
-    using     T_B_KQ    = typename mma_tile_sizes<DV, ncols>::T_B_KQ;
-    using     T_C_KQ    = typename mma_tile_sizes<DV, ncols>::T_C_KQ;
-    using     T_A_VKQ   = typename mma_tile_sizes<DV, ncols>::T_A_VKQ;
-    using     T_B_VKQ   = typename mma_tile_sizes<DV, ncols>::T_B_VKQ;
-    using     T_C_VKQ   = typename mma_tile_sizes<DV, ncols>::T_C_VKQ;
+    using     T_A_KQ    = typename mma_tile_sizes<DKQ, ncols>::T_A_KQ;
+    using     T_B_KQ    = typename mma_tile_sizes<DKQ, ncols>::T_B_KQ;
+    using     T_C_KQ    = typename mma_tile_sizes<DKQ, ncols>::T_C_KQ;
+    using     T_A_VKQ   = typename mma_tile_sizes<DKQ, ncols>::T_A_VKQ;
+    using     T_B_VKQ   = typename mma_tile_sizes<DKQ, ncols>::T_B_VKQ;
+    using     T_C_VKQ   = typename mma_tile_sizes<DKQ, ncols>::T_C_VKQ;
 
     constexpr int  cols_per_warp   = T_B_KQ::I;
     constexpr int  cols_per_thread = get_cols_per_thread();
@@ -1760,7 +1786,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
 static constexpr __host__ __device__ bool ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(
         const int DKQ, const int DV, const int ncols1, const int ncols2) {
     return (DKQ == 512 && DV == 512 && ncols1 == 1 && ncols2 == 8) ||
-           (DKQ == 576 && DV == 512 && ncols1 == 1 && ncols2 == 16);
+           (DKQ == 576 && DV == 512 && ncols1 == 1 && ncols2 == 16) ||
+           (DKQ == 256 && DV == 256 && ncols1 == 1 && ncols2 == 8) ||
+           (DKQ == 256 && DV == 256 && ncols1 == 8 && ncols2 == 8);
 }
 
 template<int DKQ, int DV, int ncols1, int ncols2, bool use_logit_softcap, bool V_is_K_view, bool use_sparse>
@@ -1794,8 +1822,9 @@ static __global__ void flash_attn_ext_f16(
     const char * GGML_CUDA_RESTRICT V              = V_ptr;
     const char * GGML_CUDA_RESTRICT mask           = mask_ptr;
     const char * GGML_CUDA_RESTRICT sinks          = sinks_ptr;
-    const int  * GGML_CUDA_RESTRICT KV_max         = use_sparse ? nullptr : KV_max_ptr;
+    // sparse: one index list per (sequence, query tile), the live count of each list follows the lists
     const int  * GGML_CUDA_RESTRICT sparse_indices = use_sparse ? KV_max_ptr : nullptr;
+    const int  * GGML_CUDA_RESTRICT KV_max         = KV_max_ptr;
     float      * GGML_CUDA_RESTRICT dst            = dst_ptr;
     float2     * GGML_CUDA_RESTRICT dst_meta       = dst_meta_ptr;
 
@@ -1860,6 +1889,10 @@ static __global__ void flash_attn_ext_f16(
     const int iter_j     = (ne01.z    + (ncols1    - 1)) / ncols1;
     const int iter_z_gqa = (gqa_ratio + (ncols2    - 1)) / ncols2;
 
+    if (use_sparse) {
+        KV_max = KV_max_ptr + int64_t(iter_j)*ne33*ne11;
+    }
+
     // kbc == k block continuous, current index in continuous ijk space.
     int       kbc      = int64_t(blockIdx.x + 0)*(iter_k*iter_j*iter_z_gqa*ne12*ne03) / gridDim.x;
     const int kbc_stop = int64_t(blockIdx.x + 1)*(iter_k*iter_j*iter_z_gqa*ne12*ne03) / gridDim.x;
@@ -1889,11 +1922,13 @@ static __global__ void flash_attn_ext_f16(
 
         const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV);
         const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr;
-        const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*ne31 + jt*ncols1)*ne11 : nullptr;
+        const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 : nullptr;
 
         const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f;
 
-        if (KV_max) {
+        if (use_sparse) {
+            kb0_stop = min(kb0_stop, (KV_max[(sequence % ne33)*iter_j + jt] + nbatch_fa - 1) / nbatch_fa);
+        } else if (KV_max) {
             kb0_stop = min(kb0_stop, KV_max[sequence*iter_j + jt] / nbatch_fa);
         }
         constexpr bool is_fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer.
@@ -1936,11 +1971,13 @@ static __global__ void flash_attn_ext_f16(
 
     const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV);
     const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr;
-    const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*ne31 + jt*ncols1)*ne11 : nullptr;
+    const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 : nullptr;
 
     const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f;
 
-    if (KV_max) {
+    if (use_sparse) {
+        kb0_stop = min(kb0_stop, (KV_max[(sequence % ne33)*iter_j + jt] + nbatch_fa - 1) / nbatch_fa);
+    } else if (KV_max) {
         kb0_stop = min(kb0_stop, KV_max[sequence*iter_j + jt] / nbatch_fa);
     }
 
@@ -1963,7 +2000,7 @@ static __global__ void flash_attn_ext_f16(
 #endif // defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE))
 }
 
-bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_tensor * dst, const int ncols1);
 
 template <int DKQ, int DV, int ncols1, int ncols2>
 void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
@@ -2016,7 +2053,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
         constexpr bool use_logit_softcap = false;
 #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
         if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, ncols1, ncols2)) {
-            if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ctx, dst)) {
+            if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, ncols1)) {
                 constexpr bool use_sparse_kernel = true;
                 fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;
                 use_sparse = true;

ggml/src/ggml-cuda/fattn.cu

@@ -6,10 +6,11 @@
 #include "fattn.cuh"
 
 #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+// one list per group of ncols1 queries: a column is selected if any query of the group can see it
 __launch_bounds__(256, 1)
 static __global__ void flash_attn_mask_to_sparse_indices(
-        const half * mask_ptr, int32_t * indices_ptr, const int ne30, const int n_kv_max,
-        const int64_t s31, const int64_t s33) {
+        const half * mask_ptr, int32_t * indices_ptr, int32_t * counts_ptr, const int ne30, const int n_queries,
+        const int ncols1, const int n_kv_max, const int64_t s31, const int64_t s33) {
     ggml_cuda_pdl_sync();
 
     constexpr int values_per_lane = 8;
@@ -17,10 +18,13 @@ static __global__ void flash_attn_mask_to_sparse_indices(
     const int warp     = tid / WARP_SIZE;
     const int lane     = tid % WARP_SIZE;
     const int sequence = blockIdx.y;
-    const int query    = blockIdx.x;
+    const int group    = blockIdx.x;
 
-    const half * mask = mask_ptr + sequence*s33 + query*s31;
-    int32_t * indices = indices_ptr + (int64_t(sequence)*gridDim.x + query)*n_kv_max;
+    const int q0 = group*ncols1;
+    const int q1 = min(q0 + ncols1, n_queries);
+
+    const half * mask = mask_ptr + sequence*s33 + q0*s31;
+    int32_t * indices = indices_ptr + (int64_t(sequence)*gridDim.x + group)*n_kv_max;
 
     __shared__ int warp_offsets[256/WARP_SIZE];
     __shared__ int row_count;
@@ -37,7 +41,10 @@ static __global__ void flash_attn_mask_to_sparse_indices(
 #pragma unroll
         for (int item = 0; item < values_per_lane; ++item) {
             const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane;
-            const bool selected = i < ne30 && isfinite(__half2float(mask[i]));
+            bool selected = false;
+            for (int q = 0; q < q1 - q0 && !selected; ++q) {
+                selected = i < ne30 && isfinite(__half2float(mask[q*s31 + i]));
+            }
             selected_warp[item] = __ballot_sync(0xFFFFFFFF, selected);
             warp_count += __popc(selected_warp[item]);
         }
@@ -78,10 +85,13 @@ static __global__ void flash_attn_mask_to_sparse_indices(
         __syncthreads();
     }
 
-    const int count = row_count;
+    const int count = min(row_count, n_kv_max);
     for (int i = count + tid; i < n_kv_max; i += blockDim.x) {
         indices[i] = -1;
     }
+    if (tid == 0) {
+        counts_ptr[int64_t(sequence)*gridDim.x + group] = count;
+    }
     __syncthreads();
 
     // the dependent grid reads indices, signal once the row is complete
@@ -90,31 +100,30 @@ static __global__ void flash_attn_mask_to_sparse_indices(
 #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
 
 void ggml_cuda_flash_attn_ext_compact_mask(
-        const ggml_tensor * mask, int32_t * indices, int32_t n_kv_max, cudaStream_t stream) {
+        const ggml_tensor * mask, int32_t * indices, int32_t * counts, int32_t n_queries, int32_t ncols1, int32_t n_kv_max, cudaStream_t stream) {
 #if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
-    GGML_UNUSED_VARS(mask, indices, n_kv_max, stream);
+    GGML_UNUSED_VARS(mask, indices, counts, n_queries, ncols1, n_kv_max, stream);
     GGML_ABORT("sparse flash attention is only supported on NVIDIA CUDA");
 #else
     const int64_t s31 = mask->nb[1] / sizeof(half);
     const int64_t s33 = mask->nb[3] / sizeof(half);
-    const dim3 blocks_num(mask->ne[1], mask->ne[3], 1);
+    const dim3 blocks_num((n_queries + ncols1 - 1)/ncols1, mask->ne[3], 1);
     const dim3 block_dim(256, 1, 1);
     const ggml_cuda_kernel_launch_params launch_params(blocks_num, block_dim, 0, stream);
     ggml_cuda_kernel_launch(flash_attn_mask_to_sparse_indices, launch_params,
-        (const half *) mask->data, indices, int(mask->ne[0]), n_kv_max, s31, s33);
+        (const half *) mask->data, indices, counts, int(mask->ne[0]), n_queries, ncols1, n_kv_max, s31, s33);
     CUDA_CHECK(cudaGetLastError());
 #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
 }
 
-bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_tensor * dst, const int ncols1) {
 #if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
-    GGML_UNUSED_VARS(ctx, dst);
+    GGML_UNUSED_VARS(cc, dst, ncols1);
     return false;
 #else
     const ggml_tensor * Q    = dst->src[0];
     const ggml_tensor * K    = dst->src[1];
     const ggml_tensor * mask = dst->src[3];
-    const int cc = ggml_cuda_info().devices[ctx.device].cc;
 
     float max_bias = 0.0f;
     float logit_softcap = 0.0f;
@@ -122,10 +131,13 @@ bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context
     memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float));
 
     const int32_t n_kv_max = ggml_get_op_params_i32(dst, 4);
+
+    const int64_t n_gather = (ncols1 == 1 ? Q->ne[1] : ncols1) * (int64_t) n_kv_max;
+
     return GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) &&
         mask != nullptr && n_kv_max > 0 && max_bias == 0.0f && logit_softcap == 0.0f &&
         mask->ne[0] == K->ne[1] && mask->ne[1] >= Q->ne[1] && mask->ne[2] == 1 &&
-        K->ne[1] >= std::max<int64_t>(4096, 2LL*n_kv_max);
+        K->ne[1] >= std::max<int64_t>(4096, 2*n_gather);
 #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
 }
 
@@ -136,7 +148,7 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con
 
 #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
     if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, 1, ncols2)) {
-        if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ctx, dst)) {
+        if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, 1)) {
             ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 1, ncols2>(ctx, dst);
             return;
         }
@@ -614,7 +626,12 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
     if (turing_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) {
         if (can_use_vector_kernel) {
             if (!ggml_is_quantized(K->type) && !ggml_is_quantized(V->type)) {
-                if (cc >= GGML_CUDA_CC_ADA_LOVELACE && Q->ne[1] == 1 && Q->ne[3] == 1 && !(gqa_ratio > 4 && K->ne[1] >= 8192)) {
+                // the sparse gather exists only in the MMA kernel: (DKQ, DV, 1, 8) with GQA > 4
+                const bool sparse_decode = gqa_opt_applies && gqa_ratio > 4 &&
+                    ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(K->ne[0], V->ne[0], 1, 8) &&
+                    ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, 1);
+                if (!sparse_decode && cc >= GGML_CUDA_CC_ADA_LOVELACE && Q->ne[1] == 1 && Q->ne[3] == 1 &&
+                        !(gqa_ratio > 4 && (Q->ne[0] >= 256 || K->ne[1] >= 8192))) {
                     return BEST_FATTN_KERNEL_VEC;
                 }
             } else {

ggml/src/ggml-cuda/mmvq.cu

@@ -365,6 +365,22 @@ bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11) {
                 return ne11 <= MMVQ_MAX_BATCH_SIZE;
         }
     }
+    if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_VOLTA) {
+        switch (type) {
+            case GGML_TYPE_Q2_K:
+                return ne11 <= 4;
+            case GGML_TYPE_Q3_K:
+                return ne11 <= 6;
+            case GGML_TYPE_Q4_K:
+                return ne11 <= 5;
+            case GGML_TYPE_Q5_K:
+                return ne11 <= 6;
+            case GGML_TYPE_Q6_K:
+                return ne11 <= 7;
+            default:
+                return ne11 <= MMVQ_MAX_BATCH_SIZE;
+        }
+    }
     if (GGML_CUDA_CC_IS_CDNA(cc)) {
         if (GGML_CUDA_CC_IS_CDNA1(cc)) {
             switch (type) {

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

@@ -14,7 +14,11 @@
 #include "htp/matmul-ops.h"
 #include "htp/flash-attn-ops.h"
 #include "htp/unary-ops.h"
+#include "htp/binary-ops.h"
 #include "htp/allreduce-ops.h"
+#include "htp/ssm-conv.h"
+#include "htp/gated-delta-net-ops.h"
+#include "htp/softmax-ops.h"
 
 struct htp_opnode {
     ggml_tensor * node   { nullptr };
@@ -325,10 +329,6 @@ struct htp_opformat {
             } else if (type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || type == HTP_MM_KERNEL_HVX_F32_F32_VTCM ||
                        type == HTP_MM_KERNEL_HVX_QUANT_ROW    || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
                 path = "hvx-tiled";
-            } else if (type == HTP_MM_KERNEL_HVX_F16_F16_DDR  || type == HTP_MM_KERNEL_HVX_F16_F32_DDR ||
-                       type == HTP_MM_KERNEL_HVX_F32_F32_DDR  || type == HTP_MM_KERNEL_HVX_F32_F16_DDR ||
-                       type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) {
-                path = "hvx-flat";
             }
             snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
         } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) {
@@ -350,6 +350,21 @@ struct htp_opformat {
             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 if (node.opcode == HTP_OP_SSM_CONV) {
+            const auto * kparams = (const struct htp_ssm_conv_kernel_params *) node.kernel_params;
+            snprintf(str, max_size, "%s vtcm %d", kparams->n_t == 1 ? "decode" : "prefill", (int) kparams->vtcm_size);
+        } else if (node.opcode == HTP_OP_SOFTMAX) {
+            const auto * kparams = (const struct htp_softmax_kernel_params *) node.kernel_params;
+            snprintf(str, max_size, "k%d nth %d vtcm %d", (int) kparams->kernel_id, (int) kparams->n_threads, (int) kparams->vtcm_size);
+        } else if (node.opcode == HTP_OP_GATED_DELTA_NET) {
+            const auto * kparams = (const struct htp_gdn_kernel_params *) node.kernel_params;
+            snprintf(str, max_size, "%s vtcm %u",
+                     kparams->kda ? "kda" : "scalar",
+                     (unsigned int) (kparams->vtcm_size ? kparams->vtcm_size : kparams->vtcm_per_thread * kparams->n_threads));
+        } else if (node.opcode == HTP_OP_MUL || node.opcode == HTP_OP_ADD || node.opcode == HTP_OP_ADD_ID ||
+                   node.opcode == HTP_OP_SUB || node.opcode == HTP_OP_DIV) {
+            const auto * kparams = (const struct htp_binary_kernel_params *) node.kernel_params;
+            snprintf(str, max_size, "vtcm %u", (unsigned int) kparams->vtcm_size);
         } else {
             snprintf(str, max_size, "----");
         }

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

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

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

@@ -7,7 +7,7 @@
 #include <math.h>
 #include <string.h>
 
-#include "hex-dma.h"
+#include "dma-queue.h"
 #include "hvx-utils.h"
 
 #define GGML_COMMON_DECL_C
@@ -53,13 +53,24 @@
     const uint32_t nb2 = dst->nb[2];                     \
     const uint32_t nb3 = dst->nb[3];
 
+struct htp_act_context;
+
+typedef void (*glu_compute_fn_t)(const float * restrict src0,
+                                 const float * restrict src1,
+                                 float * restrict dst,
+                                 const uint32_t num_rows,
+                                 const struct htp_act_context * actx);
+
 struct htp_act_context {
     struct htp_ops_context * octx;
 
+    glu_compute_fn_t         compute;
+    const char *             op_str;
+
     // Precomputed values
-    const uint8_t *          data_src0;
-    const uint8_t *          data_src1;
-    uint8_t *                data_dst;
+    dma_addr_t               data_src0;
+    dma_addr_t               data_src1;
+    dma_addr_t               data_dst;
 
     size_t                   src0_row_size;
     size_t                   src1_row_size;
@@ -134,10 +145,10 @@ static inline void htp_act_vtcm_layout_build(struct htp_act_vtcm_layout * L,
 
 // swiglu(x) = x1 * sigmoid(x0)
 static void swiglu_f32(const float * restrict src0,
-                       const float * restrict src1,
-                       float * restrict dst,
-                       const uint32_t num_rows,
-                       const struct htp_act_context * actx) {
+                                                 const float * restrict src1,
+                                                 float * restrict dst,
+                                                 const uint32_t num_rows,
+                                                 const struct htp_act_context * actx) {
     htp_glu_op_preamble;
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -152,10 +163,10 @@ static void swiglu_f32(const float * restrict src0,
 
 // out = x * sigmoid(alpha * x) * (clamp(y, -limit, limit) + 1.f)
 static void swiglu_oai_f32(const float * restrict src0,
-                           const float * restrict src1,
-                           float * restrict dst,
-                           const uint32_t num_rows,
-                           const struct htp_act_context * actx) {
+                                                     const float * restrict src1,
+                                                     float * restrict dst,
+                                                     const uint32_t num_rows,
+                                                     const struct htp_act_context * actx) {
     htp_glu_op_preamble;
     const float alpha = ((const float *) (actx->octx->op_params))[2];
     const float limit = ((const float *) (actx->octx->op_params))[3];
@@ -181,10 +192,10 @@ static void swiglu_oai_f32(const float * restrict src0,
 }
 
 static void swiglu_clamp_f32(const float * restrict src0,
-                             const float * restrict src1,
-                             float * restrict dst,
-                             const uint32_t                 num_rows,
-                             const struct htp_act_context * actx) {
+                                                       const float * restrict src1,
+                                                       float * restrict dst,
+                                                       const uint32_t                 num_rows,
+                                                       const struct htp_act_context * actx) {
     htp_glu_op_preamble;
     const float limit = ((const float *) (actx->octx->op_params))[3];
 
@@ -312,12 +323,51 @@ static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * rest
     }
 }
 
+static inline void hvx_geglu_quick_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
+    assert((unsigned long) dst  % 128 == 0);
+    assert((unsigned long) src0 % 128 == 0);
+    assert((unsigned long) src1 % 128 == 0);
+
+    HVX_Vector * restrict vdst        = (HVX_Vector *) dst;
+    const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
+    const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
+
+    const uint32_t epv  = 128 / sizeof(float);
+    const uint32_t nvec = n / epv;
+    const uint32_t nloe = n % epv;
+
+    const HVX_Vector v_scale    = hvx_vec_splat_f32(1.702f);
+    const HVX_Vector v_one      = hvx_vec_splat_f32(1.0f);
+    const HVX_Vector v_max_exp  = hvx_vec_splat_f32(87.0f);
+    const HVX_Vector v_min_exp  = hvx_vec_splat_f32(-87.0f);
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        HVX_Vector x = vsrc0[i];
+        HVX_Vector g = vsrc1[i];
+        HVX_Vector scaled_x = hvx_vec_mul_f32_f32(x, v_scale);
+        HVX_Vector sigmoid_x = hvx_vec_fast_sigmoid_f32_guard_2it(scaled_x, v_one, v_max_exp, v_min_exp);
+        vdst[i] = hvx_vec_mul_f32_f32(hvx_vec_mul_f32_f32(x, sigmoid_x), g);
+    }
+
+    if (nloe) {
+        HVX_Vector x = vsrc0[i];
+        HVX_Vector g = vsrc1[i];
+        HVX_Vector scaled_x = hvx_vec_mul_f32_f32(x, v_scale);
+        HVX_Vector sigmoid_x = hvx_vec_fast_sigmoid_f32_guard_2it(scaled_x, v_one, v_max_exp, v_min_exp);
+        HVX_Vector result = hvx_vec_mul_f32_f32(hvx_vec_mul_f32_f32(x, sigmoid_x), g);
+        hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(float), result);
+    }
+}
+
 // geglu(x, g) = gelu(x) * g
 static void geglu_f32(const float * restrict src0,
-                      const float * restrict src1,
-                      float * restrict dst,
-                      const uint32_t num_rows,
-                      const struct htp_act_context * actx) {
+                                                const float * restrict src1,
+                                                float * restrict dst,
+                                                const uint32_t num_rows,
+                                                const struct htp_act_context * actx) {
     htp_glu_op_preamble;
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -329,110 +379,117 @@ 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 = 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);                                                                                      \
-                                                                                                                         \
+// geglu_quick(x, g) = x * sigmoid(1.702 * x) * g
+static void geglu_quick_f32(const float * restrict src0,
+                            const float * restrict src1,
+                            float * restrict dst,
+                            const uint32_t num_rows,
+                            const struct htp_act_context * actx) {
+    htp_glu_op_preamble;
+
+    for (uint32_t ib = 0; ib < num_rows; ib++) {
+        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
+        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
+        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
+
+        hvx_geglu_quick_f32_aa(dst_ptr, src0_ptr, src1_ptr, nc);
+    }
+}
+
+static void glu_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;
     }
 
-DEFINE_GLU_PER_THREAD(swiglu, "swiglu-f32", swiglu_f32(src0_spad, src1_spad, dst_spad, block_size, actx))
-DEFINE_GLU_PER_THREAD(swiglu_oai, "swiglu-oai-f32", swiglu_oai_f32(src0_spad, src1_spad, dst_spad, block_size, actx))
-DEFINE_GLU_PER_THREAD(swiglu_clamp, "swiglu-clamp-f32", swiglu_clamp_f32(src0_spad, src1_spad, dst_spad, block_size, actx))
-DEFINE_GLU_PER_THREAD(geglu, "geglu-f32", geglu_f32(src0_spad, src1_spad, dst_spad, block_size, actx))
+    const dma_addr_t data_src0 = actx->data_src0;
+    const dma_addr_t data_src1 = actx->data_src1;
+    const dma_addr_t 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, "%s : VTCM reservation %zu is too small, needed %zu\n",
+             actx->op_str, actx->vtcm_src0_size_per_thread, src0_row_size_aligned);
+        return;
+    }
+
+    dma_queue * dma_q = actx->octx->ctx->dma[ith];
+    glu_compute_fn_t compute = actx->compute;
+
+    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(dma_q,
+                       dma_make_data(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)),
+                       dst_row_size, dst_row_size_aligned, dst_row_size, 0);
+
+        dma_queue_push(dma_q,
+            dma_make_data(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_q,
+            dma_make_data(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_q).src;
+        float * src0_spad = (float *) dma_queue_pop(dma_q).dst;
+        float * src1_spad = (float *) dma_queue_pop(dma_q).dst;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir);
+        compute(src0_spad, src1_spad, dst_spad, block_size, actx);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir);
+
+        dma_queue_push(dma_q, dma_make_data(data_dst + (ir * dst_row_size), dst_spad),
+                       dst_row_size, dst_row_size_aligned, dst_row_size, 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_q, dma_make_data(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_q, dma_make_data(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_q);
+}
 
 static int execute_op_activations_f32(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
@@ -444,28 +501,33 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    worker_callback_t act_op_func;
-    const char *      op_type = NULL;
+    glu_compute_fn_t compute_fn = NULL;
+    const char *      op_type    = NULL;
 
     switch (octx->op) {
         case HTP_OP_GLU_SWIGLU:
-            act_op_func = (worker_callback_t)glu_swiglu_f32_per_thread;
-            op_type     = "swiglu-f32";
+            compute_fn = swiglu_f32;
+            op_type    = "swiglu-f32";
             break;
 
         case HTP_OP_GLU_SWIGLU_OAI:
-            act_op_func = (worker_callback_t)glu_swiglu_oai_f32_per_thread;
-            op_type     = "swiglu-oai-f32";
+            compute_fn = swiglu_oai_f32;
+            op_type    = "swiglu-oai-f32";
             break;
 
         case HTP_OP_GLU_SWIGLU_CLAMP:
-            act_op_func = (worker_callback_t) glu_swiglu_clamp_f32_per_thread;
-            op_type     = "swiglu-clamp-f32";
+            compute_fn = swiglu_clamp_f32;
+            op_type    = "swiglu-clamp-f32";
             break;
 
         case HTP_OP_GLU_GEGLU:
-            act_op_func = (worker_callback_t)glu_geglu_f32_per_thread;
-            op_type     = "geglu-f32";
+            compute_fn = geglu_f32;
+            op_type    = "geglu-f32";
+            break;
+
+        case HTP_OP_GLU_GEGLU_QUICK:
+            compute_fn = geglu_quick_f32;
+            op_type    = "geglu-quick-f32";
             break;
         default:
             FARF(ERROR, "Unsupported activations Op %u\n", octx->op);
@@ -526,13 +588,11 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
              L.src0_bytes_per_thread * n_threads, L.src1_bytes_per_thread * n_threads, L.dst_bytes_per_thread * n_threads);
     }
 
-    if ((octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) {
-        return HTP_STATUS_OK;
-    }
-
     // Prepare context
     struct htp_act_context actx;
-    actx.octx = octx;
+    actx.octx    = octx;
+    actx.compute = compute_fn;
+    actx.op_str  = op_type;
 
     actx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
 
@@ -566,11 +626,15 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
 
     actx.nc = dst->ne[0];
 
-    // Pointers and GLU logic
-    const uint8_t * data_src0 = (const uint8_t *) src0->data;
-    const uint8_t * data_src1 = src1 ? (const uint8_t *) src1->data : NULL;
+    // Addresses and GLU logic
+    dma_addr_t data_src0 = src0->data;
+    dma_addr_t data_src1 = src1 ? src1->data : 0;
 
-    if (!src1 && (octx->op == HTP_OP_GLU_SWIGLU || octx->op == HTP_OP_GLU_SWIGLU_OAI || octx->op == HTP_OP_GLU_SWIGLU_CLAMP || octx->op == HTP_OP_GLU_GEGLU)) {
+    if (!src1 && (octx->op == HTP_OP_GLU_SWIGLU ||
+                  octx->op == HTP_OP_GLU_SWIGLU_OAI ||
+                  octx->op == HTP_OP_GLU_SWIGLU_CLAMP ||
+                  octx->op == HTP_OP_GLU_GEGLU ||
+                  octx->op == HTP_OP_GLU_GEGLU_QUICK)) {
          const int32_t swapped = octx->op_params[1];
          data_src1 = data_src0;
          actx.src1_row_size = actx.src0_row_size;
@@ -585,9 +649,9 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
 
     actx.data_src0 = data_src0;
     actx.data_src1 = data_src1;
-    actx.data_dst  = (uint8_t *) dst->data;
+    actx.data_dst  = dst->data;
 
-    work_queue_run(octx->ctx->work_queue, act_op_func, &actx, n_threads);
+    work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f32_per_thread, &actx, n_threads);
     return HTP_STATUS_OK;
 }
 

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

@@ -14,7 +14,7 @@
 #include "htp-ops.h"
 #include "hvx-utils.h"
 #include "htp-tensor.h"
-#include "hex-dma.h"
+#include "dma-queue.h"
 #include "hex-profile.h"
 #include "allreduce-ops.h"
 #include "htp-fence.h"
@@ -38,97 +38,97 @@ struct htp_allreduce_context {
     uint8_t * res_spad_base;
 };
 
-#define DEFINE_ALLREDUCE_THREAD_DMA_1D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD)                                             \
-static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                       \
-    struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data;                                      \
-    struct htp_ops_context * octx = actx->octx;                                                                       \
-                                                                                                                      \
-    const uint32_t n_ranks     = actx->n_ranks;                                                                       \
-    const uint32_t n_dsts      = actx->n_dsts;                                                                        \
-    const uint32_t block_elems = actx->block_elems;                                                                   \
-                                                                                                                      \
-    const uint32_t dr  = actx->elems_per_thread;                                                                      \
-    const uint32_t ir0 = actx->rank_elem_start + dr * ith;                                                            \
-    const uint32_t ir1 = MIN(ir0 + dr, actx->rank_elem_start + actx->rank_nelem);                                     \
-    if (ir0 >= ir1) return;                                                                                           \
-                                                                                                                      \
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                            \
-    dma_queue * q = octx->ctx->dma[ith];                                                                              \
-                                                                                                                      \
-    uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];                                                                 \
-    for (uint32_t s = 0; s < n_ranks; s++) {                                                                          \
-        src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread);                               \
-    }                                                                                                                 \
-    uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread);                               \
-    uint8_t * res_spad_base = HAS_ADD ? (actx->res_spad_base + (ith * actx->vtcm_size_per_thread)) : NULL;            \
-                                                                                                                      \
-    const size_t spad_half = actx->vtcm_size_per_thread / 2;                                                          \
-    uint32_t ir_prefetch = ir0;                                                                                       \
-    int spad_idx = 0;                                                                                                 \
-                                                                                                                      \
-    for (int k = 0; k < 2 && ir_prefetch < ir1; k++) {                                                                \
-        uint32_t cur_elems = MIN(block_elems, ir1 - ir_prefetch);                                                     \
-        size_t   cur_bytes = cur_elems * sizeof(TYPE);                                                                \
-        uint8_t * d_spad = dst_spad_base + spad_idx * spad_half;                                                      \
-        for (uint32_t d = 0; d < n_dsts; d++) {                                                                       \
-            uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir_prefetch * sizeof(TYPE);                           \
-            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 0);                       \
-        }                                                                                                             \
-        for (uint32_t s = 0; s < n_ranks; s++) {                                                                      \
-            uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half;                                               \
-            const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE);                \
-            dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), cur_bytes, cur_bytes, cur_bytes, 1);                       \
-        }                                                                                                             \
-        if (HAS_ADD) {                                                                                                \
-            uint8_t * r_spad = res_spad_base + spad_idx * spad_half;                                                  \
-            const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE);      \
-            dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), cur_bytes, cur_bytes, cur_bytes, 1);                       \
-        }                                                                                                             \
-        ir_prefetch += cur_elems;                                                                                     \
-        spad_idx ^= 1;                                                                                                \
-    }                                                                                                                 \
-                                                                                                                      \
-    for (uint32_t ir = ir0; ir < ir1; ) {                                                                             \
-        uint32_t cur_elems = MIN(block_elems, ir1 - ir);                                                              \
-        size_t   cur_bytes = cur_elems * sizeof(TYPE);                                                                \
-        uint8_t * d_spad = NULL;                                                                                      \
-        for (uint32_t d = 0; d < n_dsts; d++) {                                                                       \
-            d_spad = (uint8_t *) dma_queue_pop(q).src;                                                                \
-        }                                                                                                             \
-        uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS];                                                                    \
-        for (uint32_t s = 0; s < n_ranks; s++) {                                                                      \
-            s_spad[s] = (uint8_t *) dma_queue_pop(q).dst;                                                             \
-        }                                                                                                             \
-        uint8_t * r_spad = HAS_ADD ? (uint8_t *) dma_queue_pop(q).dst : NULL;                                         \
-        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);                                             \
-        HVX_ADD_FN(d_spad, s_spad[0], s_spad[1], cur_elems);                                                          \
-        for (uint32_t s = 2; s < n_ranks; s++) {                                                                      \
-            HVX_ADD_FN(d_spad, d_spad, s_spad[s], cur_elems);                                                         \
-        }                                                                                                             \
-        if (HAS_ADD) {                                                                                                \
-            HVX_ADD_FN(d_spad, d_spad, r_spad, cur_elems);                                                            \
-        }                                                                                                             \
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);                                              \
-        for (uint32_t d = 0; d < n_dsts; d++) {                                                                       \
-            uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir * sizeof(TYPE);                                    \
-            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 1);                       \
-        }                                                                                                             \
-        if (ir_prefetch < ir1) {                                                                                      \
-            uint32_t next_elems = MIN(block_elems, ir1 - ir_prefetch);                                                \
-            size_t   next_bytes = next_elems * sizeof(TYPE);                                                          \
-            for (uint32_t s = 0; s < n_ranks; s++) {                                                                  \
-                const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE);           \
-                dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), next_bytes, next_bytes, next_bytes, 1);            \
-            }                                                                                                         \
-            if (HAS_ADD) {                                                                                            \
-                const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE); \
-                dma_queue_push(q, dma_make_ptr(r_spad, r_next), next_bytes, next_bytes, next_bytes, 1);               \
-            }                                                                                                         \
-            ir_prefetch += next_elems;                                                                                \
-        }                                                                                                             \
-        ir += cur_elems;                                                                                              \
-    }                                                                                                                 \
-    dma_queue_flush(q);                                                                                               \
+#define DEFINE_ALLREDUCE_THREAD_DMA_1D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD)                                       \
+static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                 \
+    struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data;                                \
+    struct htp_ops_context * octx = actx->octx;                                                                 \
+                                                                                                                \
+    const uint32_t n_ranks     = actx->n_ranks;                                                                 \
+    const uint32_t n_dsts      = actx->n_dsts;                                                                  \
+    const uint32_t block_elems = actx->block_elems;                                                             \
+                                                                                                                \
+    const uint32_t dr  = actx->elems_per_thread;                                                                \
+    const uint32_t ir0 = actx->rank_elem_start + dr * ith;                                                      \
+    const uint32_t ir1 = MIN(ir0 + dr, actx->rank_elem_start + actx->rank_nelem);                               \
+    if (ir0 >= ir1) return;                                                                                     \
+                                                                                                                \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                      \
+    dma_queue * dma_q = octx->ctx->dma[ith];                                                                    \
+                                                                                                                \
+    uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];                                                           \
+    for (uint32_t s = 0; s < n_ranks; s++) {                                                                    \
+        src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread);                         \
+    }                                                                                                           \
+    uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread);                         \
+    uint8_t * res_spad_base = HAS_ADD ? (actx->res_spad_base + (ith * actx->vtcm_size_per_thread)) : NULL;      \
+                                                                                                                \
+    const size_t spad_half = actx->vtcm_size_per_thread / 2;                                                    \
+    uint32_t ir_prefetch = ir0;                                                                                 \
+    int spad_idx = 0;                                                                                           \
+                                                                                                                \
+    for (int k = 0; k < 2 && ir_prefetch < ir1; k++) {                                                          \
+        uint32_t cur_elems = MIN(block_elems, ir1 - ir_prefetch);                                               \
+        size_t   cur_bytes = cur_elems * sizeof(TYPE);                                                          \
+        uint8_t * d_spad = dst_spad_base + spad_idx * spad_half;                                                \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                 \
+            dma_addr_t d_ddr = octx->dsts[d]->data + ir_prefetch * sizeof(TYPE);                                \
+            dma_queue_push(dma_q, dma_make_data(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 0);            \
+        }                                                                                                       \
+        for (uint32_t s = 0; s < n_ranks; s++) {                                                                \
+            uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half;                                         \
+            const dma_addr_t s_ddr = octx->src[s]->data + ir_prefetch * sizeof(TYPE);                           \
+            dma_queue_push(dma_q, dma_make_data(s_spad, s_ddr), cur_bytes, cur_bytes, cur_bytes, 1);            \
+        }                                                                                                       \
+        if (HAS_ADD) {                                                                                          \
+            uint8_t * r_spad = res_spad_base + spad_idx * spad_half;                                            \
+            const dma_addr_t r_ddr = octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE);                 \
+            dma_queue_push(dma_q, dma_make_data(r_spad, r_ddr), cur_bytes, cur_bytes, cur_bytes, 1);            \
+        }                                                                                                       \
+        ir_prefetch += cur_elems;                                                                               \
+        spad_idx ^= 1;                                                                                          \
+    }                                                                                                           \
+                                                                                                                \
+    for (uint32_t ir = ir0; ir < ir1; ) {                                                                       \
+        uint32_t cur_elems = MIN(block_elems, ir1 - ir);                                                        \
+        size_t   cur_bytes = cur_elems * sizeof(TYPE);                                                          \
+        uint8_t * d_spad = NULL;                                                                                \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                 \
+            d_spad = (uint8_t *) dma_queue_pop(dma_q).src;                                                      \
+        }                                                                                                       \
+        uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS];                                                              \
+        for (uint32_t s = 0; s < n_ranks; s++) {                                                                \
+            s_spad[s] = (uint8_t *) dma_queue_pop(dma_q).dst;                                                   \
+        }                                                                                                       \
+        uint8_t * r_spad = HAS_ADD ? (uint8_t *) dma_queue_pop(dma_q).dst : NULL;                               \
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);                                       \
+        HVX_ADD_FN(d_spad, s_spad[0], s_spad[1], cur_elems);                                                    \
+        for (uint32_t s = 2; s < n_ranks; s++) {                                                                \
+            HVX_ADD_FN(d_spad, d_spad, s_spad[s], cur_elems);                                                   \
+        }                                                                                                       \
+        if (HAS_ADD) {                                                                                          \
+            HVX_ADD_FN(d_spad, d_spad, r_spad, cur_elems);                                                      \
+        }                                                                                                       \
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);                                        \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                 \
+            dma_addr_t d_ddr = octx->dsts[d]->data + ir * sizeof(TYPE);                                         \
+            dma_queue_push(dma_q, dma_make_data(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 1);            \
+        }                                                                                                       \
+        if (ir_prefetch < ir1) {                                                                                \
+            uint32_t next_elems = MIN(block_elems, ir1 - ir_prefetch);                                          \
+            size_t   next_bytes = next_elems * sizeof(TYPE);                                                    \
+            for (uint32_t s = 0; s < n_ranks; s++) {                                                            \
+                const dma_addr_t s_next = octx->src[s]->data + ir_prefetch * sizeof(TYPE);                      \
+                dma_queue_push(dma_q, dma_make_data(s_spad[s], s_next), next_bytes, next_bytes, next_bytes, 1); \
+            }                                                                                                   \
+            if (HAS_ADD) {                                                                                      \
+                const dma_addr_t r_next = octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE);            \
+                dma_queue_push(dma_q, dma_make_data(r_spad, r_next), next_bytes, next_bytes, next_bytes, 1);    \
+            }                                                                                                   \
+            ir_prefetch += next_elems;                                                                          \
+        }                                                                                                       \
+        ir += cur_elems;                                                                                        \
+    }                                                                                                           \
+    dma_queue_flush(dma_q);                                                                                     \
 }
 
 DEFINE_ALLREDUCE_THREAD_DMA_1D(f16,     __fp16, hvx_add_f16_aaa, 0)
@@ -154,7 +154,7 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith,
     if (r0 >= r1) return;                                                                                                                         \
                                                                                                                                                   \
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                                                        \
-    dma_queue * q = octx->ctx->dma[ith];                                                                                                          \
+    dma_queue * dma_q = octx->ctx->dma[ith];                                                                                                      \
                                                                                                                                                   \
     uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];                                                                                             \
     for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                      \
@@ -171,18 +171,18 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith,
         uint32_t cur_rows = MIN(block_rows, r1 - r_prefetch);                                                                                     \
         uint8_t * d_spad = dst_spad_base + spad_idx * spad_half;                                                                                  \
         for (uint32_t d = 0; d < n_dsts; d++) {                                                                                                   \
-            uint8_t * d_ddr  = (uint8_t *) octx->dsts[d]->data + r_prefetch * octx->dsts[d]->nb[1];                                               \
-            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, 0);                                 \
+            dma_addr_t d_ddr = octx->dsts[d]->data + r_prefetch * octx->dsts[d]->nb[1];                                                           \
+            dma_queue_push(dma_q, dma_make_data(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, 0);                            \
         }                                                                                                                                         \
         for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                  \
             uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half;                                                                           \
-            const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1];                                      \
-            dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), row_size_aligned, octx->src[s]->nb[1], row_bytes, cur_rows);                           \
+            const dma_addr_t s_ddr = octx->src[s]->data + r_prefetch * octx->src[s]->nb[1];                                                       \
+            dma_queue_push(dma_q, dma_make_data(s_spad, s_ddr), row_size_aligned, octx->src[s]->nb[1], row_bytes, cur_rows);                      \
         }                                                                                                                                         \
         if (HAS_ADD && !IS_ROW_BCAST) {                                                                                                           \
             uint8_t * r_spad = res_spad_base + spad_idx * spad_half;                                                                              \
-            const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1];                  \
-            dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, cur_rows);                 \
+            const dma_addr_t r_ddr = octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1];                                   \
+            dma_queue_push(dma_q, dma_make_data(r_spad, r_ddr), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, cur_rows);            \
         }                                                                                                                                         \
         r_prefetch += cur_rows;                                                                                                                   \
         spad_idx ^= 1;                                                                                                                            \
@@ -192,13 +192,13 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith,
         uint32_t cur_rows = MIN(block_rows, r1 - r);                                                                                              \
         uint8_t * d_spad = NULL;                                                                                                                  \
         for (uint32_t d = 0; d < n_dsts; d++) {                                                                                                   \
-            d_spad = (uint8_t *) dma_queue_pop(q).src;                                                                                            \
+            d_spad = (uint8_t *) dma_queue_pop(dma_q).src;                                                                                        \
         }                                                                                                                                         \
         uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS];                                                                                                \
         for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                  \
-            s_spad[s] = (uint8_t *) dma_queue_pop(q).dst;                                                                                         \
+            s_spad[s] = (uint8_t *) dma_queue_pop(dma_q).dst;                                                                                     \
         }                                                                                                                                         \
-        uint8_t * r_spad = (HAS_ADD && !IS_ROW_BCAST) ? (uint8_t *) dma_queue_pop(q).dst : NULL;                                                  \
+        uint8_t * r_spad = (HAS_ADD && !IS_ROW_BCAST) ? (uint8_t *) dma_queue_pop(dma_q).dst : NULL;                                              \
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);                                                                          \
         for (uint32_t row = 0; row < cur_rows; row++) {                                                                                           \
             uint8_t * d_row = d_spad + row * row_size_aligned;                                                                                    \
@@ -216,24 +216,24 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith,
         }                                                                                                                                         \
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);                                                                           \
         for (uint32_t d = 0; d < n_dsts; d++) {                                                                                                   \
-            uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + r * octx->dsts[d]->nb[1];                                                         \
-            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, cur_rows);                          \
+            dma_addr_t d_ddr = octx->dsts[d]->data + r * octx->dsts[d]->nb[1];                                                                    \
+            dma_queue_push(dma_q, dma_make_data(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, cur_rows);                     \
         }                                                                                                                                         \
         if (r_prefetch < r1) {                                                                                                                    \
             uint32_t next_rows = MIN(block_rows, r1 - r_prefetch);                                                                                \
             for (uint32_t s = 0; s < n_ranks; s++) {                                                                                              \
-                const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1];                                 \
-                dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), row_size_aligned, octx->src[s]->nb[1], row_bytes, next_rows);                  \
+                const dma_addr_t s_next = octx->src[s]->data + r_prefetch * octx->src[s]->nb[1];                                                  \
+                dma_queue_push(dma_q, dma_make_data(s_spad[s], s_next), row_size_aligned, octx->src[s]->nb[1], row_bytes, next_rows);             \
             }                                                                                                                                     \
             if (HAS_ADD && !IS_ROW_BCAST) {                                                                                                       \
-                const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1];             \
-                dma_queue_push(q, dma_make_ptr(r_spad, r_next), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, next_rows);           \
+                const dma_addr_t r_next = octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1];                              \
+                dma_queue_push(dma_q, dma_make_data(r_spad, r_next), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, next_rows);      \
             }                                                                                                                                     \
             r_prefetch += next_rows;                                                                                                              \
         }                                                                                                                                         \
         r += cur_rows;                                                                                                                            \
     }                                                                                                                                             \
-    dma_queue_flush(q);                                                                                                                           \
+    dma_queue_flush(dma_q);                                                                                                                       \
 }
 
 DEFINE_ALLREDUCE_THREAD_DMA_2D(f16,           __fp16, hvx_add_f16_aaa, 0, 0)
@@ -406,11 +406,11 @@ int op_allreduce(struct htp_ops_context * octx) {
         }
 
         if (has_add && actx.is_row_bcast) {
-            const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data;
+            const dma_addr_t r_ddr = octx->src[2 * n_ranks]->data;
             const uint32_t row_bytes = actx.ne0 * (dst->type == HTP_TYPE_F16 ? sizeof(__fp16) : sizeof(float));
-            dma_queue * q = octx->ctx->dma[0];
-            dma_queue_push(q, dma_make_ptr(actx.res_spad_base, r_ddr), actx.row_size_aligned, 0, row_bytes, 1);
-            dma_queue_pop(q);
+            dma_queue * dma_q = octx->ctx->dma[0];
+            dma_queue_push(dma_q, dma_make_data(actx.res_spad_base, r_ddr), actx.row_size_aligned, 0, row_bytes, 1);
+            dma_queue_pop(dma_q);
         }
 
         work_queue_run(octx->ctx->work_queue, reduce_fun, &actx, n_threads);

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

@@ -9,7 +9,7 @@
 #include "ggml.h"
 
 #include "hvx-utils.h"
-#include "hex-dma.h"
+#include "dma-queue.h"
 
 #include "hex-common.h"
 #include "htp-ctx.h"
@@ -170,6 +170,21 @@ static void quicksort_values_indices_desc(float * values, int32_t * indices, int
     if (i < right) quicksort_values_indices_desc(values, indices, i, right);
 }
 
+static uint32_t top_k_max_value_index(const float * values, uint32_t n, float * value) {
+    uint32_t index = 0;
+    float max_value = values[0];
+
+    for (uint32_t i = 1; i < n; i++) {
+        if (values[i] > max_value) {
+            max_value = values[i];
+            index = i;
+        }
+    }
+
+    *value = max_value;
+    return index;
+}
+
 // LUT for ramp initialization of argsort output (first 32 members)
 int32_t argosrt_ramp_lut[32] __attribute__((aligned(VLEN))) = {
     0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
@@ -303,6 +318,129 @@ static inline void bitonic_sort_generic_hvx(uint8_t * values, uint8_t * indices,
     }
 }
 
+
+// Sorts descending; values pre-padded with -INFINITY. init_indices=true resets
+// indices to a fresh ramp (normal full-row sort); false leaves caller-supplied
+// indices in place and only permutes them (used when merging candidates, to
+// preserve their original global index).
+
+static void bitonic_sort_vtcm_desc(uint8_t * values, uint8_t * indices, uint32_t n_vec, bool init_indices) {
+    HVX_Vector zero_vec = Q6_V_vzero();
+    HVX_Vector idx_vec = *(HVX_Vector *)argosrt_ramp_lut;
+
+    HVX_VectorPred pred_all_1s = Q6_Q_vcmp_eq_VwVw(zero_vec, zero_vec);
+    HVX_VectorPred pred_all_0s = Q6_Q_not_Q(pred_all_1s);
+
+    if (init_indices) {
+        // Initialize indices ramp (values are already populated by the caller)
+        for (uint32_t v = 0; v < n_vec; v++) {
+            HVX_Vector idx = Q6_Vw_vadd_VwVw(idx_vec, Q6_V_vsplat_R(v * 32));
+            *(HVX_Vector *)(indices + v * 128) = idx;
+        }
+    }
+
+    int M = 5;
+    while ((1u << (M - 5)) < n_vec) M++;
+
+    for (int s = 1; s <= M; s++) {
+        for (int stage_d = s - 1; stage_d >= 0; stage_d--) {
+            int d = 1 << stage_d;
+            if (d >= 32) {
+                uint32_t v_dist = d / 32;
+                for (uint32_t v1 = 0; v1 < n_vec; v1++) {
+                    if ((v1 & v_dist) == 0) {
+                        uint32_t v2 = v1 + v_dist;
+                        bool asc = (s < M) ? ((((v1 * 32) >> s) % 2) == 0) : false;
+
+                        HVX_Vector Vv1 = *(HVX_Vector *)(values + v1 * 128);
+                        HVX_Vector Iv1 = *(HVX_Vector *)(indices + v1 * 128);
+                        HVX_Vector Vv2 = *(HVX_Vector *)(values + v2 * 128);
+                        HVX_Vector Iv2 = *(HVX_Vector *)(indices + v2 * 128);
+
+                        vec_cas(&Vv1, &Iv1, &Vv2, &Iv2, asc);
+
+                        *(HVX_Vector *)(values + v1 * 128)  = Vv1;
+                        *(HVX_Vector *)(indices + v1 * 128) = Iv1;
+                        *(HVX_Vector *)(values + v2 * 128)  = Vv2;
+                        *(HVX_Vector *)(indices + v2 * 128) = Iv2;
+                    }
+                }
+            } else {
+                if (s < 5) {
+                    HVX_VectorPred dir_mask = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(1 << s)), zero_vec);
+                    for (uint32_t v = 0; v < n_vec; v++) {
+                        HVX_Vector Vv = *(HVX_Vector *)(values + v * 128);
+                        HVX_Vector Iv = *(HVX_Vector *)(indices + v * 128);
+
+                        bitonic_cas_32(&Vv, &Iv, d, dir_mask, idx_vec, zero_vec);
+
+                        *(HVX_Vector *)(values + v * 128)  = Vv;
+                        *(HVX_Vector *)(indices + v * 128) = Iv;
+                    }
+                } else {
+                    for (uint32_t v = 0; v < n_vec; v++) {
+                        bool asc = (s < M) ? ((((v * 32) >> s) % 2) == 0) : false;
+                        HVX_VectorPred dir_mask = asc ? pred_all_1s : pred_all_0s;
+
+                        HVX_Vector Vv = *(HVX_Vector *)(values + v * 128);
+                        HVX_Vector Iv = *(HVX_Vector *)(indices + v * 128);
+
+                        bitonic_cas_32(&Vv, &Iv, d, dir_mask, idx_vec, zero_vec);
+
+                        *(HVX_Vector *)(values + v * 128)  = Vv;
+                        *(HVX_Vector *)(indices + v * 128) = Iv;
+                    }
+                }
+            }
+        }
+    }
+}
+
+static void top_k_select_tiled(const uint8_t * src, uint32_t n, uint32_t k,
+                               float * values_buf, int32_t * indices_buf,
+                               float * top_values, int32_t * top_indices) {
+    const uint32_t tile_elems = 1024;
+    uint32_t n_tiles            = (n + tile_elems - 1) / tile_elems;
+    uint32_t candidate_count    = n_tiles * k;
+    uint32_t merge_n_vec        = hmx_ceil_div(candidate_count, 32);
+    uint32_t merge_n_vec_pow2   = 1;
+    while (merge_n_vec_pow2 < merge_n_vec) merge_n_vec_pow2 <<= 1;
+    uint32_t merge_elems        = merge_n_vec_pow2 * 32;
+    float * candidate_values    = values_buf + tile_elems;
+    int32_t * candidate_indices = indices_buf + tile_elems;
+    uint32_t candidate_pos      = 0;
+
+    for (uint32_t offset = 0; offset < n; offset += tile_elems) {
+        uint32_t tile_count = MIN(tile_elems, n - offset);
+        hvx_copy_f32_au((uint8_t *) values_buf, src + offset * sizeof(float), tile_count);
+        if (tile_count < tile_elems) {
+            hvx_splat_f32_u((uint8_t *) (values_buf + tile_count), -INFINITY, tile_elems - tile_count);
+        }
+
+        bitonic_sort_vtcm_desc((uint8_t *) values_buf, (uint8_t *) indices_buf, tile_elems / 32, true);
+        uint32_t tile_k = MIN(k, tile_count);
+        for (uint32_t j = 0; j < tile_k; j++) {
+            candidate_values[candidate_pos] = values_buf[j];
+            candidate_indices[candidate_pos] = indices_buf[j] + (int32_t) offset;
+            candidate_pos++;
+        }
+    }
+
+    if (merge_elems > candidate_pos) {
+        hvx_splat_f32_u((uint8_t *) (candidate_values + candidate_pos), -INFINITY, merge_elems - candidate_pos);
+        for (uint32_t j = candidate_pos; j < merge_elems; j++) {
+            candidate_indices[j] = 0;
+        }
+    }
+
+    bitonic_sort_vtcm_desc((uint8_t *) candidate_values, (uint8_t *) candidate_indices, merge_n_vec_pow2, false);
+
+    for (uint32_t j = 0; j < k; j++) {
+        top_values[j] = candidate_values[j];
+        top_indices[j] = candidate_indices[j];
+    }
+}
+
 __attribute__((always_inline))
 static inline void sort32_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
     bitonic_sort_generic_hvx(values, indices, 1, order == GGML_SORT_ORDER_ASC);
@@ -453,6 +591,10 @@ int op_argsort(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
 
+    if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
     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);
 
@@ -534,3 +676,420 @@ int op_argsort(struct htp_ops_context * octx) {
 
     return HTP_STATUS_OK;
 }
+
+// ggml_compute_forward_top_k
+//
+// Reuses ARGSORT's sort kernels. Only the first `k` indices are copied
+// to dst, and there's no asc/desc param -- always largest-first.
+
+struct htp_top_k_context {
+    struct htp_ops_context * octx;
+    uint32_t                 nrows_per_thread;
+    uint32_t                 row_start;
+    uint32_t                 row_end;
+    uint8_t *                vtcm_base;
+    size_t                   vtcm_per_thread;
+    uint32_t                 k;
+};
+
+#define HTP_TOP_K_FN(ne00, sort_fn)                                                                            \
+static void htp_top_k_f32_##ne00(unsigned int n, unsigned int i, void * data) {                                \
+    struct htp_top_k_context * actx = (struct htp_top_k_context *)data;                                        \
+    struct htp_ops_context * octx = actx->octx;                                                                \
+    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 row_start = actx->row_start;                                                                      \
+    uint32_t row_end = actx->row_end;                                                                          \
+    uint32_t rows_per_thread = actx->nrows_per_thread;                                                         \
+    uint32_t start_row = row_start + rows_per_thread * i;                                                      \
+    uint32_t end_row = MIN(start_row + rows_per_thread, 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);                                                  \
+    uint32_t nb01 = src0->nb[1];                                                                               \
+    uint32_t nb1 = dst->nb[1];                                                                                 \
+    uint32_t k = actx->k;                                                                                      \
+    struct htp_thread_trace * tr = &octx->ctx->trace[i];                                                       \
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, start_row);                                              \
+    for (uint32_t r = start_row; r < end_row; r++) {                                                           \
+        uint32_t src_offset = r * nb01;                                                                        \
+        uint32_t dst_offset = r * nb1;                                                                         \
+        uint8_t * src_ptr = (uint8_t *) src0->data + src_offset;                                               \
+        uint8_t * dst_ptr = (uint8_t *) dst->data  + dst_offset;                                               \
+        hex_l2fetch(src_ptr, ne00 * sizeof(float), ne00 * sizeof(float), 1);                                   \
+        hvx_copy_f32_au((uint8_t*)values_buf, src_ptr, ne00);                                                  \
+        sort_fn((uint8_t*)values_buf, (uint8_t*)indices_buf, GGML_SORT_ORDER_DESC);                            \
+        hvx_copy_f32_ua(dst_ptr, (const uint8_t *) indices_buf, k);                                            \
+    }                                                                                                          \
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, start_row);                                               \
+}
+
+HTP_TOP_K_FN(32,   sort32_f32_hvx)
+HTP_TOP_K_FN(64,   sort64_f32_hvx)
+HTP_TOP_K_FN(128,  sort128_f32_hvx)
+HTP_TOP_K_FN(256,  sort256_f32_hvx)
+HTP_TOP_K_FN(512,  sort512_f32_hvx)
+HTP_TOP_K_FN(1024, sort1024_f32_hvx)
+
+static void htp_top_k_f32_fallback(unsigned int n, unsigned int i, void * data) {
+    struct htp_top_k_context * actx = (struct htp_top_k_context *)data;
+    struct htp_ops_context * octx = actx->octx;
+
+    // Unpack context
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * dst  = octx->dst;
+
+    // Scratchpad memory
+    uint8_t * spad = actx->vtcm_base + actx->vtcm_per_thread * i;
+
+    // Dimensions
+    uint32_t ne00 = src0->ne[0];
+    uint32_t nb01 = src0->nb[1];
+
+    uint32_t nb1 = dst->nb[1];
+
+    uint32_t k = actx->k;
+
+    // Rows to process
+    uint32_t row_start = actx->row_start;
+    uint32_t row_end = actx->row_end;
+    uint32_t rows_per_thread = actx->nrows_per_thread;
+    uint32_t start_row = row_start + rows_per_thread * i;
+    uint32_t end_row = MIN(start_row + rows_per_thread, row_end);
+
+    // Pad ne00 to n_vec*32 (n_vec a power of 2) for the bitonic network;
+    // pad with -INFINITY so it never lands in the top-k.
+    uint32_t n_vec = hmx_ceil_div(ne00, 32);
+    uint32_t n_vec_pow2 = 1;
+    while (n_vec_pow2 < n_vec) n_vec_pow2 <<= 1;
+    uint32_t ne00_padded = n_vec_pow2 * 32;
+
+    size_t values_size = hex_round_up(ne00_padded * sizeof(float), 128);
+    float * values_buf = (float *) spad;
+    int32_t * indices_buf = (int32_t *) (spad + values_size);
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[i];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, start_row);
+
+    for (uint32_t r = start_row; r < end_row; r++) {
+        uint32_t src_offset = r * nb01;
+        uint32_t dst_offset = r * nb1;
+
+        uint8_t * src_ptr = (uint8_t *) src0->data + src_offset;
+        uint8_t * dst_ptr = (uint8_t *) dst->data  + dst_offset;
+
+        hex_l2fetch(src_ptr, ne00 * sizeof(float), ne00 * sizeof(float), 1);
+
+        if (k <= 64 && ne00 > 1024) {
+            float top_values[64];
+            int32_t top_indices[64];
+            top_k_select_tiled(src_ptr, ne00, k, values_buf, indices_buf, top_values, top_indices);
+            memcpy(dst_ptr, top_indices, k * sizeof(int32_t));
+            continue;
+        }
+
+        hvx_copy_f32_au((uint8_t*)values_buf, src_ptr, ne00);
+
+        // Fills the indices ramp itself, so no init needed here.
+        if (ne00_padded > ne00) {
+            hvx_splat_f32_u((uint8_t *)(values_buf + ne00), -INFINITY, ne00_padded - ne00);
+        }
+        bitonic_sort_vtcm_desc((uint8_t*)values_buf, (uint8_t*)indices_buf, n_vec_pow2, true);
+
+        // Copy top-k indices back to DDR
+        hvx_copy_f32_ua(dst_ptr, (const uint8_t *) indices_buf, k);
+    }
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, start_row);
+}
+
+// Single row (ne01=ne02=ne03=1) + large ne00 would otherwise run on one
+// HVX thread while the rest sit idle. Split the row into n_chunks
+// power-of-two chunks, sort each in parallel with bitonic_sort_vtcm_desc,
+// then merge the n_chunks*local_k winners with one more sort that carries
+// the global index through instead of re-deriving it.
+struct htp_top_k_chunk_ctx {
+    struct htp_ops_context * octx;
+    uint8_t *                vtcm_base;
+    size_t                    phase1_slot_size;
+    uint32_t                  ne00;
+    uint32_t                  chunk_elems;
+    uint32_t                  local_k;
+    size_t                    merge_values_off;
+    size_t                    merge_indices_off;
+};
+
+static void htp_top_k_chunk_job(unsigned int n, unsigned int i, void * data) {
+    struct htp_top_k_chunk_ctx * cctx = (struct htp_top_k_chunk_ctx *) data;
+    struct htp_ops_context * octx = cctx->octx;
+    const struct htp_tensor * src0 = octx->src[0];
+
+    uint32_t chunk_elems = cctx->chunk_elems;
+    uint32_t ne00        = cctx->ne00;
+    uint32_t local_k     = cctx->local_k;
+    uint32_t chunk_base  = i * chunk_elems;
+
+    uint8_t * spad = cctx->vtcm_base + cctx->phase1_slot_size * i;
+    size_t values_size = hex_round_up(chunk_elems * sizeof(float), 128);
+    float *   values_buf  = (float *) spad;
+    int32_t * indices_buf = (int32_t *) (spad + values_size);
+
+    uint32_t real_count = (chunk_base < ne00) ? MIN(chunk_elems, ne00 - chunk_base) : 0;
+
+    if (real_count == 0) {
+        float *   merge_values  = (float *) (cctx->vtcm_base + cctx->merge_values_off);
+        int32_t * merge_indices = (int32_t *) (cctx->vtcm_base + cctx->merge_indices_off);
+        for (uint32_t j = 0; j < local_k; j++) {
+            merge_values[i * local_k + j] = -INFINITY;
+            merge_indices[i * local_k + j] = 0;
+        }
+        return;
+    }
+
+    if (local_k > 1 && local_k <= 64 && chunk_elems > 1024) {
+        uint8_t * src_ptr = (uint8_t *) src0->data + (size_t) chunk_base * sizeof(float);
+        float top_values[64];
+        int32_t top_indices[64];
+        float *   merge_values  = (float *) (cctx->vtcm_base + cctx->merge_values_off);
+        int32_t * merge_indices = (int32_t *) (cctx->vtcm_base + cctx->merge_indices_off);
+        for (uint32_t j = 0; j < local_k; j++) {
+            top_values[j] = -INFINITY;
+            top_indices[j] = 0;
+        }
+        top_k_select_tiled(src_ptr, real_count, local_k, values_buf, indices_buf, top_values, top_indices);
+
+        for (uint32_t j = 0; j < local_k; j++) {
+            merge_values[i * local_k + j] = top_values[j];
+            merge_indices[i * local_k + j] = top_indices[j] + (int32_t) chunk_base;
+        }
+        return;
+    }
+
+    if (real_count > 0) {
+        uint8_t * src_ptr = (uint8_t *) src0->data + (size_t) chunk_base * sizeof(float);
+        hex_l2fetch(src_ptr, real_count * sizeof(float), real_count * sizeof(float), 1);
+        hvx_copy_f32_au((uint8_t *) values_buf, src_ptr, real_count);
+    }
+    if (chunk_elems > real_count) {
+        hvx_splat_f32_u((uint8_t *) (values_buf + real_count), -INFINITY, chunk_elems - real_count);
+    }
+
+    if (local_k == 1 && cctx->ne00 >= 128*1024) {
+        float max_value;
+        uint32_t max_index = top_k_max_value_index(values_buf, real_count, &max_value);
+        float *   merge_values  = (float *) (cctx->vtcm_base + cctx->merge_values_off);
+        int32_t * merge_indices = (int32_t *) (cctx->vtcm_base + cctx->merge_indices_off);
+        merge_values[i] = max_value;
+        merge_indices[i] = (int32_t) (max_index + chunk_base);
+        return;
+    }
+
+    // chunk_elems is always a power-of-two multiple of 32
+    bitonic_sort_vtcm_desc((uint8_t *) values_buf, (uint8_t *) indices_buf, chunk_elems / 32, true);
+
+    float *   merge_values  = (float *)   (cctx->vtcm_base + cctx->merge_values_off);
+    int32_t * merge_indices = (int32_t *) (cctx->vtcm_base + cctx->merge_indices_off);
+
+    for (uint32_t j = 0; j < local_k; j++) {
+        merge_values[i * local_k + j]  = values_buf[j];
+        merge_indices[i * local_k + j] = indices_buf[j] + (int32_t) chunk_base;
+    }
+}
+
+struct htp_top_k_merge_ctx {
+    struct htp_ops_context * octx;
+    uint8_t *                vtcm_base;
+    size_t                    merge_values_off;
+    size_t                    merge_indices_off;
+    uint32_t                  merge_elems;
+    uint32_t                  total_candidates;
+    uint32_t                  k;
+};
+
+static void htp_top_k_merge_job(unsigned int n, unsigned int i, void * data) {
+    struct htp_top_k_merge_ctx * mctx = (struct htp_top_k_merge_ctx *) data;
+    struct htp_ops_context * octx = mctx->octx;
+    const struct htp_tensor * dst = octx->dst;
+
+    float *   merge_values  = (float *)   (mctx->vtcm_base + mctx->merge_values_off);
+    int32_t * merge_indices = (int32_t *) (mctx->vtcm_base + mctx->merge_indices_off);
+
+    if (mctx->merge_elems > mctx->total_candidates) {
+        uint32_t pad = mctx->merge_elems - mctx->total_candidates;
+        hvx_splat_f32_u((uint8_t *) (merge_values + mctx->total_candidates), -INFINITY, pad);
+        for (uint32_t j = mctx->total_candidates; j < mctx->merge_elems; j++) {
+            merge_indices[j] = 0;
+        }
+    }
+
+    // Preserve the global indices computed in phase 1 -- init_indices=false
+    // so they aren't overwritten with a local ramp.
+    bitonic_sort_vtcm_desc((uint8_t *) merge_values, (uint8_t *) merge_indices, mctx->merge_elems / 32, false);
+
+    hvx_copy_f32_ua((uint8_t *) dst->data, (const uint8_t *) merge_indices, mctx->k);
+}
+
+static int op_top_k_single_row_threaded(struct htp_ops_context * octx, uint32_t ne00, uint32_t k) {
+    uint32_t n_threads_avail = octx->n_threads;
+
+    uint32_t n_vec = hmx_ceil_div(ne00, 32);
+    uint32_t n_vec_pow2 = 1;
+    while (n_vec_pow2 < n_vec) n_vec_pow2 <<= 1;
+
+    // Largest power-of-two chunk count that both fits the available
+    // threads and evenly divides n_vec_pow2
+    uint32_t n_chunks = 1;
+    while (n_chunks * 2 <= n_threads_avail && n_chunks * 2 <= n_vec_pow2) {
+        n_chunks *= 2;
+    }
+
+    uint32_t chunk_n_vec = n_vec_pow2 / n_chunks;
+    uint32_t chunk_elems = chunk_n_vec * 32;
+    uint32_t local_k     = MIN(k, chunk_elems);
+
+    uint32_t total_candidates = n_chunks * local_k;
+    uint32_t merge_n_vec = hmx_ceil_div(total_candidates, 32);
+    uint32_t merge_n_vec_pow2 = 1;
+    while (merge_n_vec_pow2 < merge_n_vec) merge_n_vec_pow2 <<= 1;
+    uint32_t merge_elems = merge_n_vec_pow2 * 32;
+
+    size_t phase1_values_size  = hex_round_up(chunk_elems * sizeof(float), 128);
+    size_t phase1_indices_size = hex_round_up(chunk_elems * sizeof(int32_t), 128);
+    size_t phase1_slot_size    = hex_round_up(phase1_values_size + phase1_indices_size, 256);
+    size_t phase1_total_size   = phase1_slot_size * n_chunks;
+
+    size_t merge_values_size  = hex_round_up(merge_elems * sizeof(float), 128);
+    size_t merge_indices_size = hex_round_up(merge_elems * sizeof(int32_t), 128);
+    size_t merge_values_off   = phase1_total_size;
+    size_t merge_indices_off  = merge_values_off + merge_values_size;
+
+    size_t total_vtcm = phase1_total_size + merge_values_size + merge_indices_size;
+    if (octx->ctx->vtcm_size < total_vtcm) {
+        return HTP_STATUS_VTCM_TOO_SMALL;
+    }
+
+    uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
+
+    struct htp_top_k_chunk_ctx cctx;
+    cctx.octx              = octx;
+    cctx.vtcm_base         = vtcm_base;
+    cctx.phase1_slot_size  = phase1_slot_size;
+    cctx.ne00              = ne00;
+    cctx.chunk_elems       = chunk_elems;
+    cctx.local_k           = local_k;
+    cctx.merge_values_off  = merge_values_off;
+    cctx.merge_indices_off = merge_indices_off;
+
+    work_queue_run(octx->ctx->work_queue, htp_top_k_chunk_job, &cctx, n_chunks);
+
+    struct htp_top_k_merge_ctx mctx;
+    mctx.octx              = octx;
+    mctx.vtcm_base         = vtcm_base;
+    mctx.merge_values_off  = merge_values_off;
+    mctx.merge_indices_off = merge_indices_off;
+    mctx.merge_elems       = merge_elems;
+    mctx.total_candidates  = total_candidates;
+    mctx.k                 = k;
+
+    work_queue_run(octx->ctx->work_queue, htp_top_k_merge_job, &mctx, 1);
+
+    return HTP_STATUS_OK;
+}
+
+int op_top_k(struct htp_ops_context * octx) {
+    // Check supported types
+    if (octx->src[0]->type != HTP_TYPE_F32) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    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;
+    }
+
+    uint32_t ne00 = src0->ne[0];
+    uint32_t k    = dst->ne[0];
+
+    // Single row + large ne00: the per-row dispatch below would run on one
+    // HVX thread while the rest sit idle. Split the row across threads.
+    if (total_rows == 1 && ne00 > 1024) {
+        int status = op_top_k_single_row_threaded(octx, ne00, k);
+        if (status != HTP_STATUS_VTCM_TOO_SMALL) {
+            return status;
+        }
+        // else: fall through to the single-thread path below.
+    }
+
+    const uint32_t n_threads = MIN(nrows, octx->n_threads);
+
+    // Scratchpad layout: values + indices
+    // For bitonic: need padding to power-of-2 size
+    // Allocate for worst case (bitonic with padding)
+    uint32_t n_vec = hmx_ceil_div(ne00, 32);
+    uint32_t n_vec_pow2 = 1;
+    while (n_vec_pow2 < n_vec) n_vec_pow2 <<= 1;
+    uint32_t ne00_padded = n_vec_pow2 * 32;
+
+    size_t values_size  = hex_round_up(ne00_padded * sizeof(float), 128);
+    size_t indices_size = hex_round_up(ne00_padded * sizeof(int32_t), 128);
+    size_t spad_per_thread = values_size + indices_size;
+
+    // Make sure we round up to 256 for alignment requirements
+    spad_per_thread = hex_round_up(spad_per_thread, 256);
+
+    size_t total_spad_size = spad_per_thread * n_threads;
+
+    if (octx->ctx->vtcm_size < total_spad_size) {
+        FARF(ERROR, "top_k: VTCM size too small. Needed %zu, have %zu", total_spad_size, octx->ctx->vtcm_size);
+        return HTP_STATUS_VTCM_TOO_SMALL;
+    }
+
+    FARF(HIGH, "top_k: %ux%ux%ux%u -> %ux%ux%ux%u (0x%x, 0x%x)",
+         octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3],
+         octx->dst->ne[0], octx->dst->ne[1], octx->dst->ne[2], octx->dst->ne[3],
+         octx->src[0]->data, octx->dst->data);
+
+    struct htp_top_k_context actx;
+    const struct fastdiv_values n_threads_div = init_fastdiv_values(n_threads);
+    actx.octx             = octx;
+    actx.nrows_per_thread = fastdiv(nrows + n_threads - 1, &n_threads_div);
+    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;
+    actx.k                = k;
+
+    worker_callback_t job_func = htp_top_k_f32_fallback;
+    switch (ne00) {
+        case 1024: job_func = htp_top_k_f32_1024; break;
+        case 512:  job_func = htp_top_k_f32_512;  break;
+        case 256:  job_func = htp_top_k_f32_256;  break;
+        case 128:  job_func = htp_top_k_f32_128;  break;
+        case 64:   job_func = htp_top_k_f32_64;   break;
+        case 32:   job_func = htp_top_k_f32_32;   break;
+        default:   job_func = htp_top_k_f32_fallback; break;
+    }
+
+    // Run jobs
+    work_queue_run(octx->ctx->work_queue, job_func, &actx, n_threads);
+
+    return HTP_STATUS_OK;
+}

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

@@ -0,0 +1,111 @@
+#ifndef HTP_BINARY_OPS_H
+#define HTP_BINARY_OPS_H
+
+#include <stddef.h>
+#include <stdint.h>
+#include <string.h>
+
+#include "hex-common.h"
+#include "htp-ops.h"
+#include "htp-vtcm.h"
+
+enum htp_binary_kernel_type {
+    HTP_BINARY_KERNEL_SAME_SHAPE = 0,
+    HTP_BINARY_KERNEL_ROW_BCAST,
+    HTP_BINARY_KERNEL_SCALAR_DMA,
+    HTP_BINARY_KERNEL_SCALAR,
+    HTP_BINARY_KERNEL_ADD_ID,
+    HTP_BINARY_KERNEL_COMPLEX,
+    HTP_BINARY_KERNEL_REPEAT,
+};
+
+struct htp_binary_kernel_params {
+    uint32_t kernel_type;
+    uint32_t n_threads;
+    uint32_t rows_per_buffer;
+
+    uint32_t src0_row_size_aligned;
+    uint32_t src1_row_size_aligned;
+    uint32_t dst_row_size_aligned;
+
+    uint32_t src1_size;
+    uint32_t vtcm_size;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_binary_kernel_params) <= 128, "htp_binary_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_binary_kernel_params) <= 128, "htp_binary_kernel_params is too large for kernel_params blob");
+#endif
+
+struct htp_binary_vtcm_layout {
+    size_t total_bytes;
+    size_t off_src0;
+    size_t off_src1;
+    size_t off_dst;
+
+    size_t src0_bytes_per_thread;
+    size_t src1_bytes_per_thread;
+    size_t dst_bytes_per_thread;
+
+    size_t src0_spad_half_size;
+    size_t src1_spad_half_size;
+    size_t dst_spad_half_size;
+
+    size_t src1_size;
+    uint32_t rows_per_buffer;
+};
+
+static inline void htp_binary_vtcm_layout_build(
+    struct htp_binary_vtcm_layout * L,
+    const struct htp_binary_kernel_params * kparams,
+    size_t vtcm_size
+) {
+    memset(L, 0, sizeof(*L));
+
+    const uint32_t n_threads = kparams->n_threads;
+    if (n_threads == 0) {
+        return;
+    }
+
+    const size_t spad_row_total = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE)
+        ? 2 * (kparams->src0_row_size_aligned + kparams->src1_row_size_aligned + kparams->dst_row_size_aligned)
+        : 2 * (kparams->src0_row_size_aligned + kparams->dst_row_size_aligned);
+
+    if (spad_row_total == 0 || vtcm_size < kparams->src1_size) {
+        return;
+    }
+
+    const size_t rows_per_buffer = (vtcm_size - kparams->src1_size) / (n_threads * spad_row_total);
+    if (rows_per_buffer == 0) {
+        return;
+    }
+
+    L->rows_per_buffer = (uint32_t) rows_per_buffer;
+    L->src1_size = kparams->src1_size;
+
+    L->src0_bytes_per_thread = rows_per_buffer * 2 * kparams->src0_row_size_aligned;
+    L->dst_bytes_per_thread  = rows_per_buffer * 2 * kparams->dst_row_size_aligned;
+    L->src1_bytes_per_thread = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE)
+        ? rows_per_buffer * 2 * kparams->src1_row_size_aligned
+        : 0;
+
+    L->src0_spad_half_size = L->src0_bytes_per_thread / 2;
+    L->src1_spad_half_size = L->src1_bytes_per_thread / 2;
+    L->dst_spad_half_size  = L->dst_bytes_per_thread / 2;
+
+    const size_t src0_total = n_threads * L->src0_bytes_per_thread;
+    const size_t src1_total = (kparams->src1_size > 0)
+        ? kparams->src1_size
+        : n_threads * L->src1_bytes_per_thread;
+    const size_t dst_total  = n_threads * L->dst_bytes_per_thread;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, src0_total);
+    VTCM_LAYOUT_ALLOC(off, off_src1, src1_total);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  dst_total);
+
+    L->total_bytes = off;
+}
+
+#endif

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

@@ -6,7 +6,7 @@
 #include "hexagon_types.h"
 #include "hexagon_protos.h"
 #include "hvx_hexagon_protos.h"
-#include "hex-dma.h"
+#include "dma-queue.h"
 #include "htp-vtcm.h"
 #include "hvx-utils.h"
 #include "hex-fastdiv.h"
@@ -41,7 +41,7 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *
     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];
+    dma_queue * dma_q = octx->ctx->dma[ith];
 
     uint8_t * spad0_base = octx->src0_spad.data + ith * octx->src0_spad.size_per_thread;
     uint8_t * spad1_base = octx->src1_spad.data + ith * octx->src1_spad.size_per_thread;
@@ -64,14 +64,14 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *
         uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i;
 
         uint32_t src1_width_bytes = current_block_i * sizeof(float);
-        uint8_t * src1_ptr = (uint8_t *)src1->data + i * src1->nb[1];
-        dma_queue_push(q, dma_make_ptr(spad1_base, src1_ptr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0);
+        const dma_addr_t src1_addr = src1->data + i * src1->nb[1];
+        dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0);
 
         uint32_t src0_row_bytes = src0_ne0 * sizeof(float);
-        uint8_t * src0_ptr = (uint8_t *)src0->data + i * src0->nb[1];
-        dma_queue_push(q, dma_make_ptr(spad0_base, src0_ptr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i);
+        const dma_addr_t src0_addr = src0->data + i * src0->nb[1];
+        dma_queue_push(dma_q, dma_make_data(spad0_base, src0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i);
 
-        dma_queue_pop(q); // src1
+        dma_queue_pop(dma_q); // src1
 
         HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride);
 
@@ -87,12 +87,12 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *
         }
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
 
-        dma_queue_pop(q); // src0
+        dma_queue_pop(dma_q); // src0
 
-        uint8_t * dst_ptr = (uint8_t *)dst->data + i * dst->nb[1];
-        dma_queue_push(q, dma_make_ptr(dst_ptr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(float), current_block_i);
+        const dma_addr_t dst_addr = dst->data + i * dst->nb[1];
+        dma_queue_push(dma_q, dma_make_data(dst_addr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(float), current_block_i);
 
-        dma_queue_pop(q);
+        dma_queue_pop(dma_q);
     }
 }
 
@@ -112,7 +112,7 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *
     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];
+    dma_queue * dma_q = octx->ctx->dma[ith];
 
     uint8_t * spad0_base = octx->src0_spad.data + ith * octx->src0_spad.size_per_thread;
     uint8_t * spad1_base = octx->src1_spad.data + ith * octx->src1_spad.size_per_thread;
@@ -135,14 +135,14 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *
         uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i;
 
         uint32_t src1_width_bytes = current_block_i * sizeof(__fp16);
-        uint8_t * src1_ptr = (uint8_t *)src1->data + i * src1->nb[1];
-        dma_queue_push(q, dma_make_ptr(spad1_base, src1_ptr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0);
+        const dma_addr_t src1_addr = src1->data + i * src1->nb[1];
+        dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0);
 
         uint32_t src0_row_bytes = src0_ne0 * sizeof(__fp16);
-        uint8_t * src0_ptr = (uint8_t *)src0->data + i * src0->nb[1];
-        dma_queue_push(q, dma_make_ptr(spad0_base, src0_ptr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i);
+        const dma_addr_t src0_addr = src0->data + i * src0->nb[1];
+        dma_queue_push(dma_q, dma_make_data(spad0_base, src0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i);
 
-        dma_queue_pop(q); // src1
+        dma_queue_pop(dma_q); // src1
 
         HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride);
 
@@ -158,12 +158,12 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *
         }
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
 
-        dma_queue_pop(q); // src0
+        dma_queue_pop(dma_q); // src0
 
-        uint8_t * dst_ptr = (uint8_t *)dst->data + i * dst->nb[1];
-        dma_queue_push(q, dma_make_ptr(dst_ptr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(__fp16), current_block_i);
+        const dma_addr_t dst_addr = dst->data + i * dst->nb[1];
+        dma_queue_push(dma_q, dma_make_data(dst_addr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(__fp16), current_block_i);
 
-        dma_queue_pop(q);
+        dma_queue_pop(dma_q);
     }
 }
 
@@ -304,6 +304,10 @@ int op_concat(struct htp_ops_context * octx) {
             worker_func = concat_2d_f16_transposed;
         }
     } else {
+        if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(src1) || htp_tensor_is_extended(dst)) {
+            return HTP_STATUS_NO_SUPPORT;
+        }
+
         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;

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

@@ -49,6 +49,30 @@ struct htp_copy_context {
     struct fastdiv_values div_ne02_ne01_ne00;
 };
 
+static inline void cpy_dma_sametype_reshape_contig(
+    dma_queue * dma_q,
+    dma_addr_t dst,
+    dma_addr_t src0,
+    uint32_t total_bytes
+) {
+    if (total_bytes == 0) {
+        return;
+    }
+
+    const uint32_t max_chunk = DMA_SAFE_CHUNK_SIZE;
+    while (total_bytes > 0) {
+        const uint32_t chunk = MIN(total_bytes, max_chunk);
+        if (!dma_queue_push(dma_q, dma_make_data(dst, src0), chunk, chunk, chunk, /*nrows=*/ 1)) {
+            dma_queue_flush(dma_q);
+            dma_queue_push(dma_q, dma_make_data(dst, src0), chunk, chunk, chunk, /*nrows=*/ 1);
+        }
+        dst         += chunk;
+        src0        += chunk;
+        total_bytes -= chunk;
+    }
+    dma_queue_flush(dma_q);
+}
+
 #define cpy_preamble                              \
     const struct htp_tensor *src0 = octx->src[0]; \
     const struct htp_tensor *dst  = octx->dst;    \
@@ -73,129 +97,131 @@ struct htp_copy_context {
     const uint32_t  nb2 = dst->nb[2];             \
     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 = 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 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 = htp_tensor_is_contiguous(src0, ELEM_SIZE) && htp_tensor_is_contiguous(dst, ELEM_SIZE); \
+    if (contiguous) {                                                                                              \
+        dma_queue * dma_q = octx->ctx->dma[ith];                                                                   \
+        dma_addr_t dst_addr  = dst->data  + ir0 * ne00 * ELEM_SIZE;                                                \
+        dma_addr_t src0_addr = src0->data + ir0 * ne00 * ELEM_SIZE;                                                \
+        cpy_dma_sametype_reshape_contig(dma_q, dst_addr, src0_addr, (ir1 - ir0) * ne00 * ELEM_SIZE);               \
+        return;                                                                                                    \
+    }                                                                                                              \
+    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 * 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 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;                                                                  \
-                                                                                                     \
-    if (htp_tensor_is_contiguous(src0, ELEM_SIZE) && htp_tensor_is_contiguous(dst, ELEM_SIZE)) {     \
-        hvx_copy_uu((uint8_t *) dst->data + (size_t) th_start * ELEM_SIZE,                           \
-                    (const uint8_t *) src0->data + (size_t) th_start * ELEM_SIZE,                    \
-                    th_end - th_start, ELEM_SIZE);                                                   \
-        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;   \
-                                                                                                     \
-    const bool rows_contig = (nb00 == ELEM_SIZE) && (nb0 == ELEM_SIZE);                              \
-                                                                                                     \
-    while (e < th_end) {                                                                             \
-        uint32_t run = 1;                                                                            \
-        if (rows_contig) {                                                                           \
-            run = MIN(MIN(ne00 - i00, ne0 - i10), th_end - e);                                       \
-            hvx_copy_uu((uint8_t *) dst_ptr, (const uint8_t *) src0_ptr, run, ELEM_SIZE);            \
-        } else {                                                                                     \
-            *((ELEM_TYPE *) dst_ptr) = *((const ELEM_TYPE *) src0_ptr);                              \
-        }                                                                                            \
-        e += run;                                                                                    \
-                                                                                                     \
-        dst_ptr += run * nb0;                                                                        \
-        i10     += run;                                                                              \
-        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 += run * nb00;                                                                      \
-        i00      += run;                                                                             \
-        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 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;                                                                   \
+                                                                                                      \
+    if (htp_tensor_is_contiguous(src0, ELEM_SIZE) && htp_tensor_is_contiguous(dst, ELEM_SIZE)) {      \
+        dma_queue * dma_q = octx->ctx->dma[ith];                                                      \
+        dma_addr_t dst_addr  = dst->data  + th_start * ELEM_SIZE;                                     \
+        dma_addr_t src0_addr = src0->data + th_start * ELEM_SIZE;                                     \
+        cpy_dma_sametype_reshape_contig(dma_q, dst_addr, src0_addr, (th_end - th_start) * ELEM_SIZE); \
+        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;    \
+                                                                                                      \
+    const bool rows_contig = (nb00 == ELEM_SIZE) && (nb0 == ELEM_SIZE);                               \
+                                                                                                      \
+    while (e < th_end) {                                                                              \
+        uint32_t run = 1;                                                                             \
+        if (rows_contig) {                                                                            \
+            run = MIN(MIN(ne00 - i00, ne0 - i10), th_end - e);                                        \
+            hvx_copy_uu((uint8_t *) dst_ptr, (const uint8_t *) src0_ptr, run, ELEM_SIZE);             \
+        } else {                                                                                      \
+            *((ELEM_TYPE *) dst_ptr) = *((const ELEM_TYPE *) src0_ptr);                               \
+        }                                                                                             \
+        e += run;                                                                                     \
+                                                                                                      \
+        dst_ptr += run * nb0;                                                                         \
+        i10     += run;                                                                               \
+        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 += run * nb00;                                                                       \
+        i00      += run;                                                                              \
+        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)
@@ -273,6 +299,27 @@ static void cpy_thread_f32_f16_sameshape(unsigned int nth, unsigned int ith, voi
     }
 }
 
+static inline void cpy_dma_push_2d_chunked(
+    dma_queue * dma_q,
+    dma_addr_t  dst,
+    dma_addr_t  src,
+    size_t      dst_stride,
+    size_t      src_stride,
+    size_t      row_size,
+    uint32_t    nrows
+) {
+    while (nrows > 0) {
+        const uint32_t cur_rows = MIN(nrows, DMA_MAX_NROWS);
+        if (!dma_queue_push(dma_q, dma_make_data(dst, src), dst_stride, src_stride, row_size, cur_rows)) {
+            dma_queue_flush(dma_q);
+            dma_queue_push(dma_q, dma_make_data(dst, src), dst_stride, src_stride, row_size, cur_rows);
+        }
+        dst   += cur_rows * dst_stride;
+        src   += cur_rows * src_stride;
+        nrows -= cur_rows;
+    }
+}
+
 static inline void cpy_dma_sametype_sameshape(
     struct htp_ops_context * octx,
     const struct htp_tensor * dst,
@@ -282,46 +329,35 @@ static inline void cpy_dma_sametype_sameshape(
     uint32_t nb01, uint32_t nb02, uint32_t nb03,
     uint32_t  nb1, uint32_t  nb2, uint32_t nb3
 ) {
+    const bool contiguous = htp_tensor_is_contiguous(src0, elem_size) && htp_tensor_is_contiguous(dst, elem_size);
+
+    dma_queue * dma_q = octx->ctx->dma[0];
+
+    if (contiguous) {
+        cpy_dma_sametype_reshape_contig(dma_q, dst->data, src0->data, ne00 * elem_size * ne01 * ne02 * ne03);
+        return;
+    }
+
     const bool contiguous_outer =
         (ne02 == 1 || (nb02 == ne01 * nb01 && nb2 == ne01 * nb1)) &&
         (ne03 == 1 || (nb03 == ne02 * nb02 && nb3 == ne02 * nb2));
 
-    dma_queue * q = octx->ctx->dma[0];
-
     if (contiguous_outer) {
-        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);
+        uint32_t total_rows = ne01 * ne02 * ne03;
+        cpy_dma_push_2d_chunked(dma_q, dst->data, src0->data, nb1, nb01, ne00 * elem_size, total_rows);
+        dma_queue_flush(dma_q);
         return;
     }
 
     for (uint32_t i03 = 0; i03 < ne03; i03++) {
         for (uint32_t i02 = 0; i02 < ne02; i02++) {
-            uint8_t * dst_ptr  = (uint8_t *) dst->data  + i02 * nb2  + i03 * nb3;
-            uint8_t * src0_ptr = (uint8_t *) src0->data + i02 * nb02 + i03 * nb03;
-
-            if (!dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01)) {
-                dma_queue_flush(q);
-                dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01);
-            }
+            dma_addr_t dst_data  = dst->data  + i02 * nb2  + i03 * nb3;
+            dma_addr_t src0_data = src0->data + i02 * nb02 + i03 * nb03;
+            cpy_dma_push_2d_chunked(dma_q, dst_data, src0_data, nb1, nb01, ne00 * elem_size, ne01);
         }
     }
 
-    dma_queue_flush(q);
-}
-
-static inline void cpy_dma_sametype_reshape_contig(
-    struct htp_ops_context * octx,
-    const struct htp_tensor * dst,
-    const struct htp_tensor * src0,
-    uint32_t total_bytes
-) {
-    dma_queue * q = octx->ctx->dma[0];
-    dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data),
-                   total_bytes, total_bytes, total_bytes, /*nrows=*/ 1);
-    dma_queue_pop(q);
+    dma_queue_flush(dma_q);
 }
 
 static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) {
@@ -345,10 +381,6 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
-        return HTP_STATUS_OK;
-    }
-
     const bool sametype   = (src0->type == dst->type);
     const bool transposed = (nb00 > nb01) || (nb0 > nb1) ||
                             (nb00 != ct.src0_type_size) || (nb0 != ct.dst_type_size) ||
@@ -360,6 +392,15 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) {
     const bool src_is_contiguous = htp_tensor_is_contiguous(src0, ct.src0_type_size);
     const bool dst_is_contiguous = htp_tensor_is_contiguous(dst, ct.dst_type_size);
 
+    if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
+        if (!sametype) {
+            return HTP_STATUS_NO_SUPPORT;
+        }
+        if (!sameshape && !(src_is_contiguous && dst_is_contiguous && octx->ctx->mdev.count <= 1)) {
+            return HTP_STATUS_NO_SUPPORT;
+        }
+    }
+
     if (sameshape) {
         const uint32_t total_rows = ne01 * ne02 * ne03;
         const uint32_t row_size   = ne00 * ct.dst_type_size;
@@ -373,7 +414,8 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) {
         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);
+            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;
         }
@@ -386,9 +428,11 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) {
         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);
+        if (sametype && (octx->ctx->mdev.count <= 1 || htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) {
+            if (octx->ctx->mdev.idx == 0) {
+                *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) {
@@ -408,7 +452,7 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) {
 
         if (octx->ctx->mdev.count <= 1 && dst_is_contiguous && src_is_contiguous) {
             *use_dma = true;
-            cpy_dma_sametype_reshape_contig(octx, dst, src0, total_elems * ct.dst_type_size);
+            cpy_dma_sametype_reshape_contig(octx->ctx->dma[0], dst->data, src0->data, total_elems * ct.dst_type_size);
             return HTP_STATUS_OK;
         }
 
@@ -424,7 +468,8 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) {
 
         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);
+            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;
         }
@@ -461,6 +506,9 @@ int op_cpy(struct htp_ops_context * octx) {
 
         if (octx->ctx->mdev.idx == 0) {
             const struct htp_tensor * sync = octx->src[1];
+            if (htp_tensor_is_extended(sync)) {
+                return HTP_STATUS_NO_SUPPORT;
+            }
             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);

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

@@ -14,7 +14,7 @@
 #include "htp-tensor.h"
 #include "hvx-types.h"
 #include "hvx-utils.h"
-#include "hex-dma.h"
+#include "dma-queue.h"
 
 #define htp_cumsum_tensors_preamble                         \
     const struct htp_tensor * restrict src0 = octx->src[0]; \
@@ -55,7 +55,7 @@ struct htp_cumsum_context {
     struct htp_cumsum_context * cctx = (struct htp_cumsum_context *) data; \
     struct htp_ops_context *    octx = cctx->octx;                         \
     htp_cumsum_tensors_preamble;                                           \
-    dma_queue * dma_queue = octx->ctx->dma[ith];
+    dma_queue * dma_q = octx->ctx->dma[ith];
 
 // ---------------------------------------------------------------------------
 // HVX prefix scan helpers
@@ -131,47 +131,47 @@ static void cumsum_thread_f32_dma(unsigned int nth, unsigned int ith, void * dat
     const size_t src_row_size_aligned = cctx->src_row_size_aligned;
     const size_t dst_row_size_aligned = cctx->dst_row_size_aligned;
 
-    const uint8_t * src_data = (const uint8_t *) src0->data;
-    uint8_t *       dst_data = (uint8_t *) dst->data;
+    const dma_addr_t src_data = src0->data;
+    const dma_addr_t dst_data = dst->data;
 
     uint8_t * src_spad = octx->src0_spad.data + (ith * src_row_size_aligned * 2);
     uint8_t * dst_spad = octx->dst_spad.data  + (ith * dst_row_size_aligned * 2);
 
     for (uint32_t ir = ir0, spad_idx = 0; ir < ir1 && spad_idx < 2; ir++, spad_idx++) {
         // Dummy dst writeback to establish queue ordering
-        dma_queue_push_vtcm_to_ddr(dma_queue,
-                                   dma_make_ptr(dst_data, dst_spad + (spad_idx * dst_row_size_aligned)),
-                                   dst_row_size, dst_row_size_aligned, 0);
+        dma_queue_push(dma_q,
+                       dma_make_data(dst_data, dst_spad + (spad_idx * dst_row_size_aligned)),
+                       dst_row_size, dst_row_size_aligned, dst_row_size, 0);
 
-        dma_queue_push_ddr_to_vtcm(dma_queue,
-                                   dma_make_ptr(src_spad + (spad_idx * src_row_size_aligned),
-                                                src_data + (ir * src_row_size)),
-                                   src_row_size_aligned, src_row_size, 1);
+        dma_queue_push(dma_q,
+                       dma_make_data(src_spad + (spad_idx * src_row_size_aligned),
+                                     src_data + (ir * src_row_size)),
+                       src_row_size_aligned, src_row_size, 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;
+        float * dst_spad_row = (float *) dma_queue_pop(dma_q).src;
+        float * src_spad_row = (float *) dma_queue_pop(dma_q).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),
-                                   dst_row_size, dst_row_size_aligned, 1);
+        dma_queue_push(dma_q,
+                       dma_make_data(dst_data + (ir * dst_row_size), dst_spad_row),
+                       dst_row_size, dst_row_size_aligned, dst_row_size, 1);
 
         const uint32_t next_row = ir + 2;
         if (next_row < ir1) {
-            dma_queue_push_ddr_to_vtcm(dma_queue,
-                                       dma_make_ptr((uint8_t *) src_spad_row, src_data + (next_row * src_row_size)),
-                                       src_row_size_aligned, src_row_size, 1);
+            dma_queue_push(dma_q,
+                           dma_make_data(src_spad_row, src_data + (next_row * src_row_size)),
+                           src_row_size_aligned, src_row_size, src_row_size, 1);
         }
     }
 
-    dma_queue_flush(dma_queue);
+    dma_queue_flush(dma_q);
 
     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,
@@ -211,10 +211,6 @@ int op_cumsum_f32(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
 
-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
-        return HTP_STATUS_OK;
-    }
-
     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);
 
@@ -264,6 +260,9 @@ int op_cumsum_f32(struct htp_ops_context * octx) {
     };
 
     if (octx->ctx->vtcm_size < spad_per_thread * n_threads) {
+        if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
+            return HTP_STATUS_NO_SUPPORT;
+        }
         work_queue_run(octx->ctx->work_queue, cumsum_thread_f32, &cctx, n_threads);
     } else {
         work_queue_run(octx->ctx->work_queue, cumsum_thread_f32_dma, &cctx, n_threads);

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

@@ -13,7 +13,7 @@
 #include "hvx-types.h"
 #include "hex-utils.h"
 #include "hvx-copy.h"
-#include "hex-dma.h"
+#include "dma-queue.h"
 
 #define htp_diag_tensors_preamble                           \
     const struct htp_tensor * restrict src0 = octx->src[0]; \
@@ -59,7 +59,7 @@ static inline void hvx_diag_row_f32(const float * restrict src, float * restrict
 
 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];
+    dma_queue * dma_q = octx->ctx->dma[ith];
 
     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);
@@ -73,8 +73,8 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
     const size_t src_batch_size_aligned = dctx->src_batch_size_aligned;
     const size_t dst_row_size_aligned   = dctx->dst_row_size_aligned;
 
-    const uint8_t * src_data = (const uint8_t *) src0->data;
-    uint8_t *       dst_data = (uint8_t *) dst->data;
+    const dma_addr_t src_data = src0->data;
+    const dma_addr_t dst_data = dst->data;
 
     // 1 src buffer + 1 dst row buffer per thread in VTCM
     uint8_t * src_spad = octx->src0_spad.data + (ith * src_batch_size_aligned);
@@ -86,13 +86,13 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
         const uint32_t i3 = ib / ne02;
         const uint32_t i2 = ib % ne02;
 
-        const uint8_t * src_batch = src_data + i3 * nb03 + i2 * nb02;
+        const dma_addr_t src_batch = src_data + i3 * nb03 + i2 * nb02;
 
         // Fetch source vector into VTCM
-        dma_queue_push_ddr_to_vtcm(dma_queue,
-                                   dma_make_ptr(src_spad, src_batch),
-                                   src_batch_size_aligned, src_batch_size, 1);
-        dma_queue_flush(dma_queue);
+        dma_queue_push(dma_q,
+                       dma_make_data(src_spad, src_batch),
+                       src_batch_size_aligned, src_batch_size, src_batch_size, 1);
+        dma_queue_flush(dma_q);
 
         const float * src_spad_f32 = (const float *) src_spad;
         float       * dst_spad_f32 = (float *) dst_spad;
@@ -104,11 +104,11 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
             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;
-            dma_queue_push_vtcm_to_ddr(dma_queue,
-                                       dma_make_ptr(dst_row, dst_spad),
-                                       dst_row_size, dst_row_size_aligned, 1);
-            dma_queue_flush(dma_queue);
+            const dma_addr_t dst_row = dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1;
+            dma_queue_push(dma_q,
+                           dma_make_data(dst_row, dst_spad),
+                           dst_row_size, dst_row_size_aligned, dst_row_size, 1);
+            dma_queue_flush(dma_q);
         }
     }
 
@@ -156,10 +156,6 @@ int op_diag_f32(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
 
-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
-        return HTP_STATUS_OK;
-    }
-
     const uint32_t total_batches = src0->ne[2] * src0->ne[3];
     const size_t dst_batch_size  = dst->ne[1] * dst->nb[1];
 
@@ -221,6 +217,9 @@ int op_diag_f32(struct htp_ops_context * octx) {
     };
 
     if (octx->ctx->vtcm_size < spad_per_thread * n_threads) {
+        if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
+            return HTP_STATUS_NO_SUPPORT;
+        }
         work_queue_run(octx->ctx->work_queue, diag_thread_f32, &dctx, n_threads);
     } else {
         work_queue_run(octx->ctx->work_queue, diag_thread_f32_dma, &dctx, n_threads);

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

@@ -22,58 +22,69 @@ static inline uintptr_t align_up(uintptr_t addr, size_t align) {
     return (addr + align - 1) & ~(align - 1);
 }
 
-size_t dma_queue_sizeof(size_t capacity) {
+static inline size_t dma_ring_sizeof(size_t capacity) {
     capacity = pow2_ceil(capacity);
 
-    size_t size_q      = sizeof(dma_queue);
-    size_t offset_r    = align_up(size_q, HEX_L2_LINE_SIZE);
     size_t size_r      = sizeof(dma_ring);
-    size_t offset_desc = align_up(offset_r + size_r, HEX_L2_LINE_SIZE);
+    size_t offset_desc = align_up(size_r, HEX_L2_LINE_SIZE);
     size_t size_desc   = capacity * sizeof(dma_descriptor_2d);
-    size_t offset_dptr = align_up(offset_desc + size_desc, HEX_L2_LINE_SIZE);
-    size_t size_dptr   = capacity * sizeof(dma_ptr);
+    size_t offset_data = align_up(offset_desc + size_desc, HEX_L2_LINE_SIZE);
+    size_t size_data   = capacity * sizeof(dma_data);
 
-    return offset_dptr + size_dptr;
+    return offset_data + size_data;
+}
+
+static inline dma_ring * dma_ring_init(void * ptr, size_t capacity, struct htp_thread_trace * trace) {
+    capacity = pow2_ceil(capacity);
+
+    size_t size_r      = sizeof(dma_ring);
+    size_t offset_desc = align_up(size_r, HEX_L2_LINE_SIZE);
+    size_t size_desc   = capacity * sizeof(dma_descriptor_2d);
+    size_t offset_data = align_up(offset_desc + size_desc, HEX_L2_LINE_SIZE);
+
+    dma_ring * r = (dma_ring *) ptr;
+    r->trace     = trace;
+    r->capacity  = capacity;
+    r->idx_mask  = capacity - 1;
+    r->push_idx  = 0;
+    r->pop_idx   = 0;
+    r->desc      = (dma_descriptor_2d *) ((uintptr_t) ptr + offset_desc);
+    r->data      = (dma_data *) ((uintptr_t) ptr + offset_data);
+    r->tail      = &r->desc[capacity - 1];
+
+    return r;
+}
+
+size_t dma_queue_sizeof(size_t capacity) {
+    size_t size_q    = sizeof(dma_queue);
+    size_t offset_r0 = align_up(size_q, HEX_L2_LINE_SIZE);
+    size_t size_r0   = dma_ring_sizeof(capacity);
+    size_t offset_r1 = align_up(offset_r0 + size_r0, HEX_L2_LINE_SIZE);
+    size_t size_r1   = dma_ring_sizeof(DMA_FALLBACK_CAPACITY);
+
+    return offset_r1 + size_r1;
 }
 
 size_t dma_queue_alignof(void) {
     return HEX_L2_LINE_SIZE;
 }
 
-dma_queue_t dma_queue_init(void * ptr, size_t capacity, uintptr_t vtcm_base, size_t vtcm_size, struct htp_thread_trace * trace) {
-    capacity = pow2_ceil(capacity);
-
-    size_t size_q      = sizeof(dma_queue);
-    size_t offset_r    = align_up(size_q, HEX_L2_LINE_SIZE);
-    size_t size_r      = sizeof(dma_ring);
-    size_t offset_desc = align_up(offset_r + size_r, HEX_L2_LINE_SIZE);
-    size_t size_desc   = capacity * sizeof(dma_descriptor_2d);
-    size_t offset_dptr = align_up(offset_desc + size_desc, HEX_L2_LINE_SIZE);
-    size_t size_dptr   = capacity * sizeof(dma_ptr);
-
-    size_t total_size = offset_dptr + size_dptr;
+dma_queue_t dma_queue_init(void * ptr, size_t capacity, struct htp_thread_trace * trace) {
+    size_t total_size = dma_queue_sizeof(capacity);
     memset(ptr, 0, total_size);
 
     dma_queue * q = (dma_queue *) ptr;
-    dma_ring * r  = (dma_ring *) ((uintptr_t) ptr + offset_r);
 
-    q->ring    = r;
-    q->nocache = 0;
-    q->alias   = false;
+    size_t size_q    = sizeof(dma_queue);
+    size_t offset_r0 = align_up(size_q, HEX_L2_LINE_SIZE);
+    size_t size_r0   = dma_ring_sizeof(capacity);
+    size_t offset_r1 = align_up(offset_r0 + size_r0, HEX_L2_LINE_SIZE);
 
-    r->trace     = trace;
-    r->vtcm_base = vtcm_base;
-    r->vtcm_end  = vtcm_base + vtcm_size;
-    r->capacity  = capacity;
-    r->idx_mask  = capacity - 1;
-    r->push_idx  = 0;
-    r->pop_idx   = 0;
+    q->ring0 = dma_ring_init((void *) ((uintptr_t) ptr + offset_r0), capacity, trace);
+    q->ring1 = dma_ring_init((void *) ((uintptr_t) ptr + offset_r1), DMA_FALLBACK_CAPACITY, trace);
+    q->alias = false;
 
-    r->desc = (dma_descriptor_2d *) ((uintptr_t) ptr + offset_desc);
-    r->dptr = (dma_ptr *) ((uintptr_t) ptr + offset_dptr);
-    r->tail = &r->desc[capacity - 1];
-
-    FARF(HIGH, "dma-queue: capacity %u, unified memory size %zu\n", capacity, total_size);
+    FARF(HIGH, "dma-queue: capacity %u, unified memory size %zu\n", (unsigned) capacity, total_size);
 
     return q;
 }
@@ -86,13 +97,13 @@ size_t dma_queue_alias_sizeof(void) {
     return sizeof(dma_queue);
 }
 
-dma_queue_t dma_queue_alias_init(void * ptr, dma_queue_t main_q, uint8_t nocache) {
+dma_queue_t dma_queue_alias_init(void * ptr, dma_queue_t main_q) {
     dma_queue * q = (dma_queue *) ptr;
     memset(q, 0, sizeof(dma_queue));
 
-    q->ring    = main_q->ring;
-    q->nocache = nocache;
-    q->alias   = true;
+    q->ring0 = main_q->ring0;
+    q->ring1 = main_q->ring1;
+    q->alias = true;
 
     return q;
 }
@@ -101,4 +112,101 @@ void dma_queue_alias_free(dma_queue_t q) {
     (void) q;
 }
 
+bool dma_queue_push_fallback_2d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
+    dma_ring * r0 = q->ring0;
+    dma_ring * r1 = q->ring1;
 
+    if (((r0->push_idx + 1) & r0->idx_mask) == r0->pop_idx) {
+        return false;
+    }
+
+    r1->tail = r0->tail;
+
+    size_t rem_rows    = nrows;
+    dma_addr_t cur_dst = ddata.dst;
+    dma_addr_t cur_src = ddata.src;
+
+    while (rem_rows > 0) {
+        const uint32_t cur_rows = MIN(rem_rows, DMA_MAX_NROWS);
+        dma_data cur_data = dma_make_data(cur_dst, cur_src);
+        if (!dma_ring_push_single_2d(r1, cur_data, dst_stride, src_stride, row_size, cur_rows)) {
+            dma_ring_flush(r1);
+            dma_ring_push_single_2d(r1, cur_data, dst_stride, src_stride, row_size, cur_rows);
+        }
+        cur_dst  += cur_rows * dst_stride;
+        cur_src  += cur_rows * src_stride;
+        rem_rows -= cur_rows;
+    }
+
+    dma_ring_flush(r1);
+    r0->tail = r1->tail;
+
+    return dma_ring_push_single_2d(r0, ddata, 0, 0, 0, /*nrows=*/ 0);
+}
+
+bool dma_queue_push_fallback_contig(dma_queue * q, dma_data ddata, size_t total) {
+    dma_ring * r0 = q->ring0;
+    dma_ring * r1 = q->ring1;
+
+    if (((r0->push_idx + 1) & r0->idx_mask) == r0->pop_idx) {
+        return false;
+    }
+
+    r1->tail = r0->tail;
+
+    size_t rem_bytes   = total;
+    dma_addr_t cur_dst = ddata.dst;
+    dma_addr_t cur_src = ddata.src;
+
+    while (rem_bytes > 0) {
+        const uint32_t cur_bytes = MIN(rem_bytes, DMA_SAFE_CHUNK_SIZE);
+        dma_data cur_data = dma_make_data(cur_dst, cur_src);
+        if (!dma_ring_push_single_1d(r1, cur_data, cur_bytes)) {
+            dma_ring_flush(r1);
+            dma_ring_push_single_1d(r1, cur_data, cur_bytes);
+        }
+        cur_dst   += cur_bytes;
+        cur_src   += cur_bytes;
+        rem_bytes -= cur_bytes;
+    }
+
+    dma_ring_flush(r1);
+    r0->tail = r1->tail;
+
+    return dma_ring_push_single_1d(r0, ddata, /*size=*/ 0);
+}
+
+#if __HVX_ARCH__ < 75
+
+bool dma_queue_push_fallback_1d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
+    dma_ring * r0 = q->ring0;
+    dma_ring * r1 = q->ring1;
+
+    if (((r0->push_idx + 1) & r0->idx_mask) == r0->pop_idx) {
+        return false;
+    }
+
+    r1->tail = r0->tail;
+
+    size_t rem_rows    = nrows;
+    dma_addr_t cur_dst = ddata.dst;
+    dma_addr_t cur_src = ddata.src;
+
+    while (rem_rows > 0) {
+        dma_data cur_data = dma_make_data(cur_dst, cur_src);
+        if (!dma_ring_push_single_1d(r1, cur_data, row_size)) {
+            dma_ring_flush(r1);
+            dma_ring_push_single_1d(r1, cur_data, row_size);
+        }
+        cur_dst  += dst_stride;
+        cur_src  += src_stride;
+        rem_rows -= 1;
+    }
+
+    dma_ring_flush(r1);
+    r0->tail = r1->tail;
+
+    return dma_ring_push_single_1d(r0, ddata, /*size=*/ 0);
+}
+
+#endif

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

@@ -3,8 +3,10 @@
 
 #include <HAP_farf.h>
 #include <hexagon_types.h>
+#include <assert.h>
 #include <stdbool.h>
 #include <stdint.h>
+#include <string.h>
 #include "hex-utils.h"
 
 #include "hex-profile.h"
@@ -24,8 +26,8 @@ typedef struct dma_descriptor_1d_s {
     uint32_t src_bypass:1;
     uint32_t order:1;
     uint32_t done:1;
-    void *   src;
-    void *   dst;
+    uint32_t src;
+    uint32_t dst;
 } dma_descriptor_1d;
 
 #if __HVX_ARCH__ < 75
@@ -40,8 +42,8 @@ typedef struct dma_descriptor_2d_s {
     uint32_t src_bypass:1;
     uint32_t order:1;
     uint32_t done:1;
-    void *   src;
-    void *   dst;
+    uint32_t src;
+    uint32_t dst;
     uint32_t desc_type:8;
     uint32_t reserved1:24;
     uint32_t row_size:16;
@@ -64,10 +66,18 @@ typedef struct dma_descriptor_2d_s {
     uint32_t src_bypass:1;
     uint32_t order:1;
     uint32_t done:1;
-    void *   src;
-    void *   dst;
+    uint32_t src;
+    uint32_t dst;
     uint32_t desc_type:8;
+#if __HVX_ARCH__ > 79
+    uint32_t src_upper:8;
+    uint32_t dst_upper:8;
+    uint32_t allocation:2;
+    uint32_t reserved0:2;
+    uint32_t transform:4;
+#else
     uint32_t reserved0:24;
+#endif
     uint32_t row_size:24;
     uint32_t nrows_lo:8;
     uint32_t nrows_hi:8;
@@ -78,45 +88,63 @@ typedef struct dma_descriptor_2d_s {
 
 #endif
 
+#if __HVX_ARCH__ > 79
+typedef uint64_t dma_addr_t;
+#else
+typedef uint32_t dma_addr_t;
+#endif
+
 typedef struct {
-    void       *dst;
-    const void *src;
-} dma_ptr;
+    dma_addr_t dst;
+    dma_addr_t src;
+} dma_data;
+
+// Hardware descriptor field limits
+#define DMA_MAX_NROWS          0xFFFFu        // 16-bit HW descriptor limit (65535)
+#define DMA_MAX_SIZE_16B       0xFFFFu        // 16-bit HW descriptor limit for row_size (65535)
+#define DMA_MAX_STRIDE_16B     0xFFFFu        // 16-bit HW descriptor limit for strides (65535)
+#define DMA_MAX_SIZE_24B       0x00FFFFFFu    // 24-bit HW descriptor limit for row_size / 1D size (16MB - 1)
+#define DMA_MAX_STRIDE_24B     0x00FFFFFFu    // 24-bit HW descriptor limit for strides (16MB - 1)
+#define DMA_SAFE_CHUNK_SIZE    0x00F00000u    // ~15MB safe contiguous chunk size
+
+#define DMA_FALLBACK_CAPACITY  16u            // descriptors in secondary fallback ring
 
 typedef struct dma_ring_s dma_ring;
 struct dma_ring_s {
     dma_descriptor_2d * desc;      // descriptor pointers
     dma_descriptor_2d * tail;      // tail pointer
-    dma_ptr           * dptr;      // dst/src pointers
+    dma_data          * data;      // dst/src data
     uint32_t            push_idx;
     uint32_t            pop_idx;
     uint32_t            capacity;
     uint32_t            idx_mask;
     struct htp_thread_trace * trace;
-    uintptr_t           vtcm_base;
-    uintptr_t           vtcm_end;
 };
 
 typedef struct dma_queue_s dma_queue;
 typedef dma_queue * dma_queue_t;
 
 struct dma_queue_s {
-    dma_ring *          ring;      // Points to the descriptor ring state
-    uint8_t             nocache;   // Queue-specific bypass flag
+    dma_ring *          ring0;     // Main descriptor ring state
+    dma_ring *          ring1;     // Secondary fallback descriptor ring state
     bool                alias;     // When set, dma_queue_delete will not free the ring
 };
 
-
-
 size_t      dma_queue_sizeof(size_t capacity);
 size_t      dma_queue_alignof(void);
-dma_queue_t dma_queue_init(void * ptr, size_t capacity, uintptr_t vtcm_base, size_t vtcm_size, struct htp_thread_trace * trace);
+dma_queue_t dma_queue_init(void * ptr, size_t capacity, struct htp_thread_trace * trace);
 void        dma_queue_free(dma_queue_t q);
 
 size_t      dma_queue_alias_sizeof(void);
-dma_queue_t dma_queue_alias_init(void * ptr, dma_queue_t main_q, uint8_t nocache);
+dma_queue_t dma_queue_alias_init(void * ptr, dma_queue_t main_q);
 void        dma_queue_alias_free(dma_queue_t q);
 
+bool        dma_queue_push_fallback_2d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows);
+bool        dma_queue_push_fallback_contig(dma_queue * q, dma_data ddata, size_t total);
+#if __HVX_ARCH__ < 75
+bool        dma_queue_push_fallback_1d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows);
+#endif
+
 // TODO: technically we don't need these and could use Q6_dmstart/wait/etc instead
 // but those do not seem to always compiler properly.
 static inline void dmstart(void * next) {
@@ -141,36 +169,37 @@ static inline unsigned int dmwait(void) {
     return ret;
 }
 
-static inline dma_ptr dma_make_ptr(void *dst, const void *src)
+static inline dma_data dma_make_data_impl(dma_addr_t dst, dma_addr_t src)
 {
-    dma_ptr p = { dst, src };
-    return p;
+    dma_data d = { dst, src };
+    return d;
 }
 
-static inline bool dma_is_vtcm(const dma_queue * q, const void * ptr) {
-    return (uintptr_t) ptr >= q->ring->vtcm_base && (uintptr_t) ptr < q->ring->vtcm_end;
-}
+#define dma_make_data(dst, src) dma_make_data_impl((dma_addr_t) (dst), (dma_addr_t) (src))
+
+static inline bool dma_ring_push_single_1d(dma_ring * r, dma_data ddata, size_t size) {
+#if __HVX_ARCH__ > 79
+    assert(!((ddata.src | ddata.dst) >> 32) || size == 0);
+#endif
 
-static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t size) {
-    dma_ring * r = q->ring;
     if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) {
         return false;
     }
 
     dma_descriptor_1d * desc = (dma_descriptor_1d *) &r->desc[r->push_idx];
-    desc->src  = (void *) dptr.src;
-    desc->dst  = (void *) dptr.dst;
+    desc->src  = (uint32_t) ddata.src;
+    desc->dst  = (uint32_t) ddata.dst;
     desc->size = size;
 
-    r->dptr[r->push_idx] = dptr;
+    r->data[r->push_idx] = ddata;
 
     htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
 
     if (size) {
         desc->next       = NULL;
         desc->desc_size  = 0; // 1D mode
-        desc->src_bypass = dma_is_vtcm(q, dptr.src) ? 1 : q->nocache;
-        desc->dst_bypass = dma_is_vtcm(q, dptr.dst) ? 1 : q->nocache;
+        desc->src_bypass = 1;
+        desc->dst_bypass = 1;
         desc->order      = 0;
         desc->done       = 0;
 
@@ -185,8 +214,17 @@ static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t
     return true;
 }
 
-static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
-    dma_ring * r = q->ring;
+static inline bool dma_ring_push_single_2d(dma_ring * r, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
+#if __HVX_ARCH__ > 79
+    const uint32_t src_hi = (uint32_t) (ddata.src >> 32);
+    const uint32_t dst_hi = (uint32_t) (ddata.dst >> 32);
+    const bool is_ext     = (src_hi | dst_hi) != 0;
+
+    if (is_ext && ((ddata.src >> 40) || (ddata.dst >> 40))) {
+        return false;
+    }
+#endif
+
     if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) {
         return false;
     }
@@ -194,34 +232,44 @@ static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t
     dma_descriptor_2d * desc = &r->desc[r->push_idx];
 
     desc->next           = NULL;
-    desc->reserved0      = 0;
     desc->reserved1      = 0;
     desc->desc_size      = 1; // 2d mode
-    desc->src_bypass     = dma_is_vtcm(q, dptr.src) ? 1 : q->nocache;
-    desc->dst_bypass     = dma_is_vtcm(q, dptr.dst) ? 1 : q->nocache;
+    desc->src_bypass     = 1;
+    desc->dst_bypass     = 1;
     desc->src_comp       = 0;
     desc->dst_comp       = 0;
     desc->order          = 0;
     desc->done           = 0;
     desc->src_stride     = src_stride;
     desc->dst_stride     = dst_stride;
-    desc->src            = (void *) dptr.src;
-    desc->dst            = (void *) dptr.dst;
+    desc->src            = (uint32_t) ddata.src;
+    desc->dst            = (uint32_t) ddata.dst;
     desc->row_size       = row_size;
 
 #if __HVX_ARCH__ < 75
+    desc->reserved0      = 0;
     desc->desc_type      = 0; // 2d (16-bit) mode
     desc->nrows          = nrows;
     desc->src_offset     = 0;
     desc->dst_offset     = 0;
 #else
+#if __HVX_ARCH__ > 79
+    desc->src_upper      = src_hi;
+    desc->dst_upper      = dst_hi;
+    desc->allocation     = 0;
+    desc->reserved0      = 0;
+    desc->transform      = 0;
+    desc->desc_type      = is_ext ? 10 : 9; // 2d 40-bit or 24-bit mode
+#else
+    desc->reserved0      = 0;
     desc->desc_type      = 9; // 2d (24-bit) mode
+#endif
     desc->nrows_lo       = (nrows & 0xff);
     desc->nrows_hi       = (nrows >> 8);
     desc->offset         = 0;
 #endif
 
-    r->dptr[r->push_idx] = dptr;
+    r->data[r->push_idx] = ddata;
 
     htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
 
@@ -236,21 +284,20 @@ static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t
     return true;
 }
 
-static inline dma_ptr dma_queue_pop(dma_queue * q) {
-    dma_ring * r = q->ring;
-    dma_ptr dptr  = { NULL };
+static inline dma_data dma_ring_pop(dma_ring * r) {
+    dma_data ddata = { 0 };
 
     if (r->push_idx == r->pop_idx) {
-        return dptr;
+        return ddata;
     }
 
-    dptr = r->dptr[r->pop_idx];
+    ddata = r->data[r->pop_idx];
 
     volatile dma_descriptor_2d * desc = &r->desc[r->pop_idx];
 
     // Wait for desc to complete
     if (!desc->done) {
-        // FARF(ALWAYS, "dma-poll: idx %u dst %p src %p", r->pop_idx, dptr.dst, dptr.src);
+        // FARF(ALWAYS, "dma-poll: idx %u dst %p src %p", r->pop_idx, ddata.dst, ddata.src);
         while (!desc->done) {
             dmpoll();
         }
@@ -259,108 +306,133 @@ static inline dma_ptr dma_queue_pop(dma_queue * q) {
     htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);
 
     r->pop_idx = (r->pop_idx + 1) & r->idx_mask;
-    return dptr;
+    return ddata;
 }
 
-static inline dma_ptr dma_queue_pop_nowait(dma_queue * q) {
-    dma_ring * r = q->ring;
-    dma_ptr dptr  = { NULL };
+static inline dma_data dma_ring_pop_nowait(dma_ring * r) {
+    dma_data ddata = { 0 };
 
     if (r->push_idx == r->pop_idx) {
-        return dptr;
+        return ddata;
     }
 
-    dptr = r->dptr[r->pop_idx];
+    ddata = r->data[r->pop_idx];
 
     htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);
 
     r->pop_idx = (r->pop_idx + 1) & r->idx_mask;
-    return dptr;
+    return ddata;
+}
+
+static inline bool dma_ring_empty(dma_ring * r) {
+    return r->push_idx == r->pop_idx;
+}
+
+static inline void dma_ring_flush(dma_ring * r) {
+    while (!dma_ring_empty(r)) {
+        dma_ring_pop(r);
+    }
+}
+
+static inline uint32_t dma_ring_depth(dma_ring * r) {
+    return (r->push_idx - r->pop_idx) & r->idx_mask;
+}
+
+static inline uint32_t dma_ring_capacity(dma_ring * r) {
+    return r->capacity;
+}
+
+static inline bool dma_queue_push_single_1d(dma_queue * q, dma_data ddata, size_t size) {
+    return dma_ring_push_single_1d(q->ring0, ddata, size);
+}
+
+static inline bool dma_queue_push_single_2d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
+    return dma_ring_push_single_2d(q->ring0, ddata, dst_stride, src_stride, row_size, nrows);
+}
+
+static inline dma_data dma_queue_pop(dma_queue * q) {
+    return dma_ring_pop(q->ring0);
+}
+
+static inline dma_data dma_queue_pop_nowait(dma_queue * q) {
+    return dma_ring_pop_nowait(q->ring0);
 }
 
 static inline bool dma_queue_empty(dma_queue * q) {
-    return q->ring->push_idx == q->ring->pop_idx;
+    return dma_ring_empty(q->ring0);
 }
 
 static inline void dma_queue_flush(dma_queue * q) {
-    while (dma_queue_pop(q).dst != NULL) ;
+    dma_ring_flush(q->ring0);
 }
 
 static inline uint32_t dma_queue_depth(dma_queue * q) {
-    return (q->ring->push_idx - q->ring->pop_idx) & q->ring->idx_mask;
+    return dma_ring_depth(q->ring0);
 }
 
 static inline uint32_t dma_queue_capacity(dma_queue * q) {
-    return q->ring->capacity;
+    return dma_ring_capacity(q->ring0);
 }
 
 #if __HVX_ARCH__ < 75
 
-// Overflow-safe DMA push: all 2d descriptor fields (row_size, nrows, src_stride, dst_stride) are 16-bit, max 65535.
-// This version transparently handles values that exceed the 16-bit limit and submits chained DMA transtions.
-
-#define DMA_MAX_FIELD_VAL 65535u
-
-static inline bool dma_queue_push(dma_queue *q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
+static inline bool dma_queue_push(dma_queue *q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
     // Fast path: everything fits in 16 bits
     if (nrows == 0 || __builtin_expect(
-            row_size   <= DMA_MAX_FIELD_VAL &&
-            nrows      <= DMA_MAX_FIELD_VAL &&
-            src_stride <= DMA_MAX_FIELD_VAL &&
-            dst_stride <= DMA_MAX_FIELD_VAL, 1)) {
-        return dma_queue_push_single_2d(q, dptr, dst_stride, src_stride, row_size, nrows);
+            nrows      <= DMA_MAX_NROWS &&
+            row_size   <= DMA_MAX_SIZE_16B &&
+            src_stride <= DMA_MAX_STRIDE_16B &&
+            dst_stride <= DMA_MAX_STRIDE_16B, 1)) {
+        return dma_ring_push_single_2d(q->ring0, ddata, dst_stride, src_stride, row_size, nrows);
     }
 
-    // Contiguous block
-    // Use 1d DMA mode which supports sizes up to 24-bits (16MB)
+    // Contiguous block: 1D DMA mode supports up to 24-bit size (16MB)
     if (nrows == 1 || (row_size == src_stride && row_size == dst_stride)) {
         size_t total = row_size * nrows;
-        return dma_queue_push_single_1d(q, dptr, total);
+        if (total <= DMA_MAX_SIZE_24B) {
+            return dma_ring_push_single_1d(q->ring0, ddata, total);
+        }
+        return dma_queue_push_fallback_contig(q, ddata, total);
     }
 
-    // Stride overflow - fall back to row-by-row.
-    {
-        const uint8_t *src = (const uint8_t *) dptr.src;
-        uint8_t       *dst = (uint8_t *)       dptr.dst;
-        size_t r = 0;
-        while (r + 1 < nrows) {
-            dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride);
-            if (!dma_queue_push_single_1d(q, p, row_size)) {
-                dma_queue_flush(q);
-            } else {
-                r++;
-            }
-        }
-        dma_queue_flush(q);
-        dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride);
-        return dma_queue_push_single_1d(q, p, row_size);
+    // Row count overflow with 16-bit strides: chunk 2D descriptors via fallback ring
+    if (row_size <= DMA_MAX_SIZE_16B && src_stride <= DMA_MAX_STRIDE_16B && dst_stride <= DMA_MAX_STRIDE_16B) {
+        return dma_queue_push_fallback_2d(q, ddata, dst_stride, src_stride, row_size, nrows);
     }
+
+    // Stride or row_size overflow: row-by-row 1D via fallback ring
+    return dma_queue_push_fallback_1d(q, ddata, dst_stride, src_stride, row_size, nrows);
 }
 
 #else // HVX_ARCH >= 75
 
-static inline bool dma_queue_push(dma_queue *q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
-    // On v75 and up we always use 2d 24-bit mode
-    return dma_queue_push_single_2d(q, dptr, dst_stride, src_stride, row_size, nrows);
+static inline bool dma_queue_push(dma_queue *q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
+    if (nrows == 0 || __builtin_expect(
+            nrows      <= DMA_MAX_NROWS &&
+            row_size   <= DMA_MAX_SIZE_24B &&
+            src_stride <= DMA_MAX_STRIDE_24B &&
+            dst_stride <= DMA_MAX_STRIDE_24B, 1)) {
+        return dma_ring_push_single_2d(q->ring0, ddata, dst_stride, src_stride, row_size, nrows);
+    }
+
+    // Contiguous block exceeding 24 bits
+    if (nrows == 1 || (row_size == src_stride && row_size == dst_stride)) {
+        size_t total = row_size * nrows;
+        return dma_queue_push_fallback_contig(q, ddata, total);
+    }
+
+    return dma_queue_push_fallback_2d(q, ddata, dst_stride, src_stride, row_size, nrows);
 }
 
 #endif
 
-static inline bool dma_queue_push_ddr_to_vtcm(dma_queue * q, dma_ptr dptr, size_t dst_row_size, size_t src_row_size, size_t nrows) {
-    return dma_queue_push(q, dptr, dst_row_size, src_row_size, src_row_size, nrows);
-}
-
-static inline bool dma_queue_push_vtcm_to_ddr(dma_queue * q, dma_ptr dptr, size_t dst_row_size, size_t src_row_size, size_t nrows) {
-    return dma_queue_push(q, dptr, dst_row_size, src_row_size, dst_row_size, nrows);
-}
-
 #define DMA_CACHE_MAX_SIZE 256U
 
 typedef struct {
     uint8_t *base;
     uint32_t line_size;
     uint32_t capacity;
-    uint32_t src[DMA_CACHE_MAX_SIZE];
+    dma_addr_t src[DMA_CACHE_MAX_SIZE];
     uint16_t age[DMA_CACHE_MAX_SIZE];
 } dma_cache;
 
@@ -376,14 +448,14 @@ static inline void dma_cache_init(dma_cache *c, uint8_t *base, uint32_t line_siz
     }
 }
 
-static inline bool dma_cache_push(dma_queue *q, dma_cache *c, const uint8_t * src, uint32_t dst_stride, uint32_t src_stride, uint32_t row_size, uint32_t nrows)
+static inline bool dma_cache_push(dma_queue *q, dma_cache *c, dma_addr_t src_addr, uint32_t dst_stride, uint32_t src_stride, uint32_t row_size, uint32_t nrows)
 {
     uint32_t o_idx = 0;
     uint16_t o_age = 0;
     uint8_t *  dst = 0;
 
     for (unsigned i=0; i < c->capacity; i++) {
-        if (c->src[i] == (uint32_t) src) {
+        if (c->src[i] == src_addr) {
             c->age[i] = 0;
             dst = c->base + (i * c->line_size); nrows = 0; // dummy dma
         } else {
@@ -393,12 +465,12 @@ static inline bool dma_cache_push(dma_queue *q, dma_cache *c, const uint8_t * sr
     }
     if (!dst) {
         c->age[o_idx] = 0;
-        c->src[o_idx] = (uint32_t) src;
+        c->src[o_idx] = src_addr;
         dst = c->base + o_idx * c->line_size; // normal nrows dma
-        return dma_queue_push(q, dma_make_ptr(dst, src), dst_stride, src_stride, row_size, nrows);
+        return dma_queue_push(q, dma_make_data(dst, src_addr), dst_stride, src_stride, row_size, nrows);
     }
 
-    return dma_queue_push_single_1d(q, dma_make_ptr(dst, src), 0);
+    return dma_queue_push_single_1d(q, dma_make_data(dst, src_addr), 0);
 }
 
 #ifdef __cplusplus

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

@@ -88,8 +88,8 @@ int op_fill(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
-        return HTP_STATUS_OK;
+    if (htp_tensor_is_extended(dst)) {
+        return HTP_STATUS_NO_SUPPORT;
     }
 
     uint32_t row_start = 0;

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

@@ -12,7 +12,7 @@
 #include <stdint.h>
 #include <string.h>
 
-#include "hex-dma.h"
+#include "dma-queue.h"
 #include "hex-fastdiv.h"
 #include "hex-profile.h"
 #include "hmx-queue.h"
@@ -86,6 +86,7 @@ struct htp_fa_context {
     uint8_t * spad_v;
     uint8_t * spad_m;
     uint8_t * spad_a;
+    float *   spad_sinks;
 
     const struct htp_tensor * k;
     const struct htp_tensor * v;
@@ -108,6 +109,7 @@ struct hmx_fa_context {
 
     // Dimensions
     uint32_t     DK, DV;
+    uint32_t     DK_pad, DV_pad;  // head_dim rounded up to 64 for HMX tiling
     uint32_t     n_kv;        // kv_len
     uint32_t     n_kv_heads;  // number of KV heads
     uint32_t     n_heads;     // number of Q heads
@@ -148,6 +150,7 @@ struct hmx_fa_context {
     uint8_t *    vtcm_hmx_scales_qk;   // HMX output scales (qk_scale)
     __fp16 *     vtcm_mask_buf;        // VTCM mask buffer [Br * m_line], DMA'd per KV block
     __fp16 *     vtcm_slopes;          // ALiBi slopes [g_br]
+    float *      vtcm_sinks;           // Attention sinks
     size_t       row_buf_stride;       // HVX vectors per row buffer (Bc/64)
     size_t       mask_buf_row_stride;  // elements (__fp16) per row in mask buffer
     size_t       q_tile_bytes;
@@ -212,7 +215,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
 
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
-    dma_queue * dma = octx->ctx->dma[ith];
+    dma_queue * dma_q = octx->ctx->dma[ith];
 
     const uint32_t DK = nek0;
     const uint32_t DV = nev0;
@@ -242,27 +245,27 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
         const uint32_t iv3 = fastdiv(iq3, &factx->broadcast_rv3);
         const uint32_t iv2 = fastdiv(iq2, &factx->broadcast_rv2);
 
-        const __fp16 * mp_base = NULL;
+        dma_addr_t mp_base = 0;
         if (mask) {
             const uint32_t im2 = fastmodulo(iq2, mask->ne[2], &factx->src3_div2);
             const uint32_t im3 = fastmodulo(iq3, mask->ne[3], &factx->src3_div3);
-            mp_base = (const __fp16 *) ((const uint8_t *) mask->data + iq1*mask->nb[1] + im2*mask->nb[2] + im3*mask->nb[3]);
+            mp_base = mask->data + iq1*mask->nb[1] + im2*mask->nb[2] + im3*mask->nb[3];
         }
 
         // Precalculate next row variables if there is a next row
         bool has_next_ir = (ir + 1 < ir1);
         uint32_t next_ik2 = 0, next_ik3 = 0, next_iv2 = 0, next_iv3 = 0;
-        const uint8_t * next_q_row_ptr = NULL;
-        const __fp16 * next_mp_base = NULL;
+        dma_addr_t next_q_row_ptr = 0;
+        dma_addr_t next_mp_base = 0;
 
-        const uint8_t * next_k_src0 = NULL;
-        const uint8_t * next_v_src0 = NULL;
-        const uint8_t * next_m_src0 = NULL;
+        dma_addr_t next_k_src0 = 0;
+        dma_addr_t next_v_src0 = 0;
+        dma_addr_t next_m_src0 = 0;
         uint32_t next_block_size0 = 0;
 
-        const uint8_t * next_k_src1 = NULL;
-        const uint8_t * next_v_src1 = NULL;
-        const uint8_t * next_m_src1 = NULL;
+        dma_addr_t next_k_src1 = 0;
+        dma_addr_t next_v_src1 = 0;
+        dma_addr_t next_m_src1 = 0;
         uint32_t next_block_size1 = 0;
 
         if (has_next_ir) {
@@ -277,22 +280,22 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
             next_iv3 = fastdiv(next_iq3, &factx->broadcast_rv3);
             next_iv2 = fastdiv(next_iq2, &factx->broadcast_rv2);
 
-            next_q_row_ptr = (const uint8_t *) q->data + (next_iq1*nbq1 + next_iq2*nbq2 + next_iq3*nbq3);
+            next_q_row_ptr = q->data + next_iq1*nbq1 + next_iq2*nbq2 + next_iq3*nbq3;
 
             if (mask) {
                 const uint32_t next_im2 = fastmodulo(next_iq2, mask->ne[2], &factx->src3_div2);
                 const uint32_t next_im3 = fastmodulo(next_iq3, mask->ne[3], &factx->src3_div3);
-                next_mp_base = (const __fp16 *) ((const uint8_t *) mask->data + next_iq1*mask->nb[1] + next_im2*mask->nb[2] + next_im3*mask->nb[3]);
+                next_mp_base = mask->data + next_iq1*mask->nb[1] + next_im2*mask->nb[2] + next_im3*mask->nb[3];
             }
 
             // Precalculate next K/V block 0 source pointers
             {
                 const uint32_t ic_start = 0;
                 next_block_size0 = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
-                next_k_src0 = (const uint8_t *) k->data + (ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3);
-                next_v_src0 = (const uint8_t *) v->data + (ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3);
+                next_k_src0 = k->data + ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3;
+                next_v_src0 = v->data + ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3;
                 if (mask) {
-                    next_m_src0 = (const uint8_t *) (next_mp_base + ic_start);
+                    next_m_src0 = next_mp_base + ic_start * sizeof(__fp16);
                 }
             }
 
@@ -300,18 +303,18 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
             if (factx->n_blocks > 1) {
                 const uint32_t ic_start = 1 * FLASH_ATTN_BLOCK_SIZE;
                 next_block_size1 = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
-                next_k_src1 = (const uint8_t *) k->data + (ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3);
-                next_v_src1 = (const uint8_t *) v->data + (ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3);
+                next_k_src1 = k->data + ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3;
+                next_v_src1 = v->data + ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3;
                 if (mask) {
-                    next_m_src1 = (const uint8_t *) (next_mp_base + ic_start);
+                    next_m_src1 = next_mp_base + ic_start * sizeof(__fp16);
                 }
             }
         }
 
         if (ir == ir0) {
             // Fetch Q row
-            const uint8_t * q_row_ptr = (const uint8_t *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3);
-            dma_queue_push(dma, dma_make_ptr(spad_q, q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
+            const dma_addr_t q_row_ptr = q->data + iq1*nbq1 + iq2*nbq2 + iq3*nbq3;
+            dma_queue_push(dma_q, dma_make_data(spad_q, q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
 
             // Prefetch first two blocks
             for (uint32_t ib = 0; ib < MIN(factx->n_blocks, 2); ++ib) {
@@ -319,20 +322,20 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                 const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
 
                 // K
-                const uint8_t * k_src = (const uint8_t *) k->data + (ic_start*nbk1 + ik2*nbk2 + ik3*nbk3);
+                const dma_addr_t k_src = k->data + ic_start*nbk1 + ik2*nbk2 + ik3*nbk3;
                 uint8_t * k_dst = spad_k + (ib % 2) * factx->size_k_block;
-                dma_queue_push(dma, dma_make_ptr(k_dst, k_src), factx->size_k_row_padded, nbk1, size_k_row, current_block_size);
+                dma_queue_push(dma_q, dma_make_data(k_dst, k_src), factx->size_k_row_padded, nbk1, size_k_row, current_block_size);
 
                 // V
-                const uint8_t * v_src = (const uint8_t *) v->data + (ic_start*nbv1 + iv2*nbv2 + iv3*nbv3);
+                const dma_addr_t v_src = v->data + ic_start*nbv1 + iv2*nbv2 + iv3*nbv3;
                 uint8_t * v_dst = spad_v + (ib % 2) * factx->size_v_block;
-                dma_queue_push(dma, dma_make_ptr(v_dst, v_src), factx->size_v_row_padded, nbv1, size_v_row, current_block_size);
+                dma_queue_push(dma_q, dma_make_data(v_dst, v_src), factx->size_v_row_padded, nbv1, size_v_row, current_block_size);
 
                 // Mask
                 if (mask) {
-                    const uint8_t * m_src = (const uint8_t *) (mp_base + ic_start);
+                    const dma_addr_t m_src = mp_base + ic_start * sizeof(__fp16);
                     // Mask is 1D contiguous for this row
-                    dma_cache_push(dma, &m_cache, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
+                    dma_cache_push(dma_q, &m_cache, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
                 }
             }
         }
@@ -347,7 +350,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
         hvx_splat_f32_a(spad_a, 0, DV);
         float * VKQ32 = (float *) (spad_a + 0);
 
-        uint8_t * q_ptr_vtcm = dma_queue_pop(dma).dst;
+        uint8_t * q_ptr_vtcm = (void *) dma_queue_pop(dma_q).dst;
         if (factx->is_q_fp32) {
             hvx_copy_f16_f32_aa(q_ptr_vtcm, q_ptr_vtcm, DK);  // inplace convert f32 to f16
         }
@@ -364,9 +367,9 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
             const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
 
             // Wait for DMA
-            uint8_t * k_base = dma_queue_pop(dma).dst; // K
-            uint8_t * v_base = dma_queue_pop(dma).dst; // V
-            __fp16  * m_base = mask ? dma_queue_pop(dma).dst : NULL; // M
+            uint8_t * k_base = (void *) dma_queue_pop(dma_q).dst; // K
+            uint8_t * v_base = (void *) dma_queue_pop(dma_q).dst; // V
+            __fp16  * m_base = mask ? (__fp16 *) dma_queue_pop(dma_q).dst : NULL; // M
 
             if (factx->k->type == HTP_TYPE_Q8_0) {
                 htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir);
@@ -423,7 +426,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
 
             if (ib + 1 == factx->n_blocks && has_next_ir) {
                 // Queue next row's Q row!
-                dma_queue_push(dma, dma_make_ptr(spad_q, next_q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
+                dma_queue_push(dma_q, dma_make_data(spad_q, next_q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
 
                 if (factx->n_blocks % 2 == 0) {
                     // Queue next row's block 0 (into buffer slot 0)
@@ -431,14 +434,14 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                     uint8_t * v_dst = spad_v + 0 * factx->size_v_block;
 
                     // K (block 0 of next row)
-                    dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0);
+                    dma_queue_push(dma_q, dma_make_data(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0);
 
                     // V (block 0 of next row)
-                    dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0);
+                    dma_queue_push(dma_q, dma_make_data(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0);
 
                     // Mask (block 0 of next row)
                     if (mask) {
-                        dma_cache_push(dma, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1);
+                        dma_cache_push(dma_q, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1);
                     }
                 }
             }
@@ -501,17 +504,17 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                 const uint32_t next_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - next_ic_start);
 
                 // K
-                const uint8_t * k_src = (const uint8_t *) k->data + (next_ic_start*nbk1 + ik2*nbk2 + ik3*nbk3);
-                dma_queue_push(dma, dma_make_ptr(k_base, k_src), factx->size_k_row_padded, nbk1, size_k_row, next_block_size);
+                const dma_addr_t k_src = k->data + next_ic_start*nbk1 + ik2*nbk2 + ik3*nbk3;
+                dma_queue_push(dma_q, dma_make_data(k_base, k_src), factx->size_k_row_padded, nbk1, size_k_row, next_block_size);
 
                 // V
-                const uint8_t * v_src = (const uint8_t *) v->data + (next_ic_start*nbv1 + iv2*nbv2 + iv3*nbv3);
-                dma_queue_push(dma, dma_make_ptr(v_base, v_src), factx->size_v_row_padded, nbv1, size_v_row, next_block_size);
+                const dma_addr_t v_src = v->data + next_ic_start*nbv1 + iv2*nbv2 + iv3*nbv3;
+                dma_queue_push(dma_q, dma_make_data(v_base, v_src), factx->size_v_row_padded, nbv1, size_v_row, next_block_size);
 
                 // Mask
                 if (mask) {
-                    const uint8_t * m_src = (const uint8_t *) (mp_base + next_ic_start);
-                    dma_cache_push(dma, &m_cache, m_src, next_block_size * 2, next_block_size * 2, next_block_size * 2, 1);
+                    const dma_addr_t m_src = mp_base + next_ic_start * sizeof(__fp16);
+                    dma_cache_push(dma_q, &m_cache, m_src, next_block_size * 2, next_block_size * 2, next_block_size * 2, 1);
                 }
             }
         }
@@ -524,14 +527,14 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                     uint8_t * v_dst = spad_v + 1 * factx->size_v_block;
 
                     // K (block 1 of next row)
-                    dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1);
+                    dma_queue_push(dma_q, dma_make_data(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1);
 
                     // V (block 1 of next row)
-                    dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1);
+                    dma_queue_push(dma_q, dma_make_data(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1);
 
                     // Mask (block 1 of next row)
                     if (mask) {
-                        dma_cache_push(dma, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1);
+                        dma_cache_push(dma_q, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1);
                     }
                 }
             } else {
@@ -541,14 +544,14 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                     uint8_t * v_dst = spad_v + 0 * factx->size_v_block;
 
                     // K (block 0 of next row)
-                    dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0);
+                    dma_queue_push(dma_q, dma_make_data(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0);
 
                     // V (block 0 of next row)
-                    dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0);
+                    dma_queue_push(dma_q, dma_make_data(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0);
 
                     // Mask (block 0 of next row)
                     if (mask) {
-                        dma_cache_push(dma, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1);
+                        dma_cache_push(dma_q, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1);
                     }
                 }
 
@@ -558,14 +561,14 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                     uint8_t * v_dst = spad_v + 1 * factx->size_v_block;
 
                     // K (block 1 of next row)
-                    dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1);
+                    dma_queue_push(dma_q, dma_make_data(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1);
 
                     // V (block 1 of next row)
-                    dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1);
+                    dma_queue_push(dma_q, dma_make_data(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1);
 
                     // Mask (block 1 of next row)
                     if (mask) {
-                        dma_cache_push(dma, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1);
+                        dma_cache_push(dma_q, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1);
                     }
                 }
             }
@@ -577,7 +580,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
         float S = hvx_vec_get_f32(S_vec);
 
         if (sinks) {
-            const float s = ((float *)((char *) sinks->data))[h];
+            const float s = factx->spad_sinks[h];
 
             float vs = 1.0f;
 
@@ -652,7 +655,7 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data)
             hvx_dequantize_row_q8_0_f16(row_k, row_k, factx->DK);
         }
     }
-    hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles[args->buf_idx], (const __fp16 *) args->curr_k, total_rows, factx->DK,
+    hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles[args->buf_idx], (const __fp16 *) args->curr_k, total_rows, factx->DK_pad,
                              args->src_stride, start, end);
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
 }
@@ -706,7 +709,7 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data)
             hvx_dequantize_row_q8_0_f16(row_v, row_v, factx->DV);
         }
     }
-    hmx_interleave_cols_to_tiles(v_tiles_dst, (const __fp16 *) args->v_src, total_rows, factx->DV,
+    hmx_interleave_cols_to_tiles(v_tiles_dst, (const __fp16 *) args->v_src, total_rows, factx->DV_pad,
                              args->src_stride, (uint32_t) args->n_col_tiles, start, end);
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start));
 }
@@ -780,7 +783,7 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
         const size_t m_end         = hex_smin(m_start + m_bytes_per_t, col_vec_bytes);
 
         if (factx->sinks) {
-            const float * sinks_data = (const float *) (uintptr_t) factx->sinks->data;
+            const float * sinks_data = factx->vtcm_sinks;
             float *       m_vec      = (float *) factx->vtcm_m_vec;
             const size_t  r_start    = l_start / sizeof(float);
             const size_t  r_end      = l_end / sizeof(float);
@@ -832,17 +835,22 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
         const uint32_t            kv_head = args->kv_head;
         const uint32_t            ib3     = args->ib3;
 
-        assert(factx->DK == factx->DV);
-
         const bool use_q_dma = (factx->vtcm_q_dma != NULL);
 
         __fp16 * q_tiles = factx->vtcm_q_tiles;
+        const size_t DK_pad = factx->DK_pad;
         if (use_q_dma) {
             const size_t g_rows_end = hex_smin(end, n_rows_g);
             const uint32_t d_limit = factx->is_q_fp32 ? DK / 32 : DK / 64;
 
             uint8_t * q_flat  = (uint8_t *) factx->vtcm_q_dma;
-            if (factx->is_q_fp32) {
+            if (DK_pad != DK) {
+                if (factx->is_q_fp32) {
+                    hmx_fa_q_prep_fp32_pad(q_tiles, q_flat, start, end, g_rows_end, DK, DK_pad, G, args->n_rows_q, &factx->div_G, args->q_transposed);
+                } else {
+                    hmx_fa_q_prep_fp16_pad(q_tiles, q_flat, start, end, g_rows_end, DK, DK_pad, G, args->n_rows_q, &factx->div_G, args->q_transposed);
+                }
+            } else if (factx->is_q_fp32) {
                 switch (d_limit) {
                 case 2:  hmx_fa_q_prep_fp32_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break;
                 case 4:  hmx_fa_q_prep_fp32_d4(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break;
@@ -858,7 +866,7 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
         } else {
             // Fallback: direct-from-DDR/L2 path
             hmx_fa_q_prep_fallback(q_tiles, q->data, q->nb[1], q->nb[2], q->nb[3],
-                                   q_start, kv_head, ib3, start, end, n_rows_g, G, DK, factx->is_q_fp32, &factx->div_G);
+                                   q_start, kv_head, ib3, start, end, n_rows_g, G, DK, DK_pad, factx->is_q_fp32, &factx->div_G);
         }
     }
 
@@ -952,6 +960,8 @@ static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
     const uint32_t            kv_head    = args->kv_head;
     const uint32_t            ib3        = args->ib3;
 
+    const size_t DV_pad = factx->DV_pad;
+
     size_t q_idx = fastdiv(start, &factx->div_G);
     size_t h_idx = fastmodulo(start, G, &factx->div_G);
 
@@ -961,7 +971,7 @@ static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
 
         size_t         r0            = r / HMX_FP16_TILE_N_ROWS;
         size_t         r1            = r % HMX_FP16_TILE_N_ROWS;
-        const __fp16 * tile_row_base = o_tile_src + r0 * HMX_FP16_TILE_N_ROWS * DV;
+        const __fp16 * tile_row_base = o_tile_src + r0 * HMX_FP16_TILE_N_ROWS * DV_pad;
 
         for (uint32_t d = 0; d < DV / 32; ++d) {
             const HVX_Vector * in_tile = (const HVX_Vector *) (tile_row_base + d * HMX_FP16_TILE_N_ELMS);
@@ -972,6 +982,16 @@ static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
                 *(HVX_UVector *) (out + d * 32) = Q6_V_hi_W(vp);
             }
         }
+        // Ragged tail: DV not a multiple of 32 (e.g. 72 -> last 8 lanes). Partial vector-write
+        // for the remaining (DV % 32) floats.
+        const uint32_t d_tail = DV / 32;
+        const uint32_t rem    = DV - d_tail * 32;
+        if (rem) {
+            const HVX_Vector * in_tile = (const HVX_Vector *) (tile_row_base + d_tail * HMX_FP16_TILE_N_ELMS);
+            HVX_VectorPair     vp      = hvx_vec_f16_to_f32_shuff(in_tile[r1 / 2]);
+            HVX_Vector         vd      = (r1 % 2 == 0) ? Q6_V_lo_W(vp) : Q6_V_hi_W(vp);
+            hvx_vec_store_u((void *) (out + d_tail * 32), rem * sizeof(float), vd);
+        }
 
         h_idx++;
         if (h_idx == G) {
@@ -1006,6 +1026,9 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
     const uint32_t            kv_head    = args->kv_head;
     const uint32_t            ib3        = args->ib3;
 
+    // O-tiles use the padded head dim (DV_pad); dst holds the real DV lanes.
+    const size_t DV_pad = factx->DV_pad;
+
     size_t q_idx = fastdiv(start, &factx->div_G);
     size_t h_idx = fastmodulo(start, G, &factx->div_G);
 
@@ -1015,7 +1038,7 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
 
         size_t         r0            = r / HMX_FP16_TILE_N_ROWS;
         size_t         r1            = r % HMX_FP16_TILE_N_ROWS;
-        const __fp16 * tile_row_base = o_tile_src + r0 * HMX_FP16_TILE_N_ROWS * DV;
+        const __fp16 * tile_row_base = o_tile_src + r0 * HMX_FP16_TILE_N_ROWS * DV_pad;
 
         for (uint32_t d = 0; d < DV / 64; ++d) {
             const __fp16 *     in_dtile = tile_row_base + d * HMX_FP16_TILE_N_ELMS * 2;
@@ -1028,6 +1051,17 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
                 *(HVX_UVector *) (out + d * 64) = Q6_V_hi_W(vp);
             }
         }
+        // Ragged tail when DV is not a multiple of 64.
+        const uint32_t d_tail = DV / 64;
+        const uint32_t rem    = DV - d_tail * 64;
+        if (rem) {
+            const __fp16 *     in_dtile = tile_row_base + d_tail * HMX_FP16_TILE_N_ELMS * 2;
+            const HVX_Vector * pv_in0   = ((const HVX_Vector *) in_dtile) + r1 / 2;
+            const HVX_Vector * pv_in1   = pv_in0 + 16;
+            HVX_VectorPair     vp       = Q6_W_vdeal_VVR(*pv_in1, *pv_in0, -2);
+            HVX_Vector         vd       = (r1 % 2 == 0) ? Q6_V_lo_W(vp) : Q6_V_hi_W(vp);
+            hvx_vec_store_u((void *) (out + d_tail * 64), rem * sizeof(__fp16), vd);
+        }
 
         h_idx++;
         if (h_idx == G) {
@@ -1747,7 +1781,7 @@ static __attribute__((noinline)) void fa_compute_slopes(
 }
 
 static void fa_push_mask_dma_gqa(
-    dma_queue *               dma,
+    dma_queue *               dma_q,
     const struct htp_tensor * mask,
     uint32_t                  q_start,
     uint32_t                  im3,
@@ -1762,36 +1796,36 @@ static void fa_push_mask_dma_gqa(
     for (uint32_t g = 0; g < G; ++g) {
         const uint32_t h_idx = kv_head * G + g;
         const uint32_t im2 = fastmodulo(h_idx, mask->ne[2], &factx->src3_div2);
-        const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] +
-                                 im2 * mask->nb[2] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
+        const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] +
+                                  im2 * mask->nb[2] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
         uint8_t * ms_dst = (uint8_t *) factx->vtcm_mask_buf + g * m_line_bytes;
-        dma_queue_push(dma, dma_make_ptr(ms_dst, ms_src), G * m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
+        dma_queue_push(dma_q, dma_make_data(ms_dst, ms_src), G * m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
     }
 }
 
-static void fa_pop_mask_dma_gqa(dma_queue * dma, uint32_t G) {
+static void fa_pop_mask_dma_gqa(dma_queue * dma_q, uint32_t G) {
     for (uint32_t g = 0; g < G; ++g) {
-        dma_queue_pop(dma);
+        dma_queue_pop(dma_q);
     }
 }
 
-static inline void fa_prefetch_block(dma_queue * dma, const struct htp_tensor * k, const struct htp_tensor * v, const struct htp_tensor * mask,
+static inline void fa_prefetch_block(dma_queue * dma_q, const struct htp_tensor * k, const struct htp_tensor * v, const struct htp_tensor * mask,
                                      uint32_t b, size_t Bc, size_t size_k_row_padded, size_t size_k_row, size_t size_v_row_padded, size_t size_v_row,
                                      uint32_t ik2, uint32_t ik3, uint32_t iv2, uint32_t iv3, uint32_t q_start, uint32_t im3, uint32_t kv_head, uint32_t G,
                                      size_t m_line_bytes, size_t n_rows_q, size_t nek1, size_t prefetch_buf, struct hmx_fa_context * factx) {
     const uint32_t  prefetch_start = b * Bc;
     const uint32_t  prefetch_rows  = hex_smin(Bc, nek1 - prefetch_start);
-    const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
-    dma_queue_push(dma, dma_make_ptr(factx->vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
-    const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
-    dma_queue_push(dma, dma_make_ptr(factx->vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
+    const dma_addr_t k_prefetch_src = k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
+    dma_queue_push(dma_q, dma_make_data(factx->vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
+    const dma_addr_t v_prefetch_src = v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
+    dma_queue_push(dma_q, dma_make_data(factx->vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
 
     if (mask) {
         if (__builtin_expect(factx->mask_broadcast, true)) {
-            const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + prefetch_start * sizeof(__fp16);
-            dma_cache_push(dma, &factx->m_cache, ms_src, m_line_bytes, mask->nb[1], prefetch_rows * sizeof(__fp16), n_rows_q);
+            const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + prefetch_start * sizeof(__fp16);
+            dma_cache_push(dma_q, &factx->m_cache, ms_src, m_line_bytes, mask->nb[1], prefetch_rows * sizeof(__fp16), n_rows_q);
         } else {
-            fa_push_mask_dma_gqa(dma, mask, q_start, im3, prefetch_start, kv_head, G, m_line_bytes, prefetch_rows, n_rows_q, factx);
+            fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, prefetch_start, kv_head, G, m_line_bytes, prefetch_rows, n_rows_q, factx);
         }
     }
 }
@@ -1829,8 +1863,11 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     const uint32_t DK = neq0;
     const uint32_t DV = nev0;
 
-    // HMX requires head_dim to be multiple of 32
-    if (DK % 32 != 0 || DV % 32 != 0) {
+    // HMX tiles head_dim in units of 64. head_dim need not be 64- (or 32-) aligned:
+    // we can operate on DK/DV rounded up to 64 with tail lanes [D, D_pad) zero-filled.
+    const uint32_t DK_pad = hex_round_up(DK, 64);
+    const uint32_t DV_pad = hex_round_up(DV, 64);
+    if (DK == 0 || DV == 0) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
@@ -1847,6 +1884,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     factx.n_threads      = kparams->n_threads;
     factx.DK             = DK;
     factx.DV             = DV;
+    factx.DK_pad         = DK_pad;
+    factx.DV_pad         = DV_pad;
     factx.n_kv           = nek1;
     factx.n_kv_heads     = n_kv_heads;
     factx.n_heads        = neq2;
@@ -1905,16 +1944,18 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
 
     // ======== VTCM allocation (GQA-aware) ========
     // K/V row sizes drive the DMA descriptors (not the VTCM layout) and are used
-    // throughout the KV loop below.
+    // throughout the KV loop below. The DMA copies only the real DK/DV columns; the
+    // staging rows are padded to hold DK_pad/DV_pad columns (tail zero-filled below)
+    // so the HMX interleave/tile logic can operate on 64-aligned head dims.
     const size_t size_k_row        = htp_tensor_get_row_size(k->type, DK);
     const size_t size_v_row        = htp_tensor_get_row_size(v->type, DV);
-    const size_t size_k_row_padded = hex_round_up(DK * sizeof(__fp16), 128);
-    const size_t size_v_row_padded = hex_round_up(DV * sizeof(__fp16), 128);
+    const size_t size_k_row_padded = hex_round_up(DK_pad * sizeof(__fp16), 128);
+    const size_t size_v_row_padded = hex_round_up(DV_pad * sizeof(__fp16), 128);
 
     // Build the VTCM layout once (shared with the host estimator) and place every
-    // scratch buffer at its computed offset.
+    // scratch buffer at its computed offset. Padded head dims size the HMX tiles.
     struct hmx_fa_vtcm_layout L;
-    hmx_fa_vtcm_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline, factx.is_q_fp32);
+    hmx_fa_vtcm_layout_build(&L, G, DK_pad, DV_pad, Br, Bc, n_threads, pipeline, factx.is_q_fp32, factx.sinks != NULL, factx.n_heads);
 
     if (L.total_bytes > ctx->vtcm_size) {
         return HTP_STATUS_VTCM_TOO_SMALL;
@@ -1956,22 +1997,41 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     factx.col_vec_bytes       = L.col_vec_bytes;
     factx.d_tile_bytes        = L.d_tile_bytes;
     factx.vtcm_slopes         = VTCM_LAYOUT_PTR(__fp16, base, L.off_slopes);
+    factx.vtcm_sinks          = VTCM_LAYOUT_PTR_OPTIONAL(float, base, L.off_sinks, factx.sinks != NULL);
 
     const size_t m_line_bytes = L.m_line_bytes;  // used by the mask DMAs in the KV loop
 
     dma_cache_init(&factx.m_cache, (uint8_t *) factx.vtcm_mask_buf, L.m_buf_slot_bytes, HMX_FA_DMA_CACHE_SIZE);
 
+    // Head-dim padding: the K/V DMA staging buffers and the flat-Q buffer are laid out
+    // with padded row strides (size_{k,v,q}_row_padded, covering D_pad columns) but the
+    // DMA only writes the real D columns per row. Zero the whole staging buffers once up
+    // front so tail lanes [D, D_pad) stay zero for all KV blocks. No-op when already aligned.
+    if (DK_pad != DK || DV_pad != DV) {
+        const size_t k_buf_bytes = (size_t) factx.Bc * size_k_row_padded;
+        const size_t v_buf_bytes = (size_t) factx.Bc * size_v_row_padded;
+        hvx_splat_u8_a((char *) factx.vtcm_k_fp16[0], 0, k_buf_bytes);
+        hvx_splat_u8_a((char *) factx.vtcm_k_fp16[1], 0, k_buf_bytes);
+        hvx_splat_u8_a((char *) factx.vtcm_v_fp16[0], 0, v_buf_bytes);
+        hvx_splat_u8_a((char *) factx.vtcm_v_fp16[1], 0, v_buf_bytes);
+        // Flat-Q DMA scratch
+        if (factx.vtcm_q_dma) {
+            const size_t q_dma_bytes = hex_align_up(factx.g_br * DK * (factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16)), 128);
+            hvx_splat_u8_a((char *) factx.vtcm_q_dma, 0, q_dma_bytes);
+        }
+    }
+
     // ======== Initialize HMX output scales ========
     hmx_init_column_scales(factx.vtcm_hmx_scales_id, Q6_V_vsplat_R(0x3c00)); // 1.0
     hmx_init_column_scales(factx.vtcm_hmx_scales_qk, hvx_vec_splat_f16(factx.scale));
 
-    // ======== Skip compute if profiling ========
-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
-        return HTP_STATUS_OK;
-    }
-
     // ======== DMA setup ========
-    dma_queue * const dma = ctx->dma[0];
+    dma_queue * const dma_q = ctx->dma[0];
+
+    if (factx.sinks) {
+        dma_queue_push(dma_q, dma_make_data(factx.vtcm_sinks, factx.sinks->data), L.sinks_bytes, 0, factx.sinks->size, 1);
+        dma_queue_pop(dma_q);
+    }
 
     const size_t n_row_tiles_g_br = g_br / HMX_FP16_TILE_N_ROWS;
     const size_t n_tiles_per_bc   = Bc / HMX_FP16_TILE_N_COLS;
@@ -2007,32 +2067,32 @@ 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 == q_start_min && kv_head == 0) {
-                    const uint8_t * q_ptr = (const uint8_t *) q->data + q_start * q->nb[1] +
+                    const dma_addr_t q_ptr = 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);
+                    dma_queue_push(dma_q, dma_make_data(factx.vtcm_q_dma, q_ptr), q_row_bytes, hex_smax(q_src_stride, q_row_bytes), q_row_bytes, n_rows);
 
                     if (factx.n_kv_blocks > 0) {
-                        const uint8_t * k_src = (const uint8_t *) k->data + ik2 * k->nb[2] + ik3 * k->nb[3];
-                        dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
+                        const dma_addr_t k_src = k->data + ik2 * k->nb[2] + ik3 * k->nb[3];
+                        dma_queue_push(dma_q, dma_make_data(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
 
-                        const uint8_t * v_src = (const uint8_t *) v->data + iv2 * v->nb[2] + iv3 * v->nb[3];
-                        dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
+                        const dma_addr_t v_src = v->data + iv2 * v->nb[2] + iv3 * v->nb[3];
+                        dma_queue_push(dma_q, dma_make_data(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
 
                         if (factx.pipeline && mask) {
                             if (__builtin_expect(factx.mask_broadcast, true)) {
-                                const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + 0;
-                                dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), n_rows_q);
+                                const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + 0;
+                                dma_cache_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), n_rows_q);
                             } else {
-                                fa_push_mask_dma_gqa(dma, mask, q_start, im3, 0, kv_head, G, m_line_bytes, kv_rows0, n_rows_q, &factx);
+                                fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, 0, kv_head, G, m_line_bytes, kv_rows0, n_rows_q, &factx);
                             }
                         }
                     }
                 }
 
                 // 2. Pop Q DMA (blocks until Q is loaded)
-                dma_queue_pop(dma);
+                dma_queue_pop(dma_q);
 
                 // ---- Load Q block & Initialize per-block state ----
                 fa_phase_q_load(&factx, q, q_start, kv_head, ib3, n_rows_g);
@@ -2059,12 +2119,12 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
 
                     // Prefetch block 1 early if there are multiple blocks
                     if (factx.n_kv_blocks > 1) {
-                        fa_prefetch_block(dma, k, v, mask, 1, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
+                        fa_prefetch_block(dma_q, k, v, mask, 1, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
                                           ik2, ik3, iv2, iv3, q_start, im3, kv_head, G, m_line_bytes, n_rows_q, nek1, 1, &factx);
                     }
 
                     // Prep and start QK-dot(0)
-                    void * curr_k0 = dma_queue_pop(dma).dst;
+                    void * curr_k0 = (void *) dma_queue_pop(dma_q).dst;
                     fa_phase_k_interleave(&factx, kv_rows0, k_src_stride, curr_k0, 0, 0);
 
                     qk_job[0].q_tiles        = factx.vtcm_q_tiles;
@@ -2072,7 +2132,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                     qk_job[0].s_tiles        = factx.vtcm_s_tiles[0];
                     qk_job[0].n_row_tiles    = n_row_tiles;
                     qk_job[0].n_col_tiles    = hmx_ceil_div(kv_rows0, HMX_FP16_TILE_N_COLS);
-                    qk_job[0].n_dot_tiles    = DK / 32;
+                    qk_job[0].n_dot_tiles    = DK_pad / 32;
                     qk_job[0].n_tiles_per_bc = n_tiles_per_bc;
                     qk_job[0].hmx_scales     = factx.vtcm_hmx_scales_qk;
                     hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job[0]));
@@ -2083,16 +2143,16 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                         const size_t   n_col_tiles = hmx_ceil_div(kv_rows, HMX_FP16_TILE_N_COLS);
 
                         // ---- 1. Pop and run V-prep for current block ----
-                        void * curr_v = dma_queue_pop(dma).dst;
+                        void * curr_v = (void *) dma_queue_pop(dma_q).dst;
                         fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, factx.vtcm_v_tiles[buf_idx], n_tiles_per_bc, kv_start);
 
                         // ---- 2. Pop and run mask-prep for current block ----
                         __fp16 * current_mask_vtcm = NULL;
                         if (mask) {
                             if (__builtin_expect(factx.mask_broadcast, true)) {
-                                current_mask_vtcm = (__fp16 *) dma_queue_pop(dma).dst;
+                                current_mask_vtcm = (__fp16 *) dma_queue_pop(dma_q).dst;
                             } else {
-                                fa_pop_mask_dma_gqa(dma, G);
+                                fa_pop_mask_dma_gqa(dma_q, G);
                                 current_mask_vtcm = factx.vtcm_mask_buf;
                             }
                         }
@@ -2116,7 +2176,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                                 hmx_ceil_div(hex_smin(Bc, nek1 - (kv_blk - 1) * Bc), HMX_FP16_TILE_N_COLS);
                             ou_job[prev_buf].n_row_tiles_g_br = n_row_tiles_g_br;
                             ou_job[prev_buf].n_tiles_per_bc   = n_tiles_per_bc;
-                            ou_job[prev_buf].DV               = DV;
+                            ou_job[prev_buf].DV               = DV_pad;
                             hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job[prev_buf]));
                         }
 
@@ -2126,7 +2186,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                             const uint32_t next_rows  = hex_smin(Bc, nek1 - next_start);
                             const size_t   next_buf   = 1 - buf_idx;
 
-                            void * next_k = dma_queue_pop(dma).dst;
+                            void * next_k = (void *) dma_queue_pop(dma_q).dst;
                             fa_phase_k_interleave(&factx, next_rows, k_src_stride, next_k, next_start, next_buf);
 
                             qk_job[next_buf].q_tiles        = factx.vtcm_q_tiles;
@@ -2134,7 +2194,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                             qk_job[next_buf].s_tiles        = factx.vtcm_s_tiles[next_buf];
                             qk_job[next_buf].n_row_tiles    = n_row_tiles;
                             qk_job[next_buf].n_col_tiles    = hmx_ceil_div(next_rows, HMX_FP16_TILE_N_COLS);
-                            qk_job[next_buf].n_dot_tiles    = DK / 32;
+                            qk_job[next_buf].n_dot_tiles    = DK_pad / 32;
                             qk_job[next_buf].n_tiles_per_bc = n_tiles_per_bc;
                             qk_job[next_buf].hmx_scales     = factx.vtcm_hmx_scales_qk;
                             hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job[next_buf]));
@@ -2177,7 +2237,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
 
                         // Prefetch block kv_blk + 2
                         if (kv_blk + 2 < factx.n_kv_blocks) {
-                            fa_prefetch_block(dma, k, v, mask, kv_blk + 2, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
+                            fa_prefetch_block(dma_q, k, v, mask, kv_blk + 2, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
                                               ik2, ik3, iv2, iv3, q_start, im3, kv_head, G, m_line_bytes, n_rows_q, nek1, buf_idx, &factx);
                         }
 
@@ -2198,7 +2258,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                         ou_job[0].n_col_tiles      = last_cols;
                         ou_job[0].n_row_tiles_g_br = n_row_tiles_g_br;
                         ou_job[0].n_tiles_per_bc   = n_tiles_per_bc;
-                        ou_job[0].DV               = DV;
+                        ou_job[0].DV               = DV_pad;
                         hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job[0]));
 
                         // Overlapped: run HVX build diag inv L while HMX is busy executing the update
@@ -2219,10 +2279,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
 
                         if (mask) {
                             if (__builtin_expect(factx.mask_broadcast, true)) {
-                                const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
-                                dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
+                                const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
+                                dma_cache_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
                             } else {
-                                fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx);
+                                fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx);
                             }
                         }
 
@@ -2230,14 +2290,14 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                             const uint32_t  prefetch_start = (kv_blk + 1) * Bc;
                             const uint32_t  prefetch_rows  = hex_smin(Bc, nek1 - prefetch_start);
                             const size_t    prefetch_buf   = 1 - buf_idx;
-                            const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
-                            dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
-                            const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
-                            dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
+                            const dma_addr_t k_prefetch_src = k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
+                            dma_queue_push(dma_q, dma_make_data(factx.vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
+                            const dma_addr_t v_prefetch_src = v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
+                            dma_queue_push(dma_q, dma_make_data(factx.vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
                         }
 
                         // Wait for current K DMA and interleave
-                        void * curr_k = dma_queue_pop(dma).dst;
+                        void * curr_k = (void *) dma_queue_pop(dma_q).dst;
                         fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start, 0);
 
                         {
@@ -2246,7 +2306,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                             qk_job.s_tiles        = factx.vtcm_s_tiles[0];
                             qk_job.n_row_tiles    = n_row_tiles;
                             qk_job.n_col_tiles    = n_col_tiles;
-                            qk_job.n_dot_tiles    = (size_t) (DK / 32);
+                            qk_job.n_dot_tiles    = (size_t) (DK_pad / 32);
                             qk_job.n_tiles_per_bc = n_tiles_per_bc;
                             qk_job.hmx_scales     = factx.vtcm_hmx_scales_qk;
 
@@ -2255,16 +2315,16 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                         }
 
                         // Wait for current V DMA and interleave
-                        void * curr_v = dma_queue_pop(dma).dst;
+                        void * curr_v = (void *) dma_queue_pop(dma_q).dst;
                         fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, factx.vtcm_v_tiles[0], n_tiles_per_bc, kv_start);
 
                         // ---- Phase 3: softmax + build_D ----
                         __fp16 * current_mask_vtcm = NULL;
                         if (mask) {
                             if (__builtin_expect(factx.mask_broadcast, true)) {
-                                current_mask_vtcm = (__fp16 *) dma_queue_pop(dma).dst;
+                                current_mask_vtcm = (__fp16 *) dma_queue_pop(dma_q).dst;
                             } else {
-                                fa_pop_mask_dma_gqa(dma, G);
+                                fa_pop_mask_dma_gqa(dma_q, G);
                                 current_mask_vtcm = factx.vtcm_mask_buf;
                             }
                         }
@@ -2302,7 +2362,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                             ou_job.n_col_tiles      = n_col_tiles;
                             ou_job.n_row_tiles_g_br = n_row_tiles_g_br;
                             ou_job.n_tiles_per_bc   = n_tiles_per_bc;
-                            ou_job.DV               = DV;
+                            ou_job.DV               = DV_pad;
 
                             hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
                             if (kv_blk + 1 == factx.n_kv_blocks) {
@@ -2336,10 +2396,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
 
                 if (has_next) {
                     const uint32_t next_n_rows_q = hex_smin(Br, neq1 - next_q_start);
-                    const uint8_t * next_q_ptr = (const uint8_t *) q->data + next_q_start * q->nb[1] + (next_kv_head * factx.G) * q->nb[2] + next_ib3 * q->nb[3];
+                    const dma_addr_t next_q_ptr = q->data + next_q_start * q->nb[1] + (next_kv_head * factx.G) * q->nb[2] + next_ib3 * q->nb[3];
                     const size_t next_q_row_bytes = q_transposed ? next_n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
                     const size_t next_n_rows      = q_transposed ? factx.G : next_n_rows_q;
-                    dma_queue_push(dma, dma_make_ptr(factx.vtcm_q_dma, next_q_ptr), next_q_row_bytes, hex_smax(q_src_stride, next_q_row_bytes), next_q_row_bytes, next_n_rows);
+                    dma_queue_push(dma_q, dma_make_data(factx.vtcm_q_dma, next_q_ptr), next_q_row_bytes, hex_smax(q_src_stride, next_q_row_bytes), next_q_row_bytes, next_n_rows);
 
                     if (factx.n_kv_blocks > 0) {
                         const uint32_t next_ik2 = next_kv_head;
@@ -2351,11 +2411,11 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                             next_iv3 = fastdiv(next_ib3, &kparams->broadcast_rv3);
                         }
 
-                        const uint8_t * next_k_src = (const uint8_t *) k->data + next_ik2 * k->nb[2] + next_ik3 * k->nb[3];
-                        dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], next_k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
+                        const dma_addr_t next_k_src = k->data + next_ik2 * k->nb[2] + next_ik3 * k->nb[3];
+                        dma_queue_push(dma_q, dma_make_data(factx.vtcm_k_fp16[0], next_k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
 
-                        const uint8_t * next_v_src = (const uint8_t *) v->data + next_iv2 * v->nb[2] + next_iv3 * v->nb[3];
-                        dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], next_v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
+                        const dma_addr_t next_v_src = v->data + next_iv2 * v->nb[2] + next_iv3 * v->nb[3];
+                        dma_queue_push(dma_q, dma_make_data(factx.vtcm_v_fp16[0], next_v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
 
                         if (factx.pipeline && mask) {
                             uint32_t next_im3 = im3;
@@ -2363,10 +2423,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                                 next_im3 = fastmodulo(next_ib3, mask->ne[3], &factx.src3_div3);
                             }
                             if (__builtin_expect(factx.mask_broadcast, true)) {
-                                const uint8_t * ms_src = (const uint8_t *) mask->data + next_q_start * mask->nb[1] + next_im3 * mask->nb[3] + 0;
-                                dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), next_n_rows_q);
+                                const dma_addr_t ms_src = mask->data + next_q_start * mask->nb[1] + next_im3 * mask->nb[3] + 0;
+                                dma_cache_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), next_n_rows_q);
                             } else {
-                                fa_push_mask_dma_gqa(dma, mask, next_q_start, next_im3, 0, next_kv_head, G, m_line_bytes, kv_rows0, next_n_rows_q, &factx);
+                                fa_push_mask_dma_gqa(dma_q, mask, next_q_start, next_im3, 0, next_kv_head, G, m_line_bytes, kv_rows0, next_n_rows_q, &factx);
                             }
                         }
                     }
@@ -2380,7 +2440,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                     on_job.hmx_scales       = factx.vtcm_hmx_scales_id;
                     on_job.n_row_tiles      = n_row_tiles;
                     on_job.n_row_tiles_g_br = n_row_tiles_g_br;
-                    on_job.DV               = DV;
+                    on_job.DV               = DV_pad;
                     hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_norm_worker, &on_job));
                     hmx_queue_pop(ctx->hmx_queue);
                 }
@@ -2408,6 +2468,10 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
+    if (htp_tensor_is_extended(dst)) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
     const struct htp_fa_kernel_params * kparams = (const struct htp_fa_kernel_params *) octx->kernel_params;
 
     if (kparams->kernel_type == HTP_FA_KERNEL_UNSUPPORTED) {
@@ -2445,11 +2509,6 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
     factx.size_k_row_padded = kparams->u.hvx.size_k_row_padded;
     factx.size_v_row_padded = kparams->u.hvx.size_v_row_padded;
 
-    size_t size_q_block = factx.size_q_row_padded * 1; // single row for now
-    factx.size_k_block = factx.size_k_row_padded * FLASH_ATTN_BLOCK_SIZE;
-    factx.size_v_block = factx.size_v_row_padded * FLASH_ATTN_BLOCK_SIZE;
-    factx.size_m_block = hex_round_up(FLASH_ATTN_BLOCK_SIZE * sizeof(__fp16), 128);
-
     factx.n_blocks = kparams->n_kv_blocks;
 
     factx.scale = kparams->scale;
@@ -2494,26 +2553,38 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
     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
+    const bool has_mask  = (mask != NULL);
+    const bool has_sinks = (octx->src[4] != NULL);
+    struct hvx_fa_vtcm_layout L;
+    hvx_fa_vtcm_layout_build(&L, k->ne[0], v->ne[0], factx.is_q_fp32, has_mask, has_sinks, n_head, n_threads);
 
-    factx.size_q_block = size_q_block;
-    factx.size_vkq_acc = size_vkq_acc;
-
-    uint8_t * vtcm_cur = octx->ctx->vtcm_base;
-
-    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) {
+    if (L.total_bytes > 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, n_threads);
+    factx.size_q_block = L.size_q_block;
+    factx.size_k_block = L.size_k_block;
+    factx.size_v_block = L.size_v_block;
+    factx.size_m_block = L.size_m_block;
+    factx.size_vkq_acc = L.size_vkq_acc;
+
+    uint8_t * const base = octx->ctx->vtcm_base;
+
+    factx.spad_q     = VTCM_LAYOUT_PTR(uint8_t, base, L.off_q);
+    factx.spad_k     = VTCM_LAYOUT_PTR(uint8_t, base, L.off_k);
+    factx.spad_v     = VTCM_LAYOUT_PTR(uint8_t, base, L.off_v);
+    factx.spad_m     = VTCM_LAYOUT_PTR_OPTIONAL(uint8_t, base, L.off_m, has_mask);
+    factx.spad_a     = VTCM_LAYOUT_PTR(uint8_t, base, L.off_a);
+    factx.spad_sinks = VTCM_LAYOUT_PTR_OPTIONAL(float, base, L.off_sinks, has_sinks);
+
+    if (has_sinks) {
+        const struct htp_tensor * sinks = octx->src[4];
+        dma_queue * dma_q = octx->ctx->dma[0];
+        dma_queue_push(dma_q, dma_make_data(factx.spad_sinks, sinks->data), L.size_sinks, 0, sinks->size, 1);
+        dma_queue_pop(dma_q);
     }
 
+    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

@@ -121,6 +121,7 @@ struct hmx_fa_vtcm_layout {
     size_t off_hmx_scales_qk;
     size_t off_mask_buf;
     size_t off_slopes;
+    size_t off_sinks;
 
     // Region byte sizes reused by the device at runtime (not just for allocation).
     size_t q_tile_bytes;
@@ -130,6 +131,7 @@ struct hmx_fa_vtcm_layout {
     size_t m_line_bytes;       // one mask row
     size_t m_buf_slot_bytes;   // one dma_cache slot = align_up(Br * m_line_bytes, 4096)
     size_t col_vec_bytes;
+    size_t sinks_bytes;
 
     // Derived strides.
     size_t row_buf_stride;       // HVX vectors (128B) per row buffer
@@ -142,8 +144,10 @@ struct hmx_fa_vtcm_layout {
 // Build the VTCM layout.
 
 static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
-                                       size_t gqa_factor, size_t DK, size_t DV,
-                                       size_t Br, size_t Bc, size_t n_threads, bool pipeline, bool is_q_fp32) {
+                                            size_t gqa_factor, size_t DK, size_t DV,
+                                            size_t Br, size_t Bc, size_t n_threads,
+                                            bool pipeline, bool is_q_fp32,
+                                            bool has_sinks, size_t n_heads) {
     const size_t g_br         = hex_align_up(gqa_factor * Br, HMX_FP16_TILE_N_ROWS);
     const size_t q_tile_size  = hex_align_up(g_br * DK   * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
     const size_t o_tile_size  = hex_align_up(g_br * DV   * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
@@ -166,6 +170,7 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
     const size_t m_buf_slot   = hex_align_up(Br * m_line_size, 256);
     const size_t m_buf_size   = m_buf_slot * HMX_FA_DMA_CACHE_SIZE;
     const size_t slopes_size  = hex_align_up(g_br * sizeof(__fp16), 128);
+    const size_t sinks_size   = hex_round_up(n_heads * sizeof(float), 128);
 
     size_t off = 0;
 
@@ -215,6 +220,7 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
     VTCM_LAYOUT_ALLOC(off, off_hmx_scales_qk, 256);
     VTCM_LAYOUT_ALLOC(off, off_mask_buf,      m_buf_size);
     VTCM_LAYOUT_ALLOC(off, off_slopes,        slopes_size);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_sinks, sinks_size, has_sinks);
 
     L->q_tile_bytes        = q_tile_size;
     L->o_tile_bytes        = o_tile_size;
@@ -228,20 +234,43 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
     L->m_buf_slot_bytes    = m_buf_slot;
     L->row_buf_stride      = row_vec_size / 128;
     L->mask_buf_row_stride = m_line_size / sizeof(__fp16);
+    L->sinks_bytes         = has_sinks ? sinks_size : 0;
     L->pipeline            = pipeline;
     L->total_bytes         = off;
 }
 
 // Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration.
-static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline, bool is_q_fp32) {
+static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline, bool is_q_fp32, bool has_sinks, size_t n_heads) {
     struct hmx_fa_vtcm_layout L;
-    hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline, is_q_fp32);
+    hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline, is_q_fp32, has_sinks, n_heads);
     return L.total_bytes;
 }
 
 #define FA_HVX_BLOCK_SIZE 64
 
-static inline size_t hvx_fa_compute_vtcm_usage(size_t DK, size_t DV, bool is_q_fp32, bool has_mask, size_t n_threads) {
+struct hvx_fa_vtcm_layout {
+    size_t off_q;
+    size_t off_k;
+    size_t off_v;
+    size_t off_m;
+    size_t off_a;
+    size_t off_sinks;
+
+    size_t size_q_block;
+    size_t size_k_block;
+    size_t size_v_block;
+    size_t size_m_block;
+    size_t size_vkq_acc;
+    size_t size_sinks;
+
+    size_t total_bytes;
+};
+
+static inline void hvx_fa_vtcm_layout_build(struct hvx_fa_vtcm_layout * L,
+                                            size_t DK, size_t DV,
+                                            bool is_q_fp32, bool has_mask,
+                                            bool has_sinks, size_t n_heads,
+                                            size_t n_threads) {
     const size_t size_q_row_padded = hex_round_up(DK * (is_q_fp32 ? 4 : 2), 128);
     const size_t size_k_row_padded = hex_round_up(DK * sizeof(__fp16), 128);
     const size_t size_v_row_padded = hex_round_up(DV * sizeof(__fp16), 128);
@@ -251,29 +280,47 @@ static inline size_t hvx_fa_compute_vtcm_usage(size_t DK, size_t DV, bool is_q_f
     const size_t size_v_block = size_v_row_padded * FA_HVX_BLOCK_SIZE;
     const size_t size_m_block = hex_round_up(FA_HVX_BLOCK_SIZE * sizeof(__fp16), 128);
     const size_t size_vkq_acc = hex_round_up(DV * sizeof(float), 128);
+    const size_t size_sinks   = hex_round_up(n_heads * sizeof(float), 128);
 
-    const size_t size_per_thread = size_q_block * 1
-                                 + size_k_block * 2
-                                 + size_v_block * 2
-                                 + (has_mask ? size_m_block * HVX_FA_DMA_CACHE_SIZE : 0)
-                                 + size_vkq_acc;
+    size_t off = 0;
 
-    return size_per_thread * n_threads;
+    VTCM_LAYOUT_ALLOC(off, off_q, size_q_block * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_k, size_k_block * 2 * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_v, size_v_block * 2 * n_threads);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_m, size_m_block * HVX_FA_DMA_CACHE_SIZE * n_threads, has_mask);
+    VTCM_LAYOUT_ALLOC(off, off_a, size_vkq_acc * n_threads);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_sinks, size_sinks, has_sinks);
+
+    L->size_q_block = size_q_block;
+    L->size_k_block = size_k_block;
+    L->size_v_block = size_v_block;
+    L->size_m_block = size_m_block;
+    L->size_vkq_acc = size_vkq_acc;
+    L->size_sinks   = has_sinks ? size_sinks : 0;
+    L->total_bytes  = off;
+}
+
+static inline size_t hvx_fa_compute_vtcm_usage(size_t DK, size_t DV, bool is_q_fp32, bool has_mask, bool has_sinks, size_t n_heads, size_t n_threads) {
+    struct hvx_fa_vtcm_layout L;
+    hvx_fa_vtcm_layout_build(&L, DK, DV, is_q_fp32, has_mask, has_sinks, n_heads, n_threads);
+    return L.total_bytes;
 }
 
 #define FA_MIN_KV_BLOCKS 3
 
 // Cost-based (Br, Bc) search for flash attention with pipeline constraint.
 static inline int hmx_fa_find_chunk_size(size_t * Br_out,
-                                  size_t * Bc_out,
-                                  size_t   gqa_factor,
-                                  size_t   DK,
-                                  size_t   DV,
-                                  size_t   qo_len,
-                                  size_t   kv_len,
-                                  size_t   vtcm_budget,
-                                  size_t   n_threads,
-                                  bool     is_q_fp32) {
+                                         size_t * Bc_out,
+                                         size_t   gqa_factor,
+                                         size_t   DK,
+                                         size_t   DV,
+                                         size_t   qo_len,
+                                         size_t   kv_len,
+                                         size_t   vtcm_budget,
+                                         size_t   n_threads,
+                                         bool     is_q_fp32,
+                                         bool     has_sinks,
+                                         size_t   n_heads) {
     const size_t T       = HMX_FP16_TILE_N_ROWS;  // 32
     const size_t br_unit = hmx_ceil_div(T, gqa_factor);
     const size_t bc_unit = HMX_FP16_TILE_N_COLS * 2;  // 64
@@ -297,7 +344,7 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
     for (size_t Br = Br_max; Br >= br_unit; Br -= br_unit) {
         // Try all Bc candidates from Bc_limit down to bc_unit
         for (size_t Bc = Bc_limit; Bc >= bc_unit; Bc -= bc_unit) {
-            size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline, is_q_fp32);
+            size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline, is_q_fp32, has_sinks, n_heads);
             if (vtcm_needed <= vtcm_budget) {
                 // This Bc fits for this Br!
                 const size_t q_blocks       = (qo_len + Br - 1) / Br;

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

@@ -0,0 +1,297 @@
+#ifndef HTP_GATED_DELTA_NET_OPS_H
+#define HTP_GATED_DELTA_NET_OPS_H
+
+#include <stdint.h>
+#include <stddef.h>
+#include <stdbool.h>
+
+#include "hex-fastdiv.h"
+#include "hex-common.h"
+#include "htp-vtcm.h"
+
+#define HTP_GDN_MAX_SV     128
+#define HTP_GDN_CHUNK_SIZE 64
+
+#ifndef HMX_FP16_TILE_SIZE
+#define HMX_FP16_TILE_SIZE 2048
+#endif
+
+enum htp_gdn_kernel_type {
+    HTP_GDN_KERNEL_HVX_RECURRENT = 0,
+    HTP_GDN_KERNEL_HMX_CHUNKED   = 1,
+};
+
+struct htp_gdn_kernel_params {
+    uint8_t  kernel_type;
+    uint8_t  pipeline;
+    uint16_t chunk_size;
+    uint16_t n_chunks;
+    uint16_t n_heads_batch;
+
+    uint32_t n_threads;
+    uint32_t S_v;
+    uint32_t H;
+    uint32_t n_tokens;
+    uint32_t n_seqs;
+    uint32_t K;
+
+    uint32_t total_rows;
+    uint32_t row_start;
+    uint32_t nrows;
+    uint32_t rows_per_thread;
+
+    uint32_t kda;
+    uint32_t state_aligned;
+    uint32_t vtcm_per_thread;
+    uint32_t vtcm_size;
+    uint32_t state_seq_stride;
+    uint32_t state_size_per_snap;
+
+    float    scale;
+
+    struct fastdiv_values div_H;
+    struct fastdiv_values div_q1;
+    struct fastdiv_values div_k1;
+    struct fastdiv_values div_rq3;
+    struct fastdiv_values div_rk3;
+    struct fastdiv_values div_n_threads;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_gdn_kernel_params) <= 128, "htp_gdn_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_gdn_kernel_params) <= 128, "htp_gdn_kernel_params is too large for kernel_params blob");
+#endif
+
+struct htp_gdn_vtcm_layout {
+    size_t state_aligned;
+    size_t bytes_per_thread;
+    size_t total_bytes;
+};
+
+static inline void htp_gdn_vtcm_layout_build(
+    struct htp_gdn_vtcm_layout * layout,
+    uint32_t S_v,
+    uint32_t n_threads
+) {
+    size_t state_bytes = (size_t) S_v * S_v * sizeof(float);
+    layout->state_aligned = hex_round_up(state_bytes, 128);
+    layout->bytes_per_thread = 2 * layout->state_aligned;
+    layout->total_bytes = layout->bytes_per_thread * n_threads;
+}
+
+struct htp_gdn_hmx_vtcm_layout {
+    size_t off_s_state;
+    size_t off_s_f16;
+    size_t off_s_col_tiles;
+    size_t off_s_update_f32;
+    size_t off_s_update_tiles;
+
+    size_t off_q_f32[2];
+    size_t off_k_f32[2];
+    size_t off_v_f32[2];
+    size_t off_g_f32[2];
+    size_t off_b_f32[2];
+    size_t off_g_raw[2];
+    size_t off_b_raw[2];
+    size_t off_o_f32[2];
+
+    size_t off_v_inter_f32;
+    size_t off_o_inter_f32;
+    size_t off_o_intra_f32;
+    size_t off_k_f16;
+    size_t off_v_prime_f16;
+    size_t off_delta_f16;
+    size_t off_d_f16;
+
+    size_t off_q_row_tiles;
+    size_t off_q_prime_row_tiles;
+    size_t off_k_row_tiles;
+    size_t off_k_col_tiles;
+    size_t off_k_prime_row_tiles;
+    size_t off_k_col_tiles_64x128;
+    size_t off_kk_tiles;
+    size_t off_qk_tiles;
+    size_t off_v_inter_tiles;
+    size_t off_o_inter_tiles;
+    size_t off_inv_row_tiles;
+    size_t off_a_row_tiles;
+    size_t off_v_prime_col_tiles;
+    size_t off_delta_tiles;
+    size_t off_delta_col_tiles;
+    size_t off_o_intra_tiles;
+    size_t off_d_row_tiles;
+
+    size_t off_gamma;
+    size_t off_lambda_init;
+    size_t off_decay_m;
+    size_t off_decay_a;
+    size_t off_rows_kk;
+    size_t off_rows_qk;
+    size_t off_rows_inv;
+    size_t off_rows_a;
+
+    size_t off_thread_scratch;
+    size_t off_attn_rem;
+    size_t off_scales_1;
+
+    size_t state_f32_bytes;
+    size_t state_f16_bytes;
+    size_t state_tiles_bytes;
+    size_t dma_chunk_bytes;
+    size_t act_f16_bytes;
+    size_t tile_64xSv_bytes;
+    size_t tile_64x64_bytes;
+
+    uint32_t n_heads_batch;
+    uint32_t n_threads;
+    bool     pipeline;
+    size_t   total_bytes;
+};
+
+static inline void htp_gdn_hmx_vtcm_layout_build(
+    struct htp_gdn_hmx_vtcm_layout * L,
+    uint32_t S_v,
+    uint32_t chunk_size,
+    uint32_t n_heads_batch,
+    uint32_t n_threads,
+    bool     pipeline
+) {
+    memset(L, 0, sizeof(*L));
+    L->n_heads_batch = n_heads_batch;
+    L->n_threads     = n_threads;
+    L->pipeline      = pipeline;
+
+    const size_t bh = (size_t) n_heads_batch;
+    const size_t nth = (size_t) (n_threads > 0 ? n_threads : 1);
+
+    const size_t state_f32_sz   = hex_round_up(S_v * S_v * sizeof(float), 2048);
+    const size_t state_f16_sz   = hex_round_up(S_v * S_v * sizeof(__fp16), 2048);
+    const size_t n_sv_tiles     = S_v / 32;
+    const size_t state_tiles_sz = n_sv_tiles * n_sv_tiles * 2048;
+
+    const size_t dma_chunk_sz   = hex_round_up(chunk_size * S_v * sizeof(float), 2048);
+    const size_t dma_scalar_sz  = hex_round_up(chunk_size * sizeof(float), 128);
+
+    const size_t act_f16_sz     = hex_round_up(chunk_size * S_v * sizeof(__fp16), 2048);
+    const size_t tile_64xSv_sz  = 2 * n_sv_tiles * 2048;
+    const size_t tile_64x64_sz  = 4 * 2048;
+
+    const size_t decay_sz       = 64 * 64 * sizeof(__fp16);
+    const size_t row_vecs_sz    = 64 * 128;
+
+    L->state_f32_bytes   = state_f32_sz;
+    L->state_f16_bytes   = state_f16_sz;
+    L->state_tiles_bytes = state_tiles_sz;
+    L->dma_chunk_bytes   = dma_chunk_sz;
+    L->act_f16_bytes     = act_f16_sz;
+    L->tile_64xSv_bytes  = tile_64xSv_sz;
+    L->tile_64x64_bytes  = tile_64x64_sz;
+
+    size_t off = 0;
+
+    VTCM_LAYOUT_ALLOC(off, off_s_state,        bh * state_f32_sz);
+    VTCM_LAYOUT_ALLOC(off, off_s_f16,          bh * state_f16_sz);
+    VTCM_LAYOUT_ALLOC(off, off_s_col_tiles,    bh * state_tiles_sz);
+    VTCM_LAYOUT_ALLOC(off, off_s_update_f32,   bh * state_f32_sz);
+    VTCM_LAYOUT_ALLOC(off, off_s_update_tiles, bh * state_tiles_sz);
+
+    VTCM_LAYOUT_ALLOC(off, off_q_f32[0], bh * dma_chunk_sz);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_q_f32[1], bh * dma_chunk_sz, pipeline);
+    VTCM_LAYOUT_ALLOC(off, off_k_f32[0], bh * dma_chunk_sz);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_k_f32[1], bh * dma_chunk_sz, pipeline);
+    VTCM_LAYOUT_ALLOC(off, off_v_f32[0], bh * dma_chunk_sz);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_v_f32[1], bh * dma_chunk_sz, pipeline);
+    VTCM_LAYOUT_ALLOC(off, off_g_f32[0], bh * dma_scalar_sz);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_g_f32[1], bh * dma_scalar_sz, pipeline);
+    VTCM_LAYOUT_ALLOC(off, off_b_f32[0], bh * dma_scalar_sz);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_b_f32[1], bh * dma_scalar_sz, pipeline);
+    const size_t raw_gb_sz = hex_round_up(bh * chunk_size * sizeof(float), 128);
+    VTCM_LAYOUT_ALLOC(off, off_g_raw[0], raw_gb_sz);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_g_raw[1], raw_gb_sz, pipeline);
+    VTCM_LAYOUT_ALLOC(off, off_b_raw[0], raw_gb_sz);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_b_raw[1], raw_gb_sz, pipeline);
+    VTCM_LAYOUT_ALLOC(off, off_o_f32[0], bh * dma_chunk_sz);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_o_f32[1], bh * dma_chunk_sz, pipeline);
+
+    VTCM_LAYOUT_ALLOC(off, off_v_inter_f32, bh * dma_chunk_sz);
+    VTCM_LAYOUT_ALLOC(off, off_o_inter_f32, bh * dma_chunk_sz);
+    VTCM_LAYOUT_ALLOC(off, off_o_intra_f32, bh * dma_chunk_sz);
+    VTCM_LAYOUT_ALLOC(off, off_k_f16,       bh * act_f16_sz);
+    VTCM_LAYOUT_ALLOC(off, off_v_prime_f16, bh * act_f16_sz);
+    VTCM_LAYOUT_ALLOC(off, off_delta_f16,   bh * act_f16_sz);
+    VTCM_LAYOUT_ALLOC(off, off_d_f16,       bh * act_f16_sz);
+
+    VTCM_LAYOUT_ALLOC(off, off_q_row_tiles,        bh * tile_64xSv_sz);
+    VTCM_LAYOUT_ALLOC(off, off_q_prime_row_tiles,  bh * tile_64xSv_sz);
+    VTCM_LAYOUT_ALLOC(off, off_k_row_tiles,        bh * tile_64xSv_sz);
+    VTCM_LAYOUT_ALLOC(off, off_k_col_tiles,        bh * tile_64xSv_sz);
+    VTCM_LAYOUT_ALLOC(off, off_k_prime_row_tiles,  bh * tile_64xSv_sz);
+    VTCM_LAYOUT_ALLOC(off, off_k_col_tiles_64x128, bh * tile_64xSv_sz);
+    VTCM_LAYOUT_ALLOC(off, off_kk_tiles,           bh * tile_64x64_sz);
+    VTCM_LAYOUT_ALLOC(off, off_qk_tiles,           bh * tile_64x64_sz);
+    VTCM_LAYOUT_ALLOC(off, off_v_inter_tiles,      bh * tile_64xSv_sz);
+    VTCM_LAYOUT_ALLOC(off, off_o_inter_tiles,      bh * tile_64xSv_sz);
+    VTCM_LAYOUT_ALLOC(off, off_inv_row_tiles,      bh * tile_64x64_sz);
+    VTCM_LAYOUT_ALLOC(off, off_a_row_tiles,        bh * tile_64x64_sz);
+    VTCM_LAYOUT_ALLOC(off, off_v_prime_col_tiles,  bh * tile_64xSv_sz);
+    VTCM_LAYOUT_ALLOC(off, off_delta_tiles,        bh * tile_64xSv_sz);
+    VTCM_LAYOUT_ALLOC(off, off_delta_col_tiles,    bh * tile_64xSv_sz);
+    VTCM_LAYOUT_ALLOC(off, off_o_intra_tiles,      bh * tile_64xSv_sz);
+    VTCM_LAYOUT_ALLOC(off, off_d_row_tiles,        bh * tile_64xSv_sz);
+
+    VTCM_LAYOUT_ALLOC(off, off_gamma,       bh * hex_round_up(chunk_size * sizeof(float), 128));
+    VTCM_LAYOUT_ALLOC(off, off_lambda_init, bh * hex_round_up(chunk_size * sizeof(float), 128));
+    VTCM_LAYOUT_ALLOC(off, off_decay_m,     bh * decay_sz);
+    VTCM_LAYOUT_ALLOC(off, off_decay_a,     bh * decay_sz);
+    VTCM_LAYOUT_ALLOC(off, off_rows_kk,     bh * row_vecs_sz);
+    VTCM_LAYOUT_ALLOC(off, off_rows_qk,     bh * row_vecs_sz);
+    VTCM_LAYOUT_ALLOC(off, off_rows_inv,    bh * row_vecs_sz);
+    VTCM_LAYOUT_ALLOC(off, off_rows_a,      bh * row_vecs_sz);
+
+    const size_t thread_scratch_sz = 64 * 128;
+    VTCM_LAYOUT_ALLOC(off, off_thread_scratch, nth * thread_scratch_sz);
+    VTCM_LAYOUT_ALLOC(off, off_attn_rem,       nth * (128 * sizeof(float)));
+    VTCM_LAYOUT_ALLOC(off, off_scales_1,       256);
+
+    L->total_bytes = off;
+}
+
+static inline bool htp_gdn_hmx_solve_layout(
+    struct htp_gdn_hmx_vtcm_layout * layout_out,
+    uint32_t S_v,
+    uint32_t chunk_size,
+    uint32_t total_rows,
+    size_t   vtcm_budget,
+    uint32_t n_threads,
+    bool     pipeline,
+    uint32_t * n_heads_batch_out
+) {
+    uint32_t max_batch = 8;
+    if (max_batch > total_rows) {
+        max_batch = total_rows;
+    }
+    if (max_batch > n_threads) {
+        max_batch = n_threads;
+    }
+    static const uint32_t candidates[] = { 8, 6, 4, 2, 1 };
+    for (size_t i = 0; i < sizeof(candidates) / sizeof(candidates[0]); ++i) {
+        uint32_t bh = candidates[i];
+        if (bh > max_batch) {
+            continue;
+        }
+        struct htp_gdn_hmx_vtcm_layout L;
+        htp_gdn_hmx_vtcm_layout_build(&L, S_v, chunk_size, bh, n_threads, pipeline);
+        if (L.total_bytes <= vtcm_budget) {
+            *layout_out = L;
+            *n_heads_batch_out = bh;
+            return true;
+        }
+    }
+    if (pipeline) {
+        return htp_gdn_hmx_solve_layout(layout_out, S_v, chunk_size, total_rows, vtcm_budget, n_threads, false, n_heads_batch_out);
+    }
+    return false;
+}
+
+#endif // HTP_GATED_DELTA_NET_OPS_H

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

@@ -11,6 +11,7 @@
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
 #include "hex-common.h"
+#include "dma-queue.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
@@ -59,140 +60,139 @@ struct get_rows_context {
                                                \
     const uint32_t nr = ne10 * ne11 * ne12;
 
-#define GET_ROWS_THREAD_ST_FN(IDX_TYPE)                                                                                \
-static void get_rows_thread_st_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) {                            \
-    struct get_rows_context * grctx = (struct get_rows_context *)data;                                                 \
-    struct htp_ops_context * octx = grctx->octx;                                                                       \
-    const struct htp_get_rows_kernel_params * kparams = grctx->kparams;                                                \
-    get_rows_preamble;                                                                                                 \
-    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, 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) {                                                                             \
-        const uint32_t i12 = fastdiv(i, &kparams->div_ne10_ne11);                                                      \
-        const uint32_t rem = i - i12 * ne11 * ne10;                                                                    \
-        const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);                                                         \
-        const uint32_t i10 = rem - i11 * ne10;                                                                         \
-        const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12);           \
-        const uint32_t i01 = (uint32_t)*src1_ptr;                                                                      \
-        assert(i01 < ne01);                                                                                            \
-        const uint32_t q02 = fastdiv(i11, &kparams->div_ne02);                                                         \
-        const uint32_t i02 = i11 - q02 * ne02;                                                                         \
-        const uint32_t q03 = fastdiv(i12, &kparams->div_ne03);                                                         \
-        const uint32_t i03 = i12 - q03 * ne03;                                                                         \
-        const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03;                                \
-        const uintptr_t dst_ptr  = octx->dst->data    + i10*nb1  + i11*nb2  + i12*nb3;                                 \
-        while (!dma_queue_push(dma_queue, dma_make_ptr((void *)dst_ptr, (const void *)src0_ptr), nb1, nb01,            \
-                               row_size_bytes, 1)) {                                                                   \
-            dma_queue_pop(dma_queue);                                                                                  \
-        }                                                                                                              \
-    }                                                                                                                  \
-    dma_queue_flush(dma_queue);                                                                                        \
+#define GET_ROWS_THREAD_ST_FN(IDX_TYPE)                                                                                 \
+static void get_rows_thread_st_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) {                             \
+    struct get_rows_context * grctx = (struct get_rows_context *)data;                                                  \
+    struct htp_ops_context * octx = grctx->octx;                                                                        \
+    const struct htp_get_rows_kernel_params * kparams = grctx->kparams;                                                 \
+    get_rows_preamble;                                                                                                  \
+    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, grctx->task_start + grctx->tasks);                                               \
+    const uint32_t row_size_bytes = htp_tensor_get_row_size(octx->src[0]->type, ne00);                                  \
+    dma_queue * dma_q = octx->ctx->dma[ith];                                                                            \
+    for (uint32_t i = ir0; i < ir1; ++i) {                                                                              \
+        const uint32_t i12 = fastdiv(i, &kparams->div_ne10_ne11);                                                       \
+        const uint32_t rem = i - i12 * ne11 * ne10;                                                                     \
+        const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);                                                          \
+        const uint32_t i10 = rem - i11 * ne10;                                                                          \
+        const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(uintptr_t)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12); \
+        const uint32_t i01 = (uint32_t)*src1_ptr;                                                                       \
+        assert(i01 < ne01);                                                                                             \
+        const uint32_t q02 = fastdiv(i11, &kparams->div_ne02);                                                          \
+        const uint32_t i02 = i11 - q02 * ne02;                                                                          \
+        const uint32_t q03 = fastdiv(i12, &kparams->div_ne03);                                                          \
+        const uint32_t i03 = i12 - q03 * ne03;                                                                          \
+        const dma_addr_t src0_data = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03;                               \
+        const dma_addr_t dst_data  = octx->dst->data    + i10*nb1  + i11*nb2  + i12*nb3;                                \
+        while (!dma_queue_push(dma_q, dma_make_data(dst_data, src0_data), nb1, nb01,                                    \
+                               row_size_bytes, 1)) {                                                                    \
+            dma_queue_pop(dma_q);                                                                                       \
+        }                                                                                                               \
+    }                                                                                                                   \
+    dma_queue_flush(dma_q);                                                                                             \
 }
 
 GET_ROWS_THREAD_ST_FN(int32_t)
 GET_ROWS_THREAD_ST_FN(int64_t)
 
-#define GET_ROWS_THREAD_DT_FN(TYPE_NAME, SRC0_SIZE_EXPR, IDX_TYPE, COMPUTE_EXPR)                                       \
-static void get_rows_thread_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) {                 \
-    struct get_rows_context * grctx = (struct get_rows_context *)data;                                                 \
-    struct htp_ops_context * octx = grctx->octx;                                                                       \
-    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  = 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, 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];                                                                       \
-    const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout;                                         \
-    uint8_t * vtcm_src0 = grctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread;         \
-    uint8_t * vtcm_dst  = grctx->vtcm_base + vtcm_layout->off_dst  + ith * vtcm_layout->dst_bytes_per_thread;          \
-    for (uint32_t step = 0, spad_idx = 0; step < ir1 - ir0 && spad_idx < 2; ++step, spad_idx++) {                      \
-        const uint32_t i = ir0 + step;                                                                                 \
-        const uint32_t row_idx   = fastdiv(i, &kparams->div_chunks_per_row);                                           \
-        const uint32_t chunk_idx = i - row_idx * chunks_per_row;                                                       \
-        const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11);                                                \
-        const uint32_t rem = row_idx - i12 * ne11 * ne10;                                                              \
-        const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);                                                         \
-        const uint32_t i10 = rem - i11 * ne10;                                                                         \
-        const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12);           \
-        const uint32_t i01 = (uint32_t)*src1_ptr;                                                                      \
-        assert(i01 < ne01);                                                                                            \
-        const uint32_t q02 = fastdiv(i11, &kparams->div_ne02);                                                         \
-        const uint32_t i02 = i11 - q02 * ne02;                                                                         \
-        const uint32_t q03 = fastdiv(i12, &kparams->div_ne03);                                                         \
-        const uint32_t i03 = i12 - q03 * ne03;                                                                         \
-        const uint32_t offset = chunk_idx * chunk_size;                                                                \
-        const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0;                               \
-        const uint32_t cur_src0_bytes = SRC0_SIZE_EXPR(cur_elems);                                                     \
-        const uint32_t cur_dst_bytes  = cur_elems * sizeof(float);                                                     \
-        const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03 + SRC0_SIZE_EXPR(offset);       \
-        dma_queue_push(dma_queue,                                                                                      \
-                       dma_make_ptr((void *)(uintptr_t)octx->dst->data,                                                \
-                                    vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size),                            \
-                       cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 0);                              \
-        dma_queue_push(dma_queue,                                                                                      \
-                       dma_make_ptr((void *)(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size),                 \
-                                    (const void *)src0_ptr),                                                           \
-                       vtcm_layout->src0_spad_half_size, cur_src0_bytes, cur_src0_bytes, 1);                           \
-    }                                                                                                                  \
-    for (uint32_t step = 0; step < ir1 - ir0; ++step) {                                                                \
-        const uint32_t i = ir0 + step;                                                                                 \
-        void * dst_spad = (void *) dma_queue_pop(dma_queue).src;                                                       \
-        void * src_spad = (void *) dma_queue_pop(dma_queue).dst;                                                       \
-        const uint32_t row_idx   = fastdiv(i, &kparams->div_chunks_per_row);                                           \
-        const uint32_t chunk_idx = i - row_idx * chunks_per_row;                                                       \
-        const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11);                                                \
-        const uint32_t rem = row_idx - i12 * ne11 * ne10;                                                              \
-        const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);                                                         \
-        const uint32_t i10 = rem - i11 * ne10;                                                                         \
-        const uint32_t offset = chunk_idx * chunk_size;                                                                \
-        const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0;                               \
-        const uint32_t cur_dst_bytes  = cur_elems * sizeof(float);                                                     \
-        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, i);                                                          \
-        COMPUTE_EXPR;                                                                                                  \
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, i);                                                           \
-        const uintptr_t dst_ptr  = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3 + offset * sizeof(float);             \
-        dma_queue_push(dma_queue,                                                                                      \
-                       dma_make_ptr((void *)dst_ptr, (const void *)dst_spad),                                          \
-                       cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 1);                              \
-        const uint32_t next_step = step + 2;                                                                           \
-        if (next_step < ir1 - ir0) {                                                                                   \
-            const uint32_t pi = ir0 + next_step;                                                                       \
-            const uint32_t prow_idx   = fastdiv(pi, &kparams->div_chunks_per_row);                                     \
-            const uint32_t pchunk_idx = pi - prow_idx * chunks_per_row;                                                \
-            const uint32_t pi12 = fastdiv(prow_idx, &kparams->div_ne10_ne11);                                          \
-            const uint32_t prem = prow_idx - pi12 * ne11 * ne10;                                                       \
-            const uint32_t pi11 = fastdiv(prem, &kparams->div_ne10);                                                   \
-            const uint32_t pi10 = prem - pi11 * ne10;                                                                  \
-            const IDX_TYPE * psrc1_ptr = (const IDX_TYPE *)(octx->src[1]->data + pi10*nb10 + pi11*nb11 + pi12*nb12);   \
-            const uint32_t pi01 = (uint32_t)*psrc1_ptr;                                                                \
-            assert(pi01 < ne01);                                                                                       \
-            const uint32_t pq02 = fastdiv(pi11, &kparams->div_ne02);                                                   \
-            const uint32_t pi02 = pi11 - pq02 * ne02;                                                                  \
-            const uint32_t pq03 = fastdiv(pi12, &kparams->div_ne03);                                                   \
-            const uint32_t pi03 = pi12 - pq03 * ne03;                                                                  \
-            const uint32_t poffset = pchunk_idx * chunk_size;                                                          \
-            const uint32_t pcur_elems = (poffset < ne00) ? MIN(chunk_size, ne00 - poffset) : 0;                        \
-            const uint32_t pcur_src0_bytes = SRC0_SIZE_EXPR(pcur_elems);                                               \
-            const uintptr_t psrc0_ptr =                                                                                \
-                octx->src[0]->data + pi01*nb01 + pi02*nb02 + pi03*nb03 + SRC0_SIZE_EXPR(poffset);                      \
-            dma_queue_push(dma_queue,                                                                                  \
-                           dma_make_ptr((void *)src_spad, (const void *)psrc0_ptr),                                    \
-                           vtcm_layout->src0_spad_half_size, pcur_src0_bytes, pcur_src0_bytes, 1);                     \
-        }                                                                                                              \
-    }                                                                                                                  \
-    dma_queue_flush(dma_queue);                                                                                        \
+#define GET_ROWS_THREAD_DT_FN(TYPE_NAME, SRC0_SIZE_EXPR, IDX_TYPE, COMPUTE_EXPR)                                                \
+static void get_rows_thread_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) {                          \
+    struct get_rows_context * grctx = (struct get_rows_context *)data;                                                          \
+    struct htp_ops_context * octx = grctx->octx;                                                                                \
+    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  = 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, 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_q = octx->ctx->dma[ith];                                                                                    \
+    const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout;                                                  \
+    uint8_t * vtcm_src0 = grctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread;                  \
+    uint8_t * vtcm_dst  = grctx->vtcm_base + vtcm_layout->off_dst  + ith * vtcm_layout->dst_bytes_per_thread;                   \
+    for (uint32_t step = 0, spad_idx = 0; step < ir1 - ir0 && spad_idx < 2; ++step, spad_idx++) {                               \
+        const uint32_t i = ir0 + step;                                                                                          \
+        const uint32_t row_idx   = fastdiv(i, &kparams->div_chunks_per_row);                                                    \
+        const uint32_t chunk_idx = i - row_idx * chunks_per_row;                                                                \
+        const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11);                                                         \
+        const uint32_t rem = row_idx - i12 * ne11 * ne10;                                                                       \
+        const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);                                                                  \
+        const uint32_t i10 = rem - i11 * ne10;                                                                                  \
+        const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(uintptr_t)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12);         \
+        const uint32_t i01 = (uint32_t)*src1_ptr;                                                                               \
+        assert(i01 < ne01);                                                                                                     \
+        const uint32_t q02 = fastdiv(i11, &kparams->div_ne02);                                                                  \
+        const uint32_t i02 = i11 - q02 * ne02;                                                                                  \
+        const uint32_t q03 = fastdiv(i12, &kparams->div_ne03);                                                                  \
+        const uint32_t i03 = i12 - q03 * ne03;                                                                                  \
+        const uint32_t offset = chunk_idx * chunk_size;                                                                         \
+        const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0;                                        \
+        const uint32_t cur_src0_bytes = SRC0_SIZE_EXPR(cur_elems);                                                              \
+        const uint32_t cur_dst_bytes  = cur_elems * sizeof(float);                                                              \
+        const dma_addr_t src0_data = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03 + SRC0_SIZE_EXPR(offset);              \
+        dma_queue_push(dma_q,                                                                                                   \
+                       dma_make_data(octx->dst->data,                                                                           \
+                                     vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size),                                    \
+                       cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 0);                                       \
+        dma_queue_push(dma_q,                                                                                                   \
+                       dma_make_data(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size, src0_data),                       \
+                       vtcm_layout->src0_spad_half_size, cur_src0_bytes, cur_src0_bytes, 1);                                    \
+    }                                                                                                                           \
+    for (uint32_t step = 0; step < ir1 - ir0; ++step) {                                                                         \
+        const uint32_t i = ir0 + step;                                                                                          \
+        void * dst_spad = (void *) dma_queue_pop(dma_q).src;                                                                    \
+        void * src_spad = (void *) dma_queue_pop(dma_q).dst;                                                                    \
+        const uint32_t row_idx   = fastdiv(i, &kparams->div_chunks_per_row);                                                    \
+        const uint32_t chunk_idx = i - row_idx * chunks_per_row;                                                                \
+        const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11);                                                         \
+        const uint32_t rem = row_idx - i12 * ne11 * ne10;                                                                       \
+        const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);                                                                  \
+        const uint32_t i10 = rem - i11 * ne10;                                                                                  \
+        const uint32_t offset = chunk_idx * chunk_size;                                                                         \
+        const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0;                                        \
+        const uint32_t cur_dst_bytes  = cur_elems * sizeof(float);                                                              \
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, i);                                                                   \
+        COMPUTE_EXPR;                                                                                                           \
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, i);                                                                    \
+        const dma_addr_t dst_data  = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3 + offset * sizeof(float);                    \
+        dma_queue_push(dma_q,                                                                                                   \
+                       dma_make_data(dst_data, dst_spad),                                                                       \
+                       cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 1);                                       \
+        const uint32_t next_step = step + 2;                                                                                    \
+        if (next_step < ir1 - ir0) {                                                                                            \
+            const uint32_t pi = ir0 + next_step;                                                                                \
+            const uint32_t prow_idx   = fastdiv(pi, &kparams->div_chunks_per_row);                                              \
+            const uint32_t pchunk_idx = pi - prow_idx * chunks_per_row;                                                         \
+            const uint32_t pi12 = fastdiv(prow_idx, &kparams->div_ne10_ne11);                                                   \
+            const uint32_t prem = prow_idx - pi12 * ne11 * ne10;                                                                \
+            const uint32_t pi11 = fastdiv(prem, &kparams->div_ne10);                                                            \
+            const uint32_t pi10 = prem - pi11 * ne10;                                                                           \
+            const IDX_TYPE * psrc1_ptr = (const IDX_TYPE *)(uintptr_t)(octx->src[1]->data + pi10*nb10 + pi11*nb11 + pi12*nb12); \
+            const uint32_t pi01 = (uint32_t)*psrc1_ptr;                                                                         \
+            assert(pi01 < ne01);                                                                                                \
+            const uint32_t pq02 = fastdiv(pi11, &kparams->div_ne02);                                                            \
+            const uint32_t pi02 = pi11 - pq02 * ne02;                                                                           \
+            const uint32_t pq03 = fastdiv(pi12, &kparams->div_ne03);                                                            \
+            const uint32_t pi03 = pi12 - pq03 * ne03;                                                                           \
+            const uint32_t poffset = pchunk_idx * chunk_size;                                                                   \
+            const uint32_t pcur_elems = (poffset < ne00) ? MIN(chunk_size, ne00 - poffset) : 0;                                 \
+            const uint32_t pcur_src0_bytes = SRC0_SIZE_EXPR(pcur_elems);                                                        \
+            const dma_addr_t psrc0_data =                                                                                       \
+                octx->src[0]->data + pi01*nb01 + pi02*nb02 + pi03*nb03 + SRC0_SIZE_EXPR(poffset);                               \
+            dma_queue_push(dma_q,                                                                                               \
+                           dma_make_data(src_spad, psrc0_data),                                                                 \
+                           vtcm_layout->src0_spad_half_size, pcur_src0_bytes, pcur_src0_bytes, 1);                              \
+        }                                                                                                                       \
+    }                                                                                                                           \
+    dma_queue_flush(dma_q);                                                                                                     \
 }
 
 #define F32_BYTES(n)  ((n) * sizeof(float))
@@ -213,11 +213,13 @@ int op_get_rows(struct htp_ops_context * octx) {
 
     if (octx->src[0]->type != HTP_TYPE_F32 &&
         octx->src[0]->type != HTP_TYPE_F16 &&
-        octx->src[0]->type != HTP_TYPE_Q8_0) {
+        octx->src[0]->type != HTP_TYPE_Q8_0 &&
+        octx->src[0]->type != HTP_TYPE_I32) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    if (octx->dst->type != HTP_TYPE_F32) {
+    if ((octx->src[0]->type == HTP_TYPE_I32 && octx->dst->type != HTP_TYPE_I32) ||
+        (octx->src[0]->type != HTP_TYPE_I32 && octx->dst->type != HTP_TYPE_F32)) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
@@ -225,8 +227,8 @@ int op_get_rows(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
-        return HTP_STATUS_OK;
+    if (htp_tensor_is_extended(octx->src[1])) {
+        return HTP_STATUS_NO_SUPPORT;
     }
 
     const struct htp_tensor * dst = octx->dst;
@@ -275,6 +277,7 @@ int op_get_rows(struct htp_ops_context * octx) {
             case HTP_TYPE_F32:  q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_f32_int32_t  : get_rows_thread_f32_int64_t);  break;
             case HTP_TYPE_F16:  q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_f16_int32_t  : get_rows_thread_f16_int64_t);  break;
             case HTP_TYPE_Q8_0: q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_q8_0_int32_t : get_rows_thread_q8_0_int64_t); break;
+            case HTP_TYPE_I32:  q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_st_int32_t   : get_rows_thread_st_int64_t);   break;
             default:            return HTP_STATUS_NO_SUPPORT;
         }
     }

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

@@ -1,2 +0,0 @@
-#pragma once
-#include "dma-queue.h"

ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h

@@ -495,12 +495,140 @@ static inline void hmx_fa_q_prep_fp16(
 }
 
 
+// Head-dim-padded Q-prep (f32). Used when DK is not a multiple of 64.
+static inline void hmx_fa_q_prep_fp32_pad(__fp16 *                      vtcm_q_tiles,
+                                          const uint8_t *               temp_q_vtcm,
+                                          size_t                        start,
+                                          size_t                        end,
+                                          size_t                        g_rows_end,
+                                          size_t                        dk_in,
+                                          size_t                        dk_out,
+                                          size_t                        G,
+                                          size_t                        n_rows_q,
+                                          const struct fastdiv_values * div_G,
+                                          bool                          q_transposed) {
+    const uint32_t n_out_tiles = (uint32_t) (dk_out / 32);
+    for (size_t r = start; r < end; r += 2) {
+        size_t   r0       = r / HMX_FP16_TILE_N_ROWS;
+        size_t   r1       = r % HMX_FP16_TILE_N_ROWS;
+        __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * dk_out;
+
+        if (r >= g_rows_end) {
+            for (uint32_t d = 0; d < n_out_tiles; ++d) {
+                ((HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero();
+            }
+            continue;
+        }
+
+        const size_t q_idx0 = fastdiv(r + 0, div_G);
+        const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+        const size_t q_idx1 = fastdiv(r + 1, div_G);
+        const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+        const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0);
+        const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1);
+
+        const HVX_UVector * pv_in0 = (const HVX_UVector *) (temp_q_vtcm + offset0 * dk_in * sizeof(float));
+        const HVX_UVector * pv_in1 = (r + 1 < g_rows_end) ? (const HVX_UVector *) (temp_q_vtcm + offset1 * dk_in * sizeof(float)) : NULL;
+
+        for (uint32_t d = 0; d < n_out_tiles; ++d) {
+            HVX_Vector * out_tile   = (HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS);
+            const size_t base_lane  = (size_t) d * 32;
+            const size_t real_lanes = (base_lane < dk_in) ? hex_smin(32, dk_in - base_lane) : 0;
+
+            if (real_lanes == 0) {
+                out_tile[r1 / 2] = Q6_V_vzero();
+                continue;
+            }
+
+            HVX_Vector v0 = pv_in0[d];
+            HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+            if (real_lanes < 32) {
+                // Straddle tile: keep the first real_lanes floats, zero the padded tail so
+                // the packed f16 lanes beyond DK are zero.
+                const HVX_VectorPred keep = Q6_Q_vsetq_R((uint32_t) (real_lanes * sizeof(float)));
+                v0                        = Q6_V_vmux_QVV(keep, v0, Q6_V_vzero());
+                v1                        = Q6_V_vmux_QVV(keep, v1, Q6_V_vzero());
+            }
+            out_tile[r1 / 2] = hvx_vec_f32_to_f16_shuff(v0, v1);
+        }
+    }
+}
+
+// Head-dim-padded Q-prep (f16). Used when DK is not a multiple of 64.
+static inline void hmx_fa_q_prep_fp16_pad(__fp16 *                      vtcm_q_tiles,
+                                          const uint8_t *               temp_q_vtcm,
+                                          size_t                        start,
+                                          size_t                        end,
+                                          size_t                        g_rows_end,
+                                          size_t                        dk_in,
+                                          size_t                        dk_out,
+                                          size_t                        G,
+                                          size_t                        n_rows_q,
+                                          const struct fastdiv_values * div_G,
+                                          bool                          q_transposed) {
+    const uint32_t n_out_pairs = (uint32_t) (dk_out / 64);
+    for (size_t r = start; r < end; r += 2) {
+        size_t   r0       = r / HMX_FP16_TILE_N_ROWS;
+        size_t   r1       = r % HMX_FP16_TILE_N_ROWS;
+        __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * dk_out;
+
+        if (r >= g_rows_end) {
+            for (uint32_t d = 0; d < n_out_pairs; ++d) {
+                __fp16 *     out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2;
+                HVX_Vector * pv_out0   = ((HVX_Vector *) out_dtile) + r1 / 2;
+                HVX_Vector * pv_out1   = pv_out0 + 16;
+                *pv_out0               = Q6_V_vzero();
+                *pv_out1               = Q6_V_vzero();
+            }
+            continue;
+        }
+
+        const size_t q_idx0 = fastdiv(r + 0, div_G);
+        const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+        const size_t q_idx1 = fastdiv(r + 1, div_G);
+        const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+        const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0);
+        const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1);
+
+        const HVX_UVector * pv_in0 = (const HVX_UVector *) (temp_q_vtcm + offset0 * dk_in * sizeof(__fp16));
+        const HVX_UVector * pv_in1 = (r + 1 < g_rows_end) ? (const HVX_UVector *) (temp_q_vtcm + offset1 * dk_in * sizeof(__fp16)) : NULL;
+
+        for (uint32_t d = 0; d < n_out_pairs; ++d) {
+            __fp16 *     out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2;
+            HVX_Vector * pv_out0   = ((HVX_Vector *) out_dtile) + r1 / 2;
+            HVX_Vector * pv_out1   = pv_out0 + 16;
+
+            const size_t base_lane  = (size_t) d * 64;
+            const size_t real_lanes = (base_lane < dk_in) ? hex_smin(64, dk_in - base_lane) : 0;
+
+            if (real_lanes == 0) {
+                *pv_out0 = Q6_V_vzero();
+                *pv_out1 = Q6_V_vzero();
+                continue;
+            }
+
+            HVX_Vector v0 = pv_in0[d];
+            HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+            if (real_lanes < 64) {
+                const HVX_VectorPred keep = Q6_Q_vsetq_R((uint32_t) (real_lanes * sizeof(__fp16)));
+                v0                        = Q6_V_vmux_QVV(keep, v0, Q6_V_vzero());
+                v1                        = Q6_V_vmux_QVV(keep, v1, Q6_V_vzero());
+            }
+            HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+            *pv_out0          = Q6_V_lo_W(vp);
+            *pv_out1          = Q6_V_hi_W(vp);
+        }
+    }
+}
+
 static inline void hmx_fa_q_prep_fallback(
     __fp16 * vtcm_q_tiles, uintptr_t q_data,
     size_t q_nb1, size_t q_nb2, size_t q_nb3,
     uint32_t q_start, uint32_t kv_head, uint32_t ib3,
     size_t start, size_t end, size_t n_rows_g,
-    size_t G, size_t DK, bool is_q_fp32,
+    size_t G, size_t dk_in, size_t dk_out, bool is_q_fp32,
     const struct fastdiv_values * div_G
 ) {
     for (size_t r = start; r < end; r += 2) {
@@ -518,33 +646,55 @@ static inline void hmx_fa_q_prep_fallback(
 
         size_t   r0       = r / HMX_FP16_TILE_N_ROWS;
         size_t   r1       = r % HMX_FP16_TILE_N_ROWS;
-        __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+        __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * dk_out;
 
         if (is_q_fp32) {
             const HVX_UVector * pv_in0 = q_ptr0 ? (const HVX_UVector *) q_ptr0 : NULL;
             const HVX_UVector * pv_in1 = q_ptr1 ? (const HVX_UVector *) q_ptr1 : NULL;
 
-            for (uint32_t d = 0; d < DK / 32; ++d) {
-                HVX_Vector v0   = pv_in0 ? pv_in0[d] : Q6_V_vzero();
-                HVX_Vector v1   = pv_in1 ? pv_in1[d] : Q6_V_vzero();
-                HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1);
+            for (uint32_t d = 0; d < dk_out / 32; ++d) {
+                HVX_Vector * out_tile   = (HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS);
+                const size_t base_lane  = (size_t) d * 32;
+                const size_t real_lanes = (base_lane < dk_in) ? hex_smin(32, dk_in - base_lane) : 0;
 
-                HVX_Vector * out_tile = (HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS);
-                out_tile[r1 / 2]      = v_hf;
+                if (real_lanes == 0) {
+                    out_tile[r1 / 2] = Q6_V_vzero();
+                    continue;
+                }
+                HVX_Vector v0 = pv_in0 ? pv_in0[d] : Q6_V_vzero();
+                HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+                if (real_lanes < 32) {
+                    const HVX_VectorPred keep = Q6_Q_vsetq_R((uint32_t) (real_lanes * sizeof(float)));
+                    v0 = Q6_V_vmux_QVV(keep, v0, Q6_V_vzero());
+                    v1 = Q6_V_vmux_QVV(keep, v1, Q6_V_vzero());
+                }
+                out_tile[r1 / 2] = hvx_vec_f32_to_f16_shuff(v0, v1);
             }
         } else {
             const HVX_UVector * pv_in0 = q_ptr0 ? (const HVX_UVector *) q_ptr0 : NULL;
             const HVX_UVector * pv_in1 = q_ptr1 ? (const HVX_UVector *) q_ptr1 : NULL;
 
-            for (uint32_t d = 0; d < DK / 64; ++d) {
-                HVX_Vector     v0 = pv_in0 ? pv_in0[d] : Q6_V_vzero();
-                HVX_Vector     v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
-                HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
-
+            for (uint32_t d = 0; d < dk_out / 64; ++d) {
                 __fp16 *     out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2;
                 HVX_Vector * pv_out0   = ((HVX_Vector *) out_dtile) + r1 / 2;
                 HVX_Vector * pv_out1   = pv_out0 + 16;
 
+                const size_t base_lane  = (size_t) d * 64;
+                const size_t real_lanes = (base_lane < dk_in) ? hex_smin(64, dk_in - base_lane) : 0;
+
+                if (real_lanes == 0) {
+                    *pv_out0 = Q6_V_vzero();
+                    *pv_out1 = Q6_V_vzero();
+                    continue;
+                }
+                HVX_Vector v0 = pv_in0 ? pv_in0[d] : Q6_V_vzero();
+                HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+                if (real_lanes < 64) {
+                    const HVX_VectorPred keep = Q6_Q_vsetq_R((uint32_t) (real_lanes * sizeof(__fp16)));
+                    v0 = Q6_V_vmux_QVV(keep, v0, Q6_V_vzero());
+                    v1 = Q6_V_vmux_QVV(keep, v1, Q6_V_vzero());
+                }
+                HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
                 *pv_out0 = Q6_V_lo_W(vp);
                 *pv_out1 = Q6_V_hi_W(vp);
             }

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

@@ -1,7 +1,7 @@
 #ifndef HTP_CTX_H
 #define HTP_CTX_H
 
-#include "hex-dma.h"
+#include "dma-queue.h"
 #include "hmx-queue.h"
 #include "htp-ops.h"
 #include "hex-profile.h"
@@ -17,16 +17,16 @@
 #ifndef HTP_MAX_NTHREADS
 #define HTP_MAX_NTHREADS 10
 #endif
-#define HTP_MAX_MMAPS    16
 
-#define HTP_MAX_DIRTY_RANGES 32
+#define HTP_MAX_MMAPS        64
+#define HTP_MAX_DIRTY_RANGES 64
 
 // Memory mapping
 struct htp_mmap {
     uint64_t size;
     uint64_t base;
     uint32_t fd;
-    uint32_t reserved;
+    uint32_t flags;
 };
 
 struct htp_dirty_range {
@@ -68,9 +68,6 @@ struct htp_ops_context {
         const struct htp_tensor * dsts[HTP_OP_MAX_OUTPUTS];
     };
 
-    dma_queue **    src_dma[HTP_OP_MAX_INPUTS];
-    dma_queue **    dst_dma[HTP_OP_MAX_OUTPUTS];
-
     // TODO convert these to an array
     struct htp_spad src0_spad;
     struct htp_spad src1_spad;
@@ -90,7 +87,6 @@ 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;
@@ -165,6 +161,7 @@ int op_get_rows(struct htp_ops_context * octx);
 int op_cpy(struct htp_ops_context * octx);
 int op_repeat(struct htp_ops_context * octx);
 int op_argsort(struct htp_ops_context * octx);
+int op_top_k(struct htp_ops_context * octx);
 int op_ssm_conv(struct htp_ops_context * octx);
 int op_cumsum(struct htp_ops_context * octx);
 int op_fill(struct htp_ops_context * octx);
@@ -175,5 +172,6 @@ int op_gated_delta_net(struct htp_ops_context * octx);
 int op_pad(struct htp_ops_context * octx);
 int op_im2col(struct htp_ops_context * octx);
 int op_allreduce(struct htp_ops_context * octx);
+int op_roll(struct htp_ops_context * octx);
 
 #endif /* HTP_CTX_H */

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

@@ -71,6 +71,7 @@ enum htp_op_code {
     HTP_OP_GLU_SWIGLU,
     HTP_OP_GLU_SWIGLU_OAI,
     HTP_OP_GLU_GEGLU,
+    HTP_OP_GLU_GEGLU_QUICK,
     HTP_OP_SOFTMAX,
     HTP_OP_ADD_ID,
     HTP_OP_ROPE,
@@ -81,6 +82,7 @@ enum htp_op_code {
     HTP_OP_CPY,
     HTP_OP_CPY_FENCE,
     HTP_OP_ARGSORT,
+    HTP_OP_TOP_K,
     HTP_OP_SQR,
     HTP_OP_SQRT,
     HTP_OP_SUM_ROWS,
@@ -104,6 +106,7 @@ enum htp_op_code {
     HTP_OP_ALLREDUCE_ADD,
     HTP_OP_GLU_SWIGLU_CLAMP,
     HTP_OP_MDEV_GROUP,
+    HTP_OP_ROLL,
 
     HTP_OP_INVALID
 };
@@ -130,10 +133,13 @@ enum htp_tensor_flags {
     HTP_TENSOR_FENCE   = (1U << 2)  // Tensor is synchronization fence (explicitly managed)
 };
 
+enum htp_buf_flags {
+    HTP_BUF_EXTENDED = (1U << 0),
+};
+
 // Tensor descriptor
 struct htp_tensor {
-    uint32_t data;                 // Buffer offset in the messages, and data pointer on the NPU
-    uint32_t reserved;             // Reserved for alignment padding (must be multiple of 8)
+    uint64_t data;                 // Buffer offset in the messages, and data pointer on the NPU
     uint32_t size;                 // Data size in bytes
     uint32_t flags;                // Buffer / tensor flags
     uint32_t type;                 // Data type
@@ -147,12 +153,12 @@ struct htp_tensor {
 struct htp_buf_desc {
     uint64_t base;     // base address
     uint64_t size;     // total size
-    uint32_t flags;    // buffer flags (unused)
+    uint32_t flags;    // HTP_BUF_*
     uint32_t fd;       // file descriptor
 };
 
 enum htp_op_flags {
-    HTP_OPFLAGS_SKIP_COMPUTE  = (1U << 0), // Skip actual computation (used for profiling)
+    HTP_OPFLAGS_STUB  = (1U << 0),
 };
 
 // Op descriptor

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

@@ -226,6 +226,10 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co
 }
 
 static void make_tensor_clean(struct htp_context * ctx, const struct htp_tensor * t) {
+    if (!t || (t->flags & (HTP_TENSOR_WEIGHT | HTP_TENSOR_FENCE))) {
+        return;
+    }
+
     uint32_t t_start = t->data;
     uint32_t t_end   = t_start + t->size;
 
@@ -236,6 +240,7 @@ static void make_tensor_clean(struct htp_context * ctx, const struct htp_tensor
         if (r->start < t_end && t_start < r->end) {
             if (t_start <= r->start && r->end <= t_end) {
                 r->start = 0;
+                r->end   = 0;
             } else if (t_start <= r->start) {
                 r->start = t_end;
             } else if (r->end <= t_end) {
@@ -246,6 +251,10 @@ static void make_tensor_clean(struct htp_context * ctx, const struct htp_tensor
 }
 
 static inline bool is_tensor_dirty(struct htp_context * ctx, const struct htp_tensor * t) {
+    if (!t || (t->flags & (HTP_TENSOR_WEIGHT | HTP_TENSOR_FENCE))) {
+        return false;
+    }
+
     uint32_t t_start = t->data;
     uint32_t t_end   = t_start + t->size;
 
@@ -327,7 +336,7 @@ void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * co
 
     for (uint32_t i = 0; i < n; i++) {
         const struct htp_tensor * t = tensors[i];
-        if (t && is_tensor_dirty(ctx, t)) {
+        if (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;

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

@@ -21,6 +21,10 @@ static inline void * htp_tensor_data(const struct htp_tensor * t) {
     return (void *) (uintptr_t) t->data;
 }
 
+static inline bool htp_tensor_is_extended(const struct htp_tensor * t) {
+    return t && (t->data >> 32) != 0;
+}
+
 static inline uint32_t * htp_tensor_flags(const struct htp_tensor * t) {
     return (uint32_t *) &t->flags;
 }
@@ -126,6 +130,7 @@ static inline uint32_t htp_tensor_get_row_size(int type, uint32_t ne00) {
         case HTP_TYPE_F32:  return ne00 * 4;
         case HTP_TYPE_F16:  return ne00 * 2;
         case HTP_TYPE_Q8_0: return (ne00 / 32) * 34;
+        case HTP_TYPE_I32:  return ne00 * 4;
         default:            return 0;
     }
 }

ggml/src/ggml-hexagon/htp/htp_iface.idl

@@ -13,7 +13,7 @@ struct htp_iface_pmu_conf {
 interface htp_iface : remote_handle64 {
     AEEResult start(in uint32 sess_id, in uint64 dsp_queue_id, in uint32 n_hvx, in uint32 n_hmx, in uint64 max_vmem);
     AEEResult stop();
-    AEEResult mmap(in uint32 fd, in uint32 size);
+    AEEResult mmap(in uint32 fd, in uint64 size);
     AEEResult munmap(in uint32 fd);
     AEEResult profiler(in uint32 mode, in htp_iface_pmu_conf pmu);
     AEEResult etm(in uint32 enable);

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

@@ -173,7 +173,7 @@ static inline void hvx_exp_f32(uint8_t * restrict dst, const uint8_t * restrict
         HVX_Vector * p_vec_in1 = (HVX_Vector *) src;
         HVX_Vector * p_vec_out = (HVX_Vector *) dst;
 
-        #pragma unroll(4)
+        #pragma unroll(2)
         for (int i = 0; i < num_elems_whole; i += VLEN_FP32) {
             if (true == negate) {
                 HVX_Vector neg_vec_in = hvx_vec_neg_f32(*p_vec_in1++);
@@ -183,7 +183,7 @@ static inline void hvx_exp_f32(uint8_t * restrict dst, const uint8_t * restrict
             }
         }
     } else {
-        #pragma unroll(4)
+        #pragma unroll(2)
         for (int i = 0; i < num_elems_whole; i += VLEN_FP32) {
             HVX_Vector in = *(HVX_UVector *) (src + i * SIZEOF_FP32);
 

ggml/src/ggml-hexagon/htp/hvx-mm-kernels-float.h

@@ -0,0 +1,382 @@
+#ifndef HVX_MM_KERNELS_FLOAT_H
+#define HVX_MM_KERNELS_FLOAT_H
+
+#include "hvx-utils.h"
+#include "htp-tensor.h"
+
+// Float activation copy/quantization kernels (DDR -> VTCM)
+
+static inline void quantize_f32_f32_kernel(
+    const uint8_t * restrict src_data,
+    uint8_t * restrict dst_data,
+    uint8_t * restrict tmp_data,
+    uint32_t ne0,
+    uint32_t nrows,
+    size_t src_stride,
+    size_t dst_stride
+) {
+    (void) tmp_data;
+    const size_t src_row_size = ne0 * sizeof(float);
+    for (uint32_t i = 0; i < nrows; ++i) {
+        hex_l2fetch(src_data, src_row_size, src_stride, 2);
+        hvx_copy_f32_au(dst_data, src_data, ne0);
+
+        dst_data += dst_stride;
+        src_data += src_stride;
+    }
+}
+
+static inline void quantize_f32_f16_kernel(
+    const uint8_t * restrict src_data,
+    uint8_t * restrict dst_data,
+    uint8_t * restrict tmp_data,
+    uint32_t ne0,
+    uint32_t nrows,
+    size_t src_stride,
+    size_t dst_stride
+) {
+    (void) tmp_data;
+    const size_t src_row_size = ne0 * sizeof(float);
+    for (uint32_t i = 0; i < nrows; ++i) {
+        hex_l2fetch(src_data, src_row_size, src_stride, 2);
+        hvx_copy_f16_f32_au(dst_data, src_data, ne0);
+
+        dst_data += dst_stride;
+        src_data += src_stride;
+    }
+}
+
+static inline void quantize_f16_f16_kernel(
+    const uint8_t * restrict src_data,
+    uint8_t * restrict dst_data,
+    uint8_t * restrict tmp_data,
+    uint32_t ne0,
+    uint32_t nrows,
+    size_t src_stride,
+    size_t dst_stride
+) {
+    (void) tmp_data;
+    const size_t src_row_size = ne0 * sizeof(float);
+    for (uint32_t i = 0; i < nrows; ++i) {
+        hex_l2fetch(src_data, src_row_size, src_stride, 2);
+        hvx_copy_f16_au(dst_data, src_data, ne0);
+
+        dst_data += dst_stride;
+        src_data += src_stride;
+    }
+}
+
+// Float dot product kernels (HVX)
+
+#if __HVX_ARCH__ < 79
+#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b))
+#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b))
+#else
+#define HVX_OP_ADD_F32(a, b) Q6_Vsf_vadd_VsfVsf(a, b)
+#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b)
+#endif
+
+static inline void vec_dot_f32_f32_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
+    const HVX_Vector * restrict x = (const HVX_Vector *) vx;
+    const HVX_Vector * restrict y = (const HVX_Vector *) vy;
+
+    uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors
+    uint32_t nloe = n % VLEN_FP32; // leftover elements
+
+    HVX_Vector rsum = Q6_V_vzero();
+
+    uint32_t i = 0;
+
+    #pragma unroll(4)
+    for (i = 0; i < nvec; i++) {
+        HVX_Vector prod = HVX_OP_MUL_F32(x[i], y[i]);
+        rsum = HVX_OP_ADD_F32(rsum, prod);
+    }
+
+    if (nloe) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+        HVX_Vector x_sf = Q6_V_vand_QV(bmask, x[i]);
+        HVX_Vector y_sf = Q6_V_vand_QV(bmask, y[i]);
+        HVX_Vector prod = HVX_OP_MUL_F32(x_sf, y_sf);
+        rsum = HVX_OP_ADD_F32(rsum, prod);
+    }
+
+    *s = hvx_vec_get_f32(hvx_vec_reduce_sum_f32(rsum));
+}
+
+static inline void vec_dot_f32_f32_aa_2x1(const uint32_t n, float * restrict s0,
+                                const void * restrict vx0, const void * restrict vx1,
+                                const void * restrict vy0) {
+    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
+    const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
+    const HVX_Vector * restrict y  = (const HVX_Vector *) vy0;
+
+    uint32_t nvec = n / VLEN_FP32;
+    uint32_t nloe = n % VLEN_FP32;
+
+    HVX_Vector rsum0 = Q6_V_vzero();
+    HVX_Vector rsum1 = Q6_V_vzero();
+
+    uint32_t i = 0;
+
+    #pragma unroll(2)
+    for (i = 0; i < nvec; i++) {
+        HVX_Vector y_sf = y[i];
+        HVX_Vector prod0 = HVX_OP_MUL_F32(x0[i], y_sf);
+        HVX_Vector prod1 = HVX_OP_MUL_F32(x1[i], y_sf);
+        rsum0 = HVX_OP_ADD_F32(rsum0, prod0);
+        rsum1 = HVX_OP_ADD_F32(rsum1, prod1);
+    }
+
+    if (nloe) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+        HVX_Vector y_sf  = Q6_V_vand_QV(bmask, y[i]);
+        HVX_Vector x0_sf = Q6_V_vand_QV(bmask, x0[i]);
+        HVX_Vector x1_sf = Q6_V_vand_QV(bmask, x1[i]);
+        HVX_Vector prod0 = HVX_OP_MUL_F32(x0_sf, y_sf);
+        HVX_Vector prod1 = HVX_OP_MUL_F32(x1_sf, y_sf);
+        rsum0 = HVX_OP_ADD_F32(rsum0, prod0);
+        rsum1 = HVX_OP_ADD_F32(rsum1, prod1);
+    }
+
+    HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1);
+    hvx_vec_store_u(s0, 8, rsum);
+}
+
+static inline void vec_dot_f32_f32_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
+                                const void * restrict vx0, const void * restrict vx1,
+                                const void * restrict vy0, const void * restrict vy1) {
+    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
+    const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
+    const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0;
+    const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1;
+
+    uint32_t nvec = n / VLEN_FP32;
+    uint32_t nloe = n % VLEN_FP32;
+
+    HVX_Vector r0_c0_sum = Q6_V_vzero();
+    HVX_Vector r0_c1_sum = Q6_V_vzero();
+    HVX_Vector r1_c0_sum = Q6_V_vzero();
+    HVX_Vector r1_c1_sum = Q6_V_vzero();
+
+    uint32_t i = 0;
+
+    #pragma unroll(2)
+    for (i = 0; i < nvec; i++) {
+        HVX_Vector r0_sf = x0[i];
+        HVX_Vector r1_sf = x1[i];
+        HVX_Vector c0_sf = y0[i];
+        HVX_Vector c1_sf = y1[i];
+
+        r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf));
+        r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf));
+        r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf));
+        r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf));
+    }
+
+    if (nloe) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+
+        HVX_Vector r0_sf = Q6_V_vand_QV(bmask, x0[i]);
+        HVX_Vector r1_sf = Q6_V_vand_QV(bmask, x1[i]);
+        HVX_Vector c0_sf = Q6_V_vand_QV(bmask, y0[i]);
+        HVX_Vector c1_sf = Q6_V_vand_QV(bmask, y1[i]);
+
+        r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf));
+        r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf));
+        r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf));
+        r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf));
+    }
+
+    // Reduce and store results
+    HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum);
+    HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum);
+
+    hvx_vec_store_u(s0, 8, r0_r1_c0_sum);
+    hvx_vec_store_u(s1, 8, r0_r1_c1_sum);
+}
+
+#undef HVX_OP_ADD_F32
+#undef HVX_OP_MUL_F32
+
+static inline void vec_dot_f16_f16_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
+    const HVX_Vector * restrict x = (const HVX_Vector *) vx;
+    const HVX_Vector * restrict y = (const HVX_Vector *) vy;
+
+    uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
+    uint32_t nloe = n % VLEN_FP16; // leftover elements
+
+    HVX_VectorPair rsum_p = Q6_W_vzero();
+
+    uint32_t i = 0;
+
+    #pragma unroll(4)
+    for (i = 0; i < nvec; i++) {
+        rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x[i], y[i]);
+    }
+
+    if (nloe) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
+        HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]);
+        HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]);
+        rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x_hf, y_hf);
+    }
+
+    HVX_Vector rsum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum_p), Q6_V_hi_W(rsum_p)));
+    hvx_vec_store_u(s, 4, hvx_vec_reduce_sum_f32(rsum));
+}
+
+static inline void vec_dot_f16_f16_aa_2x1(const uint32_t n, float * restrict s0,
+                                const void * restrict vx0, const void * restrict vx1,
+                                const void * restrict vy0) {
+    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
+    const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
+    const HVX_Vector * restrict y  = (const HVX_Vector *) vy0;
+
+    uint32_t nvec = n / VLEN_FP16;
+    uint32_t nloe = n % VLEN_FP16;
+
+    HVX_VectorPair rsum0_p = Q6_W_vzero();
+    HVX_VectorPair rsum1_p = Q6_W_vzero();
+
+    uint32_t i = 0;
+
+    #pragma unroll(2)
+    for (i = 0; i < nvec; i++) {
+        HVX_Vector y_hf = y[i];
+        rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0[i], y_hf);
+        rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1[i], y_hf);
+    }
+
+    if (nloe) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
+        HVX_Vector y_hf  = Q6_V_vand_QV(bmask, y[i]);
+        HVX_Vector x0_hf = Q6_V_vand_QV(bmask, x0[i]);
+        HVX_Vector x1_hf = Q6_V_vand_QV(bmask, x1[i]);
+        rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf);
+        rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf);
+    }
+
+    HVX_Vector rsum0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum0_p), Q6_V_hi_W(rsum0_p)));
+    HVX_Vector rsum1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum1_p), Q6_V_hi_W(rsum1_p)));
+    HVX_Vector rsum  = hvx_vec_reduce_sum_f32x2(rsum0, rsum1);
+    hvx_vec_store_u(s0, 8, rsum);
+}
+
+static inline void vec_dot_f16_f16_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
+                                const void * restrict vx0, const void * restrict vx1,
+                                const void * restrict vy0, const void * restrict vy1) {
+    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
+    const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
+    const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0;
+    const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1;
+
+    uint32_t nvec = n / VLEN_FP16;
+    uint32_t nloe = n % VLEN_FP16;
+
+    // Row sums (sf) - 4 accumulators for 2x2 tile
+    HVX_VectorPair r0_c0_sum_p = Q6_W_vzero();
+    HVX_VectorPair r0_c1_sum_p = Q6_W_vzero();
+    HVX_VectorPair r1_c0_sum_p = Q6_W_vzero();
+    HVX_VectorPair r1_c1_sum_p = Q6_W_vzero();
+
+    uint32_t i = 0;
+
+    #pragma unroll(2)
+    for (i = 0; i < nvec; i++) {
+        HVX_Vector r0_hf = x0[i];
+        HVX_Vector r1_hf = x1[i];
+        HVX_Vector c0_hf = y0[i];
+        HVX_Vector c1_hf = y1[i];
+
+        // Compute 4 dot products: r0xc0, r0xc1, r1xc0, r1xc1
+        r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf);
+        r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf);
+        r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf);
+        r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf);
+    }
+
+    if (nloe) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
+
+        HVX_Vector r0_hf = Q6_V_vand_QV(bmask, x0[i]);
+        HVX_Vector r1_hf = Q6_V_vand_QV(bmask, x1[i]);
+        HVX_Vector c0_hf = Q6_V_vand_QV(bmask, y0[i]);
+        HVX_Vector c1_hf = Q6_V_vand_QV(bmask, y1[i]);
+
+        r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf);
+        r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf);
+        r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf);
+        r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf);
+    }
+
+    HVX_Vector r0_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c0_sum_p), Q6_V_hi_W(r0_c0_sum_p)));
+    HVX_Vector r0_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c1_sum_p), Q6_V_hi_W(r0_c1_sum_p)));
+    HVX_Vector r1_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c0_sum_p), Q6_V_hi_W(r1_c0_sum_p)));
+    HVX_Vector r1_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c1_sum_p), Q6_V_hi_W(r1_c1_sum_p)));
+
+    // Reduce and store results
+    HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum);
+    HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum);
+
+    hvx_vec_store_u(&s0[0], 8, r0_r1_c0_sum);  // row0,col0 row1,col0
+    hvx_vec_store_u(&s1[0], 8, r0_r1_c1_sum);  // row0,col1 row1,col1
+}
+
+
+
+static inline void hvx_tensor_add_f32_grid(
+    const struct htp_tensor * restrict dst,
+    const struct htp_tensor * restrict src2,
+    uint32_t start_row,
+    uint32_t end_row,
+    uint32_t start_col,
+    uint32_t end_col,
+    const struct fastdiv_values * div_ne11_12,
+    const struct fastdiv_values * div_ne11
+) {
+    if (start_row >= end_row || start_col >= end_col) return;
+    const uint32_t nb1 = dst->nb[1]; // row stride in bytes
+
+    const uint32_t ne11 = dst->ne[1];
+    const uint32_t ne12 = dst->ne[2];
+    const uint32_t ne11_12 = ne11 * ne12;
+
+    const bool is_broadcast1 = (src2->ne[1] == 1);
+    const bool is_broadcast2 = (src2->ne[2] == 1);
+    const bool is_broadcast3 = (src2->ne[3] == 1);
+
+    for (uint32_t r = start_row; r < end_row; r++) {
+        float * dst_row = (float *) ((uint8_t *) dst->data + (size_t) r * nb1);
+
+        uint32_t i13 = fastdiv(r, div_ne11_12);
+        uint32_t i12 = fastdiv(r - i13 * ne11_12, div_ne11);
+        uint32_t i11 = r - i13 * ne11_12 - i12 * ne11;
+
+        uint32_t i23 = is_broadcast3 ? 0 : i13;
+        uint32_t i22 = is_broadcast2 ? 0 : i12;
+        uint32_t i21 = is_broadcast1 ? 0 : i11;
+
+        const float * src2_row = (const float *) ((const uint8_t *) src2->data +
+                                  (size_t) i21 * src2->nb[1] + (size_t) i22 * src2->nb[2] + (size_t) i23 * src2->nb[3]);
+
+        float * dst_ptr = &dst_row[start_col];
+        const float * src2_ptr = &src2_row[start_col];
+        int remaining = end_col - start_col;
+        while (remaining >= 32) {
+            HVX_Vector v_out = hvx_vmemu(dst_ptr);
+            HVX_Vector v_z   = hvx_vmemu(src2_ptr);
+            hvx_vmemu(dst_ptr) = hvx_vec_add_f32_f32(v_out, v_z);
+            dst_ptr += 32;
+            src2_ptr += 32;
+            remaining -= 32;
+        }
+        if (remaining > 0) {
+            HVX_Vector v_out = hvx_vmemu(dst_ptr);
+            HVX_Vector v_z   = hvx_vmemu(src2_ptr);
+            hvx_vec_store_u(dst_ptr, remaining * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z));
+        }
+    }
+}
+
+#endif // HVX_MM_KERNELS_FLOAT_H

ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h

@@ -48,22 +48,33 @@ static inline void quantize_block_f32_q8_1_tiled(float * restrict x, uint8_t * r
     v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 8));
     v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 16));
 
-    float vmax0[32]  __attribute__((aligned(128)));
-    float vmax1[32]  __attribute__((aligned(128)));
-    float vmax2[32]  __attribute__((aligned(128)));
-    float vmax3[32]  __attribute__((aligned(128)));
-    int32_t sums[32] __attribute__((aligned(128)));
+    const HVX_Vector v_inv127 = hvx_vec_splat_f32(1.0f / 127.0f);
+    HVX_Vector vd0_sf = hvx_vec_mul_f32_f32(vmax0_sf, v_inv127);
+    HVX_Vector vd1_sf = hvx_vec_mul_f32_f32(vmax1_sf, v_inv127);
+    HVX_Vector vd2_sf = hvx_vec_mul_f32_f32(vmax2_sf, v_inv127);
+    HVX_Vector vd3_sf = hvx_vec_mul_f32_f32(vmax3_sf, v_inv127);
 
-    hvx_vec_store_u(vmax0, 128, vmax0_sf);
-    hvx_vec_store_u(vmax1, 128, vmax1_sf);
-    hvx_vec_store_u(vmax2, 128, vmax2_sf);
-    hvx_vec_store_u(vmax3, 128, vmax3_sf);
-    hvx_vec_store_u(sums,  128, v_sums);
+    HVX_Vector v_sums_sf = Q6_Vsf_equals_Vw(v_sums);
+    HVX_Vector voff0_sf = hvx_vec_mul_f32_f32(vd0_sf, v_sums_sf);
+    HVX_Vector voff1_sf = hvx_vec_mul_f32_f32(vd1_sf, Q6_V_vror_VR(v_sums_sf, 32));
+    HVX_Vector voff2_sf = hvx_vec_mul_f32_f32(vd2_sf, Q6_V_vror_VR(v_sums_sf, 64));
+    HVX_Vector voff3_sf = hvx_vec_mul_f32_f32(vd3_sf, Q6_V_vror_VR(v_sums_sf, 96));
 
-    float d0 = vmax0[0] / 127.0f;
-    float d1 = vmax1[0] / 127.0f;
-    float d2 = vmax2[0] / 127.0f;
-    float d3 = vmax3[0] / 127.0f;
+    HVX_Vector voff01_hf = hvx_vec_f32_to_f16(voff0_sf, voff1_sf);
+    HVX_Vector voff23_hf = hvx_vec_f32_to_f16(voff2_sf, voff3_sf);
+
+    HVX_Vector r_scale[4] = {
+        hvx_vec_repl_f16(vd01_hf),
+        hvx_vec_repl_f16(Q6_V_vror_VR(vd01_hf, 64)),
+        hvx_vec_repl_f16(vd23_hf),
+        hvx_vec_repl_f16(Q6_V_vror_VR(vd23_hf, 64)),
+    };
+    HVX_Vector r_offset[4] = {
+        hvx_vec_repl_f16(voff01_hf),
+        hvx_vec_repl_f16(Q6_V_vror_VR(voff01_hf, 64)),
+        hvx_vec_repl_f16(voff23_hf),
+        hvx_vec_repl_f16(Q6_V_vror_VR(voff23_hf, 64)),
+    };
 
     static const uint8_t __attribute__((aligned(128))) repl[128] = {
         0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
@@ -89,24 +100,6 @@ static inline void quantize_block_f32_q8_1_tiled(float * restrict x, uint8_t * r
         HVX_Vector r6 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 24), v_repl_ctrl);
         HVX_Vector r7 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 28), v_repl_ctrl);
 
-        __fp16 scale_h, offset_h;
-        if (b == 0) {
-            scale_h  = (__fp16) d0;
-            offset_h = (__fp16) (sums[0] * d0);
-        } else if (b == 1) {
-            scale_h  = (__fp16) d1;
-            offset_h = (__fp16) (sums[8] * d1);
-        } else if (b == 2) {
-            scale_h  = (__fp16) d2;
-            offset_h = (__fp16) (sums[16] * d2);
-        } else {
-            scale_h  = (__fp16) d3;
-            offset_h = (__fp16) (sums[24] * d3);
-        }
-
-        HVX_Vector r_scale  = Q6_Vh_vsplat_R(*(int16_t *)&scale_h);
-        HVX_Vector r_offset = Q6_Vh_vsplat_R(*(int16_t *)&offset_h);
-
         HVX_Vector * restrict dst = (HVX_Vector *) (y_block + b * 1280);
         dst[0] = r0;
         dst[1] = r1;
@@ -116,8 +109,8 @@ static inline void quantize_block_f32_q8_1_tiled(float * restrict x, uint8_t * r
         dst[5] = r5;
         dst[6] = r6;
         dst[7] = r7;
-        dst[8] = r_scale;
-        dst[9] = r_offset;
+        dst[8] = r_scale[b];
+        dst[9] = r_offset[b];
     }
 }
 
@@ -486,51 +479,7 @@ static void tiled_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float
     HVX_Vector i8 = Q6_Vb_vsplat_R(8);
 
     uint32_t n_k_tiles = n / 32;
-    uint32_t kt = 0;
-    for (; kt + 1 < n_k_tiles; kt += 2) {
-        const HVX_Vector * restrict vptr0 = (const HVX_Vector *) (tile_ptr + (kt + 0) * 640);
-        const HVX_Vector * restrict v_act0_0 = (const HVX_Vector *) (y0_q + (kt + 0) * 1152);
-        const HVX_Vector * restrict v_act1_0 = (const HVX_Vector *) (y1_q + (kt + 0) * 1152);
-
-        const HVX_Vector * restrict vptr1 = (const HVX_Vector *) (tile_ptr + (kt + 1) * 640);
-        const HVX_Vector * restrict v_act0_1 = (const HVX_Vector *) (y0_q + (kt + 1) * 1152);
-        const HVX_Vector * restrict v_act1_1 = (const HVX_Vector *) (y1_q + (kt + 1) * 1152);
-
-        HVX_VectorPair v_sums0 = accum_4bit_32x2(vptr0, v_act0_0, v_act1_0, i8);
-        HVX_VectorPair v_sums1 = accum_4bit_32x2(vptr1, v_act0_1, v_act1_1, i8);
-
-        HVX_Vector v_sum_c0_0 = Q6_V_lo_W(v_sums0);
-        HVX_Vector v_sum_c1_0 = Q6_V_hi_W(v_sums0);
-        HVX_Vector v_sum_c0_1 = Q6_V_lo_W(v_sums1);
-        HVX_Vector v_sum_c1_1 = Q6_V_hi_W(v_sums1);
-
-        HVX_Vector v_sum_sf_c0_0 = Q6_Vsf_equals_Vw(v_sum_c0_0);
-        HVX_Vector v_sum_sf_c1_0 = Q6_Vsf_equals_Vw(v_sum_c1_0);
-        HVX_Vector v_sum_sf_c0_1 = Q6_Vsf_equals_Vw(v_sum_c0_1);
-        HVX_Vector v_sum_sf_c1_1 = Q6_Vsf_equals_Vw(v_sum_c1_1);
-
-        HVX_Vector v_scale_w0 = vptr0[4];
-        HVX_Vector v_scale_w1 = vptr1[4];
-        HVX_Vector v_scale_a_c0_0 = v_act0_0[8];
-        HVX_Vector v_scale_a_c1_0 = v_act1_0[8];
-        HVX_Vector v_scale_a_c0_1 = v_act0_1[8];
-        HVX_Vector v_scale_a_c1_1 = v_act1_1[8];
-
-        HVX_Vector v_scale_comb_c0_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c0_0);
-        HVX_Vector v_scale_comb_c1_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c1_0);
-        HVX_Vector v_scale_comb_c0_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c0_1);
-        HVX_Vector v_scale_comb_c1_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c1_1);
-
-        HVX_Vector v_sum_scaled_c0_0 = hvx_vec_mul_f32_f32(v_sum_sf_c0_0, v_scale_comb_c0_0);
-        HVX_Vector v_sum_scaled_c1_0 = hvx_vec_mul_f32_f32(v_sum_sf_c1_0, v_scale_comb_c1_0);
-        HVX_Vector v_sum_scaled_c0_1 = hvx_vec_mul_f32_f32(v_sum_sf_c0_1, v_scale_comb_c0_1);
-        HVX_Vector v_sum_scaled_c1_1 = hvx_vec_mul_f32_f32(v_sum_sf_c1_1, v_scale_comb_c1_1);
-
-        v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vec_add_f32_f32(v_sum_scaled_c0_0, v_sum_scaled_c0_1));
-        v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vec_add_f32_f32(v_sum_scaled_c1_0, v_sum_scaled_c1_1));
-    }
-
-    for (; kt < n_k_tiles; kt++) {
+    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
         const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640);
         const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1152);
         const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1152);
@@ -615,76 +564,7 @@ static void tiled_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float
     HVX_Vector v_sum_float_c1 = Q6_V_vzero();
 
     uint32_t n_k_tiles = n / 32;
-    uint32_t kt = 0;
-    for (; kt + 1 < n_k_tiles; kt += 2) {
-        const HVX_Vector * restrict vptr0 = (const HVX_Vector *) (tile_ptr + (kt + 0) * 640);
-        const HVX_Vector * restrict v_act0_0 = (const HVX_Vector *) (y0_q + (kt + 0) * 1280);
-        const HVX_Vector * restrict v_act1_0 = (const HVX_Vector *) (y1_q + (kt + 0) * 1280);
-
-        const HVX_Vector * restrict vptr1 = (const HVX_Vector *) (tile_ptr + (kt + 1) * 640);
-        const HVX_Vector * restrict v_act0_1 = (const HVX_Vector *) (y0_q + (kt + 1) * 1280);
-        const HVX_Vector * restrict v_act1_1 = (const HVX_Vector *) (y1_q + (kt + 1) * 1280);
-
-        HVX_VectorPair v_sums0 = accum_4bit_32x2(vptr0, v_act0_0, v_act1_0, Q6_V_vzero());
-        HVX_VectorPair v_sums1 = accum_4bit_32x2(vptr1, v_act0_1, v_act1_1, Q6_V_vzero());
-
-        HVX_Vector v_sum_c0_0 = Q6_V_lo_W(v_sums0);
-        HVX_Vector v_sum_c1_0 = Q6_V_hi_W(v_sums0);
-        HVX_Vector v_sum_c0_1 = Q6_V_lo_W(v_sums1);
-        HVX_Vector v_sum_c1_1 = Q6_V_hi_W(v_sums1);
-
-        HVX_Vector v_sum_sf_c0_0 = Q6_Vsf_equals_Vw(v_sum_c0_0);
-        HVX_Vector v_sum_sf_c1_0 = Q6_Vsf_equals_Vw(v_sum_c1_0);
-        HVX_Vector v_sum_sf_c0_1 = Q6_Vsf_equals_Vw(v_sum_c0_1);
-        HVX_Vector v_sum_sf_c1_1 = Q6_Vsf_equals_Vw(v_sum_c1_1);
-
-        HVX_Vector v_scale_offset0 = vptr0[4];
-        HVX_VectorPair p_deal0 = Q6_W_vdeal_VVR(v_scale_offset0, v_scale_offset0, -2);
-        HVX_Vector v_scale0 = Q6_V_lo_W(p_deal0);
-        HVX_Vector v_offset0 = Q6_V_hi_W(p_deal0);
-
-        HVX_Vector v_scale_offset1 = vptr1[4];
-        HVX_VectorPair p_deal1 = Q6_W_vdeal_VVR(v_scale_offset1, v_scale_offset1, -2);
-        HVX_Vector v_scale1 = Q6_V_lo_W(p_deal1);
-        HVX_Vector v_offset1 = Q6_V_hi_W(p_deal1);
-
-        HVX_Vector v_scale_a_c0_0 = v_act0_0[8];
-        HVX_Vector v_sum_a_c0_0   = v_act0_0[9];
-        HVX_Vector v_scale_a_c1_0 = v_act1_0[8];
-        HVX_Vector v_sum_a_c1_0   = v_act1_0[9];
-
-        HVX_Vector v_scale_a_c0_1 = v_act0_1[8];
-        HVX_Vector v_sum_a_c0_1   = v_act0_1[9];
-        HVX_Vector v_scale_a_c1_1 = v_act1_1[8];
-        HVX_Vector v_sum_a_c1_1   = v_act1_1[9];
-
-        HVX_Vector v_scale_comb_c0_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale0, v_scale_a_c0_0);
-        HVX_Vector v_offset_comb_c0_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset0, v_sum_a_c0_0);
-        HVX_Vector v_scale_comb_c1_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale0, v_scale_a_c1_0);
-        HVX_Vector v_offset_comb_c1_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset0, v_sum_a_c1_0);
-
-        HVX_Vector v_scale_comb_c0_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale1, v_scale_a_c0_1);
-        HVX_Vector v_offset_comb_c0_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset1, v_sum_a_c0_1);
-        HVX_Vector v_scale_comb_c1_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale1, v_scale_a_c1_1);
-        HVX_Vector v_offset_comb_c1_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset1, v_sum_a_c1_1);
-
-        HVX_Vector v_scaled_dot_c0_0 = hvx_vec_mul_f32_f32(v_sum_sf_c0_0, v_scale_comb_c0_0);
-        HVX_Vector v_sum_scaled_c0_0 = hvx_vec_add_f32_f32(v_scaled_dot_c0_0, v_offset_comb_c0_0);
-
-        HVX_Vector v_scaled_dot_c1_0 = hvx_vec_mul_f32_f32(v_sum_sf_c1_0, v_scale_comb_c1_0);
-        HVX_Vector v_sum_scaled_c1_0 = hvx_vec_add_f32_f32(v_scaled_dot_c1_0, v_offset_comb_c1_0);
-
-        HVX_Vector v_scaled_dot_c0_1 = hvx_vec_mul_f32_f32(v_sum_sf_c0_1, v_scale_comb_c0_1);
-        HVX_Vector v_sum_scaled_c0_1 = hvx_vec_add_f32_f32(v_scaled_dot_c0_1, v_offset_comb_c0_1);
-
-        HVX_Vector v_scaled_dot_c1_1 = hvx_vec_mul_f32_f32(v_sum_sf_c1_1, v_scale_comb_c1_1);
-        HVX_Vector v_sum_scaled_c1_1 = hvx_vec_add_f32_f32(v_scaled_dot_c1_1, v_offset_comb_c1_1);
-
-        v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vec_add_f32_f32(v_sum_scaled_c0_0, v_sum_scaled_c0_1));
-        v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vec_add_f32_f32(v_sum_scaled_c1_0, v_sum_scaled_c1_1));
-    }
-
-    for (; kt < n_k_tiles; kt++) {
+    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
         const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640);
         const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1280);
         const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1280);
@@ -771,51 +651,7 @@ static void tiled_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float
     HVX_Vector v_sum_float_c1 = Q6_V_vzero();
 
     uint32_t n_k_tiles = n / 32;
-    uint32_t kt = 0;
-    for (; kt + 1 < n_k_tiles; kt += 2) {
-        const HVX_Vector * restrict vptr0 = (const HVX_Vector *) (tile_ptr + (kt + 0) * 1152);
-        const HVX_Vector * restrict v_act0_0 = (const HVX_Vector *) (y0_q + (kt + 0) * 1152);
-        const HVX_Vector * restrict v_act1_0 = (const HVX_Vector *) (y1_q + (kt + 0) * 1152);
-
-        const HVX_Vector * restrict vptr1 = (const HVX_Vector *) (tile_ptr + (kt + 1) * 1152);
-        const HVX_Vector * restrict v_act0_1 = (const HVX_Vector *) (y0_q + (kt + 1) * 1152);
-        const HVX_Vector * restrict v_act1_1 = (const HVX_Vector *) (y1_q + (kt + 1) * 1152);
-
-        HVX_VectorPair v_sums0 = accum_q8_0_32x2(vptr0, v_act0_0, v_act1_0);
-        HVX_VectorPair v_sums1 = accum_q8_0_32x2(vptr1, v_act0_1, v_act1_1);
-
-        HVX_Vector v_sum_c0_0 = Q6_V_lo_W(v_sums0);
-        HVX_Vector v_sum_c1_0 = Q6_V_hi_W(v_sums0);
-        HVX_Vector v_sum_c0_1 = Q6_V_lo_W(v_sums1);
-        HVX_Vector v_sum_c1_1 = Q6_V_hi_W(v_sums1);
-
-        HVX_Vector v_sum_sf_c0_0 = Q6_Vsf_equals_Vw(v_sum_c0_0);
-        HVX_Vector v_sum_sf_c1_0 = Q6_Vsf_equals_Vw(v_sum_c1_0);
-        HVX_Vector v_sum_sf_c0_1 = Q6_Vsf_equals_Vw(v_sum_c0_1);
-        HVX_Vector v_sum_sf_c1_1 = Q6_Vsf_equals_Vw(v_sum_c1_1);
-
-        HVX_Vector v_scale_w0 = vptr0[8];
-        HVX_Vector v_scale_w1 = vptr1[8];
-        HVX_Vector v_scale_a_c0_0 = v_act0_0[8];
-        HVX_Vector v_scale_a_c1_0 = v_act1_0[8];
-        HVX_Vector v_scale_a_c0_1 = v_act0_1[8];
-        HVX_Vector v_scale_a_c1_1 = v_act1_1[8];
-
-        HVX_Vector v_scale_comb_c0_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c0_0);
-        HVX_Vector v_scale_comb_c1_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c1_0);
-        HVX_Vector v_scale_comb_c0_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c0_1);
-        HVX_Vector v_scale_comb_c1_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c1_1);
-
-        HVX_Vector v_sum_scaled_c0_0 = hvx_vec_mul_f32_f32(v_sum_sf_c0_0, v_scale_comb_c0_0);
-        HVX_Vector v_sum_scaled_c1_0 = hvx_vec_mul_f32_f32(v_sum_sf_c1_0, v_scale_comb_c1_0);
-        HVX_Vector v_sum_scaled_c0_1 = hvx_vec_mul_f32_f32(v_sum_sf_c0_1, v_scale_comb_c0_1);
-        HVX_Vector v_sum_scaled_c1_1 = hvx_vec_mul_f32_f32(v_sum_sf_c1_1, v_scale_comb_c1_1);
-
-        v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vec_add_f32_f32(v_sum_scaled_c0_0, v_sum_scaled_c0_1));
-        v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vec_add_f32_f32(v_sum_scaled_c1_0, v_sum_scaled_c1_1));
-    }
-
-    for (; kt < n_k_tiles; kt++) {
+    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
         const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 1152);
         const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1152);
         const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1152);
@@ -952,51 +788,7 @@ static void tiled_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, floa
     HVX_Vector lut = *(const HVX_Vector *) kvalues_iq4nl_lut;
 
     uint32_t n_k_tiles = n / 32;
-    uint32_t kt = 0;
-    for (; kt + 1 < n_k_tiles; kt += 2) {
-        const HVX_Vector * restrict vptr0 = (const HVX_Vector *) (tile_ptr + (kt + 0) * 640);
-        const HVX_Vector * restrict v_act0_0 = (const HVX_Vector *) (y0_q + (kt + 0) * 1152);
-        const HVX_Vector * restrict v_act1_0 = (const HVX_Vector *) (y1_q + (kt + 0) * 1152);
-
-        const HVX_Vector * restrict vptr1 = (const HVX_Vector *) (tile_ptr + (kt + 1) * 640);
-        const HVX_Vector * restrict v_act0_1 = (const HVX_Vector *) (y0_q + (kt + 1) * 1152);
-        const HVX_Vector * restrict v_act1_1 = (const HVX_Vector *) (y1_q + (kt + 1) * 1152);
-
-        HVX_VectorPair v_sums0 = accum_4bit_32x2_lut(vptr0, v_act0_0, v_act1_0, mask_h4, lut);
-        HVX_VectorPair v_sums1 = accum_4bit_32x2_lut(vptr1, v_act0_1, v_act1_1, mask_h4, lut);
-
-        HVX_Vector v_sum_c0_0 = Q6_V_lo_W(v_sums0);
-        HVX_Vector v_sum_c1_0 = Q6_V_hi_W(v_sums0);
-        HVX_Vector v_sum_c0_1 = Q6_V_lo_W(v_sums1);
-        HVX_Vector v_sum_c1_1 = Q6_V_hi_W(v_sums1);
-
-        HVX_Vector v_sum_sf_c0_0 = Q6_Vsf_equals_Vw(v_sum_c0_0);
-        HVX_Vector v_sum_sf_c1_0 = Q6_Vsf_equals_Vw(v_sum_c1_0);
-        HVX_Vector v_sum_sf_c0_1 = Q6_Vsf_equals_Vw(v_sum_c0_1);
-        HVX_Vector v_sum_sf_c1_1 = Q6_Vsf_equals_Vw(v_sum_c1_1);
-
-        HVX_Vector v_scale_w0 = vptr0[4];
-        HVX_Vector v_scale_w1 = vptr1[4];
-        HVX_Vector v_scale_a_c0_0 = v_act0_0[8];
-        HVX_Vector v_scale_a_c1_0 = v_act1_0[8];
-        HVX_Vector v_scale_a_c0_1 = v_act0_1[8];
-        HVX_Vector v_scale_a_c1_1 = v_act1_1[8];
-
-        HVX_Vector v_scale_comb_c0_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c0_0);
-        HVX_Vector v_scale_comb_c1_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c1_0);
-        HVX_Vector v_scale_comb_c0_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c0_1);
-        HVX_Vector v_scale_comb_c1_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c1_1);
-
-        HVX_Vector v_sum_scaled_c0_0 = hvx_vec_mul_f32_f32(v_sum_sf_c0_0, v_scale_comb_c0_0);
-        HVX_Vector v_sum_scaled_c1_0 = hvx_vec_mul_f32_f32(v_sum_sf_c1_0, v_scale_comb_c1_0);
-        HVX_Vector v_sum_scaled_c0_1 = hvx_vec_mul_f32_f32(v_sum_sf_c0_1, v_scale_comb_c0_1);
-        HVX_Vector v_sum_scaled_c1_1 = hvx_vec_mul_f32_f32(v_sum_sf_c1_1, v_scale_comb_c1_1);
-
-        v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vec_add_f32_f32(v_sum_scaled_c0_0, v_sum_scaled_c0_1));
-        v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vec_add_f32_f32(v_sum_scaled_c1_0, v_sum_scaled_c1_1));
-    }
-
-    for (; kt < n_k_tiles; kt++) {
+    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
         const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640);
         const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1152);
         const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1152);
@@ -1089,69 +881,7 @@ static void tiled_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, floa
     HVX_Vector e8m0_mask = Q6_V_vsplat_R(0x000000ff);
 
     uint32_t n_k_tiles = n / 32;
-    uint32_t kt = 0;
-    for (; kt + 1 < n_k_tiles; kt += 2) {
-        const HVX_Vector * restrict vptr0 = (const HVX_Vector *) (tile_ptr + (kt + 0) * 640);
-        const HVX_Vector * restrict v_act0_0 = (const HVX_Vector *) (y0_q + (kt + 0) * 1152);
-        const HVX_Vector * restrict v_act1_0 = (const HVX_Vector *) (y1_q + (kt + 0) * 1152);
-
-        const HVX_Vector * restrict vptr1 = (const HVX_Vector *) (tile_ptr + (kt + 1) * 640);
-        const HVX_Vector * restrict v_act0_1 = (const HVX_Vector *) (y0_q + (kt + 1) * 1152);
-        const HVX_Vector * restrict v_act1_1 = (const HVX_Vector *) (y1_q + (kt + 1) * 1152);
-
-        HVX_VectorPair v_sums0 = accum_4bit_32x2_lut(vptr0, v_act0_0, v_act1_0, mask_h4, lut);
-        HVX_VectorPair v_sums1 = accum_4bit_32x2_lut(vptr1, v_act0_1, v_act1_1, mask_h4, lut);
-
-        HVX_Vector v_sum_c0_0 = Q6_V_lo_W(v_sums0);
-        HVX_Vector v_sum_c1_0 = Q6_V_hi_W(v_sums0);
-        HVX_Vector v_sum_c0_1 = Q6_V_lo_W(v_sums1);
-        HVX_Vector v_sum_c1_1 = Q6_V_hi_W(v_sums1);
-
-        HVX_Vector v_sum_sf_c0_0 = Q6_Vsf_equals_Vw(v_sum_c0_0);
-        HVX_Vector v_sum_sf_c1_0 = Q6_Vsf_equals_Vw(v_sum_c1_0);
-        HVX_Vector v_sum_sf_c0_1 = Q6_Vsf_equals_Vw(v_sum_c0_1);
-        HVX_Vector v_sum_sf_c1_1 = Q6_Vsf_equals_Vw(v_sum_c1_1);
-
-        HVX_Vector v_scale_w0 = hvx_vmem(tile_ptr + (kt + 0) * 640 + 512);
-        HVX_Vector r0_d0 = Q6_V_vdelta_VV(v_scale_w0, expand);
-        r0_d0 = Q6_V_vand_VV(r0_d0, e8m0_mask);
-        HVX_Vector v_scale_w_f32_0 = Q6_Vw_vasl_VwR(r0_d0, 23);
-
-        HVX_Vector v_scale_w1 = hvx_vmem(tile_ptr + (kt + 1) * 640 + 512);
-        HVX_Vector r0_d1 = Q6_V_vdelta_VV(v_scale_w1, expand);
-        r0_d1 = Q6_V_vand_VV(r0_d1, e8m0_mask);
-        HVX_Vector v_scale_w_f32_1 = Q6_Vw_vasl_VwR(r0_d1, 23);
-
-        HVX_Vector v_scale_a_c0_f16_0 = v_act0_0[8];
-        HVX_Vector v_scale_a_c1_f16_0 = v_act1_0[8];
-        HVX_Vector v_scale_a_c0_f16_1 = v_act0_1[8];
-        HVX_Vector v_scale_a_c1_f16_1 = v_act1_1[8];
-
-        HVX_VectorPair p_scale_a_c0_f32_0 = hvx_vec_f16_to_f32_shuff(v_scale_a_c0_f16_0);
-        HVX_VectorPair p_scale_a_c1_f32_0 = hvx_vec_f16_to_f32_shuff(v_scale_a_c1_f16_0);
-        HVX_VectorPair p_scale_a_c0_f32_1 = hvx_vec_f16_to_f32_shuff(v_scale_a_c0_f16_1);
-        HVX_VectorPair p_scale_a_c1_f32_1 = hvx_vec_f16_to_f32_shuff(v_scale_a_c1_f16_1);
-
-        HVX_Vector v_scale_a_c0_0 = Q6_V_lo_W(p_scale_a_c0_f32_0);
-        HVX_Vector v_scale_a_c1_0 = Q6_V_lo_W(p_scale_a_c1_f32_0);
-        HVX_Vector v_scale_a_c0_1 = Q6_V_lo_W(p_scale_a_c0_f32_1);
-        HVX_Vector v_scale_a_c1_1 = Q6_V_lo_W(p_scale_a_c1_f32_1);
-
-        HVX_Vector v_scale_comb_c0_0 = hvx_vec_mul_f32_f32(v_scale_w_f32_0, v_scale_a_c0_0);
-        HVX_Vector v_scale_comb_c1_0 = hvx_vec_mul_f32_f32(v_scale_w_f32_0, v_scale_a_c1_0);
-        HVX_Vector v_scale_comb_c0_1 = hvx_vec_mul_f32_f32(v_scale_w_f32_1, v_scale_a_c0_1);
-        HVX_Vector v_scale_comb_c1_1 = hvx_vec_mul_f32_f32(v_scale_w_f32_1, v_scale_a_c1_1);
-
-        HVX_Vector v_sum_scaled_c0_0 = hvx_vec_mul_f32_f32(v_sum_sf_c0_0, v_scale_comb_c0_0);
-        HVX_Vector v_sum_scaled_c1_0 = hvx_vec_mul_f32_f32(v_sum_sf_c1_0, v_scale_comb_c1_0);
-        HVX_Vector v_sum_scaled_c0_1 = hvx_vec_mul_f32_f32(v_sum_sf_c0_1, v_scale_comb_c0_1);
-        HVX_Vector v_sum_scaled_c1_1 = hvx_vec_mul_f32_f32(v_sum_sf_c1_1, v_scale_comb_c1_1);
-
-        v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vec_add_f32_f32(v_sum_scaled_c0_0, v_sum_scaled_c0_1));
-        v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vec_add_f32_f32(v_sum_scaled_c1_0, v_sum_scaled_c1_1));
-    }
-
-    for (; kt < n_k_tiles; kt++) {
+    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
         const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640);
         const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1152);
         const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1152);

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

@@ -14,7 +14,7 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "hvx-utils.h"
-#include "hex-dma.h"
+#include "dma-queue.h"
 #include "hex-profile.h"
 #include "htp-vtcm.h"
 #include "htp-tensor.h"
@@ -25,17 +25,20 @@ struct htp_im2col_context {
     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_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
+    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
     uint8_t * pe_vtcm_dst;                         // base of the 2x dst buffers region
     uint32_t  pe_src_size_per_thread;              // 2 * pe_src_row_bytes
     uint32_t  pe_dst_size_per_thread;              // 2 * pe_dst_row_bytes
+
+    uint32_t pe_owb;                               // output-col block size
+    uint32_t pe_wb;                                // staged source window width
 };
 
 // Per-op VTCM layout for the patch-embed DMA path
@@ -59,146 +62,324 @@ static inline void htp_im2col_vtcm_layout_build(struct htp_im2col_vtcm_layout *
     L->total_bytes = L->off_dst + L->dst_bytes_per_thread * n_threads;
 }
 
-#define IM2COL_PATCHEMBED_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG)                        \
-    static void FNAME(unsigned int nth, unsigned int ith, void * data) {                                  \
-        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], 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 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_stop; p++) {                                             \
-            const uint32_t iow             = p % OW;                                                      \
-            const uint32_t ioh             = (p / OW) % OH;                                               \
-            const uint32_t in              = p / (OW * OH);                                               \
-            DST_CTYPE * restrict dst_patch = dst_data + (uint64_t) p * patch_stride;                      \
-            for (uint32_t iic = 0; iic < IC; iic++) {                                                     \
-                const float * restrict src_plane = src_data + ((uint64_t) in * IC + iic) * IH * IW;       \
-                for (uint32_t ikh = 0; ikh < KH; ikh++) {                                                 \
-                    const int32_t iih            = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1;          \
-                    DST_CTYPE * restrict out_run = dst_patch + iic * (KH * KW) + ikh * KW;                \
-                    if (iih < 0 || iih >= (int32_t) IH) {                                                 \
-                        SPLAT_FN(out_run, 0.0f, KW);                                                      \
-                        continue;                                                                         \
-                    }                                                                                     \
-                    const int32_t iiw0             = (int32_t) iow * s0 - p0;                             \
-                    const float * restrict src_run = src_plane + (uint64_t) iih * IW + iiw0;              \
-                    if (d0 == 1) {                                                                        \
-                        /* contiguous source run: [lo,hi) is in-bounds, tails are zero pad */             \
-                        const int32_t lo = iiw0 < 0 ? -iiw0 : 0;                                          \
-                        int32_t       hi = (int32_t) IW - iiw0;                                           \
-                        if (hi > (int32_t) KW) {                                                          \
-                            hi = (int32_t) KW;                                                            \
-                        }                                                                                 \
-                        if (hi <= lo) {                                                                   \
-                            SPLAT_FN(out_run, 0.0f, KW);                                                  \
-                        } else {                                                                          \
-                            if (lo > 0) {                                                                 \
-                                SPLAT_FN(out_run, 0.0f, (uint32_t) lo);                                   \
-                            }                                                                             \
-                            COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (src_run + lo),         \
-                                    (uint32_t) (hi - lo));                                                \
-                            if (hi < (int32_t) KW) {                                                      \
-                                SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi));                       \
-                            }                                                                             \
-                        }                                                                                 \
-                        continue;                                                                         \
-                    }                                                                                     \
-                    for (uint32_t ikw = 0; ikw < KW; ikw++) {                                             \
-                        const int32_t iiw = (int32_t) iow * s0 + (int32_t) ikw * d0 - p0;                 \
-                        out_run[ikw]      = (iiw < 0 || iiw >= (int32_t) IW) ?                            \
-                                                (DST_CTYPE) 0.0f :                                        \
-                                                (DST_CTYPE) src_plane[(uint64_t) iih * IW + iiw];         \
-                    }                                                                                     \
-                }                                                                                         \
-            }                                                                                             \
-        }                                                                                                 \
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, patch_start);                                    \
+#define IM2COL_PATCHEMBED_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG)                         \
+    static void FNAME(unsigned int nth, unsigned int ith, void * data) {                                   \
+        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 int32_t  is_2D                    = octx->op_params[6] == 1;                                 \
+        const uint32_t N                        = is_2D ? src1->ne[3] : src1->ne[2];                       \
+        const uint32_t IC                       = is_2D ? src1->ne[2] : src1->ne[1];                       \
+        const uint32_t IH                       = is_2D ? src1->ne[1] : 1;                                 \
+        const uint32_t IW                       = src1->ne[0];                                             \
+        const uint32_t KH                       = is_2D ? src0->ne[1] : 1;                                 \
+        const uint32_t KW                       = src0->ne[0];                                             \
+        const uint32_t OH                       = is_2D ? dst->ne[2] : 1;                                  \
+        const uint32_t OW                       = dst->ne[1];                                              \
+        const uint32_t patch_stride             = IC * KH * KW;                                            \
+        const float * restrict src_data         = (const float *) (uintptr_t) src1->data;                  \
+        DST_CTYPE * restrict dst_data           = (DST_CTYPE *) (uintptr_t) dst->data;                     \
+        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_stop; p++) {                                              \
+            const uint32_t iow             = p % OW;                                                       \
+            const uint32_t ioh             = (p / OW) % OH;                                                \
+            const uint32_t in              = p / (OW * OH);                                                \
+            DST_CTYPE * restrict dst_patch = dst_data + (uint64_t) p * patch_stride;                       \
+            for (uint32_t iic = 0; iic < IC; iic++) {                                                      \
+                const float * restrict src_plane = src_data + ((uint64_t) in * IC + iic) * IH * IW;        \
+                for (uint32_t ikh = 0; ikh < KH; ikh++) {                                                  \
+                    const int32_t iih            = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1;           \
+                    DST_CTYPE * restrict out_run = dst_patch + iic * (KH * KW) + ikh * KW;                 \
+                    if (iih < 0 || iih >= (int32_t) IH) {                                                  \
+                        SPLAT_FN(out_run, 0.0f, KW);                                                       \
+                        continue;                                                                          \
+                    }                                                                                      \
+                    const int32_t iiw0             = (int32_t) iow * s0 - p0;                              \
+                    const float * restrict src_run = src_plane + (uint64_t) iih * IW + iiw0;               \
+                    if (d0 == 1) {                                                                         \
+                        /* contiguous source run: [lo,hi) is in-bounds, tails are zero pad */              \
+                        const int32_t lo = iiw0 < 0 ? -iiw0 : 0;                                           \
+                        int32_t       hi = (int32_t) IW - iiw0;                                            \
+                        if (hi > (int32_t) KW) {                                                           \
+                            hi = (int32_t) KW;                                                             \
+                        }                                                                                  \
+                        if (hi <= lo) {                                                                    \
+                            SPLAT_FN(out_run, 0.0f, KW);                                                   \
+                        } else {                                                                           \
+                            if (lo > 0) {                                                                  \
+                                SPLAT_FN(out_run, 0.0f, (uint32_t) lo);                                    \
+                            }                                                                              \
+                            COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (src_run + lo),          \
+                                    (uint32_t) (hi - lo));                                                 \
+                            if (hi < (int32_t) KW) {                                                       \
+                                SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi));                        \
+                            }                                                                              \
+                        }                                                                                  \
+                        continue;                                                                          \
+                    }                                                                                      \
+                    for (uint32_t ikw = 0; ikw < KW; ikw++) {                                              \
+                        const int32_t iiw = (int32_t) iow * s0 + (int32_t) ikw * d0 - p0;                  \
+                        out_run[ikw]      = (iiw < 0 || iiw >= (int32_t) IW) ?                             \
+                                                (DST_CTYPE) 0.0f :                                         \
+                                                (DST_CTYPE) src_plane[(uint64_t) iih * IW + iiw];          \
+                    }                                                                                      \
+                }                                                                                          \
+            }                                                                                              \
+        }                                                                                                  \
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, patch_start);                                     \
     }
 
 IM2COL_PATCHEMBED_BODY(im2col_patchembed_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "f32-f16")
 IM2COL_PATCHEMBED_BODY(im2col_patchembed_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "f32-f32")
 
-#define IM2COL_PATCHEMBED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG)                               \
-    static void FNAME(unsigned int nth, unsigned int ith, void * data) {                                             \
-        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 src1 = octx->src[1];                                                      \
-        const struct htp_tensor * restrict dst  = octx->dst;                                                         \
-        const uint32_t N = src1->ne[3], IC = src1->ne[2], IH = src1->ne[1], IW = src1->ne[0];                        \
-        const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];                                           \
-        const uint32_t OH = dst->ne[2], 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;                                                   \
-        dma_queue *    dmaq             = octx->ctx->dma[ith];                                                       \
-        uint8_t *      src_base         = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread;                    \
-        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 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        = 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++) {                                                             \
-            const uint32_t in  = r / OH;                                                                             \
-            const uint32_t ioh = r % OH;                                                                             \
-            for (uint32_t ikh = 0; ikh < KH; ikh++) {                                                                \
-                int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh;                                          \
-                int     ok  = (iih >= 0 && iih < (int32_t) IH);                                                      \
-                for (uint32_t iic = 0; iic < IC; iic++) {                                                            \
-                    float *       vdst = srcb + ((uint64_t) (iic * KH + ikh)) * IW;                                  \
-                    const float * _vsrc =                                                                            \
-                        ok ? (src_data + ((uint64_t) (in * IC + iic) * IH + iih) * IW) : (const float *) vdst;       \
-                    dma_queue_push_ddr_to_vtcm(                                                                      \
-                        dmaq, dma_make_ptr((uint8_t *) vdst, ok ? (const uint8_t *) _vsrc : (const uint8_t *) vdst), \
-                        IW * sizeof(float), IW * sizeof(float), ok ? 1 : 0);                                         \
-                }                                                                                                    \
-            }                                                                                                        \
-            for (uint32_t i = 0; i < IC * KH; i++)                                                                   \
-                dma_queue_pop(dmaq);                                                                                 \
-            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r);                                                    \
-            for (uint32_t iow = 0; iow < OW; iow++) {                                                                \
-                DST_CTYPE * dst_patch = dstb + (uint64_t) iow * patch_stride;                                        \
-                for (uint32_t ikh = 0; ikh < KH; ikh++) {                                                            \
-                    int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh;                                      \
-                    for (uint32_t iic = 0; iic < IC; iic++) {                                                        \
-                        DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW;                                \
-                        if (iih < 0 || iih >= (int32_t) IH) {                                                        \
-                            SPLAT_FN(out_run, 0.0f, KW);                                                             \
-                            continue;                                                                                \
-                        }                                                                                            \
-                        const float * src_run = srcb + ((uint64_t) (iic * KH + ikh)) * IW + (uint64_t) iow * KW;     \
-                        COPY_FN((uint8_t *) out_run, (const uint8_t *) src_run, KW);                                 \
-                    }                                                                                                \
-                }                                                                                                    \
-            }                                                                                                        \
-            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r);                                                     \
-            DST_CTYPE * ddr_row = dst_data + ((uint64_t) (in * OH + ioh) * OW) * patch_stride;                       \
-            dma_queue_push_vtcm_to_ddr(dmaq, dma_make_ptr((uint8_t *) ddr_row, (uint8_t *) dstb),                    \
-                                       OW * patch_stride * (DST_ELEM), OW * patch_stride * (DST_ELEM), 1);           \
-            dma_queue_flush(dmaq);                                                                                   \
-        }                                                                                                            \
+// Software-pipelined 2-deep: while HVX computes block bi from buffer slot
+// (bi&1), the DMA engine stages block bi+1 into the other slot concurrently.
+// A single dma_queue_flush per iteration (after issuing the next stage-in and
+// this block's store-out) waits for both - safe because the ring is strict
+// FIFO and each buffer slot is only reused after its prior consumer (compute
+// or store-out) already finished in program order.
+#define IM2COL_BLOCKED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG)                                            \
+    static void FNAME(unsigned int nth, unsigned int ith, void * data) {                                                       \
+        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 src1 = octx->src[1];                                                                \
+        const struct htp_tensor * restrict dst  = octx->dst;                                                                   \
+        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 int32_t  is_2D = octx->op_params[6] == 1;                                                                        \
+        const uint32_t N     = is_2D ? src1->ne[3] : src1->ne[2];                                                              \
+        const uint32_t IC    = is_2D ? src1->ne[2] : src1->ne[1];                                                              \
+        const uint32_t IH    = is_2D ? src1->ne[1] : 1;                                                                        \
+        const uint32_t IW    = src1->ne[0];                                                                                    \
+        const uint32_t KH    = is_2D ? octx->src[0]->ne[1] : 1;                                                                \
+        const uint32_t KW    = octx->src[0]->ne[0];                                                                            \
+        const uint32_t OH    = is_2D ? dst->ne[2] : 1;                                                                         \
+        const uint32_t OW    = dst->ne[1];                                                                                     \
+        const uint32_t owb = ictx->pe_owb, Wb = ictx->pe_wb;                                                                   \
+        const uint32_t patch_stride     = IC * KH * KW;                                                                        \
+        const dma_addr_t src_data       = src1->data;                                                                          \
+        const dma_addr_t dst_data       = dst->data;                                                                           \
+        dma_queue *    dmaq             = octx->ctx->dma[ith];                                                                 \
+        uint8_t *      srcb_base        = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread;                              \
+        uint8_t *      dstb_base        = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread;                              \
+        float *        srcb2[2]         = { (float *) srcb_base, (float *) (srcb_base + ictx->pe_src_row_bytes) };             \
+        DST_CTYPE *    dstb2[2]   = { (DST_CTYPE *) dstb_base, (DST_CTYPE *) (dstb_base + ictx->pe_dst_row_bytes) };           \
+        const uint32_t nrows      = N * OH;                                                                                    \
+        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);                                                        \
+        if (row_start >= row_end)                                                                                              \
+            return;                                                                                                            \
+        const uint32_t nbpr         = (OW + owb - 1) / owb;                                                                    \
+        const uint32_t nrows_local  = row_end - row_start;                                                                     \
+        const uint32_t total_blocks = nrows_local * nbpr;                                                                      \
+        for (uint32_t bi = 0; bi < total_blocks; bi++) {                                                                       \
+            const uint32_t buf  = bi & 1u;                                                                                     \
+            float *        srcb = srcb2[buf];                                                                                  \
+            DST_CTYPE *    dstb = dstb2[buf];                                                                                  \
+            const uint32_t r    = row_start + bi / nbpr;                                                                       \
+            const uint32_t in   = r / OH;                                                                                      \
+            const uint32_t ioh  = r % OH;                                                                                      \
+            const uint32_t c0   = (bi % nbpr) * owb;                                                                           \
+            const uint32_t nb   = MIN(owb, OW - c0);                                                                           \
+            const int32_t  win0 = (int32_t) c0 * s0 - p0;                                                                      \
+            if (bi == 0) {                                                                                                     \
+                /* prologue: stage block 0 and wait - nothing to overlap with yet */                                           \
+                for (uint32_t ikh = 0; ikh < KH; ikh++) {                                                                      \
+                    const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1;                                          \
+                    if (iih < 0 || iih >= (int32_t) IH)                                                                        \
+                        continue;                                                                                              \
+                    const int32_t lo = win0 < 0 ? -win0 : 0;                                                                   \
+                    int32_t       hi = (int32_t) IW - win0;                                                                    \
+                    if (hi > (int32_t) Wb)                                                                                     \
+                        hi = (int32_t) Wb;                                                                                     \
+                    if (hi <= lo)                                                                                              \
+                        continue;                                                                                              \
+                    const uint32_t cpw  = (uint32_t) (hi - lo);                                                                \
+                    float *        vdst = srcb + (size_t) ikh * Wb + lo;                                                       \
+                    const dma_addr_t vsrc = src_data + (size_t) (((in * IC) * IH + iih) * IW + (win0 + lo)) * sizeof(float);   \
+                    while (!dma_queue_push(dmaq, dma_make_data(vdst, vsrc),                                                    \
+                                           (size_t) KH * Wb * sizeof(float), (size_t) IH * IW * sizeof(float),                 \
+                                           cpw * sizeof(float), IC)) {                                                         \
+                        dma_queue_pop(dmaq);                                                                                   \
+                    }                                                                                                          \
+                }                                                                                                              \
+                dma_queue_flush(dmaq);                                                                                         \
+            }                                                                                                                  \
+            if (bi + 1 < total_blocks) {                                                                                       \
+                /* prefetch: stage block bi+1 into the other slot; overlaps with this block's compute below */                 \
+                const uint32_t nbuf  = 1u - buf;                                                                               \
+                float *        nsrcb = srcb2[nbuf];                                                                            \
+                const uint32_t nr    = row_start + (bi + 1) / nbpr;                                                            \
+                const uint32_t nin   = nr / OH;                                                                                \
+                const uint32_t nioh  = nr % OH;                                                                                \
+                const uint32_t nc0   = ((bi + 1) % nbpr) * owb;                                                                \
+                const int32_t  nwin0 = (int32_t) nc0 * s0 - p0;                                                                \
+                for (uint32_t ikh = 0; ikh < KH; ikh++) {                                                                      \
+                    const int32_t iih = (int32_t) nioh * s1 + (int32_t) ikh * d1 - p1;                                         \
+                    if (iih < 0 || iih >= (int32_t) IH)                                                                        \
+                        continue;                                                                                              \
+                    const int32_t lo = nwin0 < 0 ? -nwin0 : 0;                                                                 \
+                    int32_t       hi = (int32_t) IW - nwin0;                                                                   \
+                    if (hi > (int32_t) Wb)                                                                                     \
+                        hi = (int32_t) Wb;                                                                                     \
+                    if (hi <= lo)                                                                                              \
+                        continue;                                                                                              \
+                    const uint32_t cpw  = (uint32_t) (hi - lo);                                                                \
+                    float *        vdst = nsrcb + (size_t) ikh * Wb + lo;                                                      \
+                    const dma_addr_t vsrc = src_data + (size_t) (((nin * IC) * IH + iih) * IW + (nwin0 + lo)) * sizeof(float); \
+                    while (!dma_queue_push(dmaq, dma_make_data(vdst, vsrc),                                                    \
+                                           (size_t) KH * Wb * sizeof(float), (size_t) IH * IW * sizeof(float),                 \
+                                           cpw * sizeof(float), IC)) {                                                         \
+                        dma_queue_pop(dmaq);                                                                                   \
+                    }                                                                                                          \
+                }                                                                                                              \
+            }                                                                                                                  \
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r);                                                              \
+            for (uint32_t j = 0; j < nb; j++) {                                                                                \
+                const uint32_t iow       = c0 + j;                                                                             \
+                DST_CTYPE *    dst_patch = dstb + (uint64_t) j * patch_stride;                                                 \
+                const int32_t  iiw0      = (int32_t) iow * s0 - p0;                                                            \
+                for (uint32_t ikh = 0; ikh < KH; ikh++) {                                                                      \
+                    const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1;                                          \
+                    const int     okh = (iih >= 0 && iih < (int32_t) IH);                                                      \
+                    for (uint32_t iic = 0; iic < IC; iic++) {                                                                  \
+                        DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW;                                          \
+                        if (!okh) {                                                                                            \
+                            SPLAT_FN(out_run, 0.0f, KW);                                                                       \
+                            continue;                                                                                          \
+                        }                                                                                                      \
+                        const float * vrow = srcb + ((uint64_t) (iic * KH + ikh)) * Wb; /* col win0 at idx 0*/                 \
+                        if (d0 == 1) {                                                                                         \
+                            /* contiguous run within the staged window: [lo,hi) in-bounds, tails zero pad */                   \
+                            const int32_t lo = iiw0 < 0 ? -iiw0 : 0;                                                           \
+                            int32_t       hi = (int32_t) IW - iiw0;                                                            \
+                            if (hi > (int32_t) KW) {                                                                           \
+                                hi = (int32_t) KW;                                                                             \
+                            }                                                                                                  \
+                            if (hi <= lo) {                                                                                    \
+                                SPLAT_FN(out_run, 0.0f, KW);                                                                   \
+                            } else {                                                                                           \
+                                if (lo > 0) {                                                                                  \
+                                    SPLAT_FN(out_run, 0.0f, (uint32_t) lo);                                                    \
+                                }                                                                                              \
+                                COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (vrow + (iiw0 + lo - win0)),             \
+                                        (uint32_t) (hi - lo));                                                                 \
+                                if (hi < (int32_t) KW) {                                                                       \
+                                    SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi));                                        \
+                                }                                                                                              \
+                            }                                                                                                  \
+                            continue;                                                                                          \
+                        }                                                                                                      \
+                        for (uint32_t ikw = 0; ikw < KW; ikw++) {                                                              \
+                            const int32_t iiw = iiw0 + (int32_t) ikw * d0;                                                     \
+                            out_run[ikw] =                                                                                     \
+                                (iiw < 0 || iiw >= (int32_t) IW) ? (DST_CTYPE) 0.0f : (DST_CTYPE) vrow[iiw - win0];            \
+                        }                                                                                                      \
+                    }                                                                                                          \
+                }                                                                                                              \
+            }                                                                                                                  \
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r);                                                               \
+            const dma_addr_t ddr = dst_data + (size_t) ((in * OH + ioh) * OW + c0) * patch_stride * (DST_ELEM);                \
+            dma_queue_push(dmaq, dma_make_data(ddr, dstb),                                                                     \
+                           nb * patch_stride * (DST_ELEM), nb * patch_stride * (DST_ELEM),                                     \
+                           nb * patch_stride * (DST_ELEM), 1);                                                                 \
+            dma_queue_flush(dmaq);                                                                                             \
+        }                                                                                                                      \
+    }
+IM2COL_BLOCKED_DMA_BODY(im2col_blocked_dma_thread,     __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "blk-dma-f16")
+IM2COL_BLOCKED_DMA_BODY(im2col_blocked_dma_f32_thread, float,  hvx_copy_f32_uu,     hvx_splat_f32_u, sizeof(float),  "blk-dma-f32")
+
+// Exact-tiling patch-embed DMA fast path (s0==KW, p0=0, d0=1; and 2D s1==KH,
+// p1=0, d1=1). Intentionally reads no stride/pad/dilation params so the inner
+// copy stays tight and fully hoisted - do NOT graft the general gather in here.
+#define IM2COL_PATCHEMBED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG)                           \
+    static void FNAME(unsigned int nth, unsigned int ith, void * data) {                                         \
+        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 src1 = octx->src[1];                                                  \
+        const struct htp_tensor * restrict dst  = octx->dst;                                                     \
+        const int32_t  is_2D                    = octx->op_params[6] == 1;                                       \
+        const uint32_t N                        = is_2D ? src1->ne[3] : src1->ne[2];                             \
+        const uint32_t IC                       = is_2D ? src1->ne[2] : src1->ne[1];                             \
+        const uint32_t IH                       = is_2D ? src1->ne[1] : 1;                                       \
+        const uint32_t IW                       = src1->ne[0];                                                   \
+        const uint32_t KH                       = is_2D ? octx->src[0]->ne[1] : 1;                               \
+        const uint32_t KW                       = octx->src[0]->ne[0];                                           \
+        const uint32_t OH                       = is_2D ? dst->ne[2] : 1;                                        \
+        const uint32_t OW                       = dst->ne[1];                                                    \
+        const uint32_t patch_stride             = IC * KH * KW;                                                  \
+        const dma_addr_t src_data       = src1->data;                                                            \
+        const dma_addr_t dst_data       = dst->data;                                                             \
+        dma_queue *    dma_q            = octx->ctx->dma[ith];                                                   \
+        uint8_t *      src_base         = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread;                \
+        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 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                = 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++) {                                                         \
+            const uint32_t in  = r / OH;                                                                         \
+            const uint32_t ioh = r % OH;                                                                         \
+            for (uint32_t ikh = 0; ikh < KH; ikh++) {                                                            \
+                int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh;                                      \
+                int     ok  = (iih >= 0 && iih < (int32_t) IH);                                                  \
+                for (uint32_t iic = 0; iic < IC; iic++) {                                                        \
+                    float *       vdst = srcb + (size_t) (iic * KH + ikh) * IW;                                  \
+                    const dma_addr_t vsrc = ok                                                                   \
+                        ? (src_data + (size_t) ((in * IC + iic) * IH + iih) * IW * sizeof(float))                \
+                        : src_data;                                                                              \
+                    dma_queue_push(dma_q, dma_make_data(vdst, vsrc),                                             \
+                                   IW * sizeof(float), IW * sizeof(float), IW * sizeof(float), ok ? 1 : 0);      \
+                }                                                                                                \
+            }                                                                                                    \
+            for (uint32_t i = 0; i < IC * KH; i++)                                                               \
+                dma_queue_pop(dma_q);                                                                            \
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r);                                                \
+            for (uint32_t iow = 0; iow < OW; iow++) {                                                            \
+                DST_CTYPE * dst_patch = dstb + (uint64_t) iow * patch_stride;                                    \
+                for (uint32_t ikh = 0; ikh < KH; ikh++) {                                                        \
+                    int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh;                                  \
+                    for (uint32_t iic = 0; iic < IC; iic++) {                                                    \
+                        DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW;                            \
+                        if (iih < 0 || iih >= (int32_t) IH) {                                                    \
+                            SPLAT_FN(out_run, 0.0f, KW);                                                         \
+                            continue;                                                                            \
+                        }                                                                                        \
+                        const float * src_run = srcb + ((uint64_t) (iic * KH + ikh)) * IW + (uint64_t) iow * KW; \
+                        COPY_FN((uint8_t *) out_run, (const uint8_t *) src_run, KW);                             \
+                    }                                                                                            \
+                }                                                                                                \
+            }                                                                                                    \
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r);                                                 \
+            const dma_addr_t ddr_row = dst_data + (size_t) (in * OH + ioh) * OW * patch_stride * (DST_ELEM);     \
+            dma_queue_push(dma_q, dma_make_data(ddr_row, dstb),                                                  \
+                           OW * patch_stride * (DST_ELEM), OW * patch_stride * (DST_ELEM),                       \
+                           OW * patch_stride * (DST_ELEM), 1);                                                   \
+            dma_queue_flush(dma_q);                                                                              \
+        }                                                                                                        \
     }
 
 IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_thread,     __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "pe-dma-f16")
@@ -209,21 +390,30 @@ static bool im2col_use_patchembed_dma(const struct htp_ops_context * octx) {
     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 int     is_2D = octx->op_params[6] == 1;
-    if (!is_2D) {
-        return false;
-    }
     if (octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_F32) {
         return false;
     }
-    const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];
-    if (s0 != (int32_t) KW || s1 != (int32_t) KH) {
-        return false;  // non-overlapping
+    const uint32_t KH = is_2D ? octx->src[0]->ne[1] : 1;
+    const uint32_t KW = octx->src[0]->ne[0];
+    if (s0 != (int32_t) KW) {
+        return false;  // non-overlapping (width)
     }
-    if (p0 != 0 || p1 != 0) {
-        return false;  // no padding
+    if (p0 != 0) {
+        return false;  // no padding (width)
     }
-    if (d0 != 1 || d1 != 1) {
-        return false;  // no dilation
+    if (d0 != 1) {
+        return false;  // no dilation (width)
+    }
+    if (is_2D) {
+        if (s1 != (int32_t) KH) {
+            return false;  // non-overlapping (height)
+        }
+        if (p1 != 0) {
+            return false;  // no padding (height)
+        }
+        if (d1 != 1) {
+            return false;  // no dilation (height)
+        }
     }
     return true;
 }
@@ -233,8 +423,11 @@ static bool im2col_use_patchembed_dma(const struct htp_ops_context * octx) {
 static bool im2col_patchembed_dma_fits(struct htp_ops_context *    octx,
                                        struct htp_im2col_context * ictx,
                                        uint32_t                    n_threads) {
-    const uint32_t IC = octx->src[1]->ne[2], IW = octx->src[1]->ne[0];
-    const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];
+    const int32_t  is_2D = octx->op_params[6] == 1;
+    const uint32_t IC = is_2D ? octx->src[1]->ne[2] : octx->src[1]->ne[1];
+    const uint32_t IW = octx->src[1]->ne[0];
+    const uint32_t KH = is_2D ? octx->src[0]->ne[1] : 1;
+    const uint32_t KW = octx->src[0]->ne[0];
     const uint32_t OW           = octx->dst->ne[1];
     const uint32_t patch_stride = IC * KH * KW;
 
@@ -257,6 +450,45 @@ static bool im2col_patchembed_dma_fits(struct htp_ops_context *    octx,
     return true;
 }
 
+// Sizes a per-thread 2x(src,dst) VTCM ping-pong for the blocked general kernel.
+// Stages Wb=(owb-1)*s0+(KW-1)*d0+1 source cols per (iic,ikh) row and owb patches
+// of dst. Picks the largest owb that fits; returns false if even owb=1 does not.
+static bool im2col_blocked_dma_fits(struct htp_ops_context *    octx,
+                                    struct htp_im2col_context * ictx,
+                                    uint32_t                    n_threads) {
+    const int32_t  is_2D = octx->op_params[6] == 1;
+    const int32_t  s0 = octx->op_params[0];
+    const int32_t  d0 = octx->op_params[4];
+    const uint32_t IC = is_2D ? octx->src[1]->ne[2] : octx->src[1]->ne[1];
+    const uint32_t KH = is_2D ? octx->src[0]->ne[1] : 1;
+    const uint32_t KW = octx->src[0]->ne[0];
+    const uint32_t OW = octx->dst->ne[1];
+    const uint32_t patch_stride = IC * KH * KW;
+    const uint32_t dst_elem = (octx->dst->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float);
+
+    for (uint32_t owb = (OW < 256 ? OW : 256); owb >= 1; owb--) {
+        const uint32_t Wb = (owb - 1) * (uint32_t) s0 + (KW - 1) * (uint32_t) d0 + 1;
+        const uint32_t src_row_bytes = hex_round_up(IC * KH * Wb * sizeof(float), 256);
+        const uint32_t dst_row_bytes = hex_round_up(owb * patch_stride * dst_elem, 256);
+        struct htp_im2col_vtcm_layout L;
+        htp_im2col_vtcm_layout_build(&L, src_row_bytes, dst_row_bytes, n_threads);
+        if (L.total_bytes <= octx->ctx->vtcm_size) {
+            uint8_t * const base = octx->ctx->vtcm_base;
+            ictx->pe_owb                 = owb;
+            ictx->pe_wb                  = Wb;
+            ictx->pe_src_row_bytes       = src_row_bytes;
+            ictx->pe_dst_row_bytes       = dst_row_bytes;
+            ictx->pe_vtcm_src            = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src);
+            ictx->pe_vtcm_dst            = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
+            ictx->pe_src_size_per_thread = (uint32_t) L.src_bytes_per_thread;
+            ictx->pe_dst_size_per_thread = (uint32_t) L.dst_bytes_per_thread;
+            return true;
+        }
+        if (owb == 1) break;  // avoid unsigned underflow
+    }
+    return false;
+}
+
 int op_im2col(struct htp_ops_context * octx) {
     const struct htp_tensor * src1 = octx->src[1];
     const struct htp_tensor * dst  = octx->dst;
@@ -266,12 +498,9 @@ int op_im2col(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    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 int32_t  is_2D         = octx->op_params[6] == 1;
+    const uint32_t N             = is_2D ? src1->ne[3] : src1->ne[2];
+    const uint32_t OH            = is_2D ? dst->ne[2] : 1;
     const uint32_t OW            = dst->ne[1];
     const uint32_t total_patches = N * OH * OW;
     const uint32_t total_rows    = N * OH;
@@ -280,8 +509,11 @@ int op_im2col(struct htp_ops_context * octx) {
     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);
+        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;
     }
@@ -290,8 +522,11 @@ int op_im2col(struct htp_ops_context * octx) {
     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);
+        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;
     }
@@ -309,27 +544,39 @@ int op_im2col(struct htp_ops_context * octx) {
     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) && nrows > 0) {
+    // everything else (padding/dilation/stride edges) -> blocked-staging DMA
+    // kernel; if neither fits VTCM -> pure-DDR kernel.
+    if (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);
-            } else {
-                work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_f32_thread, &ictx, pth);
+        if (pth > 0) {
+            ictx.pe_row_base = row_base;
+            ictx.pe_nrows    = nrows;
+            const bool exact = im2col_use_patchembed_dma(octx);
+            if (exact && im2col_patchembed_dma_fits(octx, &ictx, pth)) {
+                ictx.pe_rows_per_thread = (nrows + pth - 1) / pth;
+                work_queue_run(octx->ctx->work_queue,
+                    dst->type == HTP_TYPE_F16 ? im2col_patchembed_dma_thread
+                                              : im2col_patchembed_dma_f32_thread, &ictx, pth);
+                return HTP_STATUS_OK;
+            }
+            if (!exact && im2col_blocked_dma_fits(octx, &ictx, pth)) {
+                ictx.pe_rows_per_thread = (nrows + pth - 1) / pth;
+                work_queue_run(octx->ctx->work_queue,
+                    dst->type == HTP_TYPE_F16 ? im2col_blocked_dma_thread
+                                              : im2col_blocked_dma_f32_thread, &ictx, pth);
+                return HTP_STATUS_OK;
             }
-            return HTP_STATUS_OK;
         }
-        // else: doesn't fit -> fall through to the pure-DDR kernel below.
     }
-
+    // Fall through to pure-DDR.
     if (npatches == 0) {
         return HTP_STATUS_OK;
     }
 
+    if (htp_tensor_is_extended(src1) || htp_tensor_is_extended(dst)) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
     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

@@ -21,7 +21,7 @@
 #include <stdatomic.h>
 
 #include "hex-utils.h"
-#include "hex-dma.h"
+#include "dma-queue.h"
 #include "hmx-queue.h"
 
 #define GGML_COMMON_DECL_C
@@ -49,33 +49,74 @@ struct htp_handle {
     struct htp_context * ctx;
 };
 
-static inline void * htp_mmap(uint32_t fd, uint32_t size) {
+static inline uint64_t htp_mmap(uint32_t fd, uint64_t size, uint32_t flags) {
+#if __HVX_ARCH__ > 79
+    if (flags & HTP_BUF_EXTENDED) {
+        HAP_mem_req_payload_t payload;
+        memset(&payload, 0, sizeof(payload));
+        payload.request_id = HAP_MEM_MAP;
+        payload.mmap.len   = size;
+        payload.mmap.prot  = HAP_MEM_CACHE_NON_SHARED | HAP_PROT_READ;
+        payload.mmap.flags = HAP_MEM_FLAGS_EXTENDED_MAP;
+        payload.mmap.fd    = fd;
+
+        if (HAP_mem_request(&payload) != 0) {
+            FARF(ERROR, "extended mmap failed : fd %u size %llu", fd, (unsigned long long) size);
+            return 0;
+        }
+
+        return payload.mmap.dsp_va;
+    }
+#else
+    if (flags & HTP_BUF_EXTENDED) {
+        FARF(ERROR, "extended mmap is unsupported on v%d", __HVX_ARCH__);
+        return 0;
+    }
+#endif
+
+    if (size > UINT32_MAX) {
+        FARF(ERROR, "mmap failed : size %llu exceeds 32-bit limit", (unsigned long long) size);
+        return 0;
+    }
+
     void * va = (void *)-1;
     for (int retry = 0; retry < 2; retry++) {
 #if __HVX_ARCH__ > 73
-        va = HAP_mmap2(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0);
+        va = HAP_mmap2(NULL, (size_t) size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0);
 #else
         if (size > HTP_MMAP_MAX_VMEM) {
-            FARF(ERROR, "mmap failed : size %u exceeds 2GB limit for HAP_mmap", (uint32_t) size);
+            FARF(ERROR, "mmap failed : size %llu exceeds 2GB limit for HAP_mmap", (unsigned long long) size);
             abort();
         }
-        va = HAP_mmap(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0);
+        va = HAP_mmap(NULL, (int) size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0);
 #endif
         if (va != (void *)-1 && va != NULL) {
-            return va;
+            return (uint64_t) (uintptr_t) va;
         }
         if (retry == 0) {
-            FARF(HIGH, "mmap failed first try (va %p fd %u size %u), retrying...", va, fd, size);
+            FARF(HIGH, "mmap failed first try (va %p fd %u size %llu), retrying...", va, fd, (unsigned long long) size);
         }
     }
-    return NULL;
+    return 0;
 }
 
-static inline void htp_munmap(void * va, uint32_t size) {
+static inline void htp_munmap(uint64_t va, uint64_t size, uint32_t flags) {
+#if __HVX_ARCH__ > 79
+    if (flags & HTP_BUF_EXTENDED) {
+        HAP_mem_req_payload_t payload;
+        memset(&payload, 0, sizeof(payload));
+        payload.request_id   = HAP_MEM_UNMAP;
+        payload.munmap.dsp_va = va;
+        payload.munmap.len    = size;
+        HAP_mem_request(&payload);
+        return;
+    }
+#endif
+
 #if __HVX_ARCH__ > 73
-    HAP_munmap2(va, size);
+    HAP_munmap2((void *) (uintptr_t) va, (size_t) size);
 #else
-    HAP_munmap(va, size);
+    HAP_munmap((void *) (uintptr_t) va, (int) size);
 #endif
 }
 
@@ -160,10 +201,11 @@ AEEResult htp_iface_close(remote_handle64 handle) {
         // release the mmaps (if any)
         for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
             if (ctx->mmap[i].size) {
-                htp_munmap((void *) ctx->mmap[i].base, ctx->mmap[i].size);
+                htp_munmap(ctx->mmap[i].base, ctx->mmap[i].size, ctx->mmap[i].flags);
                 ctx->mmap[i].size = 0;
-                ctx->mmap[i].base = NULL;
+                ctx->mmap[i].base = 0;
                 ctx->mmap[i].fd   = -1;
+                ctx->mmap[i].flags = 0;
             }
         }
 
@@ -184,7 +226,7 @@ AEEResult htp_iface_close(remote_handle64 handle) {
     return AEE_SUCCESS;
 }
 
-AEEResult htp_iface_mmap(remote_handle64 handle, uint32_t fd, uint32_t size) {
+AEEResult htp_iface_mmap(remote_handle64 handle, uint32_t fd, uint64_t size) {
     struct htp_handle * h = (struct htp_handle *) handle;
     if (!h || !h->ctx) {
         return AEE_EBADPARM;
@@ -203,16 +245,17 @@ AEEResult htp_iface_mmap(remote_handle64 handle, uint32_t fd, uint32_t size) {
     for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
         struct htp_mmap *m = &ctx->mmap[i];
         if (!m->size) {
-            FARF(HIGH, "mmap : fd %u size %u", fd, size);
-            void *va = htp_mmap(fd, size);
-            if (va == NULL) {
-                FARF(ERROR, "mmap failed : fd %u size %u", fd, (uint32_t) size);
+            FARF(HIGH, "mmap : fd %u size %llu", fd, (unsigned long long) size);
+            uint64_t va = htp_mmap(fd, size, 0);
+            if (va == 0) {
+                FARF(ERROR, "mmap failed : fd %u size %llu", fd, (unsigned long long) size);
                 return AEE_EFAILED;
             }
 
-            m->base   = (uint64_t) va;
+            m->base   = va;
             m->fd     = fd;
             m->size   = size;
+            m->flags  = 0;
 
             return AEE_SUCCESS;
         }
@@ -231,11 +274,12 @@ AEEResult htp_iface_munmap(remote_handle64 handle, uint32 fd) {
     for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
         struct htp_mmap *m = &ctx->mmap[i];
         if (fd < 0 || m->fd == fd) {
-            FARF(HIGH, "unmmap : base %p fd %u size %u", (void*) m->base, m->fd, (uint32_t) m->size);
-            htp_munmap((void *) m->base, m->size);
+            FARF(HIGH, "unmmap : base 0x%llx fd %u size %llu", (unsigned long long) m->base, m->fd, (unsigned long long) m->size);
+            htp_munmap(m->base, m->size, m->flags);
             m->size   = 0;
             m->base   = NULL;
             m->fd     = -1;
+            m->flags  = 0;
         }
     }
 
@@ -394,8 +438,6 @@ AEEResult htp_iface_start(remote_handle64 handle, uint32_t sess_id, uint64_t dsp
     for (uint32_t i = 0; i < n_hvx; i++) {
         size_dma  = hex_align_up(size_dma, dma_queue_alignof());
         size_dma += dma_queue_sizeof(256);
-        size_dma  = hex_align_up(size_dma, dma_queue_alignof());
-        size_dma += dma_queue_alias_sizeof();
     }
     offset = offset_dma + size_dma;
 
@@ -538,16 +580,11 @@ AEEResult htp_iface_start(remote_handle64 handle, uint32_t sess_id, uint64_t dsp
     // Initialize DMA queues
     uint8_t * dma_ptr_curr = (uint8_t *) ((uintptr_t) block + offset_dma);
     size_t size_dma_q = dma_queue_sizeof(256);
-    size_t size_dma_alias = dma_queue_alias_sizeof();
 
     for (int i = 0; i < ctx->n_threads; i++) {
         dma_ptr_curr = (uint8_t *) hex_align_up((uintptr_t) dma_ptr_curr, dma_queue_alignof());
-        ctx->dma_cached[i] = dma_queue_init(dma_ptr_curr, 256, (uintptr_t) ctx->vtcm_base, ctx->vtcm_size, &ctx->trace[i]);
+        ctx->dma[i] = dma_queue_init(dma_ptr_curr, 256, &ctx->trace[i]);
         dma_ptr_curr += size_dma_q;
-
-        dma_ptr_curr = (uint8_t *) hex_align_up((uintptr_t) dma_ptr_curr, dma_queue_alignof());
-        ctx->dma[i] = dma_queue_alias_init(dma_ptr_curr, ctx->dma_cached[i], 1);
-        dma_ptr_curr += size_dma_alias;
     }
 
     ctx->ddr_spad_size = 512 * 1024; // 512 KB
@@ -608,8 +645,7 @@ AEEResult htp_iface_stop(remote_handle64 handle) {
     work_queue_free(ctx->work_queue);
 
     for (int i = 0; i < ctx->n_threads; i++) {
-        dma_queue_alias_free(ctx->dma[i]);
-        dma_queue_free(ctx->dma_cached[i]);
+        dma_queue_free(ctx->dma[i]);
     }
 
     if (ctx->hmx_queue) {
@@ -828,6 +864,7 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_GLU_SWIGLU_OAI:
         case HTP_OP_GLU_SWIGLU_CLAMP:
         case HTP_OP_GLU_GEGLU:
+        case HTP_OP_GLU_GEGLU_QUICK:
             return op_activations(octx);
 
         case HTP_OP_SOFTMAX:
@@ -858,6 +895,9 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_ARGSORT:
             return op_argsort(octx);
 
+        case HTP_OP_TOP_K:
+            return op_top_k(octx);
+
         case HTP_OP_SSM_CONV:
             return op_ssm_conv(octx);
 
@@ -879,6 +919,9 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_IM2COL:
             return op_im2col(octx);
 
+        case HTP_OP_ROLL:
+            return op_roll(octx);
+
         case HTP_OP_CONCAT:
             return op_concat(octx);
 
@@ -901,7 +944,7 @@ static inline bool reuse_buf(struct htp_context *ctx, uint32_t *m_reuse, struct
 
     for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
         struct htp_mmap *m = ctx->mmap + i;
-        if (m->size && m->fd == b->fd) {
+        if (m->size && m->fd == b->fd && m->flags == b->flags) {
             b->base   = m->base;
             *m_reuse |= (1 << i);
             return true;
@@ -913,11 +956,12 @@ static inline bool reuse_buf(struct htp_context *ctx, uint32_t *m_reuse, struct
 
 static inline void drop_mmap(struct htp_context *ctx, struct htp_mmap *m) {
     if (m->size) {
-        FARF(ALWAYS, "unmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size);
-        htp_munmap((void *) m->base, m->size);
+        FARF(ALWAYS, "unmap : fd %u base 0x%llx size %llu", m->fd, (unsigned long long) m->base, (unsigned long long) m->size);
+        htp_munmap(m->base, m->size, m->flags);
         m->size = 0;
         m->base = 0;
         m->fd   = -1;
+        m->flags = 0;
     }
 }
 
@@ -928,17 +972,18 @@ static inline bool mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) {
     for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) {
         struct htp_mmap *m = &ctx->mmap[i];
         if (!m->size) {
-            void *va = htp_mmap(b->fd, b->size);
-            if (va == NULL) {
-                FARF(HIGH, "mmap failed (will attempt defrag) : fd %u size %u", b->fd, (uint32_t) b->size);
+            uint64_t va = htp_mmap(b->fd, b->size, b->flags);
+            if (va == 0) {
+                FARF(HIGH, "mmap failed (will attempt defrag) : fd %u size %llu", b->fd, (unsigned long long) b->size);
                 return false;
             }
 
-            m->base   = b->base = (uint64_t) va;
+            m->base   = b->base = va;
             m->fd     = b->fd;
             m->size   = b->size;
+            m->flags  = b->flags;
 
-            FARF(ALWAYS, "mmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size);
+            FARF(ALWAYS, "mmap : fd %u base 0x%llx size %llu flags 0x%x", m->fd, (unsigned long long) m->base, (unsigned long long) m->size, m->flags);
             return true;
         }
     }
@@ -957,14 +1002,22 @@ static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uin
     // See what we can reuse
     for (uint32_t i=0; i < n_bufs; i++) {
         struct htp_buf_desc *b = bufs + i;
-        if (reuse_buf(ctx, &m_reuse, b)) { b_reuse++; } else { e_vmem += b->size; }
-        FARF(HIGH, "prep-buf #%u : pass0 fd %u base %p size %u flags 0x%x", i, b->fd, (void*) b->base, (uint32_t) b->size, b->flags);
+        if (reuse_buf(ctx, &m_reuse, b)) {
+            b_reuse++;
+        } else if (!(b->flags & HTP_BUF_EXTENDED)) {
+            e_vmem += b->size;
+        }
+        FARF(HIGH, "prep-buf #%u : pass0 fd %u base 0x%llx size %llu flags 0x%x", i, b->fd, (unsigned long long) b->base, (unsigned long long) b->size, b->flags);
     }
 
     if (b_reuse == n_bufs) return; // all bufs reuse existing mappings
 
     // See how much vmem we have mmaped right now
-    for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) { m_vmem += ctx->mmap[i].size; }
+    for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
+        if (!(ctx->mmap[i].flags & HTP_BUF_EXTENDED)) {
+            m_vmem += ctx->mmap[i].size;
+        }
+    }
 
     FARF(HIGH, "prep-bufs : pass1 mmap-vmem %zu extra-vmem %zu max-vmem %zu : n-bufs %u b-reuse %u",
             (size_t) m_vmem, (size_t) e_vmem, (size_t) ctx->max_vmem, n_bufs, b_reuse);
@@ -973,7 +1026,9 @@ static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uin
         // Drop unused mappings
         for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) {
             bool used = m_reuse & (1<<i);
-            if (!used) { drop_mmap(ctx, ctx->mmap + i); }
+            if (!used && !(ctx->mmap[i].flags & HTP_BUF_EXTENDED)) {
+                drop_mmap(ctx, ctx->mmap + i);
+            }
         }
     }
 
@@ -985,35 +1040,40 @@ static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uin
             mmap_ok = false;
             break;
         }
-        FARF(HIGH, "prep-buf #%u : pass1 fd %u base %p size %u flags 0x%x", i, b->fd, (void*) b->base, (uint32_t) b->size, b->flags);
+        FARF(HIGH, "prep-buf #%u : pass1 fd %u base 0x%llx size %llu flags 0x%x", i, b->fd, (unsigned long long) b->base, (unsigned long long) b->size, b->flags);
     }
 
     if (!mmap_ok) {
-        // Attempt clean defragmentation: drop all mappings and remap (pass 2)
-        FARF(HIGH, "prep-bufs : dropping all mappings to defragment address space");
-        for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) { drop_mmap(ctx, ctx->mmap + i); }
+        // Attempt defragmentation: drop 32-bit mappings and remap (pass 2)
+        FARF(HIGH, "prep-bufs : dropping 32-bit mappings to defragment address space");
+        for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) {
+            if (!(ctx->mmap[i].flags & HTP_BUF_EXTENDED)) {
+                drop_mmap(ctx, ctx->mmap + i);
+            }
+        }
 
         for (uint32_t i=0; i < n_bufs; i++) {
             struct htp_buf_desc *b = bufs + i;
-            b->base = 0;
+            if (!(b->flags & HTP_BUF_EXTENDED)) {
+                b->base = 0;
+            }
             if (!mmap_buf(ctx, b)) {
-                FARF(ERROR, "prep-bufs : mmap failed after defragmentation (fd %u size %u)", b->fd, (uint32_t) b->size);
+                FARF(ERROR, "prep-bufs : mmap failed after defragmentation (fd %u size %llu)", b->fd, (unsigned long long) b->size);
                 abort();
             }
-            FARF(HIGH, "prep-buf #%u : pass2 fd %u base %p size %u flags 0x%x", i, b->fd, (void*) b->base, (uint32_t) b->size, b->flags);
+            FARF(HIGH, "prep-buf #%u : pass2 fd %u base 0x%llx size %llu flags 0x%x", i, b->fd, (unsigned long long) b->base, (unsigned long long) b->size, b->flags);
         }
     }
 }
 
 static void prep_tensor(struct htp_context *ctx, struct htp_buf_desc *bufs, struct htp_tensor *tens, uint32_t idx, struct htp_tensor *t) {
-    uint32_t offset = t->data;
-    uint32_t size   = t->size;
+    uint64_t offset = t->data;
     uint32_t bi     = t->bi;
 
-    t->data  = (uint32_t) (bufs[bi].base + offset);  // update data to the actual pointer
+    t->data  = bufs[bi].base + offset;  // update data to the actual pointer
 
-    FARF(HIGH, "prep-tensor #%u: bi %u offset %u size %u data %p : %u:%u:%u:%u", idx, t->bi, offset, t->size, (void*) t->data,
-        t->ne[0], t->ne[1], t->ne[3], t->ne[3]);
+    FARF(HIGH, "prep-tensor #%u: bi %u offset %llu size %u data 0x%llx : %u:%u:%u:%u", idx, t->bi, (unsigned long long) offset, t->size, (unsigned long long) t->data,
+        t->ne[0], t->ne[1], t->ne[2], t->ne[3]);
 }
 
 static void prep_tensors(struct htp_context *ctx, struct htp_buf_desc *bufs, struct htp_tensor *tens, uint32_t n_tens) {
@@ -1043,15 +1103,13 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_buf_desc * bufs
         uint16_t src_idx = op->src[i];
         if (src_idx == 0xffff) {
             octx->src[i]     = NULL;
-            octx->src_dma[i] = NULL;
             continue;
         }
 
         struct htp_tensor *src = tens + src_idx;
         octx->src[i]     = src;
-        octx->src_dma[i] = octx->ctx->dma; // FIXME: ? octx->ctx->dma_cached : octx->ctx->dma;
 
-        FARF(HIGH, "prep-src #%u: data %p size %u : %u:%u:%u:%u", op->src[i], (void*) src->data, src->size,
+        FARF(HIGH, "prep-src #%u: data 0x%llx size %u : %u:%u:%u:%u", op->src[i], (unsigned long long) src->data, src->size,
             src->ne[0], src->ne[1], src->ne[2], src->ne[3]);
     }
 
@@ -1062,14 +1120,12 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_buf_desc * bufs
         uint16_t dst_idx = op->dst[i];
         if (dst_idx == 0xffff) {
             octx->dsts[i]    = NULL;
-            octx->dst_dma[i] = NULL;
             continue;
         }
         struct htp_tensor *dst = tens + dst_idx;
         octx->dsts[i]    = dst;
-        octx->dst_dma[i] = octx->ctx->dma; // FIXME: ? octx->ctx->dma_cached : octx->ctx->dma;
 
-        FARF(HIGH, "prep-dst[%u] #%u: data %p size %u : %u:%u:%u:%u", i, dst_idx, (void*) dst->data, dst->size,
+        FARF(HIGH, "prep-dst[%u] #%u: data 0x%llx size %u : %u:%u:%u:%u", i, dst_idx, (unsigned long long) dst->data, dst->size,
             dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
     }
 

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

@@ -62,17 +62,11 @@ enum htp_mm_kernel_type {
 
     // HVX floating-point paths
     HTP_MM_KERNEL_HVX_F16_F16_VTCM,
-    HTP_MM_KERNEL_HVX_F16_F16_DDR,
-    HTP_MM_KERNEL_HVX_F16_F32_DDR,
-
     HTP_MM_KERNEL_HVX_F32_F32_VTCM,
-    HTP_MM_KERNEL_HVX_F32_F32_DDR,
-    HTP_MM_KERNEL_HVX_F32_F16_DDR,
 
     // HVX quantized paths
     HTP_MM_KERNEL_HVX_QUANT_ROW,      // standard row-wise parallel quantization
     HTP_MM_KERNEL_HVX_QUANT_BLOCK,    // parallel block-wise quantization
-    HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, // row-wise fallback flat quantization
 };
 
 // Op-specific struct for precomputed matmul params
@@ -101,7 +95,7 @@ struct htp_mm_kernel_params {
     struct fastdiv_values div_ne1;
     struct fastdiv_values div_r2;
     struct fastdiv_values div_r3;
-    struct fastdiv_values div_ne11;
+    struct fastdiv_values div_ne12;
     struct fastdiv_values div_n_act_threads;
     struct fastdiv_values div_ne00_padded;
 };
@@ -246,20 +240,6 @@ static inline size_t htp_mm_q8_1_tiled_row_size(uint32_t ne) {
     return nb_32 * HTP_MM_ACT_TILE_SIZE_Q8_1;
 }
 
-static inline size_t htp_mm_q8_0_flat_row_size(uint32_t ne) {
-    const uint32_t quants_size = hex_align_up(ne, 128);
-    const uint32_t num_scales = (ne + 31) / 32;
-    const uint32_t scales_size = hex_align_up(num_scales * 2, 128);
-    return quants_size + scales_size;
-}
-
-static inline size_t htp_mm_q8_1_flat_row_size(uint32_t ne) {
-    const uint32_t quants_size = hex_align_up(ne, 128);
-    const uint32_t num_scales = (ne + 31) / 32;
-    const uint32_t scales_size = hex_align_up(num_scales * 4, 128);
-    return quants_size + scales_size;
-}
-
 static inline size_t htp_mm_get_tiled_row_stride(int weight_type, uint32_t k) {
     uint32_t nb = (k + QK_Q4_0_TILED - 1) / QK_Q4_0_TILED;
     switch (weight_type) {
@@ -331,6 +311,7 @@ struct htp_mm_hmx_vtcm_layout {
     size_t off_dst[2];        // [1] is only used when pipelined
     size_t off_scratch[2];    // dequantization scratch pads
     size_t off_scales;        // HMX scales (256 bytes)
+    size_t off_src2;          // src2 bias in VTCM
 
     // Cached sizes of regions for HMX kernel use
     size_t weight_area_bytes;
@@ -339,6 +320,7 @@ struct htp_mm_hmx_vtcm_layout {
     size_t output_area_bytes;
     size_t scratch_bytes[2];
     size_t act_head_stride;
+    size_t src2_bytes;
 
     size_t total_bytes;
 };
@@ -372,7 +354,8 @@ static inline void htp_mm_hmx_vtcm_layout_build(
     bool use_dma_activation,
     bool pipeline,
     uint32_t act_threads,
-    uint32_t aligned_tile_size
+    uint32_t aligned_tile_size,
+    size_t src2_size
 ) {
     size_t off = 0;
 
@@ -390,6 +373,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
         size_t off_group_a = 0;
         VTCM_LAYOUT_ALLOC(off_group_a, off_act, activation_area_size);
         VTCM_LAYOUT_ALLOC(off_group_a, off_scales, HTP_MM_HMX_TILE_SIZE); // Padded to 2K for alignment and future persistent data
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src2, hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE), src2_size > 0);
 
         // Group B: Compute-only buffers (starts at off_group_a)
         size_t off_group_b = off_group_a;
@@ -418,6 +402,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
         L->scratch_bytes[0]  = scratch_area_size;
         L->scratch_bytes[1]  = scratch_area_size;
         L->act_head_stride   = act_head_stride;
+        L->src2_bytes        = src2_size;
 
         off = off_group_a + hex_smax(group_b_size, group_c_size);
     } else {
@@ -441,6 +426,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
         size_t off_group_a = 0;
         VTCM_LAYOUT_ALLOC(off_group_a, off_scales, HTP_MM_HMX_TILE_SIZE); // Padded to 2K for alignment and future persistent data
         VTCM_LAYOUT_ALLOC(off_group_a, off_act, act_area_size);
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src2, hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE), src2_size > 0);
 
         // Group B: Compute-only buffers (starts at off_group_a)
         size_t off_group_b = off_group_a;
@@ -468,6 +454,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
         L->scratch_bytes[0]  = scratch0_size;
         L->scratch_bytes[1]  = scratch1_size;
         L->act_head_stride   = 0;
+        L->src2_bytes        = src2_size;
 
         off = off_group_a + hex_smax(group_b_size, group_c_size);
     }
@@ -490,6 +477,7 @@ static inline void htp_mm_hvx_vtcm_layout_build(
     bool is_matmul_id,
     bool is_fused_nx
 ) {
+    (void)src1_row_size;
     size_t src0_sz = 0;
     size_t src1_sz = 0;
     size_t src2_sz = src2_row_size > 0 ? htp_mm_round_up(src2_row_size, 128) : 0;
@@ -517,12 +505,8 @@ static inline void htp_mm_hvx_vtcm_layout_build(
             weight_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128);
         }
 
-        size_t flat_act_row_size  = (wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q4_K) ? htp_mm_q8_1_flat_row_size(ne10)  : htp_mm_q8_0_flat_row_size(ne10);
         size_t tiled_act_row_size = (wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q4_K) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
-
-        size_t act_sz = (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT)
-            ? hex_round_up(flat_act_row_size  * src1_nrows, 128)
-            : hex_round_up(tiled_act_row_size * src1_nrows, 128);
+        size_t act_sz = hex_round_up(tiled_act_row_size * src1_nrows, 128);
 
         src0_sz = weight_sz_per_thread * n_threads; // shared single-weight prefetch buffer
         src1_sz = act_sz;                           // quantized activation buffer
@@ -547,6 +531,8 @@ static inline void htp_mm_hvx_vtcm_layout_build(
 
         src0_sz = src0_sz_per_thread * n_threads;
         dst_sz  = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads;
+        src2_sz = 0;
+        src3_sz = 0;
     } else {
         const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128);
         const size_t dst_nrows = (src1_nrows > 1) ? 0 : 1;
@@ -559,15 +545,6 @@ static inline void htp_mm_hvx_vtcm_layout_build(
                 dst_sz  = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
                 break;
             }
-            case HTP_MM_KERNEL_HVX_F16_F32_DDR:
-            case HTP_MM_KERNEL_HVX_F16_F16_DDR:
-            case HTP_MM_KERNEL_HVX_F32_F32_DDR:
-            case HTP_MM_KERNEL_HVX_F32_F16_DDR: {
-                src0_sz = htp_mm_round_up(n_prefetch * src0_row_size, 256) * n_threads;
-                src1_sz = htp_mm_round_up(n_prefetch * src1_row_size, 256) * n_threads;
-                dst_sz  = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
-                break;
-            }
             case HTP_MM_KERNEL_HVX_F32_F32_VTCM: {
                 size_t f32_src1_row_size = htp_mm_round_up(ne10 * 4, 128);
                 src1_sz = htp_mm_round_up(f32_src1_row_size * src1_nrows, 256);
@@ -598,28 +575,6 @@ static inline void htp_mm_hvx_vtcm_layout_build(
                 dst_sz = dst_size_per_thread * n_threads;
                 break;
             }
-            case HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT: {
-                size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q4_K) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
-
-                src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
-                src1_sz = htp_mm_round_up(q_src1_row_size * src1_nrows, 256);
-
-                src0_sz = src0_sz * n_threads;
-
-                if (is_repack) {
-                    uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
-                    uint32_t n_k_tiles = ne10 / 32;
-                    uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
-                    size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
-                    src0_sz = repacked_vtcm_size * n_threads;
-                }
-
-                size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
-                size_t dst_slice_per_thread = dst_nrows > 0 ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0;
-                size_t dst_size_per_thread = (dst_slice_per_thread > quant_scratch_size_per_thread) ? dst_slice_per_thread : quant_scratch_size_per_thread;
-                dst_sz = dst_size_per_thread * n_threads;
-                break;
-            }
             default:
                 break;
         }
@@ -640,19 +595,99 @@ static inline void htp_mm_hvx_vtcm_layout_build(
     L->total_bytes = off;
 }
 
+static inline bool htp_mm_hvx_solve_vtcm_params(
+    int kernel_type,
+    int wtype,
+    uint32_t ne10,
+    uint32_t src1_nrows,
+    uint32_t n_threads,
+    size_t dst_row_size,
+    size_t src0_row_size,
+    size_t src1_row_size,
+    size_t src2_row_size,
+    uint32_t n_prefetch,
+    size_t vtcm_budget,
+    struct htp_mm_hvx_vtcm_layout * L_out,
+    uint32_t * m_chunk_out
+) {
+    struct htp_mm_hvx_vtcm_layout L;
+    htp_mm_hvx_vtcm_layout_build(
+        &L, kernel_type, wtype, ne10, src1_nrows, n_threads,
+        dst_row_size, src0_row_size, src1_row_size, src2_row_size, n_prefetch, false, false
+    );
+
+    if (L.total_bytes <= vtcm_budget) {
+        *L_out = L;
+        *m_chunk_out = src1_nrows;
+        return true;
+    }
+
+    const size_t fixed_bytes = L.src0_bytes + L.src2_bytes + L.dst_bytes;
+    if (vtcm_budget <= fixed_bytes) {
+        return false;
+    }
+
+    const size_t avail_act = vtcm_budget - fixed_bytes;
+    size_t row_size = 0;
+    if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
+        row_size = (wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q4_K)
+                 ? htp_mm_q8_1_tiled_row_size(ne10)
+                 : htp_mm_q8_0_tiled_row_size(ne10);
+    } else if (kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM) {
+        row_size = hex_round_up(ne10 * 2, 128);
+    } else {
+        row_size = hex_round_up(ne10 * 4, 128);
+    }
+    if (row_size == 0) {
+        return false;
+    }
+
+    uint32_t m_chunk = (uint32_t) (avail_act / row_size);
+    if (m_chunk > 1) {
+        m_chunk &= ~1U;
+    }
+    if (m_chunk > src1_nrows) {
+        m_chunk = src1_nrows;
+    }
+    if (m_chunk < 1) {
+        return false;
+    }
+
+    htp_mm_hvx_vtcm_layout_build(
+        &L, kernel_type, wtype, ne10, m_chunk, n_threads,
+        dst_row_size, src0_row_size, src1_row_size, src2_row_size, n_prefetch, false, false
+    );
+
+    while (m_chunk > 2 && L.total_bytes > vtcm_budget) {
+        m_chunk -= 2;
+        htp_mm_hvx_vtcm_layout_build(
+            &L, kernel_type, wtype, ne10, m_chunk, n_threads,
+            dst_row_size, src0_row_size, src1_row_size, src2_row_size, n_prefetch, false, false
+        );
+    }
+
+    if (L.total_bytes <= vtcm_budget) {
+        *L_out = L;
+        *m_chunk_out = m_chunk;
+        return true;
+    }
+
+    return false;
+}
+
 static inline size_t htp_mm_hmx_get_2d_vtcm_size(
-    int wtype, uint32_t k, size_t mc, size_t nc, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size
+    int wtype, uint32_t k, size_t mc, size_t nc, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size, size_t src2_size
 ) {
     struct htp_mm_hmx_vtcm_layout L;
-    htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, wtype, k, mc, nc, 1, false, pipeline, act_threads, aligned_tile_size);
+    htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, wtype, k, mc, nc, 1, false, pipeline, act_threads, aligned_tile_size, src2_size);
     return L.total_bytes;
 }
 
 static inline size_t htp_mm_hmx_get_batched_vtcm_size(
-    int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool use_dma_activation, bool pipeline, uint32_t act_threads) {
+    int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool use_dma_activation, bool pipeline, uint32_t act_threads, size_t src2_size) {
     (void)pipeline;
     struct htp_mm_hmx_vtcm_layout L;
-    htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, use_dma_activation, false, act_threads, 0);
+    htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, use_dma_activation, false, act_threads, 0, src2_size);
     return L.total_bytes;
 }
 
@@ -665,6 +700,7 @@ static inline bool htp_mm_hmx_solve_batched_params(
     bool use_dma_activation,
     int n_threads,
     bool pipeline,
+    size_t src2_size,
     size_t vtcm_budget,
     size_t * m_chunk_out,
     size_t * n_chunk_out,
@@ -679,7 +715,7 @@ static inline bool htp_mm_hmx_solve_batched_params(
 
     int act_threads = n_threads;
     while (act_threads >= 1) {
-        size_t group_overhead = htp_mm_hmx_get_batched_overhead();
+        size_t group_overhead = htp_mm_hmx_get_batched_overhead() + (src2_size > 0 ? hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE) : 0);
         size_t group_size_per_n, group_size_per_m, group_size_per_mn;
         htp_mm_hmx_get_batched_chunk_costs(k, group_size, &group_size_per_n, &group_size_per_m, &group_size_per_mn);
 
@@ -690,7 +726,7 @@ static inline bool htp_mm_hmx_solve_batched_params(
         if (htp_mm_hmx_compute_chunks(vtcm_budget, group_overhead, group_size_per_n, group_size_per_m, group_size_per_mn, hex_align_up(ne11, 32), ne01_padded,
                                (size_t) ne01_padded * HTP_MM_HMX_COST_W_DEQUANT, (size_t) ne11 * HTP_MM_HMX_COST_A_CONVERT,
                                &m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) {
-            size_t exact_size = htp_mm_hmx_get_batched_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, group_size, use_dma_activation, pipeline, act_threads);
+            size_t exact_size = htp_mm_hmx_get_batched_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, group_size, use_dma_activation, pipeline, act_threads, src2_size);
             if (exact_size <= vtcm_budget) {
                 size_t mblocks = ((size_t) ne11 + m_chunk_candidate - 1) / m_chunk_candidate;
                 if (mblocks < best_mblocks || (mblocks == best_mblocks && act_threads > best_act_threads)) {
@@ -730,6 +766,7 @@ static inline bool htp_mm_hmx_solve_2d_params(
     bool pipeline,
     bool is_matmul_id,
     uint32_t aligned_tile_size,
+    size_t src2_size,
     size_t vtcm_budget,
     size_t * m_chunk_out,
     size_t * n_chunk_out,
@@ -746,7 +783,7 @@ static inline bool htp_mm_hmx_solve_2d_params(
 
     int act_threads = n_threads;
     while (act_threads >= 1) {
-        size_t simple_2d_overhead = htp_mm_hmx_get_2d_overhead(pipeline, is_matmul_id);
+        size_t simple_2d_overhead = htp_mm_hmx_get_2d_overhead(pipeline, is_matmul_id) + (src2_size > 0 ? hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE) : 0);
         size_t simple_2d_size_per_n, simple_2d_size_per_m, simple_2d_size_per_mn;
         htp_mm_hmx_get_2d_chunk_costs(wtype, k, pipeline, aligned_tile_size, &simple_2d_size_per_n, &simple_2d_size_per_m, &simple_2d_size_per_mn);
 
@@ -757,7 +794,7 @@ static inline bool htp_mm_hmx_solve_2d_params(
         if (htp_mm_hmx_compute_chunks(vtcm_budget, simple_2d_overhead, simple_2d_size_per_n, simple_2d_size_per_m, simple_2d_size_per_mn, m_for_chunks, ne01_padded,
                                (size_t) ne01_padded * HTP_MM_HMX_COST_W_DEQUANT, (size_t) m_for_cost * HTP_MM_HMX_COST_A_CONVERT,
                                &m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) {
-            size_t exact_size = htp_mm_hmx_get_2d_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, pipeline, is_matmul_id ? 0 : act_threads, aligned_tile_size);
+            size_t exact_size = htp_mm_hmx_get_2d_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, pipeline, is_matmul_id ? 0 : act_threads, aligned_tile_size, src2_size);
             if (exact_size <= vtcm_budget) {
                 size_t mblocks = ((size_t) m_for_cost + m_chunk_candidate - 1) / m_chunk_candidate;
                 if (mblocks < best_mblocks || (mblocks == best_mblocks && act_threads > best_act_threads)) {

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

@@ -7,7 +7,7 @@
 
 #include <string.h>
 
-#include "hex-dma.h"
+#include "dma-queue.h"
 #include "hvx-utils.h"
 
 #define GGML_COMMON_DECL_C
@@ -51,6 +51,15 @@ static inline const uint8_t * pad_src_row_ptr(const struct htp_tensor * src,
         + (i3 - (uint32_t)lp3) * src->nb[3];
 }
 
+static inline dma_addr_t pad_src_row_data(const struct htp_tensor * src,
+                                          uint32_t i1, uint32_t i2, uint32_t i3,
+                                          int32_t lp1, int32_t lp2, int32_t lp3) {
+    return src->data
+        + (i1 - (uint32_t)lp1) * src->nb[1]
+        + (i2 - (uint32_t)lp2) * src->nb[2]
+        + (i3 - (uint32_t)lp3) * src->nb[3];
+}
+
 /* Compute the DDR src row pointer for a circular row (wrap-around indexing) */
 static inline const uint8_t * pad_circ_src_row_ptr(const struct htp_tensor * src,
                                                     uint32_t i1, uint32_t i2, uint32_t i3,
@@ -61,6 +70,15 @@ static inline const uint8_t * pad_circ_src_row_ptr(const struct htp_tensor * src
         + wrap_around((int32_t)i3 - lp3, src->ne[3]) * src->nb[3];
 }
 
+static inline dma_addr_t pad_circ_src_row_data(const struct htp_tensor * src,
+                                               uint32_t i1, uint32_t i2, uint32_t i3,
+                                               int32_t lp1, int32_t lp2, int32_t lp3) {
+    return src->data
+        + wrap_around((int32_t)i1 - lp1, src->ne[1]) * src->nb[1]
+        + wrap_around((int32_t)i2 - lp2, src->ne[2]) * src->nb[2]
+        + wrap_around((int32_t)i3 - lp3, src->ne[3]) * src->nb[3];
+}
+
 struct htp_pad_context {
     struct htp_ops_context * octx;
 
@@ -118,7 +136,7 @@ struct htp_pad_context {
     uint8_t * src_spad_base = octx->src0_spad.data + ith * octx->src0_spad.size_per_thread; \
     uint8_t * dst_spad_base = octx->dst_spad.data  + ith * octx->dst_spad.size_per_thread;  \
                                                                                             \
-    dma_queue * dma = octx->ctx->dma[ith];
+    dma_queue * dma_q = octx->ctx->dma[ith];
 
 // ---------------------------------------------------------------------------
 // HVX vectorized PAD kernel
@@ -196,9 +214,9 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void
         uint8_t * src_spad_cur = src_spad_base + spad_idx * src_row_size_aligned;
         uint8_t * dst_spad_cur = dst_spad_base + spad_idx * dst_row_size_aligned;
 
-        dma_queue_push_vtcm_to_ddr(dma,
-            dma_make_ptr((uint8_t *)dst->data, dst_spad_cur),
-            dst_row_size, dst_row_size_aligned, 0);
+        dma_queue_push(dma_q,
+            dma_make_data(dst->data, dst_spad_cur),
+            dst_row_size, dst_row_size_aligned, dst_row_size, 0);
 
         uint32_t i1, i2, i3;
         pad_decompose_row(ir, ne1, ne2, &i1, &i2, &i3);
@@ -207,15 +225,14 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void
                                              lp2, rp2, ne2,
                                              lp3, rp3, ne3);
 
-        const uint8_t * src_ptr = interior
-            ? pad_src_row_ptr(src, i1, i2, i3, lp1, lp2, lp3) : NULL;
+        const dma_addr_t src_data = interior
+            ? pad_src_row_data(src, i1, i2, i3, lp1, lp2, lp3) : src->data;
 
         // Interior row: real DMA (1 row) from DDR to VTCM.
         // Border row: null DMA (nrows=0)
-        dma_queue_push_ddr_to_vtcm(dma,
-            dma_make_ptr(src_spad_cur,
-                         src_ptr ? src_ptr : (const uint8_t *)src_spad_cur),
-            src_row_size_aligned, src_row_size, src_ptr ? 1 : 0);
+        dma_queue_push(dma_q,
+            dma_make_data(src_spad_cur, src_data),
+            src_row_size_aligned, src_row_size, src_row_size, interior ? 1 : 0);
     }
 
     // -----------------------------------------------------------------------
@@ -225,13 +242,13 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
     for (uint32_t ir = row_start; ir < row_end; ir++) {
-        uint8_t * dst_spad_cur = (uint8_t *) dma_queue_pop(dma).src;
-        uint8_t * src_spad_cur = (uint8_t *) dma_queue_pop(dma).dst;
+        uint8_t * dst_spad_cur = (uint8_t *) dma_queue_pop(dma_q).src;
+        uint8_t * src_spad_cur = (uint8_t *) dma_queue_pop(dma_q).dst;
 
         uint32_t i1, i2, i3;
         pad_decompose_row(ir, ne1, ne2, &i1, &i2, &i3);
 
-        uint8_t * dst_ptr = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
+        const dma_addr_t dst_data = dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
 
         const int interior = pad_is_interior(i1, i2, i3,
                                              lp1, rp1, ne1,
@@ -254,9 +271,9 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void
         }
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
 
-        dma_queue_push_vtcm_to_ddr(dma,
-            dma_make_ptr(dst_ptr, dst_spad_cur),
-            dst_row_size, dst_row_size_aligned, 1);
+        dma_queue_push(dma_q,
+            dma_make_data(dst_data, dst_spad_cur),
+            dst_row_size, dst_row_size_aligned, dst_row_size, 1);
 
         const uint32_t next_row = ir + 2;
         if (next_row < row_end) {
@@ -266,17 +283,16 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void
                                                       lp1, rp1, ne1,
                                                       lp2, rp2, ne2,
                                                       lp3, rp3, ne3);
-            const uint8_t * next_src_ptr = next_interior
-                ? pad_src_row_ptr(src, ni1, ni2, ni3, lp1, lp2, lp3) : NULL;
+            const dma_addr_t next_src_data = next_interior
+                ? pad_src_row_data(src, ni1, ni2, ni3, lp1, lp2, lp3) : src->data;
 
-            dma_queue_push_ddr_to_vtcm(dma,
-                dma_make_ptr(src_spad_cur,
-                             next_src_ptr ? next_src_ptr : (const uint8_t *)src_spad_cur),
-                src_row_size_aligned, src_row_size, next_src_ptr ? 1 : 0);
+            dma_queue_push(dma_q,
+                dma_make_data(src_spad_cur, next_src_data),
+                src_row_size_aligned, src_row_size, src_row_size, next_interior ? 1 : 0);
         }
     }
 
-    dma_queue_flush(dma);
+    dma_queue_flush(dma_q);
 
     FARF(HIGH, "pad-hvx-dma %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u\n",
          ith, nth,
@@ -372,15 +388,16 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i
         uint8_t * src_spad_cur = src_spad_base + spad_idx * src_row_size_aligned;
         uint8_t * dst_spad_cur = dst_spad_base + spad_idx * dst_row_size_aligned;
 
-        dma_queue_push_vtcm_to_ddr(dma,
-            dma_make_ptr((uint8_t *)dst->data, dst_spad_cur),
-            dst_row_size, dst_row_size_aligned, 0);
+        dma_queue_push(dma_q,
+            dma_make_data(dst->data, dst_spad_cur),
+            dst_row_size, dst_row_size_aligned, dst_row_size, 0);
 
         uint32_t pi1, pi2, pi3;
         pad_decompose_row(ir, ne1, ne2, &pi1, &pi2, &pi3);
-        dma_queue_push_ddr_to_vtcm(dma,
-            dma_make_ptr(src_spad_cur, pad_circ_src_row_ptr(src, pi1, pi2, pi3, lp1, lp2, lp3)),
-            src_row_size_aligned, src_row_size, 1);
+        const dma_addr_t src_data = pad_circ_src_row_data(src, pi1, pi2, pi3, lp1, lp2, lp3);
+        dma_queue_push(dma_q,
+            dma_make_data(src_spad_cur, src_data),
+            src_row_size_aligned, src_row_size, src_row_size, 1);
     }
 
     // -----------------------------------------------------------------------
@@ -390,12 +407,12 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
     for (uint32_t ir = row_start; ir < row_end; ir++) {
-        uint8_t * dst_spad_cur = (uint8_t *) dma_queue_pop(dma).src;
-        uint8_t * src_spad_cur = (uint8_t *) dma_queue_pop(dma).dst;
+        uint8_t * dst_spad_cur = (uint8_t *) dma_queue_pop(dma_q).src;
+        uint8_t * src_spad_cur = (uint8_t *) dma_queue_pop(dma_q).dst;
 
         uint32_t i1, i2, i3;
         pad_decompose_row(ir, ne1, ne2, &i1, &i2, &i3);
-        uint8_t * dst_ptr = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
+        const dma_addr_t dst_data = dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
 
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
         if (lp0 > 0) {
@@ -431,22 +448,22 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i
         }
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
 
-        dma_queue_push_vtcm_to_ddr(dma,
-            dma_make_ptr(dst_ptr, dst_spad_cur),
-            dst_row_size, dst_row_size_aligned, 1);
+        dma_queue_push(dma_q,
+            dma_make_data(dst_data, dst_spad_cur),
+            dst_row_size, dst_row_size_aligned, dst_row_size, 1);
 
         const uint32_t next_row = ir + 2;
         if (next_row < row_end) {
             uint32_t nri1, nri2, nri3;
             pad_decompose_row(next_row, ne1, ne2, &nri1, &nri2, &nri3);
-            dma_queue_push_ddr_to_vtcm(dma,
-                dma_make_ptr(src_spad_cur,
-                             pad_circ_src_row_ptr(src, nri1, nri2, nri3, lp1, lp2, lp3)),
-                src_row_size_aligned, src_row_size, 1);
+            const dma_addr_t next_src_data = pad_circ_src_row_data(src, nri1, nri2, nri3, lp1, lp2, lp3);
+            dma_queue_push(dma_q,
+                dma_make_data(src_spad_cur, next_src_data),
+                src_row_size_aligned, src_row_size, src_row_size, 1);
         }
     }
 
-    dma_queue_flush(dma);
+    dma_queue_flush(dma_q);
 
     FARF(HIGH, "pad-hvx-circ-dma %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u\n",
          ith, nth,
@@ -468,10 +485,6 @@ int op_pad(struct htp_ops_context * octx) {
             return HTP_STATUS_NO_SUPPORT;
     }
 
-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
-        return HTP_STATUS_OK;
-    }
-
     const int32_t lp0 = octx->op_params[0];
     const int32_t rp0 = octx->op_params[1];
     const int32_t lp1 = octx->op_params[2];
@@ -515,6 +528,10 @@ int op_pad(struct htp_ops_context * octx) {
     const int use_dma = (src0->nb[0] == (uint32_t)type_size) && (ne00 >= 512) &&
                         (octx->ctx->vtcm_size >= vtcm_needed);
 
+    if (!use_dma && (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
     if (use_dma) {
         octx->src0_spad.size_per_thread = 2 * src_row_size_aligned;
         octx->dst_spad.size_per_thread  = 2 * dst_row_size_aligned;

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

@@ -122,8 +122,8 @@ int op_repeat(struct htp_ops_context * octx) {
             return HTP_STATUS_NO_SUPPORT;
     }
 
-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
-        return HTP_STATUS_OK;
+    if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
+        return HTP_STATUS_NO_SUPPORT;
     }
 
     const uint32_t total_dst_rows  = dst->ne[1] * dst->ne[2] * dst->ne[3];

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

@@ -0,0 +1,316 @@
+#pragma clang diagnostic ignored "-Wunused-variable"
+#pragma clang diagnostic ignored "-Wunused-function"
+#pragma clang diagnostic ignored "-Wunused-but-set-variable"
+
+#include <HAP_farf.h>
+#include <HAP_perf.h>
+
+#include <string.h>
+
+#include "dma-queue.h"
+#include "hvx-utils.h"
+
+#define GGML_COMMON_DECL_C
+#include "ggml-common.h"
+#include "htp-ctx.h"
+#include "hex-common.h"
+#include "hex-profile.h"
+#include "htp-ops.h"
+#include "htp-tensor.h"
+
+struct htp_roll_context {
+    struct htp_ops_context * octx;
+
+    uint32_t row_start;
+    uint32_t nrows;
+    uint32_t nrows_per_thread;
+
+    struct fastdiv_values div_ne1;
+    struct fastdiv_values div_ne2_ne1;
+};
+
+static inline uint32_t htp_roll_wrap(int32_t i, uint32_t ne) {
+    if (i < 0) {
+        return (uint32_t) (i + (int32_t) ne);
+    }
+    if ((uint32_t) i >= ne) {
+        return (uint32_t) i - ne;
+    }
+    return (uint32_t) i;
+}
+
+#define htp_roll_preamble                             \
+    const struct htp_tensor * src0 = octx->src[0];    \
+    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 nb01 = src0->nb[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];                  \
+                                                      \
+    const int32_t s0 = octx->op_params[0];            \
+    const int32_t s1 = octx->op_params[1];            \
+    const int32_t s2 = octx->op_params[2];            \
+    const int32_t s3 = octx->op_params[3];            \
+                                                      \
+    const uint32_t i0_src0 = htp_roll_wrap(-s0, ne0); \
+    const uint32_t n0      = ne0 - i0_src0;
+
+#define htp_roll_dma_preamble dma_queue * q = octx->ctx->dma[0];
+
+static inline void roll_dma_push(dma_queue * q,
+                                 dma_addr_t  dst,
+                                 dma_addr_t  src,
+                                 uint32_t    dst_stride,
+                                 uint32_t    src_stride,
+                                 uint32_t    bytes,
+                                 uint32_t    nrows) {
+    if (bytes == 0 || nrows == 0) {
+        return;
+    }
+
+    if (!dma_queue_push(q, dma_make_data(dst, src), dst_stride, src_stride, bytes, nrows)) {
+        dma_queue_flush(q);
+        dma_queue_push(q, dma_make_data(dst, src),
+                       dst_stride, src_stride, bytes, nrows);
+    }
+}
+
+static inline void roll_dma_push_rows(dma_queue *               q,
+                                      const struct htp_tensor * dst,
+                                      const struct htp_tensor * src0,
+                                      uint32_t                  dst_row,
+                                      uint32_t                  src_row,
+                                      uint32_t                  nrows,
+                                      uint32_t                  row_size,
+                                      uint32_t                  i0_src0) {
+    const dma_addr_t dst_base = dst->data + (size_t) dst_row * row_size;
+    const dma_addr_t src_base = src0->data + (size_t) src_row * row_size;
+    const uint32_t   n0       = src0->ne[0] - i0_src0;
+
+    roll_dma_push(q, dst_base, src_base + (size_t) i0_src0 * sizeof(float),
+                  row_size, row_size, n0 * sizeof(float), nrows);
+    roll_dma_push(q, dst_base + (size_t) n0 * sizeof(float), src_base,
+                  row_size, row_size, i0_src0 * sizeof(float), nrows);
+}
+
+// Same row-wrap split as roll_dma_push_rows, but addressed with explicit byte strides so it
+// also works for a src0 that is row-contiguous only (e.g. a permuted view) rather than fully packed.
+static inline void roll_dma_push_range(dma_queue * q,
+                                       dma_addr_t  dst_row,
+                                       dma_addr_t  src_row,
+                                       uint32_t    dst_stride,
+                                       uint32_t    src_stride,
+                                       uint32_t    nrows,
+                                       uint32_t    i0_src0,
+                                       uint32_t    n0) {
+    roll_dma_push(q, dst_row, src_row + (size_t) i0_src0 * sizeof(float),
+                  dst_stride, src_stride, n0 * sizeof(float), nrows);
+    roll_dma_push(q, dst_row + (size_t) n0 * sizeof(float), src_row,
+                  dst_stride, src_stride, i0_src0 * sizeof(float), nrows);
+}
+
+static int roll_dma_f32_contiguous(struct htp_ops_context * octx) {
+    htp_roll_preamble;
+    htp_roll_dma_preamble;
+
+    const uint32_t row_size = ne0 * sizeof(float);
+
+    if (s1 == 0 && s2 == 0 && s3 == 0) {
+        roll_dma_push_rows(q, dst, src0, 0, 0, ne1 * ne2 * ne3, row_size, i0_src0);
+        dma_queue_flush(q);
+        return HTP_STATUS_OK;
+    }
+
+    if (s1 == 0) {
+        const uint32_t i2_src0 = htp_roll_wrap(-s2, ne2);
+        for (uint32_t i3 = 0; i3 < ne3; i3++) {
+            const uint32_t i03 = htp_roll_wrap((int32_t) i3 - s3, ne3);
+            const uint32_t dst_row0 = i3 * ne2 * ne1;
+            const uint32_t src_row0 = (i03 * ne2 + i2_src0) * ne1;
+            const uint32_t n2_first = ne2 - i2_src0;
+
+            roll_dma_push_rows(q, dst, src0, dst_row0, src_row0, n2_first * ne1,
+                               row_size, i0_src0);
+            roll_dma_push_rows(q, dst, src0, dst_row0 + n2_first * ne1, i03 * ne2 * ne1,
+                               i2_src0 * ne1, row_size, i0_src0);
+        }
+
+        dma_queue_flush(q);
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t i1_src0 = htp_roll_wrap(-s1, ne1);
+    const uint32_t n1_first = ne1 - i1_src0;
+
+    for (uint32_t i3 = 0; i3 < ne3; i3++) {
+        const uint32_t i03 = htp_roll_wrap((int32_t) i3 - s3, ne3);
+        for (uint32_t i2 = 0; i2 < ne2; i2++) {
+            const uint32_t i02 = htp_roll_wrap((int32_t) i2 - s2, ne2);
+            const uint32_t dst_row0 = (i3 * ne2 + i2) * ne1;
+            const uint32_t src_row0 = (i03 * ne2 + i02) * ne1;
+
+            roll_dma_push_rows(q, dst, src0, dst_row0, src_row0 + i1_src0,
+                               n1_first, row_size, i0_src0);
+            roll_dma_push_rows(q, dst, src0, dst_row0 + n1_first, src_row0,
+                               i1_src0, row_size, i0_src0);
+        }
+    }
+
+    dma_queue_flush(q);
+    return HTP_STATUS_OK;
+}
+
+// DMA path for a row-contiguous but otherwise arbitrarily strided src0 (e.g. a permuted view).
+// Same row-wrap split as above, one DMA push per (i2,i3), addressed via the real nb01/nb02/nb03
+// instead of assuming a packed layout.
+static int roll_dma_f32_strided(struct htp_ops_context * octx) {
+    htp_roll_preamble;
+    htp_roll_dma_preamble;
+
+    const uint32_t i1_src0  = htp_roll_wrap(-s1, ne1);
+    const uint32_t n1_first = ne1 - i1_src0;
+
+    for (uint32_t i3 = 0; i3 < ne3; i3++) {
+        const uint32_t i03 = htp_roll_wrap((int32_t) i3 - s3, ne3);
+        for (uint32_t i2 = 0; i2 < ne2; i2++) {
+            const uint32_t i02 = htp_roll_wrap((int32_t) i2 - s2, ne2);
+
+            const dma_addr_t dst_row0 = dst->data  + (size_t) i2  * nb2  + (size_t) i3  * nb3;
+            const dma_addr_t src_row0 = src0->data + (size_t) i02 * nb02 + (size_t) i03 * nb03;
+
+            roll_dma_push_range(q, dst_row0, src_row0 + (size_t) i1_src0 * nb01,
+                                nb1, nb01, n1_first, i0_src0, n0);
+            roll_dma_push_range(q, dst_row0 + (size_t) n1_first * nb1, src_row0,
+                                nb1, nb01, i1_src0, i0_src0, n0);
+        }
+    }
+
+    dma_queue_flush(q);
+    return HTP_STATUS_OK;
+}
+
+static void roll_thread_f32(unsigned int nth, unsigned int ith, void * data) {
+    struct htp_roll_context * rctx = (struct htp_roll_context *) data;
+    struct htp_ops_context * octx = rctx->octx;
+
+    htp_roll_preamble;
+
+    const uint32_t row_start = rctx->row_start + rctx->nrows_per_thread * ith;
+    const uint32_t row_end   = MIN(row_start + rctx->nrows_per_thread, rctx->row_start + rctx->nrows);
+    if (row_start >= row_end) {
+        return;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, row_start);
+
+    for (uint32_t row = row_start; row < row_end; row++) {
+        const uint32_t i3  = fastdiv(row, &rctx->div_ne2_ne1);
+        const uint32_t rem = row - i3 * ne2 * ne1;
+        const uint32_t i2  = fastdiv(rem, &rctx->div_ne1);
+        const uint32_t i1  = rem - i2 * ne1;
+
+        const uint32_t i01 = htp_roll_wrap((int32_t) i1 - s1, ne1);
+        const uint32_t i02 = htp_roll_wrap((int32_t) i2 - s2, ne2);
+        const uint32_t i03 = htp_roll_wrap((int32_t) i3 - s3, ne3);
+
+        const uint8_t * src_row = (const uint8_t *) (uintptr_t) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
+        uint8_t * dst_row = (uint8_t *) (uintptr_t) dst->data + i1*nb1 + i2*nb2 + i3*nb3;
+
+        hex_l2fetch(src_row + i0_src0 * sizeof(float), n0 * sizeof(float), ne0 * sizeof(float), 1);
+        hvx_copy_uu(dst_row, src_row + i0_src0 * sizeof(float), n0, sizeof(float));
+
+        if (i0_src0 != 0) {
+            hex_l2fetch(src_row, i0_src0 * sizeof(float), ne0 * sizeof(float), 1);
+            hvx_copy_uu(dst_row + n0 * sizeof(float), src_row, i0_src0, sizeof(float));
+        }
+    }
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, row_start);
+
+    FARF(HIGH, "roll %d/%d: (%ux%ux%ux%u) rows %u:%u shift=(%d,%d,%d,%d)\n",
+         ith, nth, ne0, ne1, ne2, ne3,
+         row_start, row_end, s0, s1, s2, s3);
+}
+
+int execute_op_roll_f32(struct htp_ops_context * octx) {
+    htp_roll_preamble;
+
+    if (src0->type != HTP_TYPE_F32 || dst->type != HTP_TYPE_F32) {
+        FARF(ERROR, "roll: unsupported type %u -> %u\n", src0->type, dst->type);
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    if (src0->nb[0] != sizeof(float) || dst->nb[0] != sizeof(float)) {
+        FARF(ERROR, "roll: unsupported nb0 %u -> %u\n", src0->nb[0], dst->nb[0]);
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    if (src0->ne[0] != ne0 || src0->ne[1] != ne1 ||
+        src0->ne[2] != ne2 || src0->ne[3] != ne3) {
+        FARF(ERROR, "roll: shape mismatch\n");
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    const uint32_t total_rows = ne1 * ne2 * ne3;
+    const size_t dst_row_size = ne0 * 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_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;
+    }
+
+    if (octx->ctx->mdev.count <= 1) {
+        if (htp_tensor_is_contiguous(src0, sizeof(float)) && htp_tensor_is_contiguous(dst, sizeof(float))) {
+            return roll_dma_f32_contiguous(octx);
+        }
+        return roll_dma_f32_strided(octx);
+    }
+
+    if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+    struct htp_roll_context rctx = {
+        .octx             = octx,
+        .row_start        = row_start,
+        .nrows            = nrows,
+        .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
+        .div_ne1          = init_fastdiv_values(dst->ne[1]),
+        .div_ne2_ne1      = init_fastdiv_values(dst->ne[2] * dst->ne[1]),
+    };
+
+    work_queue_run(octx->ctx->work_queue, roll_thread_f32, &rctx, n_threads);
+
+    return HTP_STATUS_OK;
+}
+
+int op_roll(struct htp_ops_context * octx) {
+    switch (octx->src[0]->type) {
+        case HTP_TYPE_F32:
+            return execute_op_roll_f32(octx);
+
+        default:
+            return HTP_STATUS_NO_SUPPORT;
+    }
+}

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

@@ -9,7 +9,7 @@
 #include <string.h>
 #include <stdlib.h>
 
-#include "hex-dma.h"
+#include "dma-queue.h"
 #include "hvx-utils.h"
 #include "hex-fastdiv.h"
 
@@ -85,6 +85,8 @@ struct htp_rope_context {
 
     struct fastdiv_values div_ne2_ne1;
     struct fastdiv_values div_ne1;
+
+    const float * freq_factors;
 };
 
 static float rope_yarn_ramp(const float low, const float high, const int i0) {
@@ -562,10 +564,10 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) {
     float *   theta_cache    = (float *) (src0_spad_base);
               src0_spad_base = src0_spad_base + rctx->theta_cache_offset;
 
-    dma_queue * dma_queue = octx->ctx->dma[ith];
+    dma_queue * dma_q = octx->ctx->dma[ith];
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
-    const int32_t * pos = (const int32_t *) src1->data;
-    const float * freq_factors = src2 ? (const float *) src2->data : NULL;
+    const int32_t * pos = (const int32_t *) (uintptr_t) src1->data;
+    const float * freq_factors = rctx->freq_factors;
 
     const uint32_t i3_start = fastdiv(src0_start_row, &rctx->div_ne2_ne1);
     const uint32_t rem      = fastmodulo(src0_start_row, ne2 * ne1, &rctx->div_ne2_ne1);
@@ -587,7 +589,7 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) {
                 const uint32_t nrows = MIN(src0_end_row - ir, ne1 - i1);
 
                 // Depth before prefetch
-                const uint32_t dma_depth = dma_queue_depth(dma_queue);
+                const uint32_t dma_depth = dma_queue_depth(dma_q);
 
                 // Prefetch up to 2 blocks
                 const uint32_t p_nrows = MIN(nrows, 2 * HTP_ROPE_SPAD_BLOCK);
@@ -595,12 +597,12 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) {
                     const uint32_t pnr = MIN(nrows - pr, HTP_ROPE_SPAD_BLOCK);
                     const uint32_t slot = (cur_slot + pr / HTP_ROPE_SPAD_BLOCK) % HTP_ROPE_SPAD_NSLOTS;
                     uint8_t * spad_slot = rope_spad_slot(src0_spad_base, slot, rctx->src0_row_size_aligned);
-                    const uint8_t * src_addr = (const uint8_t *) src0->data + i3 * nb03 + i2 * nb02 + (i1 + pr) * nb01;
+                    const dma_addr_t src0_data = src0->data + i3 * nb03 + i2 * nb02 + (i1 + pr) * nb01;
 
                     // Dummy DMA transaction for sequencing (interleaving wr, rd, wr, rd, ...)
-                    dma_queue_push(dma_queue, dma_make_ptr((void *) dst->data, spad_slot), 0, 0, 0, 0);
+                    dma_queue_push(dma_q, dma_make_data(dst->data, spad_slot), 0, 0, 0, 0);
 
-                    dma_queue_push(dma_queue, dma_make_ptr(spad_slot, src_addr),
+                    dma_queue_push(dma_q, dma_make_data(spad_slot, src0_data),
                         rctx->src0_row_size_aligned, rctx->src0_row_stride, rctx->src0_row_size, pnr);
                 }
 
@@ -634,7 +636,7 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) {
                 }
 
                 // Skip output DMA transactions from prev block (if any)
-                for (uint32_t d = 0; d < dma_depth; d++) { dma_queue_pop_nowait(dma_queue); }
+                for (uint32_t d = 0; d < dma_depth; d++) { dma_queue_pop_nowait(dma_q); }
 
                 // Compute loop
                 const uint32_t ne = is_vision ? ne0 : rctx->n_dims;
@@ -647,8 +649,8 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) {
                     const uint32_t cur_ir   = base_ir + cr;
                     const uint32_t cur_i1   = base_i1 + cr;
 
-                    dma_queue_pop(dma_queue);
-                    uint8_t * cur_spad = (uint8_t *) dma_queue_pop(dma_queue).dst;
+                    dma_queue_pop(dma_q);
+                    uint8_t * cur_spad = (uint8_t *) dma_queue_pop(dma_q).dst;
 
                     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, cur_ir);
                     if (is_neox || is_vision) {
@@ -658,8 +660,8 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) {
                     }
                     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, cur_ir);
 
-                    uint8_t * dst_addr = (uint8_t *) dst->data + i3 * nb3 + i2 * nb2 + cur_i1 * nb1;
-                    dma_queue_push(dma_queue, dma_make_ptr(dst_addr, cur_spad),
+                    const dma_addr_t dst_data = dst->data + i3 * nb3 + i2 * nb2 + cur_i1 * nb1;
+                    dma_queue_push(dma_q, dma_make_data(dst_data, cur_spad),
                         rctx->dst_row_stride, rctx->src0_row_size_aligned, rctx->dst_row_size, cnr);
 
                     // Prefetch 2 blocks ahead into the slot just freed
@@ -668,9 +670,9 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) {
                         const uint32_t pnr    = MIN(nrows - p_cr, HTP_ROPE_SPAD_BLOCK);
                         const uint32_t p_slot = (cur_slot + p_cr / HTP_ROPE_SPAD_BLOCK) % HTP_ROPE_SPAD_NSLOTS;
                         uint8_t * p_spad      = rope_spad_slot(src0_spad_base, p_slot, rctx->src0_row_size_aligned);
-                        const uint8_t * src_addr = (const uint8_t *) src0->data + i3 * nb03 + i2 * nb02 + (base_i1 + p_cr) * nb01;
+                        const dma_addr_t p_src0_data = src0->data + i3 * nb03 + i2 * nb02 + (base_i1 + p_cr) * nb01;
 
-                        dma_queue_push(dma_queue, dma_make_ptr(p_spad, src_addr),
+                        dma_queue_push(dma_q, dma_make_data(p_spad, p_src0_data),
                             rctx->src0_row_size_aligned, rctx->src0_row_stride, rctx->src0_row_size, pnr);
                     }
                 }
@@ -685,7 +687,7 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) {
     }
 
 done:
-    dma_queue_flush(dma_queue);
+    dma_queue_flush(dma_q);
 
     FARF(HIGH, "rope-f32: %d/%d: (%u:%u)\n", ith, nth, src0_start_row, src0_end_row);
 }
@@ -713,6 +715,10 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) {
     }
     assert(octx->ctx->vtcm_size >= kparams->vtcm_size);
 
+    if (htp_tensor_is_extended(src1)) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
     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);
 
@@ -748,6 +754,16 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) {
     rctx.spad_per_thread       = kparams->spad_per_thread;
     rctx.theta_cache_offset    = kparams->theta_cache_offset;
 
+    if (src2) {
+        dma_queue * dma_q = octx->ctx->dma[0];
+        const size_t ff_size = src2->ne[0] * sizeof(float);
+        float * vtcm_freq_factors = (float *) (rctx.vtcm_base + kparams->freq_factors_offset);
+        dma_queue_push(dma_q, dma_make_data(vtcm_freq_factors, src2->data),
+                       kparams->freq_factors_size, 0, ff_size, 1);
+        dma_queue_pop(dma_q);
+        rctx.freq_factors = vtcm_freq_factors;
+    }
+
     const int32_t * op_params = &octx->op_params[0];
     rctx.n_dims     = ((const int32_t *) op_params)[1];
     rctx.mode       = ((const int32_t *) op_params)[2];

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

@@ -16,6 +16,8 @@ struct htp_rope_kernel_params {
     uint32_t spad_per_thread;
     uint32_t theta_cache_offset;
     uint32_t src0_row_size_aligned;
+    uint32_t freq_factors_offset;
+    uint32_t freq_factors_size;
 
     struct fastdiv_values div_ne2_ne1;
     struct fastdiv_values div_ne1;
@@ -32,21 +34,25 @@ struct htp_rope_vtcm_layout {
     size_t bytes_per_thread;
     size_t theta_cache_size_aligned;
     size_t src0_row_size_aligned;
+    size_t freq_factors_size_aligned;
 };
 
 static inline void htp_rope_vtcm_layout_build(
     struct htp_rope_vtcm_layout * layout,
     uint32_t ne00,
-    uint32_t n_threads
+    uint32_t n_threads,
+    uint32_t n_freq_factors
 ) {
-    const size_t src0_row_size            = ne00 * sizeof(float);
-    const size_t src0_row_size_aligned    = hex_round_up((uint32_t) src0_row_size, 128);
-    const size_t theta_cache_size_aligned = hex_round_up((uint32_t) src0_row_size, 256);
+    const size_t src0_row_size             = ne00 * sizeof(float);
+    const size_t src0_row_size_aligned     = hex_round_up((uint32_t) src0_row_size, 128);
+    const size_t theta_cache_size_aligned  = hex_round_up((uint32_t) src0_row_size, 256);
+    const size_t freq_factors_size_aligned = hex_round_up(n_freq_factors * sizeof(float), 256);
 
-    layout->src0_row_size_aligned    = src0_row_size_aligned;
-    layout->theta_cache_size_aligned = theta_cache_size_aligned;
-    layout->bytes_per_thread         = theta_cache_size_aligned + HTP_ROPE_SPAD_NROWS * src0_row_size_aligned;
-    layout->total_bytes              = layout->bytes_per_thread * n_threads;
+    layout->src0_row_size_aligned     = src0_row_size_aligned;
+    layout->theta_cache_size_aligned  = theta_cache_size_aligned;
+    layout->freq_factors_size_aligned = freq_factors_size_aligned;
+    layout->bytes_per_thread          = theta_cache_size_aligned + HTP_ROPE_SPAD_NROWS * src0_row_size_aligned;
+    layout->total_bytes               = layout->bytes_per_thread * n_threads + freq_factors_size_aligned;
 }
 
 static inline uint8_t * rope_spad_slot(uint8_t * base, uint32_t slot, size_t row_size_aligned) {

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

@@ -77,7 +77,7 @@ static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsig
         return;                                                                                                  \
     }                                                                                                            \
     const uint32_t ir1 = MIN(ir0 + dr, srctx->task_start + srctx->tasks);                                        \
-    dma_queue * dma_queue = octx->ctx->dma[ith];                                                                 \
+    dma_queue * dma_q = octx->ctx->dma[ith];                                                                     \
     const struct htp_set_rows_vtcm_layout * vtcm_layout = &srctx->vtcm_layout;                                   \
     uint8_t * vtcm_src0 = srctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread;   \
     uint8_t * vtcm_dst  = srctx->vtcm_base + vtcm_layout->off_dst  + ith * vtcm_layout->dst_bytes_per_thread;    \
@@ -90,14 +90,14 @@ static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsig
     uint32_t pi03 = 0;                                                                                           \
     for (uint32_t step = 0, spad_idx = 0; step < total_steps && spad_idx < 2; ++step, spad_idx++) {              \
         uint32_t i = ir0 + pi_step;                                                                              \
-        const uintptr_t src0_ptr = octx->src[0]->data + i*nb01 + pi02*nb02 + pi03*nb03;                          \
-        dma_queue_push(dma_queue,                                                                                \
-                       dma_make_ptr((void *)octx->dst->data,                                                     \
-                                    vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size),                      \
+        const dma_addr_t src0_data = octx->src[0]->data + i*nb01 + pi02*nb02 + pi03*nb03;                        \
+        dma_queue_push(dma_q,                                                                                    \
+                       dma_make_data(octx->dst->data,                                                            \
+                                     vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size),                     \
                        dst_row_size, vtcm_layout->dst_spad_half_size, dst_row_size, 0);                          \
-        dma_queue_push(dma_queue,                                                                                \
-                       dma_make_ptr((void *)(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size),           \
-                                    (const void *)src0_ptr),                                                     \
+        dma_queue_push(dma_q,                                                                                    \
+                       dma_make_data(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size,                    \
+                                     src0_data),                                                                 \
                        vtcm_layout->src0_spad_half_size, src0_row_size, src0_row_size, 1);                       \
         pi_step++;                                                                                               \
         if (pi_step == nrows_per_thread) {                                                                       \
@@ -115,8 +115,8 @@ static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsig
     uint32_t ci11_base = 0;                                                                                      \
     uint32_t ci12_base = 0;                                                                                      \
     for (uint32_t step = 0; step < total_steps; ++step) {                                                        \
-        void * dst_spad = (void *) dma_queue_pop(dma_queue).src;                                                 \
-        void * src_spad = (void *) dma_queue_pop(dma_queue).dst;                                                 \
+        void * dst_spad = (void *) dma_queue_pop(dma_q).src;                                                     \
+        void * src_spad = (void *) dma_queue_pop(dma_q).dst;                                                     \
         uint32_t i = ir0 + ci_step;                                                                              \
         const uintptr_t src1_addr = octx->src[1]->data + i*nb10 + ci11_base*nb11 + ci12_base*nb12;               \
         const IDX_TYPE i1 = *(const IDX_TYPE *)src1_addr;                                                        \
@@ -128,21 +128,21 @@ static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsig
         }                                                                                                        \
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, step);                                                  \
         if (valid_i1) {                                                                                          \
-            const uintptr_t dst_ptr = octx->dst->data + target_i1*nb1 + ci02*nb2 + ci03*nb3;                     \
-            dma_queue_push(dma_queue,                                                                            \
-                           dma_make_ptr((void *)dst_ptr, (const void *)dst_spad),                                \
+            const dma_addr_t dst_data = octx->dst->data + target_i1*nb1 + ci02*nb2 + ci03*nb3;                   \
+            dma_queue_push(dma_q,                                                                                \
+                           dma_make_data(dst_data, dst_spad),                                                    \
                            dst_row_size, vtcm_layout->dst_spad_half_size, dst_row_size, 1);                      \
         } else {                                                                                                 \
-            dma_queue_push(dma_queue,                                                                            \
-                           dma_make_ptr((void *)octx->dst->data, (const void *)dst_spad),                        \
+            dma_queue_push(dma_q,                                                                                \
+                           dma_make_data(octx->dst->data, dst_spad),                                             \
                            dst_row_size, vtcm_layout->dst_spad_half_size, dst_row_size, 0);                      \
         }                                                                                                        \
         const uint32_t next_step = step + 2;                                                                     \
         if (next_step < total_steps) {                                                                           \
             uint32_t ni = ir0 + pi_step;                                                                         \
-            const uintptr_t psrc0_ptr = octx->src[0]->data + ni*nb01 + pi02*nb02 + pi03*nb03;                    \
-            dma_queue_push(dma_queue,                                                                            \
-                           dma_make_ptr((void *)src_spad, (const void *)psrc0_ptr),                              \
+            const dma_addr_t psrc0_data = octx->src[0]->data + ni*nb01 + pi02*nb02 + pi03*nb03;                  \
+            dma_queue_push(dma_q,                                                                                \
+                           dma_make_data(src_spad, psrc0_data),                                                  \
                            vtcm_layout->src0_spad_half_size, src0_row_size, src0_row_size, 1);                   \
             pi_step++;                                                                                           \
             if (pi_step == nrows_per_thread) {                                                                   \
@@ -172,7 +172,7 @@ static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsig
             }                                                                                                    \
         }                                                                                                        \
     }                                                                                                            \
-    dma_queue_flush(dma_queue);                                                                                  \
+    dma_queue_flush(dma_q);                                                                                      \
 }
 
 SET_ROWS_THREAD_DMA_FN(f32,  int32_t, { hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, ne00); })
@@ -196,8 +196,8 @@ int op_set_rows(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
-        return HTP_STATUS_OK;
+    if (htp_tensor_is_extended(octx->src[1])) {
+        return HTP_STATUS_NO_SUPPORT;
     }
 
     const struct htp_tensor * dst = octx->dst;

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

@@ -8,54 +8,49 @@
 #include <math.h>
 #include <string.h>
 
-#include "hex-dma.h"
+#include "dma-queue.h"
+#include "work-queue.h"
 #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 "hex-common.h"
-#include "hex-profile.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
-
-#define htp_softmax_preamble3                     \
-    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 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 ne10 = src1 ? src1->ne[0] : 1; \
-    const uint32_t ne11 = src1 ? src1->ne[1] : 1; \
-    const uint32_t ne12 = src1 ? src1->ne[2] : 1; \
-    const uint32_t ne13 = src1 ? src1->ne[3] : 1; \
-                                                  \
-    const uint32_t nb10 = src1 ? src1->nb[0] : 1; \
-    const uint32_t nb11 = src1 ? src1->nb[1] : 1; \
-    const uint32_t nb12 = src1 ? src1->nb[2] : 1; \
-    const uint32_t nb13 = src1 ? src1->nb[3] : 1; \
-                                                  \
-    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 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];
+#include "htp-vtcm.h"
+#include "htp/softmax-ops.h"
+#include "hvx-flash-attn.h"
 
 struct htp_softmax_context {
     struct htp_ops_context * octx;
+    const struct htp_softmax_kernel_params * kparams;
+
+    void *   compute;
+
+    dma_addr_t data_src0;
+    dma_addr_t data_src1;
+    dma_addr_t data_dst;
+
+    uint8_t * vtcm_src0;
+    uint8_t * vtcm_src1;
+    uint8_t * vtcm_dst;
+
+    uint32_t  vtcm_src0_size_per_thread;
+    uint32_t  vtcm_src1_size_per_thread;
+    uint32_t  vtcm_dst_size_per_thread;
+
+    uint32_t  src0_spad_half_size;
+    uint32_t  src1_spad_half_size;
+    uint32_t  dst_spad_half_size;
+
+    uint32_t  src0_row_size_aligned;
+    uint32_t  src1_row_size_aligned;
+    uint32_t  dst_row_size_aligned;
 
     bool     use_f16;
-    bool     use_src1;
 
     uint32_t n_head;
     uint32_t n_head_log2;
@@ -65,127 +60,148 @@ struct htp_softmax_context {
     float    m0;
     float    m1;
 
-    struct fastdiv_values fastdiv_ne01;
-    struct fastdiv_values fastdiv_ne02;
-    struct fastdiv_values fastdiv_ne12; // For mask broadcasting
-    struct fastdiv_values fastdiv_ne13; // For mask broadcasting
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne02;
+    struct fastdiv_values div_ne12;
+    struct fastdiv_values div_ne13;
 
     uint32_t src0_nrows_per_thread;
     uint32_t row_start;
     uint32_t nrows;
+
+    float    slopes[512] __attribute__((aligned(128)));
 };
 
-static void apply_mask(float * restrict wp0,
-                       const float * restrict mp_f32,
-                       const __fp16 * restrict mp_f16,
-                       uint32_t ne00,
-                       float slope,
-                       bool use_f16) {
-    if (!mp_f32) {
-        return;
+typedef void (*softmax_compute_fn_t)(
+    void * restrict dst,
+    const void * restrict src0,
+    const void * restrict mask,
+    uint32_t ne00,
+    float scale,
+    float slope
+);
+
+static void hvx_fast_softmax_prep_f16(const uint8_t * restrict src,
+                                      uint8_t * restrict dst,
+                                      const int num_elems,
+                                      float     scale,
+                                      const uint8_t * restrict mask,
+                                      float slope) {
+    const HVX_Vector * restrict v_src  = (const HVX_Vector *) src;
+    HVX_Vector * restrict v_dst        = (HVX_Vector *) dst;
+    const HVX_Vector * restrict v_mask = (const HVX_Vector *) mask;
+
+    HVX_Vector scale_vec = hvx_vec_splat_f32(scale);
+    HVX_Vector slope_vec = hvx_vec_splat_f32(slope);
+
+    const int nvec_64 = num_elems / VLEN_FP16;
+    const int nloe_64 = num_elems % VLEN_FP16;
+
+    #pragma unroll(2)
+    for (int i = 0; i < nvec_64; i++) {
+        HVX_VectorPair p = hvx_vec_f16_to_f32(v_mask[i]);
+        HVX_Vector m0 = Q6_V_lo_W(p);
+        HVX_Vector m1 = Q6_V_hi_W(p);
+
+        HVX_Vector s0 = v_src[2 * i];
+        HVX_Vector s1 = v_src[2 * i + 1];
+
+        HVX_Vector v0 = Q6_Vqf32_vadd_Vqf32Vqf32(Q6_Vqf32_vmpy_VsfVsf(s0, scale_vec), Q6_Vqf32_vmpy_VsfVsf(m0, slope_vec));
+        HVX_Vector v1 = Q6_Vqf32_vadd_Vqf32Vqf32(Q6_Vqf32_vmpy_VsfVsf(s1, scale_vec), Q6_Vqf32_vmpy_VsfVsf(m1, slope_vec));
+
+        v_dst[2 * i]     = Q6_Vsf_equals_Vqf32(v0);
+        v_dst[2 * i + 1] = Q6_Vsf_equals_Vqf32(v1);
     }
-    if (use_f16) {
-        for (uint32_t i = 0; i < ne00; ++i) {
-            wp0[i] += slope * (float) mp_f16[i];
-        }
-    } else {
-        for (uint32_t i = 0; i < ne00; ++i) {
-            wp0[i] += slope * mp_f32[i];
+
+    if (nloe_64 > 0) {
+        HVX_VectorPair p = hvx_vec_f16_to_f32(v_mask[nvec_64]);
+        HVX_Vector m0 = Q6_V_lo_W(p);
+
+        HVX_Vector s0 = v_src[2 * nvec_64];
+        HVX_Vector v0 = Q6_Vqf32_vadd_Vqf32Vqf32(Q6_Vqf32_vmpy_VsfVsf(s0, scale_vec), Q6_Vqf32_vmpy_VsfVsf(m0, slope_vec));
+
+        if (nloe_64 <= VLEN_FP32) {
+            hvx_vec_store_a(&v_dst[2 * nvec_64], nloe_64 * sizeof(float), Q6_Vsf_equals_Vqf32(v0));
+        } else {
+            v_dst[2 * nvec_64] = Q6_Vsf_equals_Vqf32(v0);
+
+            HVX_Vector m1 = Q6_V_hi_W(p);
+            HVX_Vector s1 = v_src[2 * nvec_64 + 1];
+            HVX_Vector v1 = Q6_Vqf32_vadd_Vqf32Vqf32(Q6_Vqf32_vmpy_VsfVsf(s1, scale_vec), Q6_Vqf32_vmpy_VsfVsf(m1, slope_vec));
+
+            hvx_vec_store_a(&v_dst[2 * nvec_64 + 1], (nloe_64 - VLEN_FP32) * sizeof(float), Q6_Vsf_equals_Vqf32(v1));
         }
     }
 }
 
-static void init_softmax_ctx(struct htp_softmax_context * smctx, struct htp_ops_context * octx) {
-    const struct htp_tensor * src0 = octx->src[0];
-    const struct htp_tensor * src1 = octx->src[1];
-
-    memset(smctx, 0, sizeof(struct htp_softmax_context));
-
-    memcpy(&smctx->scale,    (float *) octx->op_params,     sizeof(float));
-    memcpy(&smctx->max_bias, (float *) octx->op_params + 1, sizeof(float));
-
-    smctx->n_head      = src0->ne[2];
-    smctx->n_head_log2 = 1u << (uint32_t) floor(log2(smctx->n_head));
-
-    smctx->m0 = powf(2.0f, -(smctx->max_bias) / smctx->n_head_log2);
-    smctx->m1 = powf(2.0f, -(smctx->max_bias / 2.0f) / smctx->n_head_log2);
-
-    smctx->use_src1 = (src1 != 0);
-    smctx->use_f16  = (src1 != 0) && (src1->type == HTP_TYPE_F16);
-
-    smctx->octx = octx;
-
-    // Initialize fastdiv values
-    const uint32_t ne01 = src0->ne[1];
-    const uint32_t ne02 = src0->ne[2];
-
-    if (ne01 > 0) smctx->fastdiv_ne01 = init_fastdiv_values(ne01);
-    if (ne02 > 0) smctx->fastdiv_ne02 = init_fastdiv_values(ne02);
-
-    const uint32_t ne12 = src1 ? src1->ne[2] : 1;
-    const uint32_t ne13 = src1 ? src1->ne[3] : 1;
-
-    if (ne12 > 0) smctx->fastdiv_ne12 = init_fastdiv_values(ne12);
-    if (ne13 > 0) smctx->fastdiv_ne13 = init_fastdiv_values(ne13);
-}
-
 static void hvx_fast_softmax_prep_f32(const uint8_t * restrict src,
                                       uint8_t * restrict dst,
                                       const int num_elems,
                                       float     scale,
                                       const uint8_t * restrict mask,
                                       float slope) {
-    const uint8_t * restrict src_curr  = src;
-    uint8_t * restrict dst_curr        = dst;
-    const uint8_t * restrict mask_curr = mask;
+    const HVX_Vector * restrict v_src  = (const HVX_Vector *) src;
+    HVX_Vector * restrict v_dst        = (HVX_Vector *) dst;
+    const HVX_Vector * restrict v_mask = (const HVX_Vector *) mask;
 
     HVX_Vector scale_vec = hvx_vec_splat_f32(scale);
     HVX_Vector slope_vec = hvx_vec_splat_f32(slope);
 
-    int step_of_1 = num_elems >> 5;
+    const int nvec = num_elems / VLEN_FP32;
+    const int nloe = num_elems % VLEN_FP32;
 
     #pragma unroll(4)
-    for (int i = 0; i < step_of_1; i++) {
-        HVX_Vector v1 = *(HVX_Vector *) src_curr;
-
-        HVX_Vector v3 = *(HVX_Vector *) mask_curr;
+    for (int i = 0; i < nvec; i++) {
+        HVX_Vector v1 = v_src[i];
+        HVX_Vector v3 = v_mask[i];
 
         HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_vec);
-
         HVX_Vector v4 = Q6_Vqf32_vmpy_VsfVsf(v3, slope_vec);
-
         HVX_Vector v5 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, v4);
 
-        *(HVX_Vector *) dst_curr = Q6_Vsf_equals_Vqf32(v5);
+        v_dst[i] = Q6_Vsf_equals_Vqf32(v5);
+    }
 
-        src_curr += VLEN;
-        dst_curr += VLEN;
-        mask_curr += VLEN;
+    if (nloe > 0) {
+        HVX_Vector v1 = v_src[nvec];
+        HVX_Vector v3 = v_mask[nvec];
+
+        HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_vec);
+        HVX_Vector v4 = Q6_Vqf32_vmpy_VsfVsf(v3, slope_vec);
+        HVX_Vector v5 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, v4);
+
+        hvx_vec_store_a(&v_dst[nvec], nloe * sizeof(float), Q6_Vsf_equals_Vqf32(v5));
     }
 }
 
-static void hvx_fast_softmax_f32(const uint8_t * restrict src, uint8_t * restrict dst, uint8_t * restrict pad, const int num_elems) {
-    const HVX_Vector * restrict v_src = (HVX_Vector *) src;
-    HVX_Vector * restrict v_pad       = (HVX_Vector *) pad;
+static void hvx_fast_softmax_f32(const uint8_t * restrict src, uint8_t * restrict dst, const int num_elems) {
+    const HVX_Vector * restrict v_src = (const HVX_Vector *) src;
     HVX_Vector * restrict v_dst       = (HVX_Vector *) dst;
 
-    HVX_Vector sum_vec = Q6_V_vsplat_R(0x00000000);
+    const int nvec = num_elems / VLEN_FP32;
+    const int nloe = num_elems % VLEN_FP32;
+
     HVX_Vector max_vec = hvx_vec_splat_f32(((const float *) src)[0]);
-    HVX_Vector zero_v  = Q6_V_vzero();
-    HVX_Vector one_v   = hvx_vec_splat_f32(1.0);
 
-    int step_of_1 = num_elems >> 5;
-
-    #pragma unroll(4)
-    for (int i = 0; i < step_of_1; i++) {
+    #pragma unroll(2)
+    for (int i = 0; i < nvec; i++) {
         HVX_Vector v1 = v_src[i];
         max_vec       = Q6_Vsf_vmax_VsfVsf(max_vec, v1);
     }
 
-    max_vec = hvx_vec_reduce_max_f32(max_vec); // replicated over all lanes
+    if (nloe > 0) {
+        HVX_VectorPred q_mask = Q6_Q_vsetq_R(nloe * sizeof(float));
+        HVX_Vector neg_inf    = hvx_vec_splat_f32(-INFINITY);
+        HVX_Vector v_tail     = Q6_V_vmux_QVV(q_mask, v_src[nvec], neg_inf);
+        max_vec               = Q6_Vsf_vmax_VsfVsf(max_vec, v_tail);
+    }
 
-    #pragma unroll(4)
-    for (int i = 0; i < step_of_1; i++) {
+    max_vec = hvx_vec_reduce_max_f32(max_vec);
+
+    HVX_Vector sum_vec = Q6_V_vsplat_R(0x00000000);
+
+    #pragma unroll(2)
+    for (int i = 0; i < nvec; i++) {
         HVX_Vector v1 = v_src[i];
         HVX_Vector v2 = Q6_Vqf32_vsub_VsfVsf(v1, max_vec);
 
@@ -193,39 +209,91 @@ static void hvx_fast_softmax_f32(const uint8_t * restrict src, uint8_t * restric
 
         sum_vec = Q6_Vqf32_vadd_VsfVsf(Q6_Vsf_equals_Vqf32(sum_vec), v3);
 
-        v_pad[i] = v3;
+        v_dst[i] = v3;
     }
 
-    sum_vec = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_vec)); // replicated over all lanes
+    if (nloe > 0) {
+        HVX_VectorPred q_mask = Q6_Q_vsetq_R(nloe * sizeof(float));
+        HVX_Vector v1     = v_src[nvec];
+        HVX_Vector v2     = Q6_Vqf32_vsub_VsfVsf(v1, max_vec);
+        HVX_Vector v3     = hvx_vec_exp_f32(Q6_Vsf_equals_Vqf32(v2));
+        HVX_Vector v3_pad = Q6_V_vmux_QVV(q_mask, v3, Q6_V_vzero());
 
-    HVX_VectorPred pos_sum   = Q6_Q_vcmp_gt_VwVw(sum_vec, zero_v);
+        sum_vec     = Q6_Vqf32_vadd_VsfVsf(Q6_Vsf_equals_Vqf32(sum_vec), v3_pad);
+        v_dst[nvec] = v3_pad;
+    }
+
+    sum_vec = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_vec));
+
+    HVX_VectorPred pos_sum   = Q6_Q_vcmp_gt_VwVw(sum_vec, Q6_V_vzero());
     HVX_Vector     v4        = hvx_vec_inverse_f32(sum_vec);
-    HVX_Vector     scale_vec = Q6_V_vmux_QVV(pos_sum, v4, one_v);
+    HVX_Vector     scale_vec = Q6_V_vmux_QVV(pos_sum, v4, hvx_vec_splat_f32(1.0f));
 
-    #pragma unroll(4)
-    for (int i = 0; i < step_of_1; i++) {
-        HVX_Vector v1 = v_pad[i];
+    #pragma unroll(2)
+    for (int i = 0; i < nvec; i++) {
+        HVX_Vector v1 = v_dst[i];
         HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_vec);
         v_dst[i]      = Q6_Vsf_equals_Vqf32(v2);
     }
+
+    if (nloe > 0) {
+        HVX_Vector v1 = v_dst[nvec];
+        HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_vec);
+        hvx_vec_store_a(&v_dst[nvec], nloe * sizeof(float), Q6_Vsf_equals_Vqf32(v2));
+    }
 }
 
-static float hvx_softmax_f32(const uint8_t * restrict src, uint8_t * restrict dst, uint8_t * restrict spad, const int  num_elems, const float max) {
-    hvx_sub_scalar_f32(spad, src, max, num_elems);
-
-    hvx_exp_f32(dst, spad, num_elems, false);
-    return hvx_reduce_sum_f32(dst, num_elems);
+static void compute_fast_softmax_f32_nomask(
+    void * restrict dst,
+    const void * restrict src0,
+    const void * restrict mask,
+    uint32_t ne00,
+    float scale,
+    float slope
+) {
+    (void) mask;
+    (void) slope;
+    hvx_scale_f32((uint8_t *) dst, (const uint8_t *) src0, ne00, scale);
+    hvx_fast_softmax_f32((const uint8_t *) dst, (uint8_t *) dst, ne00);
 }
 
-static void softmax_job_f32(unsigned int nth, unsigned int ith, void * data) {
-    struct htp_softmax_context * smctx = (struct htp_softmax_context *) data;
+static void compute_fast_softmax_f32_mask_f32(
+    void * restrict dst,
+    const void * restrict src0,
+    const void * restrict mask,
+    uint32_t ne00,
+    float scale,
+    float slope
+) {
+    hvx_fast_softmax_prep_f32((const uint8_t *) src0, (uint8_t *) dst, ne00, scale, (const uint8_t *) mask, slope);
+    hvx_fast_softmax_f32((const uint8_t *) dst, (uint8_t *) dst, ne00);
+}
+
+static void compute_fast_softmax_f32_mask_f16(
+    void * restrict dst,
+    const void * restrict src0,
+    const void * restrict mask,
+    uint32_t ne00,
+    float scale,
+    float slope
+) {
+    hvx_fast_softmax_prep_f16((const uint8_t *) src0, (uint8_t *) dst, ne00, scale, (const uint8_t *) mask, slope);
+    hvx_fast_softmax_f32((const uint8_t *) dst, (uint8_t *) dst, ne00);
+}
+
+static const softmax_compute_fn_t softmax_kernels[HTP_SOFTMAX_KERNEL_COUNT] = {
+    [HTP_SOFTMAX_KERNEL_NOMASK]   = compute_fast_softmax_f32_nomask,
+    [HTP_SOFTMAX_KERNEL_MASK_F32] = compute_fast_softmax_f32_mask_f32,
+    [HTP_SOFTMAX_KERNEL_MASK_F16] = compute_fast_softmax_f32_mask_f16,
+};
+
+static void softmax_thread_dma(unsigned int nth, unsigned int ith, void * data) {
+    (void) nth;
+    const struct htp_softmax_context * smctx = (const struct htp_softmax_context *) data;
     struct htp_ops_context * octx = smctx->octx;
-
     const struct htp_tensor * src0 = octx->src[0];
-    const struct htp_tensor * src1 = octx->src[1];
     const struct htp_tensor * dst  = octx->dst;
-
-    htp_softmax_preamble3;
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
     const uint32_t src0_nrows            = smctx->nrows;
     const uint32_t src0_nrows_per_thread = smctx->src0_nrows_per_thread;
@@ -233,122 +301,213 @@ static void softmax_job_f32(unsigned int nth, unsigned int ith, void * data) {
     const uint32_t src0_start_row = smctx->row_start + src0_nrows_per_thread * ith;
     const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, smctx->row_start + src0_nrows);
 
-    // no work for this thread
     if (src0_start_row >= src0_end_row) {
         return;
     }
 
-    int is_aligned = 1;
-    int opt_path   = 0;
+    const dma_addr_t data_src0 = smctx->data_src0;
+    const dma_addr_t data_dst  = smctx->data_dst;
 
-    if (!hex_is_aligned((void *) src0->data, VLEN) || !hex_is_aligned((void *) dst->data, VLEN)) {
-        is_aligned = 0;
-        FARF(HIGH, "softmax-f32: unaligned addresses in elementwise op, possibly slower execution\n");
+    const size_t src0_row_size = src0->ne[0] * sizeof(float);
+    const size_t dst_row_size  = src0->ne[0] * sizeof(float);
+
+    uint8_t * src0_vtcm_base = smctx->vtcm_src0 + (ith * smctx->vtcm_src0_size_per_thread);
+    uint8_t * dst_vtcm_base  = smctx->vtcm_dst  + (ith * smctx->vtcm_dst_size_per_thread);
+
+    const size_t src0_vtcm_half = smctx->src0_spad_half_size;
+    const size_t dst_vtcm_half  = smctx->dst_spad_half_size;
+
+    dma_queue * dma_q = octx->ctx->dma[ith];
+
+    for (uint32_t r = src0_start_row, idx = 0; r < src0_end_row && idx < 2; r++, idx++) {
+        dma_addr_t cur_dst  = data_dst  + r * dst_row_size;
+        dma_addr_t cur_src0 = data_src0 + r * src0_row_size;
+        void * d_spad = dst_vtcm_base  + idx * dst_vtcm_half;
+        void * s_spad = src0_vtcm_base + idx * src0_vtcm_half;
+
+        dma_queue_push(dma_q, dma_make_data(cur_dst, d_spad),
+                       dst_row_size, smctx->dst_row_size_aligned, dst_row_size, 0);
+        dma_queue_push(dma_q, dma_make_data(s_spad, cur_src0),
+                       smctx->src0_row_size_aligned, src0_row_size, src0_row_size, 1);
     }
 
-    // Only use the fast path when aligned AND row size is multiple of VLEN (128 bytes)
-    // The fast path (hvx_fast_softmax_f32) doesn't handle tail elements
-    // The non-opt path uses hvx_softmax_f32 which properly handles all sizes via its helper functions
-    if ((1 == is_aligned) && !(nb01 & (VLEN - 1))) {
-        opt_path = 1;
-    }
-
-    uint8_t * src0_spad_data = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread);
-    uint8_t * src1_spad_data = octx->src1_spad.data + (ith * octx->src1_spad.size_per_thread);
-    uint8_t * dst_spad_data  = octx->dst_spad.data  + (ith * octx->dst_spad.size_per_thread);
-
-    float * wp0 = (float *) src0_spad_data;
-    float * wp1 = (float *) src1_spad_data;
-    float * wp2 = (float *) dst_spad_data;
-
-    uint32_t prev_i2 = (uint32_t)-1;
-    float slope = 1.0f;
-
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
-    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, src0_start_row);
+    softmax_compute_fn_t compute = (softmax_compute_fn_t) smctx->compute;
+    const uint32_t ne00 = src0->ne[0];
 
     for (uint32_t r = src0_start_row; r < src0_end_row; ++r) {
-        uint32_t i1 = fastmodulo(r, ne01, &smctx->fastdiv_ne01);
-        uint32_t r_div_ne01 = fastdiv(r, &smctx->fastdiv_ne01);
-        uint32_t i2 = fastmodulo(r_div_ne01, ne02, &smctx->fastdiv_ne02);
-        uint32_t i3 = fastdiv(r_div_ne01, &smctx->fastdiv_ne02);
+        void * d_spad = (void *) dma_queue_pop(dma_q).src;
+        void * s_spad = (void *) dma_queue_pop(dma_q).dst;
 
-        // Map to original logic indices
-        // i01 = i1
-        // i02 = i2
-        // i03 = i3
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r);
+        compute(d_spad, s_spad, NULL, ne00, smctx->scale, 1.0f);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r);
 
-        const uint32_t i11 = i1;
-        // const uint32_t i12 = i2 % ne12;
-        // const uint32_t i13 = i3 % ne13;
+        dma_addr_t cur_dst = data_dst + r * dst_row_size;
+        dma_queue_push(dma_q, dma_make_data(cur_dst, d_spad),
+                       dst_row_size, smctx->dst_row_size_aligned, dst_row_size, 1);
 
-        uint32_t i12, i13;
-        if (ne12 == ne02) {
-             i12 = i2;
-        } else {
-             i12 = fastmodulo(i2, ne12, &smctx->fastdiv_ne12);
-        }
-
-        if (ne13 == ne03) {
-             i13 = i3;
-        } else {
-             i13 = fastmodulo(i3, ne13, &smctx->fastdiv_ne13);
-        }
-
-        // ALiBi
-        if (i2 != prev_i2) {
-            const uint32_t h = i2;  // head
-            slope = (smctx->max_bias > 0.0f) ? h < smctx->n_head_log2 ? powf(smctx->m0, h + 1) : powf(smctx->m1, 2 * (h - smctx->n_head_log2) + 1) : 1.0f;
-            prev_i2 = i2;
-        }
-
-        float * sp = (float *) ((char *) src0->data + i1 * nb01 + i2 * nb02 + i3 * nb03);
-        float * dp = (float *) ((char *) dst->data  + i1 * nb1  + i2 * nb2  + i3 * nb3);
-
-        // broadcast the mask across rows
-        __fp16 * mp_f16 = (smctx->use_src1) ? (__fp16 *) ((char *) src1->data + i11 * nb11 + i12 * nb12 + i13 * nb13) : NULL;
-        float *  mp_f32 = (smctx->use_src1) ? (float *)  ((char *) src1->data + i11 * nb11 + i12 * nb12 + i13 * nb13) : NULL;
-
-        if ((1 == opt_path) && (mp_f32) && !(smctx->use_f16)) {
-            hvx_fast_softmax_prep_f32((const uint8_t *) sp, (uint8_t *) wp0, ne00, smctx->scale, (const uint8_t *) mp_f32, slope);
-            hvx_fast_softmax_f32((const uint8_t *) wp0, (uint8_t *) dp, (uint8_t *) wp1, ne00);
-        } else if (1 == opt_path) {
-            hvx_scale_f32((uint8_t *) wp0, (const uint8_t *) sp, ne00, smctx->scale);
-            apply_mask(wp0, mp_f32, mp_f16, ne00, slope, smctx->use_f16);
-            hvx_fast_softmax_f32((const uint8_t *) wp0, (uint8_t *) dp, (uint8_t *) wp1, ne00);
-        } else {
-            // Non-optimized path: uses HVX helper functions that properly handle all tensor sizes
-            // including non-multiples of 32 (the HVX vector lane count for f32)
-            hvx_scale_f32((uint8_t *) wp0, (const uint8_t *) sp, ne00, smctx->scale);
-            apply_mask(wp0, mp_f32, mp_f16, ne00, slope, smctx->use_f16);
-            float max = hvx_reduce_max_f32((const uint8_t *) wp0, ne00);
-            float sum = hvx_softmax_f32((const uint8_t *) wp0, (uint8_t *) wp2, (uint8_t *) wp1, ne00, max);
-            sum       = sum > 0.0 ? (1.0 / sum) : 1;
-            hvx_scale_f32((uint8_t *) dp, (const uint8_t *) wp2, ne00, sum);
+        const uint32_t next_r = r + 2;
+        if (next_r < src0_end_row) {
+            dma_addr_t next_src0 = data_src0 + next_r * src0_row_size;
+            dma_queue_push(dma_q, dma_make_data(s_spad, next_src0),
+                           smctx->src0_row_size_aligned, src0_row_size, src0_row_size, 1);
         }
     }
 
-    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, src0_start_row);
-
-    FARF(HIGH, "softmax-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u : opt %u f16 %u\n", ith, nth,
-         ne00, ne01, ne02, ne03, src0_start_row, src0_end_row, ne10, ne11, ne12, ne13,
-         ne0, ne1, ne2, ne3, opt_path, smctx->use_f16);
+    dma_queue_flush(dma_q);
 }
 
-static int execute_op_softmax_f32(struct htp_ops_context * octx) {
-    int err = HTP_STATUS_OK;
-
+static void softmax_thread_mask_dma(unsigned int nth, unsigned int ith, void * data) {
+    (void) nth;
+    const struct htp_softmax_context * smctx = (const struct htp_softmax_context *) data;
+    struct htp_ops_context * octx = smctx->octx;
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * src1 = octx->src[1];
     const struct htp_tensor * dst  = octx->dst;
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
+    const uint32_t src0_nrows            = smctx->nrows;
+    const uint32_t src0_nrows_per_thread = smctx->src0_nrows_per_thread;
+
+    const uint32_t src0_start_row = smctx->row_start + src0_nrows_per_thread * ith;
+    const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, smctx->row_start + src0_nrows);
+
+    if (src0_start_row >= src0_end_row) {
+        return;
+    }
+
+    const dma_addr_t data_src0 = smctx->data_src0;
+    const dma_addr_t data_src1 = smctx->data_src1;
+    const dma_addr_t data_dst  = smctx->data_dst;
+
+    const size_t src0_row_size = src0->ne[0] * sizeof(float);
+    const size_t dst_row_size  = src0->ne[0] * sizeof(float);
+    const size_t mask_row_size = smctx->use_f16 ? (src1->ne[0] * sizeof(__fp16)) : (src1->ne[0] * sizeof(float));
+
+    uint8_t * src0_vtcm_base = smctx->vtcm_src0 + (ith * smctx->vtcm_src0_size_per_thread);
+    uint8_t * src1_vtcm_base = smctx->vtcm_src1 + (ith * smctx->vtcm_src1_size_per_thread);
+    uint8_t * dst_vtcm_base  = smctx->vtcm_dst  + (ith * smctx->vtcm_dst_size_per_thread);
+
+    const size_t src0_vtcm_half = smctx->src0_spad_half_size;
+    const size_t src1_vtcm_half = smctx->src1_spad_half_size;
+    const size_t dst_vtcm_half  = smctx->dst_spad_half_size;
+
+    const uint32_t nb11 = src1->nb[1];
+    const uint32_t nb12 = src1->nb[2];
+    const uint32_t nb13 = src1->nb[3];
+
+    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 ne12 = src1->ne[2];
+    const uint32_t ne13 = src1->ne[3];
+
+    const struct fastdiv_values * div_ne01 = &smctx->div_ne01;
+    const struct fastdiv_values * div_ne02 = &smctx->div_ne02;
+    const struct fastdiv_values * div_ne12 = &smctx->div_ne12;
+    const struct fastdiv_values * div_ne13 = &smctx->div_ne13;
+
+    dma_queue * dma_q = octx->ctx->dma[ith];
+
+    for (uint32_t r = src0_start_row, idx = 0; r < src0_end_row && idx < 2; r++, idx++) {
+        dma_addr_t cur_dst  = data_dst  + r * dst_row_size;
+        dma_addr_t cur_src0 = data_src0 + r * src0_row_size;
+
+        uint32_t i1 = fastmodulo(r, ne01, div_ne01);
+        uint32_t r_div_ne01 = fastdiv(r, div_ne01);
+        uint32_t i2 = fastmodulo(r_div_ne01, ne02, div_ne02);
+        uint32_t i3 = fastdiv(r_div_ne01, div_ne02);
+        uint32_t i12 = (ne12 == ne02) ? i2 : fastmodulo(i2, ne12, div_ne12);
+        uint32_t i13 = (ne13 == ne03) ? i3 : fastmodulo(i3, ne13, div_ne13);
+        dma_addr_t cur_src1 = data_src1 + i1 * nb11 + i12 * nb12 + i13 * nb13;
+
+        void * d_spad = dst_vtcm_base  + idx * dst_vtcm_half;
+        void * s_spad = src0_vtcm_base + idx * src0_vtcm_half;
+        void * m_spad = src1_vtcm_base + idx * src1_vtcm_half;
+
+        dma_queue_push(dma_q, dma_make_data(cur_dst, d_spad),
+                       dst_row_size, smctx->dst_row_size_aligned, dst_row_size, 0);
+        dma_queue_push(dma_q, dma_make_data(s_spad, cur_src0),
+                       smctx->src0_row_size_aligned, src0_row_size, src0_row_size, 1);
+        dma_queue_push(dma_q, dma_make_data(m_spad, cur_src1),
+                       smctx->src1_row_size_aligned, mask_row_size, mask_row_size, 1);
+    }
+
+    softmax_compute_fn_t compute = (softmax_compute_fn_t) smctx->compute;
+    const bool has_bias = smctx->max_bias > 0.0f;
+    uint32_t prev_i2 = (uint32_t)-1;
+    float slope = 1.0f;
+
+    for (uint32_t r = src0_start_row; r < src0_end_row; ++r) {
+        void * d_spad = (void *) (uintptr_t) dma_queue_pop(dma_q).src;
+        void * s_spad = (void *) (uintptr_t) dma_queue_pop(dma_q).dst;
+        void * m_spad = (void *) (uintptr_t) dma_queue_pop(dma_q).dst;
+
+        if (has_bias) {
+            uint32_t r_div_ne01 = fastdiv(r, div_ne01);
+            uint32_t i2 = fastmodulo(r_div_ne01, ne02, div_ne02);
+            if (i2 != prev_i2) {
+                slope = smctx->slopes[i2];
+                prev_i2 = i2;
+            }
+        }
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r);
+        compute(d_spad, s_spad, m_spad, ne00, smctx->scale, slope);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r);
+
+        dma_addr_t cur_dst = data_dst + r * dst_row_size;
+        dma_queue_push(dma_q, dma_make_data(cur_dst, d_spad),
+                       dst_row_size, smctx->dst_row_size_aligned, dst_row_size, 1);
+
+        const uint32_t next_r = r + 2;
+        if (next_r < src0_end_row) {
+            dma_addr_t next_src0 = data_src0 + next_r * src0_row_size;
+
+            uint32_t ni1 = fastmodulo(next_r, ne01, div_ne01);
+            uint32_t nr_div_ne01 = fastdiv(next_r, div_ne01);
+            uint32_t ni2 = fastmodulo(nr_div_ne01, ne02, div_ne02);
+            uint32_t ni3 = fastdiv(nr_div_ne01, div_ne02);
+            uint32_t ni12 = (ne12 == ne02) ? ni2 : fastmodulo(ni2, ne12, div_ne12);
+            uint32_t ni13 = (ne13 == ne03) ? ni3 : fastmodulo(ni3, ne13, div_ne13);
+            dma_addr_t next_src1 = data_src1 + ni1 * nb11 + ni12 * nb12 + ni13 * nb13;
+
+            dma_queue_push(dma_q, dma_make_data(s_spad, next_src0),
+                           smctx->src0_row_size_aligned, src0_row_size, src0_row_size, 1);
+            dma_queue_push(dma_q, dma_make_data(m_spad, next_src1),
+                           smctx->src1_row_size_aligned, mask_row_size, mask_row_size, 1);
+        }
+    }
+
+    dma_queue_flush(dma_q);
+}
+
+static int execute_op_softmax_f32(struct htp_ops_context * octx) {
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * dst  = octx->dst;
 
-    struct htp_softmax_context smctx;
     const char * op_type = "softmax-f32";
 
-    init_softmax_ctx(&smctx, octx);
+    const struct htp_softmax_kernel_params * kparams =
+        (const struct htp_softmax_kernel_params *) octx->kernel_params;
+
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    if (kparams->kernel_id >= HTP_SOFTMAX_KERNEL_COUNT) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    if (octx->ctx->vtcm_size < (size_t) kparams->vtcm_size) {
+        FARF(ERROR, "%s : current VTCM reservation %zu is too small, needed %u\n",
+             op_type, octx->ctx->vtcm_size, kparams->vtcm_size);
+        return HTP_STATUS_VTCM_TOO_SMALL;
+    }
 
     const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
-    const size_t elem_size = sizeof(float);
+    const size_t elem_size    = sizeof(float);
     const size_t dst_row_size = dst->nb[1];
 
     uint32_t row_start = 0;
@@ -357,9 +516,13 @@ static int execute_op_softmax_f32(struct htp_ops_context * octx) {
     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);
+        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 < octx->n_threads) {
+            htp_ops_context_set_n_threads(octx, nrows ? nrows : 1);
+        }
     }
 
     if (nrows == 0) {
@@ -367,50 +530,71 @@ static int execute_op_softmax_f32(struct htp_ops_context * octx) {
     }
 
     const uint32_t n_threads = octx->n_threads;
+    uint8_t * const vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
 
-    smctx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
-    smctx.row_start             = row_start;
-    smctx.nrows                 = nrows;
+    const uint32_t off_src0 = 0;
+    const uint32_t off_dst  = off_src0 + kparams->vtcm_src0_size_per_thread * kparams->n_threads;
+    const uint32_t off_src1 = off_dst  + kparams->vtcm_dst_size_per_thread  * kparams->n_threads;
 
-    const size_t src0_row_size = src0->nb[1];
-    const size_t src1_row_size = src0_row_size;
+    struct htp_softmax_context smctx = {
+        .octx                  = octx,
+        .kparams               = kparams,
+        .compute               = (void *) softmax_kernels[kparams->kernel_id],
 
-    // VTCM scratchpads for all tensors
-    // 4 rows per thread, padded to HVX vector size
-    octx->src0_spad.size_per_thread = hex_round_up(4 * src0_row_size, 128);
-    octx->src1_spad.size_per_thread = hex_round_up(4 * src1_row_size, 128);
-    octx->dst_spad.size_per_thread  = hex_round_up(4 * dst_row_size, 128);
+        .data_src0             = src0->data,
+        .data_src1             = kparams->use_src1 ? octx->src[1]->data : 0,
+        .data_dst              = dst->data,
 
-    octx->src0_spad.size = octx->src0_spad.size_per_thread * n_threads;
-    octx->src1_spad.size = octx->src1_spad.size_per_thread * n_threads;
-    octx->dst_spad.size  = octx->dst_spad.size_per_thread  * n_threads;
+        .vtcm_src0             = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, off_src0),
+        .vtcm_dst              = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, off_dst),
+        .vtcm_src1             = VTCM_LAYOUT_PTR_OPTIONAL(uint8_t, vtcm_base, off_src1, kparams->use_src1),
 
-    size_t spad_size = octx->src0_spad.size + octx->src1_spad.size + octx->dst_spad.size;
+        .vtcm_src0_size_per_thread = kparams->vtcm_src0_size_per_thread,
+        .vtcm_src1_size_per_thread = kparams->vtcm_src1_size_per_thread,
+        .vtcm_dst_size_per_thread  = kparams->vtcm_dst_size_per_thread,
 
-    if (src1) {
-        FARF(HIGH, "%s: %ux%ux%ux%u x %ux%ux%ux%u -> %ux%ux%ux%u : src0-spad-size %u src1-spad-size %u dst-spad-size %u\n",
-             op_type, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2],
-             src1->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], octx->src0_spad.size, octx->src1_spad.size,
-             octx->dst_spad.size);
-    } else {
-        FARF(HIGH, "%s: %ux%ux%ux%u -> %ux%ux%ux%u : src0-spad-size %u src1-spad-size %u dst-spad-size %u\n", op_type,
-             src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-             octx->src0_spad.size, octx->src1_spad.size, octx->dst_spad.size);
+        .src0_spad_half_size   = kparams->src0_spad_half_size,
+        .src1_spad_half_size   = kparams->src1_spad_half_size,
+        .dst_spad_half_size    = kparams->dst_spad_half_size,
+
+        .src0_row_size_aligned = kparams->src0_row_size_aligned,
+        .src1_row_size_aligned = kparams->src1_row_size_aligned,
+        .dst_row_size_aligned  = kparams->dst_row_size_aligned,
+
+        .use_f16               = kparams->use_f16 != 0,
+
+        .n_head                = kparams->n_head,
+        .n_head_log2           = kparams->n_head_log2,
+
+        .scale                 = kparams->scale,
+        .max_bias              = kparams->max_bias,
+        .m0                    = kparams->m0,
+        .m1                    = kparams->m1,
+
+        .div_ne01              = kparams->div_ne01,
+        .div_ne02              = kparams->div_ne02,
+        .div_ne12              = kparams->div_ne12,
+        .div_ne13              = kparams->div_ne13,
+
+        .src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
+        .row_start             = row_start,
+        .nrows                 = nrows,
+    };
+
+    if (kparams->max_bias > 0.0f && kparams->use_src1) {
+        if (kparams->n_head > 512) {
+            return HTP_STATUS_INVAL_PARAMS;
+        }
+        for (uint32_t h = 0; h < kparams->n_head; h += 32) {
+            HVX_Vector v_slopes = hvx_alibi_slopes(h, 1, kparams->n_head_log2, kparams->m0, kparams->m1);
+            hvx_vmem(&smctx.slopes[h]) = v_slopes;
+        }
     }
 
-    // Make sure the reserved vtcm size is sufficient
-    if (octx->ctx->vtcm_size < spad_size) {
-        FARF(ERROR, "%s : current VTCM reservation %zu is too small, needed %zu\n", op_type, octx->ctx->vtcm_size, spad_size);
-        return HTP_STATUS_VTCM_TOO_SMALL;
-    }
+    work_queue_func_t task_func = kparams->use_src1 ? softmax_thread_mask_dma : softmax_thread_dma;
+    work_queue_run(octx->ctx->work_queue, task_func, &smctx, n_threads);
 
-    octx->src0_spad.data = octx->ctx->vtcm_base;                        octx->src0_spad.src = NULL;
-    octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size; octx->src1_spad.src = NULL;
-    octx->dst_spad.data  = octx->src1_spad.data + octx->src1_spad.size; octx->dst_spad.src  = NULL;
-
-    work_queue_run(octx->ctx->work_queue, softmax_job_f32, &smctx, n_threads);
-
-    return err;
+    return HTP_STATUS_OK;
 }
 
 int op_softmax(struct htp_ops_context * octx) {

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

@@ -0,0 +1,106 @@
+#ifndef HTP_SOFTMAX_OPS_H
+#define HTP_SOFTMAX_OPS_H
+
+#include <stdint.h>
+#include <stddef.h>
+#include <stdbool.h>
+#include <math.h>
+#include "hex-fastdiv.h"
+#include "hex-common.h"
+
+enum htp_softmax_kernel_id {
+    HTP_SOFTMAX_KERNEL_NOMASK = 0,
+    HTP_SOFTMAX_KERNEL_MASK_F32,
+    HTP_SOFTMAX_KERNEL_MASK_F16,
+    HTP_SOFTMAX_KERNEL_COUNT,
+};
+
+struct htp_softmax_kernel_params {
+    uint32_t n_threads;
+    uint32_t src0_nrows;
+    uint32_t src0_nrows_per_thread;
+    uint32_t vtcm_size;
+
+    uint32_t vtcm_src0_size_per_thread;
+    uint32_t vtcm_src1_size_per_thread;
+    uint32_t vtcm_dst_size_per_thread;
+
+    uint32_t src0_row_size_aligned;
+    uint32_t src1_row_size_aligned;
+    uint32_t dst_row_size_aligned;
+
+    uint32_t src0_spad_half_size;
+    uint32_t src1_spad_half_size;
+    uint32_t dst_spad_half_size;
+
+    uint32_t n_head;
+    uint32_t n_head_log2;
+    uint32_t use_src1;
+    uint32_t use_f16;
+    uint32_t kernel_id;
+
+    float    scale;
+    float    max_bias;
+    float    m0;
+    float    m1;
+
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne02;
+    struct fastdiv_values div_ne12;
+    struct fastdiv_values div_ne13;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_softmax_kernel_params) <= 128, "htp_softmax_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_softmax_kernel_params) <= 128, "htp_softmax_kernel_params is too large for kernel_params blob");
+#endif
+
+struct htp_softmax_vtcm_layout {
+    size_t total_bytes;
+    size_t off_src0;
+    size_t off_dst;
+    size_t off_src1;
+
+    size_t src0_bytes_per_thread;
+    size_t dst_bytes_per_thread;
+    size_t src1_bytes_per_thread;
+
+    size_t src0_spad_half_size;
+    size_t dst_spad_half_size;
+    size_t src1_spad_half_size;
+};
+
+static inline void htp_softmax_vtcm_layout_build(
+    struct htp_softmax_vtcm_layout * layout,
+    uint32_t ne00,
+    uint32_t ne10,
+    bool use_src1,
+    bool use_f16,
+    uint32_t n_threads
+) {
+    size_t src0_row_size = ne00 * sizeof(float);
+    size_t dst_row_size  = ne00 * sizeof(float);
+    size_t src1_row_size = use_src1 ? (ne10 * (use_f16 ? 2 : 4)) : 0;
+
+    size_t src0_row_size_aligned = hex_round_up(src0_row_size, 128);
+    size_t dst_row_size_aligned  = hex_round_up(dst_row_size,  128);
+    size_t src1_row_size_aligned = use_src1 ? hex_round_up(src1_row_size, 128) : 0;
+
+    layout->src0_spad_half_size = src0_row_size_aligned;
+    layout->dst_spad_half_size  = dst_row_size_aligned;
+    layout->src1_spad_half_size = src1_row_size_aligned;
+
+    // Double buffering: 2 half-buffers per thread
+    layout->src0_bytes_per_thread = src0_row_size_aligned * 2;
+    layout->dst_bytes_per_thread  = dst_row_size_aligned  * 2;
+    layout->src1_bytes_per_thread = src1_row_size_aligned * 2;
+
+    layout->off_src0 = 0;
+    layout->off_dst  = layout->off_src0 + layout->src0_bytes_per_thread * n_threads;
+    layout->off_src1 = layout->off_dst  + layout->dst_bytes_per_thread  * n_threads;
+
+    layout->total_bytes = layout->off_src1 + layout->src1_bytes_per_thread * n_threads;
+}
+
+#endif // HTP_SOFTMAX_OPS_H