CANN: missing const qualifiers added for consistency (#30933 )

sycl : pass device into mul_mat_uses_onednn (#31014 )
Extend probabilistic sampling for DFlash2 drafts (#30663 )
2026-10-23 05:03:45 -05:00 · 2026-10-23 15:34:20 +08:00 · 2026-10-23 08:37:53 +02:00 · 2026-10-23 11:15:36 +05:30 · 2026-10-23 07:31:30 +02:00 · 2026-10-22 21:05:44 -07:00
343 changed files with 32263 additions and 17108 deletions
@@ -5,8 +5,8 @@ exclude =
    # Do not traverse examples and tools
    examples,
    tools,
-    # Reference scripts for skills
-    skills,
+    # Do not traverse .venv directories
+    .venv*,
    # Do not include package initializers
    __init__.py,
    # No need to traverse our git directory
@@ -35,7 +35,7 @@ jobs:
  android-ndk-snapdragon:
    runs-on: ubuntu-24.04  # previously ubuntu-latest
    container:
-      image: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.7'
+      image: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.8'
    defaults:
      run:
        shell: bash
@@ -69,7 +69,7 @@ jobs:
  linux-iot-snapdragon:
    runs-on: ubuntu-24.04  # previously ubuntu-latest
    container:
-      image: 'ghcr.io/snapdragon-toolchain/arm64-linux:v0.7'
+      image: 'ghcr.io/snapdragon-toolchain/arm64-linux:v0.8'
    defaults:
      run:
        shell: bash
@@ -0,0 +1,142 @@
+name: CI (wasi)
+
+on:
+  workflow_dispatch: # allows manual triggering
+  push:
+    branches:
+      - master
+    paths: [
+      '.github/workflows/build-wasi.yml',
+      '**/CMakeLists.txt',
+      '**/.cmake',
+      '**/*.h',
+      '**/*.hpp',
+      '**/*.c',
+      '**/*.cpp',
+      '**/*.wgsl',
+      '**/*.tmpl',
+      'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py'
+    ]
+
+  pull_request:
+    types: [opened, synchronize, reopened]
+    paths: [
+      '.github/workflows/build-wasi.yml',
+      '**/CMakeLists.txt',
+      '**/.cmake',
+      '**/*.h',
+      '**/*.hpp',
+      '**/*.c',
+      '**/*.cpp',
+      '**/*.wgsl',
+      '**/*.tmpl',
+      'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py'
+    ]
+
+cache-mode: none
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
+  cancel-in-progress: true
+
+env:
+  WASI_SDK_VERSION: 34
+  DAWN_VERSION: v20260908.214631
+  DAWN_COMMIT: 94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655
+  WASI_WEBGPU_HEADERS_COMMIT: 917765edd7cb8f8cf95bda1217514c78f71b582b
+
+jobs:
+  ubuntu-webgpu:
+    runs-on: ubuntu-24.04-arm
+
+    steps:
+      - name: Clone
+        id: checkout
+        uses: actions/checkout@v6
+
+      - name: ccache
+        uses: ggml-org/ccache-action@v1.2.24
+        with:
+          restore: false
+          save: false
+
+      - name: Install wasi-sdk
+        uses: bytecodealliance/setup-wasi-sdk-action@77d09f3df1e8da344d8403cf8631af8977938b4c # v1.0.0
+        with:
+          version: ${{ env.WASI_SDK_VERSION }}
+          add-to-path: false
+
+      - name: wasi:webgpu Dependency
+        run: |
+          mkdir -p wasi-webgpu/dawn-headers
+          curl -fL "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/dawn-headers-${DAWN_COMMIT}.tar.gz" |
+            tar -xz -C wasi-webgpu/dawn-headers --strip-components=1
+
+          git init -q wasi-webgpu/wasi-webgpu-headers
+          cd wasi-webgpu/wasi-webgpu-headers
+          git fetch -q --depth 1 https://github.com/wasi-gfx/wasi-webgpu-headers.git ${WASI_WEBGPU_HEADERS_COMMIT}
+          git checkout -q FETCH_HEAD
+          git submodule update -q --init --depth 1
+          cd ../..
+
+          # ggml-webgpu does find_package(Dawn) and links dawn::webgpu_dawn. There is no Dawn on WASI,
+          # so this target instead compiles wasi-webgpu-headers (webgpu.h over wasi:webgpu imports)
+          # and exposes Dawn's webgpu_cpp.h on top of it.
+          cat > wasi-webgpu/DawnConfig.cmake << 'EOF'
+          set(_wasi_webgpu "${CMAKE_CURRENT_LIST_DIR}/wasi-webgpu-headers")
+          add_library(dawn::webgpu_dawn INTERFACE IMPORTED)
+          set_target_properties(dawn::webgpu_dawn PROPERTIES
+              INTERFACE_SOURCES "${_wasi_webgpu}/webgpu.c;${_wasi_webgpu}/async_futures.c;${_wasi_webgpu}/imports.c;${_wasi_webgpu}/imports_component_type.o"
+              INTERFACE_INCLUDE_DIRECTORIES "${CMAKE_CURRENT_LIST_DIR}/dawn-headers/include/webgpu_upstream;${_wasi_webgpu}")
+          EOF
+
+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
+        with:
+          key: webgpu-ubuntu-24.04-arm-wasi
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache
+
+      # llama-common does not build for WASI, so the executables are linked by hand
+      - name: Build
+        id: cmake_build
+        run: |
+          WASI_FLAGS="-msimd128 -D_WASI_EMULATED_SIGNAL -fwasm-exceptions -mllvm -wasm-use-legacy-eh=false"
+          cmake -B build-wasi \
+            -G "Ninja" \
+            -DCMAKE_TOOLCHAIN_FILE=${WASI_SDK_PATH}/share/cmake/wasi-sdk-p3.cmake \
+            -DCMAKE_BUILD_TYPE=Release \
+            -DCMAKE_C_FLAGS="${WASI_FLAGS}" \
+            -DCMAKE_CXX_FLAGS="${WASI_FLAGS}" \
+            -DCMAKE_HAVE_LIBC_PTHREAD=1 \
+            -DBUILD_SHARED_LIBS=OFF \
+            -DGGML_NATIVE=OFF \
+            -DGGML_OPENMP=OFF \
+            -DGGML_WEBGPU=ON \
+            -DDawn_DIR=$PWD/wasi-webgpu \
+            -DLLAMA_BUILD_COMMON=OFF
+
+          time cmake --build build-wasi --config Release --target llama -j $(nproc)
+
+          LIBS="build-wasi/ggml/src/libggml.a build-wasi/ggml/src/ggml-webgpu/libggml-webgpu.a build-wasi/ggml/src/libggml-cpu.a build-wasi/ggml/src/libggml-base.a"
+          LDFLAGS="-lunwind -ldl -Wl,-z,stack-size=8388608"
+          CXX="${WASI_SDK_PATH}/bin/clang++ --target=wasm32-wasip3 -O3 -std=c++17 ${WASI_FLAGS} -I include -I ggml/include"
+
+          $CXX examples/simple/simple.cpp build-wasi/src/libllama.a $LIBS $LDFLAGS -o build-wasi/llama-simple.wasm
+          $CXX tests/test-backend-ops.cpp $LIBS $LDFLAGS -o build-wasi/test-backend-ops.wasm
+
+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: webgpu-ubuntu-24.04-arm-wasi
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true
@@ -7,15 +7,7 @@ on:
      - master
    paths: [
      '.github/workflows/build-webgpu.yml',
-      '**/CMakeLists.txt',
-      '**/.cmake',
-      '**/*.h',
-      '**/*.hpp',
-      '**/*.c',
-      '**/*.cpp',
-      '**/*.wgsl',
-      '**/*.tmpl',
-      'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py'
+      'ggml/src/ggml-webgpu/**'
    ]

  pull_request:
@@ -63,135 +55,3 @@ jobs:
            -type f \( -name '*.cpp' -o -name '*.hpp' -o -name '*.h' \) \
            -print0 |
            xargs -0 clang-format-22 --dry-run --Werror
-
-  macos:
-    runs-on: macos-latest
-
-    steps:
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      - name: ccache
-        uses: ggml-org/ccache-action@v1.2.24
-        with:
-          restore: false
-          save: false
-
-      - name: Dawn Dependency
-        id: dawn-depends
-        run: |
-          DAWN_VERSION="v20260908.214631"
-          DAWN_OWNER="google"
-          DAWN_REPO="dawn"
-          DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-macos-latest-Release"
-          echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
-          curl -L -o artifact.tar.gz \
-            "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
-          mkdir dawn
-          tar -xvf artifact.tar.gz -C dawn --strip-components=1
-
-      - name: ccache-buckets-restore
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
-        with:
-          key: webgpu-macos-latest
-          folder: llama.cpp
-          hf_bucket: ggml-org/cache
-
-      - name: Build
-        id: cmake_build
-        run: |
-          export CMAKE_PREFIX_PATH=dawn
-          cmake -B build -G "Ninja" -DCMAKE_BUILD_TYPE=Release -DGGML_WEBGPU=ON -DGGML_METAL=OFF -DGGML_BLAS=OFF
-          time cmake --build build --config Release -j $(sysctl -n hw.logicalcpu)
-
-      - name: ccache-buckets-save
-        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
-        with:
-          key: webgpu-macos-latest
-          folder: llama.cpp
-          evict-old-files: 1d
-          hf_bucket: ggml-org/cache
-          save: true
-
-      - name: Test
-        id: cmake_test
-        run: |
-          cd build
-          ctest -L main --verbose --timeout 900
-
-  ubuntu:
-    runs-on: ubuntu-24.04
-
-    steps:
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      - name: ccache
-        uses: ggml-org/ccache-action@v1.2.24
-        with:
-          restore: false
-          save: false
-
-      - name: Dependencies
-        id: depends
-        run: |
-          sudo add-apt-repository -y ppa:kisak/kisak-mesa
-          sudo apt-get update -y
-          sudo apt-get install -y build-essential mesa-vulkan-drivers \
-            libxcb-xinput0 libxcb-xinerama0 libxcb-cursor-dev libssl-dev
-
-      - name: Dawn Dependency
-        id: dawn-depends
-        run: |
-          sudo apt-get install -y libxrandr-dev libxinerama-dev libxcursor-dev mesa-common-dev libx11-xcb-dev libxi-dev
-          DAWN_VERSION="v20260908.214631"
-          DAWN_OWNER="google"
-          DAWN_REPO="dawn"
-          DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-ubuntu-latest-Release"
-          echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
-          curl -L -o artifact.tar.gz \
-            "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
-          mkdir dawn
-          tar -xvf artifact.tar.gz -C dawn --strip-components=1
-
-      - name: ccache-buckets-restore
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
-        with:
-          key: webgpu-ubuntu-24.04
-          folder: llama.cpp
-          hf_bucket: ggml-org/cache
-
-      - name: Build
-        id: cmake_build
-        run: |
-          export Dawn_DIR=dawn/lib64/cmake/Dawn
-          cmake -B build \
-            -DGGML_WEBGPU=ON
-          time cmake --build build --config Release -j $(nproc)
-
-      - name: ccache-buckets-save
-        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
-        with:
-          key: webgpu-ubuntu-24.04
-          folder: llama.cpp
-          evict-old-files: 1d
-          hf_bucket: ggml-org/cache
-          save: true
-
-      - name: Test
-        id: cmake_test
-        run: |
-          cd build
-          ctest -L main --verbose --timeout 900
@@ -35,7 +35,7 @@ permissions:
 env:
  GH_TOKEN: ${{ github.token }}
  BRANCH_NAME: ${{ github.head_ref || github.ref_name }}
-  CMAKE_ARGS: "-DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_TOOLS=ON -DLLAMA_BUILD_SERVER=ON -DGGML_RPC=ON"
+  CMAKE_ARGS: "-DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_TOOLS=ON -DLLAMA_BUILD_SERVER=ON -DGGML_RPC=ON -DLLAMA_CONNECT=ON"

 # note: run this workflow one at a time for better cache reuse
 concurrency:
@@ -547,7 +547,7 @@ jobs:
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    runs-on: ubuntu-latest
-    container: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.7'
+    container: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.8'

    defaults:
      run:
@@ -600,7 +600,7 @@ jobs:
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    runs-on: ubuntu-latest
-    container: 'ghcr.io/snapdragon-toolchain/arm64-linux:v0.7'
+    container: 'ghcr.io/snapdragon-toolchain/arm64-linux:v0.8'

    defaults:
      run:
@@ -1478,6 +1478,7 @@ jobs:
            -DCMAKE_INSTALL_RPATH='$ORIGIN' \
            -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
            -DGGML_NATIVE=OFF \
+            -DLLAMA_CONNECT=ON \
            -DGGML_SYCL_F16=${{ matrix.fp16 }}
          time cmake --build build --config Release -j $(nproc)

@@ -136,6 +136,7 @@ option(LLAMA_BUILD_SERVER    "llama: build server example"
 option(LLAMA_BUILD_APP       "llama: build the unified binary"                                                   ${LLAMA_STANDALONE})
 option(LLAMA_BUILD_UI        "llama: build the embedded Web UI for server"                                       OFF)
 option(LLAMA_USE_PREBUILT_UI "llama: use prebuilt UI from HF Bucket when available" ON)
+option(LLAMA_CONNECT         "llama: include the llama-connect tool (downloads a prebuilt binary)"               OFF)

 option(LLAMA_TOOLS_INSTALL "llama: install tools" ${LLAMA_TOOLS_INSTALL_DEFAULT})
 option(LLAMA_TESTS_INSTALL "llama: install tests" ON)
@@ -447,10 +447,14 @@ function gg_run_qwen3_0_6b {

    (time ./bin/llama-imatrix --model ${model_f16} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-imatrix.log

-    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa off --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
-    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa on  --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
-    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa off                ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
-    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa on                 ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+    # TODO: these tests are extremely slow on WebGPU - re-enable when fixed
+    #       https://github.com/ggml-org/llama.cpp/actions/runs/39639985732
+    if [ -z ${GG_BUILD_WEBGPU} ]; then
+        (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa off --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+        (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa on  --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+        (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa off                ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+        (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa on                 ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+    fi

    function check_ppl {
        qnt="$1"
@@ -2,7 +2,7 @@ include("ggml/cmake/common.cmake")

 function(llama_add_compile_flags)
    if (LLAMA_FATAL_WARNINGS)
-        if (CMAKE_CXX_COMPILER_ID MATCHES "GNU" OR CMAKE_CXX_COMPILER_ID MATCHES "Clang")
+        if (CMAKE_CXX_COMPILER_ID MATCHES "GNU" OR CMAKE_CXX_COMPILER_ID MATCHES "Clang" OR CMAKE_CXX_COMPILER_ID MATCHES "IntelLLVM")
            list(APPEND C_FLAGS   -Werror)
            list(APPEND CXX_FLAGS -Werror)
        elseif (CMAKE_CXX_COMPILER_ID STREQUAL "MSVC")
@@ -3494,6 +3494,26 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
            params.ui = value;
        }
    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_UI"));
+    add_opt(common_arg(
+        {"--connect"},
+        string_format("(experimental) open a peer-to-peer tunnel for server using llama-connect (default: %s)", params.server_connect ? "enabled" : "disabled"),
+        [](common_params & params) {
+            params.server_connect = true;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CONNECT"));
+    add_opt(common_arg(
+        {"--connect-code"}, "CODE",
+        "manually specify a 40-character code for --connect (default: generate a new one for each run)",
+        [](common_params & params, const std::string & value) {
+            std::string val = value;
+            string_replace_all(val, "-", "");
+            string_replace_all(val, " ", "");
+            if (val.size() != 40) {
+                throw std::invalid_argument(string_format("error: invalid connect code '%s', must be 40 characters\n", value.c_str()));
+            }
+            params.server_connect_code = val;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CONNECT_CODE"));
    add_opt(common_arg(
        {"--embedding", "--embeddings"},
        string_format("restrict to only support embedding use case; use only with dedicated embedding models (default: %s)", params.embedding ? "enabled" : "disabled"),
@@ -1175,6 +1175,7 @@ static const std::map<common_decision_type, std::string> COMMON_DECISION_TYPE_NA
    { COMMON_DECISION_TYPE_LFM2_D1,        "lfm2-d1"       },
    { COMMON_DECISION_TYPE_LFM2_D1_OMNI,   "lfm2-d1-omni"  },
    { COMMON_DECISION_TYPE_GRANITE_SWITCH, "granite-switch"},
+    { COMMON_DECISION_TYPE_H2O_LIGHTNING,  "h2o-lightning" },
    // for future models, ALWAYS add org name prefix to avoid name collisions, e.g. "orgname-modelname"
 };

@@ -1234,6 +1235,9 @@ common_gguf_info common_get_gguf_info(const std::string & fname) {
        info.n_ctx_train = gguf_get_val_u32(gguf_ctx.get(), ctx_id);
    }

+    // same key as llama_model_chat_template(model, nullptr)
+    info.has_chat_tmpl = gguf_find_key(gguf_ctx.get(), "tokenizer.chat_template") >= 0;
+
    return info;
 }

@@ -642,6 +642,10 @@ struct common_params {
    int32_t cache_ram_mib       = 8192;  // -1 = no limit, 0 - disable, 1 = 1 MiB, etc.
    bool    cache_ram_explicit  = false; // set by --cache-ram, decision models disable the cache otherwise

+    // llama-connect params
+    bool        server_connect      = false;
+    std::string server_connect_code = "";
+
    std::string public_path   = "";                                                                         // NOLINT
    std::string api_prefix    = "";                                                                         // NOLINT
    std::string chat_template = "";                                                                         // NOLINT
@@ -971,6 +975,7 @@ enum common_decision_type {
    COMMON_DECISION_TYPE_LFM2_D1, // same as openjev, the labels depend on the question type
    COMMON_DECISION_TYPE_LFM2_D1_OMNI, // same as laya, other prompt layout
    COMMON_DECISION_TYPE_GRANITE_SWITCH, // same as openjev, run by the decision aLoRA of granite switch
+    COMMON_DECISION_TYPE_H2O_LIGHTNING, // same as openjev, the labels are read after an answer prefill
    COMMON_DECISION_TYPE_UNKNOWN, // a decision model of a type that is not supported
 };

@@ -980,6 +985,7 @@ common_decision_type common_get_decision_type(const struct llama_model * model);
 struct common_gguf_info {
    common_decision_type decision_type = COMMON_DECISION_TYPE_UNKNOWN; // UNKNOWN if the file is missing, unreadable, or invalid
    uint32_t             n_ctx_train   = 0;                            // 0 if unknown
+    bool                 has_chat_tmpl = false;                        // has the default tokenizer.chat_template
 };

 common_gguf_info common_get_gguf_info(const std::string & fname);
@@ -1,7 +1,9 @@
 #include "lexer.h"
 #include "runtime.h"
+#include "unicode.h"

 #include <cctype>
+#include <charconv>
 #include <functional>
 #include <map>
 #include <string>
@@ -75,6 +77,24 @@ lexer_result lexer::tokenize(const std::string & source) {
                }
                // add escaped char
                char escaped_char = src[pos++];
+                if (escaped_char == 'x') {
+                    if (pos + 1 >= src.size()) {
+                        throw lexer_exception("incomplete hexadecimal escape", source, pos);
+                    }
+
+                    unsigned int value = 0;
+                    const char * first = &src[pos];
+                    const char * last  = first + 2;
+
+                    auto [ptr, ec] = std::from_chars(first, last, value, 16);
+                    if (ec != std::errc{} || ptr != last) {
+                        throw lexer_exception("invalid hexadecimal escape", source, pos);
+                    }
+
+                    str += common_unicode_cpt_to_utf8(value);
+                    pos += 2;
+                    continue;
+                }
                if (escape_chars.find(escaped_char) == escape_chars.end()) {
                    throw lexer_exception(std::string("unknown escape character \\") + escaped_char, source, pos);
                }
@@ -425,7 +425,12 @@ const func_builtins & global_builtins() {
            throw raised_exception("Jinja Exception: " + msg);
        }},
        {"dict", toobject},
-        {"namespace", toobject},
+        {"namespace", [](const func_args & args) -> value {
+            // unlike a dict, a namespace has no methods: ns.items is the attribute, as in jinja2
+            value out = toobject(args);
+            cast_val<value_object>(out)->has_builtins = false;
+            return out;
+        }},
        {"strftime_now", [](const func_args & args) -> value {
            args.ensure_vals<value_string>();
            std::string format = args.get_pos(0)->as_string().str();
@@ -1028,10 +1033,6 @@ const func_builtins & value_bool_t::get_builtins() const {
    return builtins;
 }

-[[noreturn]] static value array_unique_not_implemented(const func_args &) {
-    throw not_implemented_exception("Array unique builtin not implemented");
-}
-
 const func_builtins & value_array_t::get_builtins() const {
    static const func_builtins builtins = {
        {"default", default_value},
@@ -1257,7 +1258,36 @@ const func_builtins & value_array_t::get_builtins() const {
            }
            return result;
        }},
-        {"unique", array_unique_not_implemented},
+        {"unique", [](const func_args & args) -> value {
+            args.ensure_count(1, 3);
+            args.ensure_vals<value_array>();
+            value val       = args.get_pos(0);
+            value val_case  = args.get_kwarg_or_pos("case_sensitive", 1);
+            value attribute = args.get_kwarg_or_pos("attribute",      2);
+            // FIXME: unique is currently always case sensitive
+            //const bool case_sensitive = val_case->as_bool(); // undefined == false
+            value undef = mk_val<value_undefined>();
+            std::vector<value> result;
+            std::vector<value> seen;
+            for (const auto & item : val->as_array()) {
+                value key = item;
+                if (!attribute->is_undefined()) {
+                    key = get_attribute(item, attribute, undef);
+                }
+                bool dup = false;
+                for (const auto & s : seen) {
+                    if (value_compare(s, key, value_compare_op::eq)) {
+                        dup = true;
+                        break;
+                    }
+                }
+                if (!dup) {
+                    seen.push_back(key);
+                    result.push_back(item);
+                }
+            }
+            return mk_val<value_array>(std::move(result));
+        }},
    };
    return builtins;
 }
@@ -16,8 +16,10 @@
 #include <cassert>
 #include <cmath>
 #include <cstring>
+#include <fstream>
 #include <iomanip>
 #include <map>
+#include <random>
 #include <cinttypes>

 #define SPC_DBG(fmt, ...) LOG_DBG("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)
@@ -1009,6 +1011,10 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
    bool    is_mrope       = false;
    int32_t selector_top_k = 0;

+    // DFlash2 probabilistic drafting: one generator per sequence and the seed it was last seeded with (-1: reseed)
+    std::vector<std::mt19937> dflash2_rng;
+    std::vector<int64_t>      dflash2_seed;
+
    // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout
    const bool is_dspark;

@@ -1059,6 +1065,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {

        selector_top_k = llama_model_dflash_selector_top_k(model_dft);
        is_dflash2     = selector_top_k > 0;
+        dflash2_rng.resize(n_seq);
+        dflash2_seed.assign(n_seq, -1);
        char lilicorr_meta[32] = {};
        lilicorr_top_k = llama_model_meta_val_str(model_dft, "lilicorr.candidate_topk", lilicorr_meta, sizeof(lilicorr_meta)) >= 0
                ? std::atoi(lilicorr_meta) : 0;
@@ -1157,6 +1165,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
            return;
        }

+        // reseed per request rather than per round, or two identical requests differ
+        dflash2_seed[seq_id] = -1;
+
        const int32_t N = (int32_t) prompt.size();
        if (N <= 0) {
            return;
@@ -1337,6 +1348,21 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
                const float * lattice = llama_get_embeddings_nextn(ctx_dft);
                GGML_ASSERT(lattice && "DFlash2 selector produced no lattice");

+                // DFlash2 can instead sample each step and hand q to the verifier for rejection sampling
+                const bool sample = is_dflash2 && params.probabilistic && dp.result_q;
+
+                // greedy keeps temperature 1, sampling uses the target's
+                const float temp = sample ? dp.temp : 1.0f;
+
+                if (sample && dflash2_seed[seq_id] != dp.seed) {
+                    dflash2_seed[seq_id] = dp.seed;
+                    // the default seed means random; mixed so it differs from the target's, as in spec_retune
+                    dflash2_rng[seq_id].seed(dp.seed == LLAMA_DEFAULT_SEED ? std::random_device{}() : dp.seed ^ 0x85ebca6bu);
+                }
+
+                std::vector<float>   probs;
+                std::vector<int32_t> sorted_idx;
+
                int32_t predecessor = 0;
                for (int32_t i = 1; i < (is_dflash2 ? n_block_tokens : params.n_max + 1); ++i) {
                    const float * row = lattice + (size_t) (beg + i) * n_embd_dec;
@@ -1345,21 +1371,53 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {

                    predecessor = (int32_t) std::distance(scores,
                            std::max_element(scores, scores + top_k));
-                    if (params.p_min > 0.0f) {
-                        // softmax(scores) at the argmax, i.e. 1 / sum(exp(s_k - s_max))
+                    if (params.p_min > 0.0f || sample) {
+                        // softmax(scores / temp) at the argmax, i.e. 1 / sum(exp((s_k - s_max) / temp))
+                        probs.resize(top_k);
                        float sum = 0.0f;
                        for (int32_t k = 0; k < top_k; ++k) {
-                            sum += std::exp(scores[k] - scores[predecessor]);
+                            probs[k] = std::exp((scores[k] - scores[predecessor]) / temp);
+                            sum += probs[k];
                        }
                        if (1.0f / sum < params.p_min) {
                            break;
                        }
+                        for (int32_t k = 0; k < top_k; ++k) {
+                            probs[k] /= sum;
+                        }
+                    }
+                    if (sample) {
+                        // ggml_top_k leaves candidates unordered, so walk them by token id to keep a seed's draft stable
+                        sorted_idx.resize(top_k);
+                        for (int32_t k = 0; k < top_k; ++k) {
+                            sorted_idx[k] = k;
+                        }
+                        std::sort(sorted_idx.begin(), sorted_idx.end(), [row](int32_t a, int32_t b) { return row[a] < row[b]; });
+
+                        // a round-off overshoot keeps the argmax, which has q > 0
+                        double u = std::uniform_real_distribution<double>(0.0, 1.0)(dflash2_rng[seq_id]);
+                        for (const int32_t k : sorted_idx) {
+                            u -= probs[k];
+                            if (u < 0.0) {
+                                predecessor = k;
+                                break;
+                            }
+                        }
+
+                        auto & q = dp.result_q->emplace_back();
+                        q.reserve(top_k);
+                        for (int32_t k = 0; k < top_k; ++k) {
+                            q.push_back({ (llama_token) row[k], scores[k], probs[k] });
+                        }
                    }
                    result.push_back((llama_token) row[predecessor]);
                }

                if (result.size() < (size_t) params.n_min) {
                    result.clear();
+                    if (sample) {
+                        dp.result_q->clear();
+                    }
                }
                continue;
            }
@@ -2152,9 +2210,6 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {

    uint16_t n_draft;

-    bool save_dynamic;
-    bool save_static;
-
    struct seq_info {
        size_t cache_size = 0; // number of tokens in n-gram cache

@@ -2168,54 +2223,72 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
    common_speculative_impl_ngram_cache(
            const common_params_speculative & params,
            uint32_t n_seq,
-            uint16_t n_draft,
-            const std::string & path_static,
-            const std::string & path_dynamic,
-            bool save_dynamic,
-            bool save_static)
+            uint16_t n_draft)
        : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, n_seq, n_draft)
        , params(params.ngram_cache)
        , n_draft(n_draft)
-        , save_dynamic(save_dynamic)
-        , save_static(save_static)
    {
+        const auto & ngram = params.ngram_cache;
+
        SPC_TRC("%s", "adding speculative implementation 'ngram-cache'\n");
        SPC_TRC("- n_draft=%d, cache_static=%s, cache_dynamic=%s\n",
                n_draft,
-                path_static.empty() ? "none" : path_static.c_str(),
-                path_dynamic.empty() ? "none" : path_dynamic.c_str());
+                ngram.lookup_cache_static.empty() ? "none" : ngram.lookup_cache_static.c_str(),
+                ngram.lookup_cache_dynamic.empty() ? "none" : ngram.lookup_cache_dynamic.c_str());

        sinfos.resize(n_seq);

-        if (!path_static.empty()) {
+        if (!ngram.lookup_cache_static.empty()) {
            try {
-                auto ngram_cache_static = common_ngram_cache_load(path_static);
+                auto ngram_cache_static = common_ngram_cache_load(ngram.lookup_cache_static);

                for (auto & sinfo : sinfos) {
                    sinfo.ngram_cache_static = ngram_cache_static;
                }
-            } catch (...) {
-                SPC_ERR("failed to open static lookup cache: %s", path_static.c_str());
+            } catch (const std::ifstream::failure &) {
+                SPC_ERR("failed to open static lookup cache: %s\n", ngram.lookup_cache_static.c_str());
                GGML_ABORT("Couldn't read static lookup cache");
            }
        }

-        if (!path_dynamic.empty()) {
+        if (!ngram.lookup_cache_dynamic.empty()) {
            try {
-                auto ngram_cache_dynamic = common_ngram_cache_load(path_dynamic);
+                auto ngram_cache_dynamic = common_ngram_cache_load(ngram.lookup_cache_dynamic);

                for (auto & sinfo : sinfos) {
                    sinfo.ngram_cache_dynamic = ngram_cache_dynamic;
                }
-            } catch (...) {
-                SPC_ERR("failed to open dynamic lookup cache: %s", path_dynamic.c_str());
-                GGML_ABORT("Couldn't read dynamic lookup cache");
+            } catch (const std::ifstream::failure &) {
+                std::error_code ec;
+                const bool exists = std::filesystem::exists(ngram.lookup_cache_dynamic, ec);
+                if (!exists && !ec) {
+                    // missing file is created on destruction
+                    SPC_WRN("dynamic lookup cache file does not exist: %s, starting with empty cache\n",
+                            ngram.lookup_cache_dynamic.c_str());
+                } else {
+                    SPC_ERR("failed to open dynamic lookup cache: %s\n", ngram.lookup_cache_dynamic.c_str());
+                    GGML_ABORT("Couldn't read dynamic lookup cache");
+                }
            }
        }
    }

-    void begin(llama_seq_id /*seq_id*/, const llama_tokens & /*prompt*/) override {
-        // noop
+    ~common_speculative_impl_ngram_cache() override {
+        if (params.lookup_cache_dynamic.empty() || sinfos.empty()) {
+            return;
+        }
+
+        auto & dynamic_cache = sinfos.front().ngram_cache_dynamic;
+        for (auto & sinfo : sinfos) {
+            common_ngram_cache_merge(dynamic_cache, sinfo.ngram_cache_context);
+        }
+
+        SPC_INF("saving dynamic lookup cache at %s\n", params.lookup_cache_dynamic.c_str());
+        common_ngram_cache_save(dynamic_cache, params.lookup_cache_dynamic);
+    }
+
+    void begin(llama_seq_id seq_id, const llama_tokens & /*prompt*/) override {
+        sinfos[seq_id].cache_size = 0;
    }

    void draft_one(
@@ -2226,22 +2299,6 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {

        const auto & prompt = *dparams.prompt;

-        if (sinfo.cache_size < prompt.size() + 1) {
-            llama_tokens tokens_new;
-            tokens_new.reserve(prompt.size() + 1 - sinfo.cache_size);
-            for (size_t j = sinfo.cache_size; j < prompt.size(); ++j) {
-                tokens_new.push_back(prompt[j]);
-            }
-            tokens_new.push_back(dparams.id_last); // add the last token
-
-            // Update context ngram cache with new dparams.prompt:
-            common_ngram_cache_update(
-                    sinfo.ngram_cache_context,
-                    LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX,
-                    tokens_new, tokens_new.size(), false);
-            sinfo.cache_size = prompt.size() + 1;
-        }
-
        llama_tokens inp;
        inp.reserve(prompt.size() + 1);
        for (size_t j = 0; j < prompt.size(); ++j) {
@@ -2249,6 +2306,16 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
        }
        inp.push_back(dparams.id_last);

+        if (sinfo.cache_size < inp.size()) {
+            // Update context ngram cache with the full history so n-grams
+            // ending at the new tokens have their preceding context
+            common_ngram_cache_update(
+                    sinfo.ngram_cache_context,
+                    LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX,
+                    inp, inp.size() - sinfo.cache_size, false);
+            sinfo.cache_size = inp.size();
+        }
+
        result.push_back(dparams.id_last);

        common_ngram_cache_draft(
@@ -2309,20 +2376,12 @@ static common_ngram_map get_common_ngram_map(
    return common_ngram_map(size_key, size_value, key_only, min_hits);
 }

-static common_speculative_impl_ngram_cache create_state_ngram_cache(
+static std::unique_ptr<common_speculative_impl_ngram_cache> create_state_ngram_cache(
        const common_speculative_config & config,
-        uint32_t n_seq,
-        const std::string & path_static,
-        const std::string & path_dynamic) {
+        uint32_t n_seq) {
    uint16_t n_draft = 8; // TODO get from config?

-    // TODO bool param in common/common.h to set save_static/save_dynamic?
-    bool save_static = false;
-    bool save_dynamic = false;
-
-    common_speculative_impl_ngram_cache state(config.params, n_seq, n_draft, path_static, path_dynamic, save_static, save_dynamic);
-
-    return state;
+    return std::make_unique<common_speculative_impl_ngram_cache>(config.params, n_seq, n_draft);
 }

 std::string common_speculative_type_name_str(const std::vector<common_speculative_type> & types) {
@@ -2822,11 +2881,7 @@ common_speculative * common_speculative_init(common_params_speculative & params,
                break;
            }
            case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: {
-                auto state = create_state_ngram_cache(
-                        config, n_seq,
-                        params.ngram_cache.lookup_cache_static,
-                        params.ngram_cache.lookup_cache_dynamic);
-                impls.push_back(std::make_unique<common_speculative_impl_ngram_cache>(state));
+                impls.push_back(create_state_ngram_cache(config, n_seq));
                break;
            }
            default:
@@ -73,7 +73,7 @@ struct common_speculative_draft_params {
    // the generated draft from the last _draft() call
    llama_tokens * result;

-    // candidate distribution per drafted token; set it to make draft-simple and draft-mtp sample
+    // candidate distribution per drafted token; set it to make draft-simple, draft-mtp and draft-dflash (DFlash2 only) sample
    std::vector<std::vector<llama_token_data>> * result_q = nullptr;

    // the target's temp and seed, read only when the drafter samples probabilistically
@@ -17,8 +17,10 @@ __all__ = [


 TEXT_MODEL_MAP: dict[str, str] = {
+    "NEOChatModel": "sensenova_u1",
    "AfmoeForCausalLM": "afmoe",
    "LagunaForCausalLM": "laguna",
+    "FlintForCausalLM": "flint",
    "ApertusForCausalLM": "llama",
    "ArceeForCausalLM": "llama",
    "ArcticForCausalLM": "arctic",
@@ -194,13 +196,16 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "MiniCPMForCausalLM": "minicpm",
    "MiniCPMV4_6ForConditionalGeneration": "minicpm",
    "MiniCPMV4_7ForConditionalGeneration": "minicpm",
+    "MiniCPMOcrForConditionalGeneration": "minicpm",
    "MiniMaxText01ForCausalLM": "minimax",
    "MiniMaxM1ForCausalLM": "minimax",
    "MiniMaxM2ForCausalLM": "minimax",
    "MiniMaxM3SparseForCausalLM": "minimax",
    "MiniMaxM3SparseForConditionalGeneration": "minimax",
    "Ministral3ForCausalLM": "mistral3",
+    "Ministral3Model": "mistral3",
    "Mistral3ForConditionalGeneration": "mistral3",
+    "Mistral3Model": "mistral3",
    "VoxtralRealtimeForConditionalGeneration": "mistral3",
    "MistralForCausalLM": "llama",
    "MixtralForCausalLM": "llama",
@@ -226,6 +231,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "MuseSparkForConditionalGeneration": "muse_spark",
    "OpenELMForCausalLM": "openelm",
    "OpenJevModel": "qwen",
+    "H2OLightningModel": "h2o_lightning",
    "OrionForCausalLM": "orion",
    "PLMForCausalLM": "plm",
    "PLaMo2ForCausalLM": "plamo",
@@ -356,13 +362,16 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "MiniMaxM3SparseForConditionalGeneration": "minimax",
    "MiniCPMV4_6ForConditionalGeneration": "minicpm",
    "MiniCPMV4_7ForConditionalGeneration": "minicpm",
+    "MiniCPMOcrForConditionalGeneration": "minicpm",
    "Mistral3ForConditionalGeneration": "llava",
+    "Mistral3Model": "llava",
    "NemotronH_Omni_Reasoning_V3": "nemotron",
    "NemotronH_Nano_VL_V2": "nemotron",
    "MuseGlimmerForConditionalGeneration": "muse_glimmer",
    "MuseSparkForConditionalGeneration": "muse_spark",
    "PaddleOCRVisionModel": "ernie",
    "Phi4ForCausalLMV": "phi",
+    "NEOVisionModel": "sensenova_u1",
    "Qwen2AudioForConditionalGeneration": "ultravox",
    "Qwen2VLForConditionalGeneration": "qwenvl",
    "Qwen2VLModel": "qwenvl",
@@ -375,6 +384,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "Qwen3VLForConditionalGeneration": "qwen3vl",
    "Qwen3VLMoeForConditionalGeneration": "qwen3vl",
    "OpenJevModel": "qwen3vl",
+    "H2OLightningModel": "h2o_lightning",
    "Qwen3_5ForConditionalGeneration": "qwen3vl",
    "Qwen3_5MoeForConditionalGeneration": "qwen3vl",
    "Qwen4ExpForConditionalGeneration": "qwen4exp",
@@ -10,6 +10,7 @@ import json
 import os
 import re
 import sys
+import tempfile
 from enum import IntEnum
 from pathlib import Path
 from hashlib import sha256
@@ -1700,7 +1701,7 @@ class TextModel(ModelBase):
            self.gguf_writer.add_expert_group_used_count(n_group_used)
            logger.info(f"gguf: expert groups used count = {n_group_used}")

-        if (score_func := self.find_hparam(["score_function", "scoring_func", "score_func", "moe_router_activation", "moe_router_activation_func", "expert_selection_fn", "router_score_func"], optional=True)) is not None:
+        if (score_func := self.find_hparam(["score_function", "scoring_func", "score_func", "moe_router_activation", "moe_router_activation_func", "expert_selection_fn", "router_score_func", "router_score_function"], optional=True)) is not None:
            if score_func == "sigmoid":
                self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)
            elif score_func == "softmax":
@@ -1836,6 +1837,9 @@ class TextModel(ModelBase):
        if chkhsh == "1431a23e583c97432bc230bff598d103ddb5a1f89960c8f1d1051aaa944d0b35":
            # ref: https://huggingface.co/sapienzanlp/Minerva-7B-base-v1.0
            res = "minerva-7b"
+        if chkhsh == "6e4d72c6e5260b90606e26a0931fa15e4b42ec6e97719d12919f2b1f93478262":
+            # ref: https://huggingface.co/domyn/Domyn-Small-v1.0
+            res = "domyn-small"
        if chkhsh == "7e57df22b1fe23a7b1e1c7f3dc4e3f96d43a4eb0836d0c6bdc3436d7b2f1c664":
            # ref: https://huggingface.co/tencent/Hunyuan-A13B-Instruct
            res = "hunyuan"
@@ -1938,6 +1942,9 @@ class TextModel(ModelBase):
        if chkhsh == "52df12b4c8d4176e7481aab4b6e8454d1fd0a210a04a574f6d4e067d10e23c3e":
            # ref: https://huggingface.co/CohereLabs/North-Mini-Code-1.0
            res = "cohere2moe"
+        if chkhsh == "d5420e601e979078c670bc503d9ac411c311fa76b76a76267c31e0d6ba816637":
+            # ref: https://huggingface.co/RecursiveSI/flint
+            res = "flint"
        if chkhsh == "e636dc30a262dcc0d8c323492e32ae2b70728f4df7dfe9737d9f920a282b8aea":
            # ref: https://huggingface.co/Qwen/Qwen1.5-7B
            res = "qwen2"
@@ -2169,20 +2176,34 @@ class TextModel(ModelBase):
        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
        special_vocab.add_to_gguf(self.gguf_writer)

+    def _is_hybriddna(self) -> bool:
+        tokenizer_config_file = self.dir_model / "tokenizer_config.json"
+        if not tokenizer_config_file.is_file():
+            return False
+        with open(tokenizer_config_file, "r", encoding="utf-8") as f:
+            return json.load(f).get("tokenizer_class") == "HybridDNATokenizer"
+
    def _set_vocab_hybriddna(self):
        from transformers import AutoTokenizer
        tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True)
-        vocab_size = self.hparams.get("vocab_size", len(tokenizer.vocab))  # ty: ignore[unresolved-attribute]
-        assert max(tokenizer.vocab.values()) < vocab_size  # ty: ignore[unresolved-attribute]
+        vocab = tokenizer.get_vocab()
+        vocab_size = self.hparams.get("vocab_size", len(vocab))
+        assert max(vocab.values()) < vocab_size

-        reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in tokenizer.vocab.items()}  # ty: ignore[unresolved-attribute]
-        # k-mers can share text with a base-vocab BPE token (e.g. CCCCCC) and get
-        # dropped by get_vocab(); a reserved marker suffix (U+E000) keeps each
-        # k-mer's own id (llama.cpp strips it on detokenization)
-        for kmer in tokenizer.kmers:  # ty: ignore[unresolved-attribute]
+        reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in vocab.items()}
+        # k-mers can share text with a base-vocab BPE token (e.g. CCCCCC) and get dropped by get_vocab()
+        # a reserved marker suffix (U+E000) keeps each k-mer's own id (llama.cpp strips it on detokenization)
+        # tail="pad_n" (Carbon-14B) adds N-padded tail tokens ("CT" -> "CTNNNN") that get the same marker
+        # llama.cpp detects them to pad with N and strips the padding on detokenization, tail="char" is unsupported
+        tail = getattr(tokenizer, "tail", "pad_a")
+        if tokenizer.k != 6 or tail not in ("pad_a", "pad_n"):
+            raise NotImplementedError(f"HybridDNATokenizer with k={tokenizer.k}, tail={tail!r}")
+        for kmer in [*tokenizer.kmers, *getattr(tokenizer, "npad_tokens", [])]:
            reverse_vocab[tokenizer.dna_token_to_id[kmer]] = kmer + "\ue000"  # ty: ignore[unresolved-attribute]
-        added_vocab = tokenizer.get_added_vocab()  # ty: ignore[unresolved-attribute]
-        added_tokens_decoder = tokenizer.added_tokens_decoder  # ty: ignore[unresolved-attribute]
+        # older HybridDNATokenizer versions (Carbon-500M/3B/8B) do not register the base tokenizer's added tokens (<|im_start|>, <think>, ...) as their own
+        # but they are still encoded as special tokens: take their types from the base tokenizer
+        base_tokenizer = getattr(tokenizer, "_base_tokenizer", tokenizer)
+        added_tokens_decoder = {**base_tokenizer.added_tokens_decoder, **tokenizer.added_tokens_decoder}  # ty: ignore[unresolved-attribute]

        tokens: list[str] = []
        toktypes: list[int] = []
@@ -2192,7 +2213,7 @@ class TextModel(ModelBase):
                toktypes.append(gguf.TokenType.UNUSED)
            else:
                token: str = reverse_vocab[i]
-                if token in added_vocab:
+                if i in added_tokens_decoder:
                    if added_tokens_decoder[i].special or self.does_token_look_special(token):
                        toktypes.append(gguf.TokenType.CONTROL)
                    else:
@@ -2204,10 +2225,19 @@ class TextModel(ModelBase):
        tokpre = self.get_vocab_base_pre(tokenizer)
        self.gguf_writer.add_tokenizer_model("hybriddna")
        self.gguf_writer.add_tokenizer_pre(tokpre)
+        self.gguf_writer.add_dna_auto_tags(getattr(tokenizer, "auto_dna_tags", False))
        self.gguf_writer.add_token_list(tokens)
        self.gguf_writer.add_token_types(toktypes)

        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
+        if not special_vocab.merges:
+            # Carbon-500M/3B/8B ship only tokenizer.py, no tokenizer.json
+            # take the merges and special token ids from the base BPE tokenizer, keeping the model's own chat template (if any)
+            with tempfile.TemporaryDirectory() as tmp:
+                base_tokenizer.save_pretrained(tmp)
+                base_special_vocab = gguf.SpecialVocab(Path(tmp), load_merges=True)
+            base_special_vocab.chat_template = special_vocab.chat_template
+            special_vocab = base_special_vocab
        special_vocab.add_to_gguf(self.gguf_writer)

    def _set_vocab_qwen(self):
@@ -3119,7 +3149,7 @@ def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> st
    # Kimi-K3's text_config reports "KimiLinearForCausalLM", which is the older
    # Kimi-Linear-48B architecture and cannot load K3 (no attention residuals,
    # latent MoE, situ, ...). Route on the top-level architecture instead.
-    if model_type == ModelType.TEXT and arch in ("StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration"):
+    if model_type == ModelType.TEXT and arch in ("NEOChatModel", "StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration"):
        return arch

    # if "architectures" is found in the sub-config, use that instead
@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from .base import ModelBase, gguf
+from .qwen import Qwen2MoeModel
+
+
+@ModelBase.register("FlintForCausalLM")
+@ModelBase.example("RecursiveSI/flint")
+class FlintModel(Qwen2MoeModel):
+    model_arch = gguf.MODEL_ARCH.FLINT
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        dense = self.hparams["mlp_only_layers"]
+        if dense != list(range(len(dense))):
+            raise ValueError("Flint dense layers must precede the MoE layers")
+        self.gguf_writer.add_leading_dense_block_count(len(dense))
+        self.gguf_writer.add_expert_shared_count(
+            self.hparams["shared_expert_intermediate_size"] // self.hparams["moe_intermediate_size"]
+        )
+        self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])
+        self.gguf_writer.add_expert_weights_scale(1.0)
+        self.gguf_writer.add_sliding_window(self.hparams["sliding_window"])
+        self.gguf_writer.add_sliding_window_pattern([
+            kind == "sliding_attention" for kind in self.hparams["layer_types"]
+        ])
+        self.gguf_writer.add_rope_pattern([
+            i not in self.hparams["no_rope_layer_indices"] for i in range(self.block_count)
+        ])
+
+    def modify_tensors(self, data_torch, name, bid):
+        if name.endswith(".softmax_offset"):
+            name += ".weight"
+        if name.endswith(".gate.expert_bias"):
+            if not self.hparams["router_expert_bias"]:
+                return
+            name += ".bias"
+        yield from super().modify_tensors(data_torch, name, bid)
@@ -999,6 +999,8 @@ class Gemma4VisionAudioModel(MmprojModel):
        assert self.hparams_vision is not None
        self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.GEMMA4V)
        self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams_vision.get("layer_norm_eps", 1e-6))
+        assert self.hparams_vision.get("hidden_activation", "gelu_pytorch_tanh") == "gelu_pytorch_tanh"
+        self.gguf_writer.add_vision_use_gelu(True)

        # audio params
        if self.has_audio_encoder:
@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import json
+
+from pathlib import Path
+from typing import Any
+
+from .base import ModelBase, gguf, logger
+from .qwen import Qwen3_5TextModel
+from .qwen3vl import Qwen3VLVisionModel
+
+
+def _is_h2o_lightning_checkpoint(dir_model: Path) -> bool:
+    return (dir_model / "h2o_lightning_shim.py").is_file() and (dir_model / "serve_config.json").is_file()
+
+
+@ModelBase.register_hparams_loader(_is_h2o_lightning_checkpoint)
+def _load_h2o_lightning_hparams(dir_model: Path) -> dict[str, Any]:
+    logger.info("gguf: detected H2O-Lightning checkpoint")
+    hparams = ModelBase.load_hparams(dir_model, False, guess=False)
+    hparams["architectures"] = ["H2OLightningModel"]
+    return hparams
+
+
+@ModelBase.register("H2OLightningModel")
+@ModelBase.example("h2oai/h2o-lightning-4b")
+class H2OLightningModel(Qwen3_5TextModel):
+    model_arch = gguf.MODEL_ARCH.QWEN35
+
+    # prompt and calibration follow serve_config.json and h2o_lightning_shim.py of the model repo
+    # the state budget is hard-coded in server-decision.cpp
+    _MAX_STATE_TOKENS = 32000
+    _DEFAULT_INSTRUCTIONS = "Answer the question below."
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        with open(self.dir_model / "serve_config.json", "r", encoding="utf-8") as f:
+            self._serve_config = json.load(f)
+        prompt = self._serve_config["prompt"]
+        expected = {
+            "state_compact":    (prompt.get("state_compact", True), True),
+            "max_state_tokens": (prompt.get("max_state_tokens"), self._MAX_STATE_TOKENS),
+            "layout":           (prompt.get("layout", "plain"), "plain"),
+            "chat_template":    (prompt.get("chat_template", True), True),
+            "enable_thinking":  (prompt.get("enable_thinking", False), False),
+        }
+        for name, (got, want) in expected.items():
+            if got != want:
+                raise ValueError(f"serve_config.json: {name} is {got!r}, the h2o-lightning decision type expects {want!r}")
+
+    def get_vocab_base_pre(self, tokenizer) -> str:
+        # transformers loads this checkpoint as Qwen2Tokenizer: the model was trained and is served with the Qwen2 split
+        pre = super().get_vocab_base_pre(tokenizer)
+        if pre != "qwen35":
+            raise ValueError(f"expected the Qwen3.5 tokenizer, got pre-tokenizer {pre!r}")
+        return "qwen2"
+
+    def set_vocab(self):
+        super().set_vocab()
+        self.gguf_writer.add_chat_template([{"name": "systemone", "template": self._systemone_template()}])
+
+    def _systemone_template(self) -> str:
+        # the server gives the template texts rendered as the shim renders them, it only lays them out
+        prompt = self._serve_config["prompt"]
+        nl = "{{ '\\n' }}"
+        description = (
+            "{% if type == 'noul' %}"
+            "{% if o.description %}{{ o.description }}{% elif o.key == 'true' %}the statement holds{% else %}it does not{% endif %}"
+            "{% elif o.description is none %}{{ o.key }}{% else %}{{ o.description }}{% endif %}"
+        )
+        user = (
+            "{% set user %}record: {{ state }}" + nl
+            + "question: {{ instructions if instructions else '" + self._DEFAULT_INSTRUCTIONS + "' }}" + nl
+            + "options:{% for o in options %}" + nl + "{{ o.label | trim }}) {{ o.key }}: " + description + "{% endfor %}{% endset %}"
+        )
+        return (
+            "<|im_start|>system" + nl + prompt["system"].strip() + "<|im_end|>" + nl
+            + "<|im_start|>user" + nl + "{% for image in images %}{{ image }}{% endfor %}" + user + "{{ user | trim }}<|im_end|>" + nl
+            + "<|im_start|>assistant" + nl + "<think>{{ '\\n\\n' }}</think>{{ '\\n\\n' }}" + prompt["prefill"]
+        )
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_decision_type(gguf.DecisionType.H2O_LIGHTNING)
+        # the answer follows "Answer:", so the token of a label has a leading space
+        cfg = self._serve_config
+        self.gguf_writer.add_decision_labels([" " + label for label in cfg["labels"][:cfg.get("max_options", len(cfg["labels"]))]])
+        for name in ("choice", "noul", "score"):
+            self.gguf_writer.add_decision_temperature(name, self._serve_config["temperature"][name])
+        # a request with images has this temperature for all question types
+        self.gguf_writer.add_decision_temperature("image", self._serve_config["image"]["temperature"])
+
+
+@ModelBase.register("H2OLightningModel")
+class H2OLightningVisionModel(Qwen3VLVisionModel):
+    pass
@@ -374,6 +374,16 @@ class LFM2VLModel(MmprojModel):
            if resize_algo is None:
                raise ValueError(f"unsupported resample: {resample}")
            self.gguf_writer.add_vision_image_resize_algo(resize_algo)
+        # tiling params from processor_config.json (do_image_splitting=False -> single tile, i.e. no tiling)
+        preproc = self.preprocessor_config
+        if preproc.get("do_image_splitting", True):
+            self.gguf_writer.add_vision_preproc_min_tiles(preproc.get("min_tiles", 2))
+            self.gguf_writer.add_vision_preproc_max_tiles(preproc.get("max_tiles", 10))
+        else:
+            self.gguf_writer.add_vision_preproc_min_tiles(1)
+            self.gguf_writer.add_vision_preproc_max_tiles(1)
+        self.gguf_writer.add_vision_preproc_tile_size(preproc.get("tile_size", 512))
+        self.gguf_writer.add_vision_preproc_max_pixels_tolerance(preproc.get("max_pixels_tolerance", 2.0))

    @classmethod
    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
@@ -129,8 +129,9 @@ class LlamaModel(TextModel):
                tokenizer_config_json = json.load(f)
                if (add_prefix_space := tokenizer_config_json.get("add_prefix_space")) is not None:
                    self.gguf_writer.add_add_space_prefix(add_prefix_space)
-                if tokenizer_config_json.get("tokenizer_class") == "HybridDNATokenizer":
-                    return self._set_vocab_hybriddna()
+
+        if self._is_hybriddna():
+            return self._set_vocab_hybriddna()

        try:
            self._set_vocab_sentencepiece()
@@ -2,7 +2,7 @@ from __future__ import annotations

 import json

-from typing import Iterable, TYPE_CHECKING
+from typing import Callable, Iterable, TYPE_CHECKING

 if TYPE_CHECKING:
    from torch import Tensor
@@ -128,3 +128,24 @@ class LlavaVisionModel(MmprojModel):
            yield from super().modify_tensors(img_break_embd, name, bid)

        return # skip other tensors
+
+
+@ModelBase.register("Mistral3Model")
+@ModelBase.example("nvidia/Nemotron-3.5-Embed-1B-BF16")
+class Mistral3EmbeddingVisionModel(LlavaVisionModel):
+    def __init__(self, dir_model, *args, **kwargs):
+        hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False)
+        text_config = hparams.get("text_config", hparams)
+        if text_config.get("is_causal") is not False:
+            raise NotImplementedError("Expected a Nemotron Embed model for Mistral3Model, but is_causal is not False")
+        super().__init__(dir_model, *args, **kwargs)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        if item[0].endswith("embed_tokens.weight"):
+            return ModelBase.filter_tensors(item)
+        return super().filter_tensors(item)
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_vision_preproc_image_size(self.preprocessor_config["size"]["longest_edge"])
@@ -272,3 +272,18 @@ class MiniCPMV4_7VisionModel(MiniCPMV4_6VisionModel):
            yield self.format_tensor_name(gguf.MODEL_TENSOR.V_TOK_EMBD_SEP, suffix=""), data_torch[ids]
            return
        yield from super().modify_tensors(data_torch, name, bid)
+
+
+# MiniCPM-OCR 1.0: the 4.6 stack with model_type minicpm_ocr.
+
+@ModelBase.register("MiniCPMOcrForConditionalGeneration")
+@ModelBase.example("openbmb/MiniCPM-OCR-1.0")
+class MiniCPMOCR1_0TextModel(MiniCPMV4_6TextModel):
+    model_arch = gguf.MODEL_ARCH.QWEN35
+    no_mtp = True
+
+
+@ModelBase.register("MiniCPMOcrForConditionalGeneration")
+@ModelBase.example("openbmb/MiniCPM-OCR-1.0")
+class MiniCPMOCR1_0VisionModel(MiniCPMV4_6VisionModel):
+    projector_type = gguf.VisionProjectorType.MINICPM_OCR_1_0
@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING
 if TYPE_CHECKING:
    from torch import Tensor

-from .base import ModelBase, TextModel, gguf
+from .base import ModelBase, TextModel, get_model_architecture, gguf

 from .deepseek import DeepseekV2Model
 from .llama import LlamaModel
@@ -13,9 +13,11 @@ from .llama import LlamaModel

@ModelBase.register(
    "Mistral3ForConditionalGeneration",
+    "Mistral3Model",
    "Ministral3ForCausalLM",
+    "Ministral3Model",
 )
-@ModelBase.example("mistralai/Mistral-Small-3.1-24B-Instruct-2503", "hf-tiny-v2/tiny-random-Ministral3ForCausalLM")
+@ModelBase.example("mistralai/Mistral-Small-3.1-24B-Instruct-2503", "hf-tiny-v2/tiny-random-Ministral3ForCausalLM", "nvidia/Nemotron-3-Embed-1B-BF16", "nvidia/Nemotron-3.5-Embed-1B-BF16")
 class Mistral3Model(TextModel):
    class Ministral3Model(LlamaModel):
        model_arch = gguf.MODEL_ARCH.MISTRAL3
@@ -29,6 +31,23 @@ class Mistral3Model(TextModel):
                self.gguf_writer.add_rope_scaling_yarn_log_mul(rope_params["mscale_all_dim"])
                self.gguf_writer.add_attn_temperature_scale(rope_params["llama_4_scaling_beta"])

+    class Mistral3EmbedNemotronModel(Ministral3Model):
+        model_arch = gguf.MODEL_ARCH.MISTRAL3
+
+        def set_vocab(self):
+            super().set_vocab()
+            self.gguf_writer.add_add_bos_token(False)
+            self.gguf_writer.add_add_eos_token(False)
+
+        def get_vocab_base_pre(self, tokenizer) -> str:
+            if self.origin_hf_arch == "Mistral3Model":
+                return "nemotron-3.5-embed"
+            return super().get_vocab_base_pre(tokenizer)
+
+        def set_gguf_parameters(self):
+            super().set_gguf_parameters()
+            self._try_set_pooling_type()
+
    class Mistral4Model(DeepseekV2Model):
        model_arch = gguf.MODEL_ARCH.MISTRAL4
        skip_mtp = False # model contains no MTP layers, so no need to skip
@@ -42,12 +61,20 @@ class Mistral3Model(TextModel):
    model_arch = gguf.MODEL_ARCH.MISTRAL3 # unused
    impl: TextModel

-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self, dir_model, *args, **kwargs):
+        hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False)
+        hf_arch = get_model_architecture(hparams, self.model_type)
+        if hf_arch in ("Ministral3Model", "Mistral3Model"):
+            text_config = hparams.get("text_config", hparams)
+            if text_config.get("is_causal") is not False:
+                raise NotImplementedError(f"Expected a Nemotron Embed model for {hf_arch}, but is_causal is not False")
+        super().__init__(dir_model, *args, **kwargs)
        if self.hparams.get("model_type") == "mistral4":
-            self.impl = Mistral3Model.Mistral4Model(*args, **kwargs)
+            self.impl = Mistral3Model.Mistral4Model(dir_model, *args, **kwargs)
+        elif self.hf_arch in ("Ministral3Model", "Mistral3Model"):
+            self.impl = Mistral3Model.Mistral3EmbedNemotronModel(dir_model, *args, **kwargs)
        else:
-            self.impl = Mistral3Model.Ministral3Model(*args, **kwargs)
+            self.impl = Mistral3Model.Ministral3Model(dir_model, *args, **kwargs)

    def set_vocab(self):
        self.impl.set_vocab()
@@ -1,5 +1,6 @@
 from __future__ import annotations

+import re
 from typing import Any, Callable, Iterable, Sequence, TYPE_CHECKING

 import numpy as np
@@ -367,9 +368,21 @@ class NemotronModel(TextModel):
    model_arch = gguf.MODEL_ARCH.NEMOTRON

    def set_vocab(self):
-        self._set_vocab_sentencepiece()
-        self.gguf_writer.add_pad_token_id(0)
-        self.gguf_writer.add_unk_token_id(1)
+        if (self.dir_model / "tokenizer.model").exists():
+            self._set_vocab_sentencepiece()
+            self.gguf_writer.add_pad_token_id(0)
+            self.gguf_writer.add_unk_token_id(1)
+        else:
+            # SPM-style BPE (Metaspace + byte fallback): mark <0xXX> tokens as BYTE so they decode to raw bytes
+            tokens, toktypes, tokpre = self.get_vocab_base()
+            toktypes = [gguf.TokenType.BYTE if re.fullmatch(r"<0x[0-9A-F]{2}>", tok) else typ for tok, typ in zip(tokens, toktypes)]
+            self.gguf_writer.add_tokenizer_model("gpt2")
+            self.gguf_writer.add_tokenizer_pre(tokpre)
+            self.gguf_writer.add_token_list(tokens)
+            self.gguf_writer.add_token_types(toktypes)
+
+            special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
+            special_vocab.add_to_gguf(self.gguf_writer)

    def set_gguf_parameters(self):
        super().set_gguf_parameters()
@@ -383,7 +396,10 @@ class NemotronModel(TextModel):
        rot_pct = self.rope_parameters["partial_rotary_factor"]
        n_embd = self.find_hparam(["hidden_size", "n_embd"])
        n_head = self.find_hparam(["num_attention_heads", "n_head"])
-        self.gguf_writer.add_rope_dimension_count(int(rot_pct * n_embd) // n_head)
+        # Use explicit head_dim when available; fallback keeps standard Nemotron working.
+        # Without this, models where n_head * head_dim != n_embd (e.g. 48 heads, head_dim=128) produce the wrong rope dimension count.
+        head_dim = self.hparams.get("head_dim") or n_embd // n_head
+        self.gguf_writer.add_rope_dimension_count(int(rot_pct * head_dim))

        # * RopeScaling for Nemotron
        factor = self.hparams.get("factor") or self.rope_parameters.get("factor")
@@ -59,6 +59,10 @@ class Qwen2Model(TextModel):
    model_arch = gguf.MODEL_ARCH.QWEN2

    def set_vocab(self):
+        # HuggingFaceBio/Carbon-14B: Qwen3 with the HybridDNATokenizer (Qwen3 BPE plus DNA 6-mers)
+        if self._is_hybriddna():
+            self._set_vocab_hybriddna()
+            return
        try:
            self._set_vocab_sentencepiece()
        except FileNotFoundError:
@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+from .base import MmprojModel, ModelBase, gguf
+from .qwen import Qwen2Model
+
+
+@ModelBase.example("sensenova/SenseNova-U1.5-8B-MoT")
+@ModelBase.register("NEOChatModel")
+class SenseNovaU1Model(Qwen2Model):
+    model_arch = gguf.MODEL_ARCH.SENSENOVA_U1
+
+    @classmethod
+    def filter_tensors(cls, item):
+        name, _ = item
+        if not name.startswith("language_model.") or "_mot_gen" in name:
+            return None
+        return super().filter_tensors(item)
+
+    def set_gguf_parameters(self):
+        if int(self.hparams.get("num_experts") or 0) > 1:
+            raise ValueError("SenseNova U1 conversion currently supports the dense understanding branch")
+        if self.rope_parameters.get("rope_type", "default") not in ("default", None):
+            raise ValueError("SenseNova U1 scaled RoPE is not supported")
+        if self.hparams.get("attention_bias", False):
+            raise ValueError("SenseNova U1 attention bias is not supported")
+        if self.hparams.get("use_sliding_window", False):
+            raise ValueError("SenseNova U1 sliding attention is not supported")
+        super().set_gguf_parameters()
+        self.gguf_writer.add_rope_freq_base_spatial(self.hparams["rope_theta_hw"])
+
+    def get_tensors(self):
+        for name, gen in self.model_tensors.items():
+            if name.endswith((".q_norm_hw.weight", ".k_norm_hw.weight")):
+                continue
+            data = gen()
+            if name.endswith((".q_norm.weight", ".k_norm.weight")):
+                spatial = self.model_tensors[name.replace("_norm.weight", "_norm_hw.weight")]()
+                # Keep one weight vector; the graph normalizes its two halves independently.
+                data = torch.cat((data, spatial), dim=0)
+            yield name, data
+
+
+@ModelBase.register("NEOVisionModel")
+@ModelBase.example("sensenova/SenseNova-U1.5-8B-MoT")
+class SenseNovaU1VisionModel(MmprojModel):
+    def __init__(self, dir_model: Path, *args, **kwargs):
+        hparams = kwargs.pop("hparams", None)
+        if hparams is None:
+            hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False)
+        else:
+            hparams = dict(hparams)
+
+        vision_config = dict(hparams["vision_config"])
+        # NEOVisionModel has no transformer blocks
+        # these defaults satisfy the generic mmproj metadata contract without replacing checkpoint values
+        vision_config.setdefault("image_size", 0)
+        vision_config.setdefault("intermediate_size", 0)
+        vision_config.setdefault("num_hidden_layers", 0)
+        vision_config.setdefault("num_attention_heads", 1)
+        hparams["vision_config"] = vision_config
+        super().__init__(dir_model, *args, hparams=hparams, **kwargs)
+
+        self.preprocessor_config.setdefault("image_mean", [0.485, 0.456, 0.406])
+        self.preprocessor_config.setdefault("image_std", [0.229, 0.224, 0.225])
+
+    @classmethod
+    def filter_tensors(cls, item):
+        name, gen = item
+        if not name.startswith("vision_model."):
+            return None
+        return super().filter_tensors((name, gen))
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_clip_projector_type("sensenova_u1")
+        vision_config = self.hparams["vision_config"]
+        # These defaults match the known U1.5 checkpoint when older configs omit fields.
+        hidden_size = int(vision_config.get("hidden_size", 1024))
+        min_pixels = int(vision_config.get("min_pixels", 65536))
+        max_pixels = int(vision_config.get("max_pixels", 16777216))
+        downsample_ratio = float(vision_config.get("downsample_ratio", 0.5))
+        if downsample_ratio <= 0:
+            raise ValueError(f"SenseNova U1 downsample_ratio must be positive, got {downsample_ratio}")
+        merge_size = round(1.0 / downsample_ratio)
+        if abs((1.0 / downsample_ratio) - merge_size) > 1e-6:
+            raise ValueError(f"SenseNova U1 downsample_ratio must have an integer reciprocal, got {downsample_ratio}")
+        norm_eps = float(vision_config.get("layer_norm_eps", vision_config.get("rms_norm_eps", 1e-6)))
+        self.gguf_writer.add_vision_head_dim(hidden_size)
+        self.gguf_writer.add_vision_attention_layernorm_eps(norm_eps)
+        self.gguf_writer.add_vision_min_pixels(min_pixels)
+        self.gguf_writer.add_vision_max_pixels(max_pixels)
+        self.gguf_writer.add_vision_spatial_merge_size(merge_size)
+
+    def modify_tensors(self, data_torch, name, bid):
+        tensor_names = {
+            "vision_model.embeddings.patch_embedding.weight": "v.patch_embd.weight",
+            "vision_model.embeddings.patch_embedding.bias":   "v.patch_embd.bias",
+            "vision_model.embeddings.dense_embedding.weight": self.format_tensor_name(gguf.MODEL_TENSOR.V_MMPROJ, 0),
+            "vision_model.embeddings.dense_embedding.bias": self.format_tensor_name(
+                gguf.MODEL_TENSOR.V_MMPROJ, 0, suffix=".bias"
+            ),
+        }
+        if name not in tensor_names:
+            raise ValueError(f"Unexpected SenseNova U1 vision tensor: {name!r}")
+        return [(tensor_names[name], data_torch)]
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if new_name == self.format_tensor_name(gguf.MODEL_TENSOR.V_MMPROJ, 0):
+            return gguf.GGMLQuantizationType.F16 if self.ftype == gguf.LlamaFileType.MOSTLY_F16 else gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
@@ -101,6 +101,7 @@ models = [
    {"name": "command-r",        "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/CohereForAI/c4ai-command-r-v01", },
    {"name": "tiny_aya",         "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/CohereLabs/tiny-aya-base", },
    {"name": "cohere2moe",       "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/CohereLabs/North-Mini-Code-1.0", },
+    {"name": "flint",            "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/RecursiveSI/flint", },
    {"name": "qwen2",            "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/Qwen/Qwen1.5-7B", },
    {"name": "olmo",             "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/allenai/OLMo-1.7-7B-hf", },
    {"name": "dbrx",             "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/databricks/dbrx-base", },
@@ -177,6 +178,7 @@ pre_computed_hashes = [
    {"name": "glm4", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/zai-org/GLM-4.5-Air", "chkhsh": "9ca2dd618e8afaf09731a7cf6e2105b373ba6a1821559f258b272fe83e6eb902"},
    {"name": "glm4", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/zai-org/GLM-4.7-Flash", "chkhsh": "cdf5f35325780597efd76153d4d1c16778f766173908894c04afc20108536267"},
    {"name": "minerva-7b", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/sapienzanlp/Minerva-7B-base-v1.0", "chkhsh": "1431a23e583c97432bc230bff598d103ddb5a1f89960c8f1d1051aaa944d0b35"},
+    {"name": "domyn-small", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/domyn/Domyn-Small-v1.0", "chkhsh": "6e4d72c6e5260b90606e26a0931fa15e4b42ec6e97719d12919f2b1f93478262"},
    {"name": "hunyuan", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/tencent/Hunyuan-A13B-Instruct", "chkhsh": "7e57df22b1fe23a7b1e1c7f3dc4e3f96d43a4eb0836d0c6bdc3436d7b2f1c664"},
    {"name": "hunyuan-dense", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/tencent/Hunyuan-4B-Instruct", "chkhsh": "bba3b3366b646dbdded5dbc42d59598b849371afc42f7beafa914afaa5b70aa6"},
    {"name": "hy_v4", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/tencent/Hy4-preview", "chkhsh": "e6ddf9c6686791c12d698d34c31ab9be1fea9af5a3d9a6909783ab382198ae1c"},
@@ -783,7 +783,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL          | ON (mandatory)                        | Enable build with SYCL code path.           |
 | GGML_SYCL_TARGET   | INTEL *(default)*                     | Set the SYCL target device type.            |
 | GGML_SYCL_DEVICE_ARCH | Optional                           | Set the SYCL device architecture. Setting the device architecture can improve the performance. See the table [--offload-arch](https://github.com/intel/llvm/blob/sycl/sycl/doc/design/OffloadDesign.md#--offload-arch) for a list of valid architectures. |
-| GGML_SYCL_F16      | OFF *(default)* \|ON *(optional)*     | Enable FP16 build with SYCL code path. (1.) |
+| GGML_SYCL_F16      | OFF *(default)* \|ON *(optional)*     | Enable FP16 build with SYCL code path. [1.](#sycl-f16-footnote) |
 | GGML_SYCL_GRAPH    | ON *(default)* \|OFF *(Optional)*     | Enable build with [SYCL Graph extension](https://github.com/intel/llvm/blob/sycl/sycl/doc/extensions/experimental/sycl_ext_oneapi_graph.asciidoc). |
 | GGML_SYCL_DNN      | ON *(default)* \|OFF *(Optional)*     | Enable build with oneDNN.                   |
 | GGML_SYCL_HOST_MEM_FALLBACK | ON *(default)* \|OFF *(Optional)* | Allow host memory fallback when device memory is full during quantized weight reorder. Enables inference to continue at reduced speed (reading over PCIe) instead of failing. Requires Linux kernel 6.8+. |
@@ -791,7 +791,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | CMAKE_C_COMPILER   | `icx` *(Linux)*, `icx/cl` *(Windows)* | Set `icx` compiler for SYCL code path.      |
 | CMAKE_CXX_COMPILER | `icpx` *(Linux)*, `icx` *(Windows)*   | Set `icpx/icx` compiler for SYCL code path. |

-1. FP32 or FP16 have different performance impact to LLM. Recommended to test them for better prompt processing performance on your models. You need to rebuild the code after change `GGML_SYCL_F16=OFF/ON`.
+1. <a id="sycl-f16-footnote"></a>FP32 or FP16 have different performance impact to LLM. Recommended to test them for better prompt processing performance on your models. You need to rebuild the code after change `GGML_SYCL_F16=OFF/ON`.

 ### Runtime

@@ -802,7 +802,8 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_DEV2DEV_MEMCPY | 0 (default), 1, 2 | Choose the method of dev2dev memory copy.<br>Value: <br>*  0: SYCL API (default), only support dGPUs.<br>* 1: L0 API -- Better performance, only support dGPUs, found to lead to abnormal crash in some case. <br>* 2: Host Forward -- Most stable method for all cases (including iGPU + dGPU*N), but with lower performance (-2% to -5%).<br>SYCL & L0 API are easy to be impacted by Intel GPU driver issue. When you meet the garbled output or crash issues in multiple GPUs case, try with this debug flag to work around or check the issue.|
 | GGML_SYCL_ENABLE_FLASH_ATTN | 1 (default) or 0| Enable Flash-Attention. It can reduce memory usage. The performance impact depends on the LLM.|
 | GGML_SYCL_ENABLE_OPT | 0 or 1 (default)| Enable optimize features for Intel GPUs. (Recommended to 0 for Intel devices older than Gen 10) |
-| GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still on development, no better performance. |
+| GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still in development, performance depends on model and device. [1.](#sycl-graph-footnote) |
+| GGML_SYCL_GRAPH_EVICTION_TIMEOUT | 10 (default) or positive integer | Seconds a cached SYCL graph can go unused before it is evicted. Checked on a sweep that runs every half of this timeout. |
 | GGML_SYCL_ENABLE_HOST_PINNED_MEM | 0 or 1 (default) | Enable host pinned memory to speed up copy data from host to device. When disable it, host memory will common malloc() on CPU. Disable it when use `--load-model mlock`.|
 | GGML_SYCL_HOST_PINNED_MEM_2G | 0 (default) or 1 | Limit the max memory allocation to be no more than 2GB when enable host pinned memory. USM allocations above 2 GiB take the relaxed/large-allocation path, which serializes H2D copies with compute and prevents copy/compute overlap. It will impact the startup time. Need more test. Depend on `GGML_SYCL_ENABLE_HOST_PINNED_MEM=1`.|
 | GGML_SYCL_UPLOAD_STAGING_SLOTS | 4 (default) or non-negative integer | Number of 8 MiB pinned host slots used to stage tensor uploads (model loading), so the host copy of one slot overlaps the transfer of the previous one. Set to 0 to use the old path: a malloc'd bounce buffer and a blocking copy per tensor. |
@@ -819,6 +820,9 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_MKL_FA_DIAG | 0 (default) or 1 | Enable output fingerprinting for MKL flash attention. Dumps the first 64 float output values for the first 6 FA calls with n_kv ≥ 1024, labeled with kernel type (MKL/TILE/VEC) for cross-kernel comparison. |
 | GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute. Unsupported types and layouts fall back to the standalone op kernels. See `ggml_sycl_can_fuse()`. |
 | GGML_SYCL_ENABLE_ESIMD | 0 or 1 (default)| Enable ESIMD kernels when available. |
+| GGML_SYCL_ENABLE_XMX_ESIMD | 0 or 1 (default) | Use the XMX (DPAS) multi-column mat-vec kernels for reordered Q2_K to Q6_K and Q8_0 weights on Xe2 GPUs. Speeds up speculative decoding and short prompt batches. |
+| GGML_SYCL_XMX_MAX_COLS | non-negative integer, 80 (default) | Most activation columns (batch size) the `GGML_SYCL_ENABLE_XMX_ESIMD` kernels take; wider batches use the other mat-mul kernels. A larger value keeps XMX for bigger batches (as repeated launches), `0` never uses it. |
+| GGML_SYCL_XMX_GLU_MAX_COLS | 0 to 16, 16 (default) | Most columns the fused gate and up XMX kernel takes; wider batches run gate and up as separate mat-vecs. 16 is the most the kernel is built for, larger values are clamped to 16 and `0` turns the fused kernel off. |
 | GGML_SYCL_XMX_GATHER_TYPES | decimal bitmask, all bits set (default) | Weight formats that may use the XMX dequant-GEMM paths, which dequantize weights straight into the XMX tiles. This speeds up prompt processing of MoE models on GPUs with XMX units (Arc A- and B-series, Arc Pro, Data Center GPU Max), for example pp512 of Qwen3-30B-A3B UD-IQ3_XXS by about 50% on an Arc Pro B60. Bits:<br>* 1: IQ4_NL, 2: IQ3_S, 4: IQ4_XS, 8: IQ3_XXS, 16: IQ2_XXS, 32: IQ2_XS, 64: IQ2_S, 128: IQ1_S, 256: IQ1_M<br>* 512: Q8_0, 1024: Q4_K, 2048: Q5_K, 4096: Q6_K (MoE `MUL_MAT_ID` only)<br>Add values to combine them, for example `3` for IQ4_NL and IQ3_S; `0` disables the paths. A set bit does not force the path: batches of more than 64 tokens per expert or row lengths that are not a multiple of 256 (32 for IQ4_NL and Q8_0) use the library GEMM. |
 | GGML_SYCL_XMX_GATHER_SHAPES | decimal bitmask, 255 (default) | XMX `joint_matrix` combinations the paths of `GGML_SYCL_XMX_GATHER_TYPES` may use; the operand type comes from `GGML_SYCL_DYNAMIC_PRECISION` and the best supported combination is picked automatically (logged as `fg_pick_combo`). Bits:<br>* Xe2, Xe3, Xe-HPC: 1: f16 8x16x16, 2: f16 16x16x16, 4: f16 32x64x16, 8: f16 32x64x32, 32: tf32 8x16x8, 64: bf16 8x16x16<br>* Xe-HPG (Arc A770, ARL-H): 16: f16 8x8x16, 128: bf16 8x8x16<br>Clear a bit to exclude a combination, or set a single bit to force one for testing. |
 | GGML_SYCL_DYNAMIC_PRECISION | `F16` (default with `GGML_SYCL_F16=ON`), `BF16`, `TF32` or `F32` (default otherwise) | Operand type of the XMX dequant-GEMM paths (`GGML_SYCL_XMX_GATHER_TYPES`); accumulation is always f32. `F16` is the fastest, but activations above 65504 overflow. `BF16` keeps the f32 range at a 7-bit mantissa, `TF32` keeps the range and the f16 mantissa but is about 30% slower and needs Xe2, Xe3 or Xe-HPC, and `F32` turns the XMX paths off. Ops that request a higher src1 precision ([TAG_GGML_PREC]) get it regardless of this setting. |
@@ -831,6 +835,17 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS | 0 (default) or 1 | Allow SYCL/Unified Runtime Level Zero device allocations larger than 4 GiB. llama.cpp's direct Level Zero allocation path requests the relaxed maximum-size limit itself when GGML_SYCL_ENABLE_LEVEL_ZERO=1. |
 | GGML_SYCL_USM_SYSTEM | 0 (default) or 1 | Enable experimental support for [USM system allocations](https://github.khronos.org/SYCL_Reference/iface/usm_basic_concept.html#system-allocations) for large GPU buffers. This requires enough host memory for model weights and caches, an Intel Xe2+ GPU such as BMG or newer and supported on Linux only, with CONFIG_DRM_XE_GPUSVM enabled. |

+1. <a id="sycl-graph-footnote"></a>SYCL Graph support (`GGML_SYCL_ENABLE_GRAPH=1`) is limited to single-GPU workloads and may also depend on the model architecture and kernels in use.
+   Performance depends on how often a recorded graph can be reused, which will vary by workload.
+   For this reason, it is recommended to measure performance both with and without graphs enabled on your workloads.
+
+   To check whether graphs are supported for your workload, enable INFO logging (`-lv 4`) and look for `enabling SYCL graphs` or `disabling SYCL graphs`.
+   Support may differ between prefill and decode, so check the status for the phases that matter to you.
+   DEBUG logging (`-lv 5`) reports why graphs are disabled.
+
+   Then benchmark with your preferred tool, such as `llama-bench`.
+   For graph activity details (including reuse), set `GGML_SYCL_DEBUG=1` and look for `[SYCL-GRAPH]` messages.
+
 ## Compile-time Flags

 Pass these via `CXXFLAGS` or add a one-off `#define` to enable a flag on the spot.
@@ -6,8 +6,8 @@ The cross-compilation toolchain images are provided by the
 [Qualcomm Snapdragon Toolchain registry](https://github.com/snapdragon-toolchain).
 These Docker images include the Android NDK, OpenCL SDK, Hexagon SDK, CMake, and the necessary cross-compilers:

-* **Android toolchain**: `ghcr.io/snapdragon-toolchain/arm64-android:v0.7`
-* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.7`
+* **Android toolchain**: `ghcr.io/snapdragon-toolchain/arm64-android:v0.8`
+* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.8`

 The unified build utility (`scripts/snapdragon/build.py`) automatically pulls
 and orchestrates these containers to perform target compilation.
@@ -40,7 +40,7 @@ Alternatively, you can build llama.cpp manually by entering the cross-compilatio

 ```bash
 # Start the cross-compilation container manually:
-~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-android:v0.7
+~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-android:v0.8

 # Inside the container, build the project using presets:
 [d]/workspace> cp docs/backend/snapdragon/CMakeUserPresets.json .
@@ -53,7 +53,7 @@ Preset CMake variables:
  GGML_HEXAGON="ON"
  GGML_OPENCL="ON"
  GGML_OPENMP="OFF"
-  HEXAGON_SDK_ROOT="/opt/hexagon/6.6.0.0"
+  HEXAGON_SDK_ROOT="/opt/hexagon/6.7.0.0"
 ...
 -- Including OpenCL backend
 -- Including Hexagon backend
@@ -3,7 +3,7 @@
 The cross-compilation is performed using the Snapdragon Linux Docker toolchain image (see
 [github.com/snapdragon-toolchain](https://github.com/snapdragon-toolchain)):

-* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.7`
+* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.8`

 The unified build utility (`scripts/snapdragon/build.py`) automatically pulls
 and orchestrates this container to perform target compilation. You only need to
@@ -29,7 +29,7 @@ Alternatively, you can build llama.cpp manually by entering the cross-compilatio

 ```bash
 # Start the cross-compilation container manually:
-~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-linux:v0.7
+~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-linux:v0.8

 # Inside the container, build the project using presets:
 [d]/workspace> cp docs/backend/snapdragon/CMakeUserPresets.json .
@@ -8,7 +8,7 @@ Native Windows 11 arm64 builds have the following tool dependencies:
  - UCRT and Driver Kit
 - LLVM core libraries and Clang compiler (winget)
 - CMake, Git, Python (winget)
- Hexagon SDK Community Edition 6.6 or later (see below)
+- Hexagon SDK Community Edition 6.7 or later (see below)
 - OpenCL SDK 2.3 or later (see below)

 Note: The rest of the **Windows** build process assumes that you're running natively in Powershell.
@@ -54,15 +54,15 @@ c:\Qualcomm\OpenCL_SDK\2.3.2

 Either use the trimmed down version (optimized for CI) from

-    https://github.com/snapdragon-toolchain/hexagon-sdk/releases/download/v6.6.0.0/hexagon-sdk-v6.6.0.0-arm64-wos.tar.xz
+    https://github.com/snapdragon-toolchain/hexagon-sdk/releases/download/v6.7.0.0/hexagon-sdk-v6.7.0.0-arm64-wos.tar.xz

 Or download the complete official version from

-    https://softwarecenter.qualcomm.com/catalog/item/Hexagon_SDK?version=6.6.0.0
+    https://softwarecenter.qualcomm.com/catalog/item/Hexagon_SDK?version=6.7.0.0

 Unzip/untar the archive into
 ```
-c:\Qualcomm\Hexagon_SDK\6.6.0.0
+c:\Qualcomm\Hexagon_SDK\6.7.0.0
 ```

 ## Install the latest Adreno GPU driver
@@ -150,8 +150,8 @@ The overall Hexagon backend build procedure for Windows on Snapdragon is the sam
 However, additional settings are required for generating and signing HTP Ops libraries.
 ```
 > $env:OPENCL_SDK_ROOT="C:\Qualcomm\OpenCL_SDK\2.3.2"
-> $env:HEXAGON_SDK_ROOT="C:\Qualcomm\Hexagon_SDK\6.6.0.0"
-> $env:HEXAGON_TOOLS_ROOT="C:\Qualcomm\Hexagon_SDK\6.6.0.0\tools\HEXAGON_Tools\19.0.07"
+> $env:HEXAGON_SDK_ROOT="C:\Qualcomm\Hexagon_SDK\6.7.0.0"
+> $env:HEXAGON_TOOLS_ROOT="C:\Qualcomm\Hexagon_SDK\6.7.0.0\tools\HEXAGON_Tools\21.0.04"
 > $env:HEXAGON_HTP_CERT="c:\Users\MyUsers\Certs\ggml-htp-v1.pfx"
 > $env:WINDOWS_SDK_BIN="C:\Program Files (x86)\Windows Kits\10\bin\10.0.26100.0"

@@ -2,6 +2,7 @@

 > [!WARNING]
 > **Note:** zDNN is **not** the same as ZenDNN.
+>
 > - **zDNN** (this page): IBM's Deep Neural Network acceleration library for IBM Z & LinuxONE Mainframes
 > - **ZenDNN**: AMD's deep learning library for AMD EPYC CPUs ([see ZenDNN documentation](ZenDNN.md))

@@ -46,6 +47,14 @@ The IBM zDNN backend has the following CMake options that control the behaviour
 | --------------------- | ------------- | -------------------------------------------------------------------------------------- |
 | `GGML_ZDNN_HIGH_PREC` | `0`           | Set to `1` to requantize Q8_0 weights in smaller groups, for lower loss at lower speed |

+## Performance Tunables
+
+| Flag Option             | Default Value | Description                                            |
+| ----------------------- | ------------- | ------------------------------------------------------ |
+| `--flash-attn` or `-fa` | `auto`        | Set to `on` to run attention on the NNPA (Recommended) |
+
+Note: Set `-fa on` explicitly. Flash Attention fuses each layer's attention (the KQ matrix multiplication, softmax and KQV matrix multiplication) into one `FLASH_ATTN_EXT` operation, which zDNN runs on the NNPA instead of the CPU. The speedup grows with the context length. With `auto`, llama.cpp currently turns Flash Attention off, because the operation runs on a different device than the layer. zDNN runs it when the KV cache is F16 (the default), the model uses no attention sinks, ALiBi or logit softcapping, and there is one KV cache stream (one sequence, or `--kv-unified`). Otherwise it falls back to the CPU.
+
 ## 1. Install zDNN Library

 Note: A minimum of libzdnn 1.1.2 is required for the IBM zDNN backend to link correctly.
@@ -852,6 +852,36 @@ WebGPU allows cross-platform access to the GPU from supported browsers. We utili

 Follow the instructions [here](https://dawn.googlesource.com/dawn/+/refs/heads/main/src/emdawnwebgpu/) to download or build the emdawnwebgpu package (Note that it might be safer to build the emdawnwebgpu package locally, so that it stays in sync with the version of Dawn you have installed above). When building using CMake, the path to the emdawnwebgpu port file needs to be set with the flag `EMDAWNWEBGPU_DIR`.

+## WASI
+
+Requires [wasi-sdk](https://github.com/WebAssembly/wasi-sdk) 34+ with `WASI_SDK_PATH` set. Only the libraries build (`llama-common` does not support WASI yet), so link your program against them:
+
+```bash
+FLAGS="-msimd128 -D_WASI_EMULATED_SIGNAL -fwasm-exceptions -mllvm -wasm-use-legacy-eh=false"
+cmake -B build-wasi -DCMAKE_TOOLCHAIN_FILE=$WASI_SDK_PATH/share/cmake/wasi-sdk-p2.cmake \
+    -DCMAKE_C_FLAGS="$FLAGS" -DCMAKE_CXX_FLAGS="$FLAGS" -DCMAKE_HAVE_LIBC_PTHREAD=1 \
+    -DBUILD_SHARED_LIBS=OFF -DGGML_NATIVE=OFF -DGGML_OPENMP=OFF -DLLAMA_BUILD_COMMON=OFF
+cmake --build build-wasi --target llama
+
+$WASI_SDK_PATH/bin/clang++ --target=wasm32-wasip3 -O3 $FLAGS -I include -I ggml/include examples/simple/simple.cpp \
+    build-wasi/src/libllama.a build-wasi/ggml/src/libggml*.a -lunwind -ldl -Wl,-z,stack-size=8388608 -o llama-simple.wasm
+wasmtime run -W exceptions=y --dir models llama-simple.wasm -m models/model.gguf "Hello"
+```
+
+### wasi:webgpu
+
+The WebGPU backend can target [wasi:webgpu](https://github.com/WebAssembly/wasi-webgpu) via [wasi-webgpu-headers](https://github.com/wasi-gfx/wasi-webgpu-headers) (with its submodule) and Dawn's `dawn-headers` release. Point `Dawn_DIR` at a directory containing this `DawnConfig.cmake` next to both:
+
+```cmake
+set(_wasi_webgpu "${CMAKE_CURRENT_LIST_DIR}/wasi-webgpu-headers")
+add_library(dawn::webgpu_dawn INTERFACE IMPORTED)
+set_target_properties(dawn::webgpu_dawn PROPERTIES
+    INTERFACE_SOURCES "${_wasi_webgpu}/webgpu.c;${_wasi_webgpu}/async_futures.c;${_wasi_webgpu}/imports.c;${_wasi_webgpu}/imports_component_type.o"
+    INTERFACE_INCLUDE_DIRECTORIES "${CMAKE_CURRENT_LIST_DIR}/dawn-headers/include/webgpu_upstream;${_wasi_webgpu}")
+```
+
+Then build as above, but with `wasi-sdk-p3.cmake` / `--target=wasm32-wasip3`, `-DGGML_WEBGPU=ON -DDawn_DIR=...` at configure, and `build-wasi/ggml/src/ggml-webgpu/libggml-webgpu.a` added at link. See `.github/workflows/build-wasi.yml` for a more complete example.
+
 ## IBM Z & LinuxONE

 To read documentation for how to build on IBM Z & LinuxONE, [click here](./build-s390x.md)
@@ -27,7 +27,7 @@ Legend:
 |                        COL2IM_1D | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           CONCAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                             CONT | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
-|                       CONV_1D_DW | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ |
+|                       CONV_1D_DW | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ |
 |                          CONV_2D | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
 |                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          CONV_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -76,10 +76,10 @@ Legend:
 |                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
 |                 MUL_MAT_HADAMARD | ❌ | ❌ | ✅ | ❌ | ✅ | 🟡 | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 |
 |                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | 🟡 | ❌ |
-|                  MUL_MAT_ID_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-|                  MUL_MAT_ID_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-|                     MUL_MAT_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-|                     MUL_MAT_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -2,10 +2,16 @@ function(ggml_get_flags CCID CCVER)
    set(C_FLAGS "")
    set(CXX_FLAGS "")

-    if (CCID MATCHES "Clang")
+    if (CCID MATCHES "Clang" OR CCID MATCHES "IntelLLVM")
        set(C_FLAGS   -Wunreachable-code-break -Wunreachable-code-return)
        set(CXX_FLAGS -Wunreachable-code-break -Wunreachable-code-return -Wmissing-prototypes -Wextra-semi)

+        if (CCID STREQUAL "IntelLLVM")
+            # for Intel DPC++ and LLVM SYCL, using debug does not disable optimization
+            list(APPEND C_FLAGS    -Rno-debug-disables-optimization)
+            list(APPEND CXX_FLAGS  -Rno-debug-disables-optimization)
+        endif()
+
        if (
            (CCID STREQUAL "Clang"      AND CCVER VERSION_GREATER_EQUAL 3.8.0) OR
            (CCID STREQUAL "AppleClang" AND CCVER VERSION_GREATER_EQUAL 7.3.0)
@@ -8,10 +8,10 @@ extern "C" {

 #define RPC_PROTO_MAJOR_VERSION    8
 #define RPC_PROTO_MINOR_VERSION    0
-#define RPC_PROTO_PATCH_VERSION    0
+#define RPC_PROTO_PATCH_VERSION    1

 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif

 #define GGML_RPC_MAX_SERVERS       16
@@ -585,6 +585,7 @@ extern "C" {
        GGML_OP_DSV4_HC_COMB,
        GGML_OP_DSV4_HC_PRE,
        GGML_OP_DSV4_HC_POST,
+        GGML_OP_TOP_K_MASK,

        GGML_OP_UNARY,

@@ -2732,6 +2733,26 @@ extern "C" {
            struct ggml_tensor  * post,
            struct ggml_tensor  * comb);

+    // top-k sparse attention mask: keep the mask cells of the live picks and the extra cells, set the rest to -inf
+    //
+    // mask:   [n_kv,    m1,  m2,  m3]  f16 or f32, broadcast to the rows of top_k
+    // top_k:  [n_top,   ne1, ne2, ne3] i32, picks of each row
+    // score:  [n_pick,  ne1, ne2, ne3] f32, can be NULL: a pick is live when its score is above -inf, else every pick is live
+    // blocks: [n_blk,   n_pick]        i32, can be NULL: the member cells of each pick, else a pick is a cell
+    // extra:  [n_extra, ne1, ne2, ne3] i32, can be NULL: cells of each row kept without a pick
+    // res:    [n_kv,    ne1, ne2, ne3] same type as mask
+    //
+    // res[c, i] = mask[c, i % m] when c is a cell of a live pick of row i or an extra cell of row i, else -inf
+    // picks and cells out of range are skipped
+    //
+    GGML_API struct ggml_tensor * ggml_top_k_mask(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * mask,
+            struct ggml_tensor  * top_k,
+            struct ggml_tensor  * score,
+            struct ggml_tensor  * blocks,
+            struct ggml_tensor  * extra);
+
    // custom operators

    typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata);
@@ -774,6 +774,13 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
        return src_ss[0];
    };

+    auto handle_pool_1d = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+        if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) {
+            return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
+        }
+        return src_ss[0];
+    };
+
    auto handle_flash_attn_ext = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
        GGML_ASSERT(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);

@@ -990,12 +997,14 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
            case GGML_OP_CONV_3D:
            case GGML_OP_CONV_2D_DW:
            case GGML_OP_CONV_TRANSPOSE_2D:
-            case GGML_OP_POOL_1D:
            case GGML_OP_POOL_2D:
            case GGML_OP_POOL_2D_BACK:
            case GGML_OP_UPSCALE: {
                split_state = handle_generic(src_ss, /*scalar_only =*/ true);
            } break;
+            case GGML_OP_POOL_1D: {
+                split_state = handle_pool_1d(src_ss);
+            } break;
            case GGML_OP_PAD: {
                split_state = handle_pad(src_ss);
            } break;
@@ -1046,7 +1055,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
            } break;
            case GGML_OP_DSV4_HC_COMB:
            case GGML_OP_DSV4_HC_PRE:
-            case GGML_OP_DSV4_HC_POST: {
+            case GGML_OP_DSV4_HC_POST:
+            case GGML_OP_TOP_K_MASK: {
                split_state = handle_generic(src_ss, /*scalar_only =*/ true);
            } break;
            case GGML_OP_UNARY: {
@@ -2226,6 +2226,7 @@ bool ggml_op_alloc_size_may_expand(enum ggml_op op) {
        case GGML_OP_CUMSUM:
        case GGML_OP_ARGSORT:
        case GGML_OP_TOP_K:
+        case GGML_OP_ARGMAX:
            return true;
        default:
            return false;
@@ -920,7 +920,7 @@ static void * ggml_backend_cann_buffer_get_base(ggml_backend_buffer_t buffer) {
 * @param dst Pointer to the destination buffer where transformed data will be
 * stored.
 */
-static void ggml_backend_cann_transform_q4_0(ggml_tensor * tensor, const void * src, void * dst) {
+static void ggml_backend_cann_transform_q4_0(const ggml_tensor * tensor, const void * src, void * dst) {
    int64_t n_elems     = ggml_nelements(tensor);
    int64_t groups      = n_elems / QK4_0;
    size_t  quant_bytes = n_elems * sizeof(uint8_t) / 2;
@@ -1014,7 +1014,7 @@ static void ggml_backend_cann_transform_back_q4_0(const ggml_tensor * tensor, vo
 * @param dst Pointer to the destination buffer where transformed data will be
 * stored.
 */
-static void ggml_backend_cann_transform_q8_0(ggml_tensor * tensor, const void * src, void * dst) {
+static void ggml_backend_cann_transform_q8_0(const ggml_tensor * tensor, const void * src, void * dst) {
    int64_t n_elems     = ggml_nelements(tensor);
    int64_t groups      = n_elems / QK8_0;
    size_t  quant_bytes = n_elems * sizeof(uint8_t);
@@ -1075,7 +1075,7 @@ static void ggml_backend_cann_transform_back_q8_0(const ggml_tensor * tensor, co
 * @param dst Pointer to the destination buffer where transformed data will be
 * stored.
 */
-static void ggml_backend_cann_transform(ggml_tensor * tensor, const void * src, void * dst) {
+static void ggml_backend_cann_transform(const ggml_tensor * tensor, const void * src, void * dst) {
    switch (tensor->type) {
        case GGML_TYPE_Q4_0:
            ggml_backend_cann_transform_q4_0(tensor, src, dst);
@@ -650,9 +650,9 @@ function(ggml_add_cpu_backend_variant_impl tag_name)

        # Fetch KleidiAI sources:
        include(FetchContent)
-        set(KLEIDIAI_COMMIT_TAG "v1.31.0")
+        set(KLEIDIAI_COMMIT_TAG "v1.32.0")
        set(KLEIDIAI_DOWNLOAD_URL "https://github.com/ARM-software/kleidiai/releases/download/${KLEIDIAI_COMMIT_TAG}/kleidiai-${KLEIDIAI_COMMIT_TAG}-src.tar.gz")
-        set(KLEIDIAI_RELEASE_ARCHIVE_MD5  "930495f3aeb16b25dc18a2bde45924f4")
+        set(KLEIDIAI_RELEASE_ARCHIVE_MD5  "05f1953428ffe8e8c7f585dc6d508887")

        set(KLEIDIAI_FETCH_ARGS
            URL ${KLEIDIAI_DOWNLOAD_URL}
@@ -2105,6 +2105,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
            {
                ggml_compute_forward_dsv4_hc_post(params, tensor);
            } break;
+        case GGML_OP_TOP_K_MASK:
+            {
+                ggml_compute_forward_top_k_mask(params, tensor);
+            } break;
        case GGML_OP_MAP_CUSTOM1:
            {
                ggml_compute_forward_map_custom1(params, tensor);
@@ -2288,6 +2292,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
        case GGML_OP_DSV4_HC_COMB:
        case GGML_OP_DSV4_HC_PRE:
        case GGML_OP_DSV4_HC_POST:
+        case GGML_OP_TOP_K_MASK:
            {
                n_tasks = n_threads;
            } break;
@@ -201,85 +201,117 @@ static struct kai_matmul_uker_config q4_k_matmul_config() {
    return config;
 }

-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_m_step() {
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    return api.get_step(&config).m;
+static struct kai_matmul_uker_config default_matmul_config() {
+    return {};
 }

-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_n_step() {
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    return api.get_step(&config).n;
-}
+template<
+        struct kai_matmul_uker_api (*GetApi)(void),
+        struct kai_matmul_uker_config (*GetConfig)(void),
+        size_t MR, size_t NR, size_t KR, size_t SR>
+struct matmul_uker_adapter {
+    static size_t get_m_step() {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        return api.get_step(&config).m;
+    }

-template<size_t Value>
-static size_t q4_k_get_constant() {
-    return Value;
-}
+    static size_t get_n_step() {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        return api.get_step(&config).n;
+    }

-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_dst_offset(size_t m_idx, size_t n_idx, size_t dst_stride) {
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_dst_dim_args index = { m_idx, n_idx };
-    const struct kai_matmul_uker_dst_stride_args stride = { dst_stride };
-    return api.get_dst_offset(&config, &index, &stride);
-}
+    static size_t get_mr() {
+        if constexpr (MR == 0) {
+            return get_m_step();
+        }
+        return MR;
+    }

-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_lhs_packed_offset(size_t m_idx, size_t k, size_t bl) {
-    GGML_ASSERT(bl == KLEIDIAI_Q4_K_BLOCK_LENGTH);
+    static size_t get_nr() {
+        if constexpr (NR == 0) {
+            return get_n_step();
+        }
+        return NR;
+    }

-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_lhs_dim_args shape = { m_idx, k };
-    const struct kai_matmul_uker_lhs_dim_args index = { m_idx, 0 };
-    const struct kai_matmul_uker_lhs_stride_args stride = api.get_lhs_stride(&config, &shape);
-    return api.get_lhs_offset(&config, &index, &stride);
-}
+    static size_t get_kr() {
+        return KR;
+    }

-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_rhs_packed_offset(size_t n_idx, size_t k, size_t bl) {
-    GGML_ASSERT(bl == KLEIDIAI_Q4_K_BLOCK_LENGTH);
+    static size_t get_sr() {
+        return SR;
+    }

-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_rhs_dim_args shape = { n_idx, k };
-    const struct kai_matmul_uker_rhs_dim_args index = { n_idx, 0 };
-    const struct kai_matmul_uker_rhs_stride_args stride = api.get_rhs_stride(&config, &shape);
-    return api.get_rhs_offset(&config, &index, &stride);
-}
+    static size_t get_dst_offset(size_t m_idx, size_t n_idx, size_t dst_stride) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        const struct kai_matmul_uker_dst_dim_args index = { m_idx, n_idx };
+        const struct kai_matmul_uker_dst_stride_args stride = { dst_stride };
+        return api.get_dst_offset(&config, &index, &stride);
+    }

-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static void q4_k_run_matmul(
-        size_t m, size_t n, size_t k, size_t bl,
-        const void * lhs_packed, const void * rhs_packed,
-        void * dst, size_t dst_stride_row, size_t dst_stride_col,
-        float clamp_min, float clamp_max) {
-    GGML_ASSERT(bl == KLEIDIAI_Q4_K_BLOCK_LENGTH);
-    GGML_ASSERT(dst_stride_col == sizeof(float));
+    static size_t get_lhs_packed_offset(size_t m_idx, size_t k, size_t bl) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        GGML_ASSERT(bl == config.format.bl);
+        const struct kai_matmul_uker_lhs_dim_args shape = { m_idx, k };
+        const struct kai_matmul_uker_lhs_dim_args index = { m_idx, 0 };
+        const struct kai_matmul_uker_lhs_stride_args stride = api.get_lhs_stride(&config, &shape);
+        return api.get_lhs_offset(&config, &index, &stride);
+    }

-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_lhs_dim_args lhs_shape = { m, k };
-    const struct kai_matmul_uker_rhs_dim_args rhs_shape = { n, k };
+    static size_t get_rhs_packed_offset(size_t n_idx, size_t k, size_t bl) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        GGML_ASSERT(bl == config.format.bl);
+        const struct kai_matmul_uker_rhs_dim_args shape = { n_idx, k };
+        const struct kai_matmul_uker_rhs_dim_args index = { n_idx, 0 };
+        const struct kai_matmul_uker_rhs_stride_args stride = api.get_rhs_stride(&config, &shape);
+        return api.get_rhs_offset(&config, &index, &stride);
+    }

-    struct kai_matmul_uker_args args = {};
-    args.flags = KAI_MATMUL_UKER_FLAGS_ARGS_CLAMP;
-    args.shape = { m, n, k };
-    args.operand.lhs.ptr = lhs_packed;
-    args.operand.lhs.stride = api.get_lhs_stride(&config, &lhs_shape);
-    args.operand.rhs.ptr = rhs_packed;
-    args.operand.rhs.stride = api.get_rhs_stride(&config, &rhs_shape);
-    args.operand.dst.ptr = dst;
-    args.operand.dst.stride.m = dst_stride_row;
-    args.activation.clamp.min_ptr = &clamp_min;
-    args.activation.clamp.max_ptr = &clamp_max;
-    api.run(&config, &args);
-}
+    static void run(
+            size_t m, size_t n, size_t k, size_t bl,
+            const void * lhs_packed, const void * rhs_packed,
+            void * dst, size_t dst_stride_row, size_t dst_stride_col,
+            float clamp_min, float clamp_max) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        GGML_ASSERT(bl == config.format.bl);
+        GGML_ASSERT(dst_stride_col == sizeof(float));
+        const struct kai_matmul_uker_lhs_dim_args lhs_shape = { m, k };
+        const struct kai_matmul_uker_rhs_dim_args rhs_shape = { n, k };
+
+        struct kai_matmul_uker_args args = {};
+        args.flags = KAI_MATMUL_UKER_FLAGS_ARGS_CLAMP;
+        args.shape = { m, n, k };
+        args.operand.lhs.ptr = lhs_packed;
+        args.operand.lhs.stride = api.get_lhs_stride(&config, &lhs_shape);
+        args.operand.rhs.ptr = rhs_packed;
+        args.operand.rhs.stride = api.get_rhs_stride(&config, &rhs_shape);
+        args.operand.dst.ptr = dst;
+        args.operand.dst.stride.m = dst_stride_row;
+        args.activation.clamp.min_ptr = &clamp_min;
+        args.activation.clamp.max_ptr = &clamp_max;
+        api.run(&config, &args);
+    }
+};
+
+using q4_k_gemm_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa,
+    q4_k_matmul_config, 0, 0, 4, 2>;
+using q4_k_gemv_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot,
+    q4_k_matmul_config, 0, 0, 4, 2>;
+// GEMM and GEMV share the RHS layout only at SVE256 (NR=8, KR=8, SR=1).
+using sve_q8_gemm_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_qai8dxp4x8sf32_qsi8cxp4vsx8sf32bf32_16x4vs_sve_i8mm,
+    default_matmul_config, 4, 0, 8, 1>;
+using sve_q8_gemv_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_qai8dxp1x8_qsi8cxp8x8_1x8_sve_dot,
+    default_matmul_config, 1, 8, 8, 1>;

 static struct kai_matmul_pack_lhs_uker_config q4_k_lhs_pack_config() {
    struct kai_matmul_pack_lhs_uker_config config = {};
@@ -628,7 +660,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_size_ex        = */ &rhs_ps_fn5<kai_get_rhs_packed_size_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
            /* .pack_func_ex          = */ &rhs_pack_fn12<kai_run_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
        },
-        /* .required_cpu       = */ CPU_FEATURE_SVE | CPU_FEATURE_I8MM | CPU_FEATURE_DOTPROD,
+        /* .required_cpu       = */ CPU_FEATURE_SVE | CPU_FEATURE_SVEI8MM | CPU_FEATURE_DOTPROD,
        /* .lhs_type           = */ GGML_TYPE_F32,
        /* .rhs_type           = */ GGML_TYPE_Q4_0,
        /* .op_type            = */ GGML_TYPE_F32,
@@ -737,16 +769,16 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q4_k[] = {
    {
        /* SME2 GEMM */
        {
-            /* .get_m_step            = */ &q4_k_get_m_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_n_step            = */ &q4_k_get_n_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_mr                = */ &q4_k_get_m_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_nr                = */ &q4_k_get_n_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_kr                = */ &q4_k_get_constant<4>,
-            /* .get_sr                = */ &q4_k_get_constant<2>,
-            /* .get_dst_offset        = */ &q4_k_get_dst_offset<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_lhs_offset_ex     = */ &q4_k_get_lhs_packed_offset<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_rhs_packed_offset_ex = */ &q4_k_get_rhs_packed_offset<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .run_kernel_ex         = */ &q4_k_run_matmul<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
+            /* .get_m_step            = */ q4_k_gemm_adapter::get_m_step,
+            /* .get_n_step            = */ q4_k_gemm_adapter::get_n_step,
+            /* .get_mr                = */ q4_k_gemm_adapter::get_mr,
+            /* .get_nr                = */ q4_k_gemm_adapter::get_nr,
+            /* .get_kr                = */ q4_k_gemm_adapter::get_kr,
+            /* .get_sr                = */ q4_k_gemm_adapter::get_sr,
+            /* .get_dst_offset        = */ q4_k_gemm_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ q4_k_gemm_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ q4_k_gemm_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ q4_k_gemm_adapter::run,
        },
        {
            /* .get_offset            = */ kai_get_lhs_offset_lhs_pack_f16pmrx2_f32_neon,
@@ -756,16 +788,16 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q4_k[] = {
        },
        /* SME2 GEMV */
        {
-            /* .get_m_step            = */ &q4_k_get_m_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_n_step            = */ &q4_k_get_n_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_mr                = */ &q4_k_get_m_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_nr                = */ &q4_k_get_n_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_kr                = */ &q4_k_get_constant<4>,
-            /* .get_sr                = */ &q4_k_get_constant<2>,
-            /* .get_dst_offset        = */ &q4_k_get_dst_offset<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_lhs_offset_ex     = */ &q4_k_get_lhs_packed_offset<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_rhs_packed_offset_ex = */ &q4_k_get_rhs_packed_offset<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .run_kernel_ex         = */ &q4_k_run_matmul<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
+            /* .get_m_step            = */ q4_k_gemv_adapter::get_m_step,
+            /* .get_n_step            = */ q4_k_gemv_adapter::get_n_step,
+            /* .get_mr                = */ q4_k_gemv_adapter::get_mr,
+            /* .get_nr                = */ q4_k_gemv_adapter::get_nr,
+            /* .get_kr                = */ q4_k_gemv_adapter::get_kr,
+            /* .get_sr                = */ q4_k_gemv_adapter::get_sr,
+            /* .get_dst_offset        = */ q4_k_gemv_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ q4_k_gemv_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ q4_k_gemv_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ q4_k_gemv_adapter::run,
        },
        {
            /* .get_offset            = */ q4_k_lhs_get_offset,
@@ -883,6 +915,54 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
        /* .rhs_type           = */ GGML_TYPE_Q8_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
+    {
+        /* SVE I8MM GEMM */
+        {
+            /* .get_m_step            = */ sve_q8_gemm_adapter::get_m_step,
+            /* .get_n_step            = */ sve_q8_gemm_adapter::get_n_step,
+            /* .get_mr                = */ sve_q8_gemm_adapter::get_mr,
+            /* .get_nr                = */ sve_q8_gemm_adapter::get_nr,
+            /* .get_kr                = */ sve_q8_gemm_adapter::get_kr,
+            /* .get_sr                = */ sve_q8_gemm_adapter::get_sr,
+            /* .get_dst_offset        = */ sve_q8_gemm_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ sve_q8_gemm_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ sve_q8_gemm_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ sve_q8_gemm_adapter::run,
+        },
+        /* .gemm_lhs_info = */ {
+            /* .get_offset            = */ kai_get_lhs_offset_lhs_quant_pack_qai8dxp_f32,
+            /* .get_packed_offset_ex  = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_quant_pack_qai8dxp_f32>,
+            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32>,
+            /* .pack_func_ex          = */ &lhs_pack_float_fn9_no_bl<kai_run_lhs_quant_pack_qai8dxp_f32>,
+        },
+        /* SVE DOT GEMV */
+        {
+            /* .get_m_step            = */ sve_q8_gemv_adapter::get_m_step,
+            /* .get_n_step            = */ sve_q8_gemv_adapter::get_n_step,
+            /* .get_mr                = */ sve_q8_gemv_adapter::get_mr,
+            /* .get_nr                = */ sve_q8_gemv_adapter::get_nr,
+            /* .get_kr                = */ sve_q8_gemv_adapter::get_kr,
+            /* .get_sr                = */ sve_q8_gemv_adapter::get_sr,
+            /* .get_dst_offset        = */ sve_q8_gemv_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ sve_q8_gemv_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ sve_q8_gemv_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ sve_q8_gemv_adapter::run,
+        },
+        /* .gemv_lhs_info = */ {
+            /* .get_offset            = */ kai_get_lhs_offset_lhs_quant_pack_qai8dxp_f32,
+            /* .get_packed_offset_ex  = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_quant_pack_qai8dxp_f32>,
+            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32>,
+            /* .pack_func_ex          = */ &lhs_pack_float_fn9_no_bl<kai_run_lhs_quant_pack_qai8dxp_f32>,
+        },
+        /* .rhs_info = */ {
+            /* .packed_size_ex        = */ &rhs_ps_fn5<kai_get_rhs_packed_size_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
+            /* .pack_func_ex          = */ &rhs_pack_scale_fn12<kai_run_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
+        },
+        /* .required_cpu       = */ CPU_FEATURE_SVE | CPU_FEATURE_SVEI8MM | CPU_FEATURE_DOTPROD,
+        /* .lhs_type           = */ GGML_TYPE_F32,
+        /* .rhs_type           = */ GGML_TYPE_Q8_0,
+        /* .op_type            = */ GGML_TYPE_F32,
+    },
    {
        /* I8MM GEMM */
        {
@@ -13,7 +13,8 @@ enum cpu_feature {
    CPU_FEATURE_SVE     = 4,
    CPU_FEATURE_SME     = 8,
    CPU_FEATURE_SME2    = 16,
-    CPU_FEATURE_FP16    = 32
+    CPU_FEATURE_FP16    = 32,
+    CPU_FEATURE_SVEI8MM = 64
 };

 inline cpu_feature& operator|=(cpu_feature& lhs, cpu_feature rhs) {
@@ -319,8 +319,10 @@ static void init_kleidiai_context(void) {

        size_t detected_smcus = 0;

+        // The fixed-width SVE kernels require 32-byte vectors, including the Q8 shared RHS layout.
        ctx.features  = (runtime_feat.has_dotprod  ? CPU_FEATURE_DOTPROD : CPU_FEATURE_NONE) |
                        (runtime_feat.has_i8mm     ? CPU_FEATURE_I8MM    : CPU_FEATURE_NONE) |
+                        (runtime_feat.has_svei8mm  ? CPU_FEATURE_SVEI8MM : CPU_FEATURE_NONE) |
                        (runtime_feat.has_fp16     ? CPU_FEATURE_FP16    : CPU_FEATURE_NONE) |
                        (runtime_feat.sve_cnt == QK8_0 ? CPU_FEATURE_SVE : CPU_FEATURE_NONE);

@@ -1616,6 +1618,11 @@ public:
            const size_t kr = kernel->get_kr();
            const size_t sr = kernel->get_sr();
            const ggml_type rhs_type = kernels->rhs_type;
+            if (rhs_type == GGML_TYPE_Q8_0 && (kernels->required_cpu & CPU_FEATURE_SVE)) {
+                GGML_ASSERT(nr == kernels->gemv.get_nr());
+                GGML_ASSERT(kr == kernels->gemv.get_kr());
+                GGML_ASSERT(sr == kernels->gemv.get_sr());
+            }
            const size_t block_len = rhs_type == GGML_TYPE_Q8_0 ? QK8_0 :
                                     rhs_type == GGML_TYPE_Q4_0 ? QK4_0 :
                                     rhs_type == GGML_TYPE_Q4_K ? GGML_KLEIDIAI_Q4_K_BL :
@@ -11736,6 +11736,92 @@ void ggml_compute_forward_dsv4_hc_post(
    }
 }

+// ggml_compute_forward_top_k_mask
+
+void ggml_compute_forward_top_k_mask(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+    const ggml_tensor * mask   = dst->src[0];
+    const ggml_tensor * top_k  = dst->src[1];
+    const ggml_tensor * score  = dst->src[2];
+    const ggml_tensor * blocks = dst->src[3];
+    const ggml_tensor * extra  = dst->src[4];
+
+    GGML_ASSERT(dst->type == mask->type);
+    GGML_ASSERT(dst->nb[0] == ggml_type_size(dst->type));
+
+    const int64_t n_kv    = mask->ne[0];
+    const int64_t ne1     = dst->ne[1];
+    const int64_t ne2     = dst->ne[2];
+    const int64_t n_rows  = ggml_nrows(dst);
+    const int64_t n_top   = top_k->ne[0];
+    const int64_t n_blk   = blocks ? blocks->ne[0] : 1;
+    const int64_t n_pick  = blocks ? blocks->ne[1] : score ? score->ne[0] : n_kv;
+    const int64_t n_extra = extra ? extra->ne[0] : 0;
+
+    const size_t ts = ggml_type_size(dst->type);
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    for (int64_t ir = ith; ir < n_rows; ir += nth) {
+        const int64_t i1 = ir % ne1;
+        const int64_t i2 = (ir / ne1) % ne2;
+        const int64_t i3 = ir / (ne1*ne2);
+
+        const char * mask_row = (const char *) mask->data
+            + (i1 % mask->ne[1])*mask->nb[1] + (i2 % mask->ne[2])*mask->nb[2] + (i3 % mask->ne[3])*mask->nb[3];
+        char * dst_row = (char *) dst->data + i1*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3];
+
+        if (dst->type == GGML_TYPE_F16) {
+            const ggml_fp16_t ninf = GGML_CPU_FP32_TO_FP16(-INFINITY);
+            for (int64_t c = 0; c < n_kv; ++c) {
+                ((ggml_fp16_t *) dst_row)[c] = ninf;
+            }
+        } else {
+            for (int64_t c = 0; c < n_kv; ++c) {
+                ((float *) dst_row)[c] = -INFINITY;
+            }
+        }
+
+        auto keep = [&](int32_t c) {
+            if (c >= 0 && c < n_kv) {
+                memcpy(dst_row + c*ts, mask_row + c*ts, ts);
+            }
+        };
+
+        const char * top_k_row = (const char *) top_k->data + i1*top_k->nb[1] + i2*top_k->nb[2] + i3*top_k->nb[3];
+
+        for (int64_t j = 0; j < n_top; ++j) {
+            const int32_t p = *(const int32_t *) (top_k_row + j*top_k->nb[0]);
+            if (p < 0 || p >= n_pick) {
+                continue;
+            }
+            if (score) {
+                const float s = *(const float *) ((const char *) score->data
+                    + p*score->nb[0] + i1*score->nb[1] + i2*score->nb[2] + i3*score->nb[3]);
+                if (!(s > -INFINITY)) {
+                    continue;
+                }
+            }
+            if (!blocks) {
+                keep(p);
+                continue;
+            }
+            for (int64_t k = 0; k < n_blk; ++k) {
+                keep(*(const int32_t *) ((const char *) blocks->data + k*blocks->nb[0] + p*blocks->nb[1]));
+            }
+        }
+
+        if (extra) {
+            const char * extra_row = (const char *) extra->data + i1*extra->nb[1] + i2*extra->nb[2] + i3*extra->nb[3];
+            for (int64_t j = 0; j < n_extra; ++j) {
+                keep(*(const int32_t *) (extra_row + j*extra->nb[0]));
+            }
+        }
+    }
+}
+
 // ggml_compute_forward_rwkv_wkv7

 static void ggml_compute_forward_rwkv_wkv7_f32(
@@ -108,6 +108,7 @@ void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * p
 void ggml_compute_forward_dsv4_hc_comb(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_dsv4_hc_pre(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_dsv4_hc_post(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_top_k_mask(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst);
@@ -42,6 +42,7 @@ extern "C" {
    #define GGML_CPU_COMPUTE_FP32_TO_FP16(x) neon_compute_fp32_to_fp16(x)

    #define GGML_CPU_FP16_TO_FP32(x) GGML_CPU_COMPUTE_FP16_TO_FP32(x)
+    #define GGML_CPU_FP32_TO_FP16(x) GGML_CPU_COMPUTE_FP32_TO_FP16(x)

    static inline float neon_compute_fp16_to_fp32(ggml_fp16_t h) {
        __fp16 tmp;
@@ -150,9 +151,15 @@ extern float ggml_table_f32_ue4m3[1 << 8];
 #define GGML_CPU_FP32_TO_FP16(x) GGML_CPU_COMPUTE_FP32_TO_FP16(x)
 #endif

-// On ARM NEON, it's quicker to directly convert x -> x instead of calling into ggml_lookup_fp16_to_fp32,
-// so we define GGML_CPU_FP16_TO_FP32 and GGML_CPU_FP32_TO_FP16 elsewhere for NEON.
-// This is also true for POWER9.
+// a branch that defines a GGML_CPU_COMPUTE_* macro must also select it, or the fallback takes over silently
+#if defined(GGML_CPU_COMPUTE_FP16_TO_FP32) && !defined(GGML_CPU_FP16_TO_FP32) && !defined(__F16C__) // see: https://github.com/ggml-org/llama.cpp/issues/30342
+#error "a branch defined GGML_CPU_COMPUTE_FP16_TO_FP32 without selecting it as GGML_CPU_FP16_TO_FP32"
+#endif
+#if defined(GGML_CPU_COMPUTE_FP32_TO_FP16) && !defined(GGML_CPU_FP32_TO_FP16)
+#error "a branch defined GGML_CPU_COMPUTE_FP32_TO_FP16 without selecting it as GGML_CPU_FP32_TO_FP16"
+#endif
+
+// fall back to the lookup table / the portable software conversion for whichever direction no branch selected
 #if !defined(GGML_CPU_FP16_TO_FP32)
 inline static float ggml_lookup_fp16_to_fp32(ggml_fp16_t f) {
    uint16_t s;
@@ -460,6 +460,92 @@ void ggml_vec_swiglu_f32(const int n, float * y, const float * x, const float *
    }
 }

+#ifndef GGML_GELU_QUICK_FP16
+// computed in f32 without the fp16 table, which loses precision on f32 input
+void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) {
+    int i = 0;
+#if defined(__AVX512F__) && defined(__AVX512DQ__)
+    for (; i + 15 < n; i += 16) {
+        _mm512_storeu_ps(y + i, ggml_v_gelu_quick(_mm512_loadu_ps(x + i)));
+    }
+#elif defined(__AVX2__) && defined(__FMA__)
+    for (; i + 7 < n; i += 8) {
+        _mm256_storeu_ps(y + i, ggml_v_gelu_quick(_mm256_loadu_ps(x + i)));
+    }
+#elif defined(__SSE2__)
+    for (; i + 3 < n; i += 4) {
+        _mm_storeu_ps(y + i, ggml_v_gelu_quick(_mm_loadu_ps(x + i)));
+    }
+#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__)
+    const int vlen = svcntw();
+    for (; i < n; i += vlen) {
+        const svbool_t pg = svwhilelt_b32_s32(i, n);
+        svst1_f32(pg, y + i, ggml_v_gelu_quick(pg, svld1_f32(pg, x + i)));
+    }
+#elif defined(__ARM_NEON) && defined(__aarch64__)
+    for (; i + 3 < n; i += 4) {
+        vst1q_f32(y + i, ggml_v_gelu_quick(vld1q_f32(x + i)));
+    }
+#elif defined(__VXE__) || defined(__VXE2__)
+    for (; i + 3 < n; i += 4) {
+        vec_xst(ggml_v_gelu_quick(vec_xl(0, x + i)), 0, y + i);
+    }
+#elif defined(__riscv_v_intrinsic)
+    for (int vl; i < n; i += vl) {
+        vl = __riscv_vsetvl_e32m2(n - i);
+        vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl);
+        vfloat32m2_t vy = ggml_v_gelu_quick_m2(vx, vl);
+        __riscv_vse32_v_f32m2(&y[i], vy, vl);
+    }
+#endif
+    for (; i < n; ++i) {
+        y[i] = ggml_gelu_quick_f32(x[i]);
+    }
+}
+
+void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) {
+    int i = 0;
+#if defined(__AVX512F__) && defined(__AVX512DQ__)
+    for (; i + 15 < n; i += 16) {
+        _mm512_storeu_ps(y + i, _mm512_mul_ps(ggml_v_gelu_quick(_mm512_loadu_ps(x + i)), _mm512_loadu_ps(g + i)));
+    }
+#elif defined(__AVX2__) && defined(__FMA__)
+    for (; i + 7 < n; i += 8) {
+        _mm256_storeu_ps(y + i, _mm256_mul_ps(ggml_v_gelu_quick(_mm256_loadu_ps(x + i)), _mm256_loadu_ps(g + i)));
+    }
+#elif defined(__SSE2__)
+    for (; i + 3 < n; i += 4) {
+        _mm_storeu_ps(y + i, _mm_mul_ps(ggml_v_gelu_quick(_mm_loadu_ps(x + i)), _mm_loadu_ps(g + i)));
+    }
+#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__)
+    const int vlen = svcntw();
+    for (; i < n; i += vlen) {
+        const svbool_t pg = svwhilelt_b32_s32(i, n);
+        svst1_f32(pg, y + i, svmul_f32_x(pg, ggml_v_gelu_quick(pg, svld1_f32(pg, x + i)), svld1_f32(pg, g + i)));
+    }
+#elif defined(__ARM_NEON) && defined(__aarch64__)
+    for (; i + 3 < n; i += 4) {
+        vst1q_f32(y + i, vmulq_f32(ggml_v_gelu_quick(vld1q_f32(x + i)), vld1q_f32(g + i)));
+    }
+#elif defined(__VXE__) || defined(__VXE2__)
+    for (; i + 3 < n; i += 4) {
+        vec_xst(vec_mul(ggml_v_gelu_quick(vec_xl(0, x + i)), vec_xl(0, g + i)), 0, y + i);
+    }
+#elif defined(__riscv_v_intrinsic)
+    for (int vl; i < n; i += vl) {
+        vl = __riscv_vsetvl_e32m2(n - i);
+        vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl);
+        vfloat32m2_t vg = __riscv_vle32_v_f32m2(&g[i], vl);
+        vfloat32m2_t vy = __riscv_vfmul_vv_f32m2(ggml_v_gelu_quick_m2(vx, vl), vg, vl);
+        __riscv_vse32_v_f32m2(&y[i], vy, vl);
+    }
+#endif
+    for (; i < n; ++i) {
+        y[i] = ggml_gelu_quick_f32(x[i]) * g[i];
+    }
+}
+#endif
+
 ggml_float ggml_vec_cvar_f32(const int n, float * y, const float * x, const float mean) {
    int i = 0;
    ggml_float sum = 0;
@@ -44,7 +44,6 @@ inline static ggml_float ggml_sve_sum_f32x2(svfloat32_t sum_lo, svfloat32_t sum_
 #endif

 #define GGML_GELU_FP16
-#define GGML_GELU_QUICK_FP16

 #define GGML_SOFT_MAX_UNROLL 4
 #define GGML_VEC_DOT_UNROLL  2
@@ -1054,11 +1053,7 @@ inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float *
    }
 }
 #else
-inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) {
-    for (int i = 0; i < n; ++i) {
-        y[i] = ggml_gelu_quick_f32(x[i]);
-    }
-}
+void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x);
 #endif

 // Sigmoid Linear Unit (SiLU) function
@@ -1143,6 +1138,13 @@ inline static svfloat32_t ggml_v_silu(svbool_t pg, svfloat32_t x) {
    return svdiv_f32_x(pg, x, one_plus_exp_neg_x);
 }

+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static svfloat32_t ggml_v_gelu_quick(svbool_t pg, svfloat32_t x) {
+    const svfloat32_t one = svdup_n_f32_x(pg, 1.0f);
+    const svfloat32_t exp_cx = ggml_v_expf(pg, svmul_n_f32_x(pg, x, GELU_QUICK_COEF));
+    return svdiv_f32_x(pg, x, svadd_f32_x(pg, one, exp_cx));
+}
+
 #elif defined(__ARM_NEON) && defined(__aarch64__)

 // adapted from arm limited optimized routine
@@ -1182,6 +1184,13 @@ inline static float32x4_t ggml_v_silu(float32x4_t x) {
    return vdivq_f32(x, one_plus_exp_neg_x);
 }

+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static float32x4_t ggml_v_gelu_quick(float32x4_t x) {
+    const float32x4_t one = vdupq_n_f32(1.0f);
+    const float32x4_t exp_cx = ggml_v_expf(vmulq_n_f32(x, GELU_QUICK_COEF));
+    return vdivq_f32(x, vaddq_f32(one, exp_cx));
+}
+
 #elif defined(__AVX512F__) && defined(__AVX512DQ__)

 // adapted from arm limited optimized routine
@@ -1225,6 +1234,13 @@ inline static __m512 ggml_v_silu(__m512 x) {
    return _mm512_div_ps(x, one_plus_exp_neg_x);
 }

+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static __m512 ggml_v_gelu_quick(__m512 x) {
+    const __m512 one = _mm512_set1_ps(1);
+    const __m512 exp_cx = ggml_v_expf(_mm512_mul_ps(x, _mm512_set1_ps(GELU_QUICK_COEF)));
+    return _mm512_div_ps(x, _mm512_add_ps(one, exp_cx));
+}
+
 #elif defined(__AVX2__) && defined(__FMA__)

 // adapted from arm limited optimized routine
@@ -1280,6 +1296,13 @@ inline static __m256 ggml_v_silu(__m256 x) {
    return _mm256_div_ps(x, one_plus_exp_neg_x);
 }

+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static __m256 ggml_v_gelu_quick(__m256 x) {
+    const __m256 one = _mm256_set1_ps(1);
+    const __m256 exp_cx = ggml_v_expf(_mm256_mul_ps(x, _mm256_set1_ps(GELU_QUICK_COEF)));
+    return _mm256_div_ps(x, _mm256_add_ps(one, exp_cx));
+}
+
 #elif defined(__SSE2__) // __AVX2__ / __ARM_NEON

 #if defined(__FMA__)
@@ -1334,6 +1357,13 @@ inline static __m128 ggml_v_silu(__m128 x) {
    return _mm_div_ps(x, one_plus_exp_neg_x);
 }

+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static __m128 ggml_v_gelu_quick(__m128 x) {
+    const __m128 one = _mm_set1_ps(1);
+    const __m128 exp_cx = ggml_v_expf(_mm_mul_ps(x, _mm_set1_ps(GELU_QUICK_COEF)));
+    return _mm_div_ps(x, _mm_add_ps(one, exp_cx));
+}
+
 #elif defined(__VXE__) || defined(__VXE2__)

 // adapted from arm limited optimized routine
@@ -1374,6 +1404,13 @@ inline static float32x4_t ggml_v_silu(float32x4_t v_x) {
    return vec_div(v_x, v_one_plus_exp_neg_x);
 }

+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static float32x4_t ggml_v_gelu_quick(float32x4_t x) {
+    const float32x4_t one = vec_splats(1.0f);
+    const float32x4_t exp_cx = ggml_v_expf(vec_mul(x, vec_splats(GELU_QUICK_COEF)));
+    return vec_div(x, vec_add(one, exp_cx));
+}
+
 #elif defined(__riscv_v_intrinsic)

 // adapted from arm limited optimized routine
@@ -1426,6 +1463,12 @@ inline static vfloat32m2_t ggml_v_silu_m2(vfloat32m2_t x, int vl) {
    return __riscv_vfdiv_vv_f32m2(x, one_plus_exp_neg_x, vl);
 }

+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static vfloat32m2_t ggml_v_gelu_quick_m2(vfloat32m2_t x, int vl) {
+    const vfloat32m2_t exp_cx = ggml_v_expf_m2(__riscv_vfmul_vf_f32m2(x, GELU_QUICK_COEF, vl), vl);
+    return __riscv_vfdiv_vv_f32m2(x, __riscv_vfadd_vf_f32m2(exp_cx, 1.0f, vl), vl);
+}
+
 #endif // __ARM_NEON / __AVX2__ / __SSE2__ / __VXE__ / __riscv_v_intrinsic

 inline static void ggml_vec_silu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
@@ -1571,11 +1614,7 @@ inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float
    }
 }
 #else
-inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) {
-    for (int i = 0; i < n; ++i) {
-        y[i] = ggml_gelu_quick_f32(x[i]) * g[i];
-    }
-}
+void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g);
 #endif

 inline static void ggml_vec_geglu_quick_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) {
@@ -1551,11 +1551,11 @@ struct ggml_backend_cuda_context {
    // pool
    std::unique_ptr<ggml_cuda_pool> pools[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS];

-    static std::unique_ptr<ggml_cuda_pool> new_pool_for_device(int device, int stream_no);
+    static std::unique_ptr<ggml_cuda_pool> new_pool_for_device(ggml_backend_cuda_context & ctx, int device, int stream_no);

    ggml_cuda_pool & pool(int device) {
        if (pools[device][curr_stream_no] == nullptr) {
-            pools[device][curr_stream_no] = new_pool_for_device(device, curr_stream_no);
+            pools[device][curr_stream_no] = new_pool_for_device(*this, device, curr_stream_no);
        }
        return *pools[device][curr_stream_no];
    }
@@ -70,6 +70,7 @@
 #include "ggml-cuda/cumsum.cuh"
 #include "ggml-cuda/fill.cuh"
 #include "ggml-cuda/lightning-indexer.cuh"
+#include "ggml-cuda/top-k-mask.cuh"
 #include "ggml.h"

 #include <algorithm>
@@ -457,6 +458,7 @@ const ggml_cuda_device_info & ggml_cuda_info() {
 struct ggml_cuda_pool_leg : public ggml_cuda_pool {
    static const int MAX_BUFFERS = 256;

+    ggml_backend_cuda_context & ctx;
    int device;
    struct ggml_cuda_buffer {
        void * ptr = nullptr;
@@ -466,7 +468,8 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
    ggml_cuda_buffer buffer_pool[MAX_BUFFERS] = {};
    size_t pool_size = 0;

-    explicit ggml_cuda_pool_leg(int device) :
+    ggml_cuda_pool_leg(ggml_backend_cuda_context & ctx, int device) :
+        ctx(ctx),
        device(device) {
    }

@@ -488,6 +491,20 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
        }
    }

+    // captured CUDA graphs can still use freed buffers, so they must be captured again
+    void invalidate_cuda_graphs() {
+#ifdef USE_CUDA_GRAPH
+        int n_captured = 0;
+        for (const auto & [key, graph] : ctx.cuda_graphs) {
+            n_captured += graph->instance != nullptr;
+        }
+        if (n_captured > 0) {
+            GGML_LOG_DEBUG(GGML_CUDA_NAME " pool[%d]: memory freed, %d captured CUDA graphs must be captured again\n", device, n_captured);
+        }
+        ctx.cuda_graphs.clear();
+#endif // USE_CUDA_GRAPH
+    }
+
    void * alloc(size_t size, size_t * actual_size) override {
 #ifdef DEBUG_CUDA_MALLOC
        int nnz = 0;
@@ -538,6 +555,7 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
                           device, look_ahead_size/1024.0/1024.0, cached_bytes/1024.0/1024.0);
            CUDA_CHECK(cudaDeviceSynchronize());
            clear_pool();
+            invalidate_cuda_graphs();
            err = ggml_cuda_device_malloc(&ptr, look_ahead_size, device);
            if (err == cudaSuccess) {
                GGML_LOG_DEBUG(GGML_CUDA_NAME " pool[%d]: retry succeeded\n", device);
@@ -566,6 +584,7 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
        ggml_cuda_set_device(device);
        CUDA_CHECK(cudaFree(ptr));
        pool_size -= size;
+        invalidate_cuda_graphs();
    }
 };

@@ -720,14 +739,15 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
 };
 #endif // defined(GGML_USE_VMM)

-std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(int                  device,
-                                                                               [[maybe_unused]] int stream_no) {
+std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(ggml_backend_cuda_context & ctx,
+                                                                               int                         device,
+                                                                               [[maybe_unused]] int        stream_no) {
 #if defined(GGML_USE_VMM)
    if (ggml_cuda_info().devices[device].vmm) {
        return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_vmm(device));
    }
 #endif // defined(GGML_USE_VMM)
-    return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_leg(device));
+    return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_leg(ctx, device));
 }

 // destroying a cuBLAS handle while a graph is being captured in a different thread can result in a CUDA error
@@ -2014,6 +2034,8 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c
    const ggml_tensor * src0 = dst->src[0];
    const ggml_tensor * src1 = dst->src[1];

+    const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
+
    if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
        return true;
    }
@@ -2032,7 +2054,7 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c
        return false;
    }

-    if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
+    if (ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
        return false;
    }

@@ -2050,6 +2072,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *
    GGML_TENSOR_BINARY_OP_LOCALS

    const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
+    const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;

    // [TAG_MUL_MAT_ID_CUDA_GRAPHS]
    if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
@@ -2074,7 +2097,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *
            return;
        }

-        if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
+        if (ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
            ggml_cuda_mul_mat_f(ctx, src0, src1, ids, dst);
            return;
        }
@@ -2548,6 +2571,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
        case GGML_OP_LIGHTNING_INDEXER:
            ggml_cuda_lightning_indexer(ctx, dst);
            break;
+        case GGML_OP_TOP_K_MASK:
+            ggml_cuda_op_top_k_mask(ctx, dst);
+            break;
        default:
            return false;
    }
@@ -5981,6 +6007,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
            return true;
        case GGML_OP_LIGHTNING_INDEXER:
            return ggml_cuda_lightning_indexer_supported(dev_ctx->device, op);
+        case GGML_OP_TOP_K_MASK:
+            return op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32;

        default:
            return false;
@@ -151,9 +151,10 @@ static void ssm_conv_f32_cuda(const float * src0, const float * src1, const floa
        case 3:  launch_kernel(std::integral_constant<int, 3 >{}); break;
        case 4:  launch_kernel(std::integral_constant<int, 4 >{}); break;
        case 5:  launch_kernel(std::integral_constant<int, 5 >{}); break;
+        case 7:  launch_kernel(std::integral_constant<int, 7 >{}); break;
        case 9:  launch_kernel(std::integral_constant<int, 9 >{}); break;
        case 15: launch_kernel(std::integral_constant<int, 15>{}); break;
-        default: GGML_ABORT("Only support kernel sizes 3, 4, 5, 9, 15 right now.");
+        default: GGML_ABORT("Only support kernel sizes 3, 4, 5, 7, 9, 15 right now.");
    }
 }

@@ -0,0 +1,141 @@
+#include "top-k-mask.cuh"
+#include "convert.cuh"
+
+#define CUDA_TOP_K_MASK_BLOCK_SIZE 512
+
+// one block per row: fill the row with -inf, then copy the mask cells of the live picks and the extra cells
+template <typename T>
+static __global__ void top_k_mask_kernel(
+        const T * mask, const int32_t * top_k, const float * score, const int32_t * blocks, const int32_t * extra, T * dst,
+        const int64_t n_kv, const int64_t ne1, const int64_t ne2,
+        const int64_t m1, const int64_t m2, const int64_t m3,
+        const int64_t n_top, const int64_t n_pick, const int64_t n_blk, const int64_t n_extra,
+        const int64_t sm1, const int64_t sm2, const int64_t sm3,
+        const int64_t sd1, const int64_t sd2, const int64_t sd3,
+        const int64_t sk0, const int64_t sk1, const int64_t sk2, const int64_t sk3,
+        const int64_t ss0, const int64_t ss1, const int64_t ss2, const int64_t ss3,
+        const int64_t sb0, const int64_t sb1,
+        const int64_t se0, const int64_t se1, const int64_t se2, const int64_t se3,
+        const bool vec) {
+    ggml_cuda_pdl_lc();
+
+    const int64_t ir = blockIdx.x;
+    const int64_t i1 = ir % ne1;
+    const int64_t i2 = (ir / ne1) % ne2;
+    const int64_t i3 = ir / (ne1*ne2);
+
+    const T       * mask_row  = mask  + (i1 % m1)*sm1 + (i2 % m2)*sm2 + (i3 % m3)*sm3;
+    T             * dst_row   = dst   + i1*sd1 + i2*sd2 + i3*sd3;
+    const int32_t * top_k_row = top_k + i1*sk1 + i2*sk2 + i3*sk3;
+
+    ggml_cuda_pdl_sync();
+
+    const T ninf = ggml_cuda_cast<T>(-INFINITY);
+    if (vec) {
+        constexpr int nv = sizeof(int4)/sizeof(T);
+        T fill[nv];
+#pragma unroll
+        for (int j = 0; j < nv; ++j) {
+            fill[j] = ninf;
+        }
+        const int4 fill4 = *(const int4 *) fill;
+        for (int64_t i = threadIdx.x; i < n_kv/nv; i += blockDim.x) {
+            ((int4 *) dst_row)[i] = fill4;
+        }
+    } else {
+        for (int64_t i = threadIdx.x; i < n_kv; i += blockDim.x) {
+            dst_row[i] = ninf;
+        }
+    }
+
+    __syncthreads();
+
+    // one thread per pick or extra cell; cells kept twice get the same value
+    for (int64_t i = threadIdx.x; i < n_top + n_extra; i += blockDim.x) {
+        if (i >= n_top) {
+            const int32_t c = extra[(i - n_top)*se0 + i1*se1 + i2*se2 + i3*se3];
+            if (c >= 0 && c < n_kv) {
+                dst_row[c] = mask_row[c];
+            }
+            continue;
+        }
+
+        const int32_t p = top_k_row[i*sk0];
+        if (p < 0 || p >= n_pick) {
+            continue;
+        }
+        if (score && !(score[p*ss0 + i1*ss1 + i2*ss2 + i3*ss3] > -INFINITY)) {
+            continue;
+        }
+        if (!blocks) {
+            if (p < n_kv) {
+                dst_row[p] = mask_row[p];
+            }
+            continue;
+        }
+        for (int64_t k = 0; k < n_blk; ++k) {
+            const int32_t c = blocks[k*sb0 + p*sb1];
+            if (c >= 0 && c < n_kv) {
+                dst_row[c] = mask_row[c];
+            }
+        }
+    }
+}
+
+template <typename T>
+static void top_k_mask_cuda(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * mask   = dst->src[0];
+    const ggml_tensor * top_k  = dst->src[1];
+    const ggml_tensor * score  = dst->src[2];
+    const ggml_tensor * blocks = dst->src[3];
+    const ggml_tensor * extra  = dst->src[4];
+
+    const int64_t n_kv = mask->ne[0];
+
+    GGML_ASSERT(dst->nb[0] == sizeof(T) && mask->nb[0] == sizeof(T));
+    GGML_ASSERT(dst->nb[1] % sizeof(T) == 0 && mask->nb[1] % sizeof(T) == 0);
+
+    // 16-byte stores need every dst row aligned
+    const bool vec = (n_kv*sizeof(T)) % sizeof(int4) == 0 &&
+        dst->nb[1] % sizeof(int4) == 0 && dst->nb[2] % sizeof(int4) == 0 && dst->nb[3] % sizeof(int4) == 0 &&
+        (uintptr_t) dst->data % sizeof(int4) == 0;
+
+    const int64_t n_pick = blocks ? blocks->ne[1] : score ? score->ne[0] : n_kv;
+
+    const size_t si = sizeof(int32_t);
+    const size_t sf = sizeof(float);
+
+    const dim3 block_dims(CUDA_TOP_K_MASK_BLOCK_SIZE, 1, 1);
+    const dim3 grid_dims(ggml_nrows(dst), 1, 1);
+    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
+
+    ggml_cuda_kernel_launch(top_k_mask_kernel<T>, launch_params,
+            (const T *) mask->data, (const int32_t *) top_k->data,
+            score  ? (const float   *) score->data  : nullptr,
+            blocks ? (const int32_t *) blocks->data : nullptr,
+            extra  ? (const int32_t *) extra->data  : nullptr,
+            (T *) dst->data,
+            n_kv, dst->ne[1], dst->ne[2],
+            mask->ne[1], mask->ne[2], mask->ne[3],
+            top_k->ne[0], n_pick, blocks ? blocks->ne[0] : 1, extra ? extra->ne[0] : 0,
+            mask->nb[1]/sizeof(T), mask->nb[2]/sizeof(T), mask->nb[3]/sizeof(T),
+            dst->nb[1]/sizeof(T),  dst->nb[2]/sizeof(T),  dst->nb[3]/sizeof(T),
+            top_k->nb[0]/si, top_k->nb[1]/si, top_k->nb[2]/si, top_k->nb[3]/si,
+            score  ? score->nb[0]/sf  : 0, score ? score->nb[1]/sf : 0, score ? score->nb[2]/sf : 0, score ? score->nb[3]/sf : 0,
+            blocks ? blocks->nb[0]/si : 0, blocks ? blocks->nb[1]/si : 0,
+            extra  ? extra->nb[0]/si  : 0, extra ? extra->nb[1]/si : 0, extra ? extra->nb[2]/si : 0, extra ? extra->nb[3]/si : 0,
+            vec);
+}
+
+void ggml_cuda_op_top_k_mask(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    switch (dst->type) {
+        case GGML_TYPE_F16:
+            top_k_mask_cuda<half>(ctx, dst);
+            break;
+        case GGML_TYPE_F32:
+            top_k_mask_cuda<float>(ctx, dst);
+            break;
+        default:
+            GGML_ABORT("unsupported type");
+    }
+}
@@ -0,0 +1,3 @@
+#include "common.cuh"
+
+void ggml_cuda_op_top_k_mask(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
@@ -22,6 +22,10 @@
 #define HWCAP2_I8MM (1ULL << 13)
 #endif

+#if !defined(HWCAP2_SVEI8MM)
+#define HWCAP2_SVEI8MM (1ULL << 9)
+#endif
+
 #if !defined(HWCAP_ASIMDDP)
 #define HWCAP_ASIMDDP (1 << 20)
 #endif
@@ -63,6 +67,10 @@
 #define PF_ARM_SVE2_INSTRUCTIONS_AVAILABLE 47
 #endif

+#if !defined(PF_ARM_SVE_I8MM_INSTRUCTIONS_AVAILABLE)
+#define PF_ARM_SVE_I8MM_INSTRUCTIONS_AVAILABLE 57
+#endif
+
 #if !defined(PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE)
 #define PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE 66
 #endif
@@ -87,6 +95,7 @@ typedef struct ggml_feats_arch64_runtime {
    bool has_sve;
    bool has_sve2;
    bool has_i8mm;
+    bool has_svei8mm;
    bool has_sme;
    bool has_sme2;
    int sve_cnt;
@@ -104,6 +113,7 @@ static inline ggml_feats_arch64_runtime_t ggml_feats_get_arch64_runtime(void) {
    runtime_feat.has_sve     = !!(hwcap & HWCAP_SVE);
    runtime_feat.has_sve2    = !!(hwcap2 & HWCAP2_SVE2);
    runtime_feat.has_i8mm    = !!(hwcap2 & HWCAP2_I8MM);
+    runtime_feat.has_svei8mm = !!(hwcap2 & HWCAP2_SVEI8MM);
    runtime_feat.has_sme     = !!(hwcap2 & HWCAP2_SME);
    runtime_feat.has_sme2    = !!(hwcap2 & HWCAP2_SME2);

@@ -153,6 +163,7 @@ static inline ggml_feats_arch64_runtime_t ggml_feats_get_arch64_runtime(void) {
    runtime_feat.has_sve     = IsProcessorFeaturePresent(PF_ARM_SVE_INSTRUCTIONS_AVAILABLE) != 0;
    runtime_feat.has_sve2    = IsProcessorFeaturePresent(PF_ARM_SVE2_INSTRUCTIONS_AVAILABLE) != 0;
    runtime_feat.has_i8mm    = IsProcessorFeaturePresent(PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE) != 0;
+    runtime_feat.has_svei8mm = IsProcessorFeaturePresent(PF_ARM_SVE_I8MM_INSTRUCTIONS_AVAILABLE) != 0;
    runtime_feat.has_sme     = IsProcessorFeaturePresent(PF_ARM_SME_INSTRUCTIONS_AVAILABLE) != 0;
    runtime_feat.has_sme2    = IsProcessorFeaturePresent(PF_ARM_SME2_INSTRUCTIONS_AVAILABLE) != 0;

@@ -21,6 +21,7 @@ message(STATUS "hexagon: using ${HEXAGON_SDK_ROOT} and ${HEXAGON_TOOLS_ROOT} for

 include(${HEXAGON_SDK_ROOT}/build/cmake/hexagon_fun.cmake)
 include(ExternalProject)
+find_package(Python3 REQUIRED COMPONENTS Interpreter)

 option(GGML_HEXAGON_HTP_DEBUG  "ggml-hexagon: enable HTP debug output" OFF)
 set(GGML_HEXAGON_HTP_BUILD_TYPE "Release" CACHE STRING "ggml-hexagon: HTP skel build type (Release, RelWithDebInfo, Debug)")
@@ -50,6 +51,13 @@ endif()
 set(TARGET_NAME ggml-hexagon)
 ggml_add_backend_library(${TARGET_NAME}
    ggml-hexagon.cpp
+    htp-common.h
+    htp-kparams.cpp
+    htp-kparams.h
+    htp-opnode.cpp
+    htp-opnode.h
+    htp-fusion.cpp
+    htp-fusion.h
    htp-repack.cpp
    htp-repack.h
    htp-drv.cpp
@@ -74,7 +82,8 @@ function(build_htp_skel V)
            -DHEXAGON_TOOLS_ROOT=${HEXAGON_TOOLS_ROOT}
            -DHEXAGON_HTP_DEBUG=${GGML_HEXAGON_HTP_DEBUG}
            -DDSP_VERSION=${V}
-            -DPREBUILT_LIB_DIR="toolv19_${V}")
+            -DPython3_EXECUTABLE=${Python3_EXECUTABLE}
+            -DPREBUILT_LIB_DIR="toolv21_${V}")
    list(APPEND HTP_SKELS ${CMAKE_CURRENT_BINARY_DIR}/libggml-htp-${V}.so)
    list(APPEND HTP_PROJECTS htp-${V})
    set(HTP_SKELS ${HTP_SKELS} PARENT_SCOPE)
@@ -85,6 +94,7 @@ build_htp_skel(v73)
 build_htp_skel(v75)
 build_htp_skel(v79)
 build_htp_skel(v81)
+build_htp_skel(v85)

 # Install Hexagon skels required at runtime
 install(FILES ${HTP_SKELS} TYPE LIB)
@@ -0,0 +1,106 @@
+#ifndef HTP_COMMON_H
+#define HTP_COMMON_H
+
+#include "htp-opnode.h"
+#include "htp-repack.h"
+
+struct htp_tensor_extra {
+    std::vector<uint8_t> shadow_buf;
+    size_t               shadow_size { 0 };
+    uint32_t             flags { 0 };
+};
+
+extern int    opt_arch;
+extern size_t opt_ndev;
+extern size_t opt_nhvx;
+extern int    opt_nhmx;
+extern size_t opt_vmem;
+extern int    opt_etm;
+extern int    opt_verbose;
+extern int    opt_profile;
+extern bool   opt_hostbuf;
+extern bool   opt_dma64;
+extern int    opt_power_mode;
+
+extern size_t opt_mbuf_dyn;
+extern size_t opt_mbuf_static;
+extern size_t opt_mbuf_total;
+
+extern int    opt_mm_select;
+extern int    opt_fa_select;
+extern int    opt_fa_head_split;
+extern int    opt_gdn_select;
+extern int    opt_ar_select;
+extern int    opt_ar_scatter;
+
+extern int    opt_opbatch;
+extern int    opt_opqueue;
+extern int    opt_optrace;
+extern int    opt_oppoll;
+extern int    opt_opfusion;
+extern int    opt_rt_repack;
+
+#ifndef HEX_VERBOSE
+#define HEX_VERBOSE(...) \
+    if (opt_verbose) GGML_LOG_DEBUG(__VA_ARGS__)
+#endif
+
+static inline bool htp_tensor_is_fuseable(const struct ggml_tensor * t) {
+    if (!t || !t->extra) return false;
+    const auto * extra = (const struct htp_tensor_extra *) t->extra;
+    return (extra->flags & HTP_TENSOR_FUSEABLE) != 0;
+}
+
+static inline bool htp_can_rt_repack(const struct ggml_tensor * t) {
+    if (opt_rt_repack == 0 || !htp_is_tiled_weight_tensor(t)) {
+        return false;
+    }
+    if (t->buffer && ggml_backend_buft_is_host(t->buffer->buft)) {
+        return false;
+    }
+    return opt_rt_repack > 1 || ggml_is_quantized(t->type);
+}
+
+static inline bool htp_tensor_is_tiled(const struct ggml_tensor * t) {
+    const auto * extra = (const struct htp_tensor_extra *) t->extra;
+    if (extra && (extra->flags & HTP_TENSOR_TILED)) {
+        return true;
+    }
+    return htp_can_rt_repack(t);
+}
+
+static inline bool htp_needs_repack(const struct ggml_tensor * t) {
+    const auto * extra = (const struct htp_tensor_extra *) t->extra;
+    if (extra && (extra->flags & HTP_TENSOR_TILED)) {
+        return false;
+    }
+    return htp_can_rt_repack(t);
+}
+
+static inline bool htp_tensors_overlap(const struct ggml_tensor * a, const struct ggml_tensor * b) {
+    const uintptr_t a0 = (uintptr_t) a->data;
+    const uintptr_t b0 = (uintptr_t) b->data;
+    const uintptr_t a1 = a0 + ggml_nbytes(a);
+    const uintptr_t b1 = b0 + ggml_nbytes(b);
+
+    return a0 < b1 && b0 < a1;
+}
+
+static inline bool htp_can_row_partition(const struct ggml_tensor * t) {
+    if (t->ne[1] > 1 && (t->nb[1] & 127) != 0) return false;
+    if (t->ne[2] > 1 && (t->nb[2] & 127) != 0) return false;
+    if (t->ne[3] > 1 && (t->nb[3] & 127) != 0) return false;
+    return true;
+}
+
+struct ggml_hexagon_session;
+
+const char * htp_session_name(const struct ggml_hexagon_session * sess);
+size_t       htp_session_vtcm_size(const struct ggml_hexagon_session * sess);
+uint32_t     htp_session_n_threads(const struct ggml_hexagon_session * sess);
+uint32_t     htp_session_n_hmx(const struct ggml_hexagon_session * sess);
+uint32_t     htp_session_mdev_count(const struct ggml_hexagon_session * sess);
+
+htp_op_code  htp_op_remap(const struct ggml_tensor * t);
+
+#endif // HTP_COMMON_H
@@ -0,0 +1,650 @@
+#include "htp-fusion.h"
+#include "htp-kparams.h"
+#include "ggml-impl.h"
+
+#include <algorithm>
+#include <cstring>
+
+static inline bool htp_same_tensor_data(const ggml_tensor * a, const ggml_tensor * b) {
+    return a && b && (a == b || a->data == b->data);
+}
+
+static inline bool htp_same_shape_and_contiguity(const ggml_tensor * a, const ggml_tensor * b) {
+    return ggml_are_same_shape(a, b) && (ggml_is_contiguous(a) == ggml_is_contiguous(b));
+}
+
+static inline bool htp_inplace_overlap_unsafe(const ggml_tensor * dst, const ggml_tensor * in) {
+    return htp_tensors_overlap(dst, in) && (dst->data != in->data || !ggml_are_same_stride(dst, in));
+}
+
+static inline bool htp_inputs_overlap_unsafe(const htp_opnode & prev, const ggml_tensor * dst, const ggml_tensor * extra = nullptr) {
+    if (extra && htp_inplace_overlap_unsafe(dst, extra)) return true;
+    for (const auto * in : prev.inputs) {
+        if (htp_inplace_overlap_unsafe(dst, in)) return true;
+    }
+    return false;
+}
+
+static inline bool htp_match_consumer(
+    const htp_opnode & prev,
+    const htp_opnode & cur,
+    bool commutative,
+    const ggml_tensor ** other_src = nullptr
+) {
+    if (prev.outputs.size() != 1) {
+        return false;
+    }
+    const ggml_tensor * prev_out = prev.dst();
+    if (!htp_tensor_is_fuseable(prev_out)) {
+        return false;
+    }
+    if (htp_same_tensor_data(cur.src0(), prev_out)) {
+        if (other_src) *other_src = cur.src1();
+        return true;
+    }
+    if (commutative && htp_same_tensor_data(cur.src1(), prev_out)) {
+        if (other_src) *other_src = cur.src0();
+        return true;
+    }
+    return false;
+}
+
+static bool is_supported_mul_mat_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams) {
+    if (kparams->n_hmx) {
+        return kparams->kernel_type == HTP_MM_KERNEL_HMX_2D;
+    }
+    return htp_tensor_is_tiled(src0) && (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW);
+}
+
+static bool is_supported_mul_mat_id_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams) {
+    if (kparams->n_hmx) {
+        return kparams->kernel_type == HTP_MM_KERNEL_HMX_2D;
+    }
+    return htp_tensor_is_tiled(src0) &&
+           (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK);
+}
+
+// --- Fusion rules ---
+
+static bool fuse_allreduce_add(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (opt_ar_select != 2 || prev.opcode != HTP_OP_ALLREDUCE) return false;
+
+    const auto * ar_kparams = prev.kparams<htp_allreduce_kernel_params>();
+    const uint32_t rank    = (uint32_t) ar_kparams->rank;
+    const uint32_t n_ranks = (uint32_t) ar_kparams->n_ranks;
+
+    const ggml_tensor * ar_local = prev.inputs[rank];
+    if (!htp_tensor_is_fuseable(ar_local)) return false;
+
+    const ggml_tensor * add_dst    = cur.dst();
+    const ggml_tensor * res_tensor = nullptr;
+    if (htp_same_tensor_data(cur.src0(), ar_local)) {
+        res_tensor = cur.src1();
+    } else if (htp_same_tensor_data(cur.src1(), ar_local)) {
+        res_tensor = cur.src0();
+    } else {
+        return false;
+    }
+
+    if (ar_local->type != res_tensor->type) return false;
+
+    const bool is_same_shape = ggml_are_same_shape(ar_local, res_tensor);
+    const bool is_row_bcast  = !is_same_shape && (ar_local->ne[0] == res_tensor->ne[0] &&
+                                                  res_tensor->ne[1] == 1 && res_tensor->ne[2] == 1 && res_tensor->ne[3] == 1);
+    if (!is_same_shape && !is_row_bcast) return false;
+
+    if (is_same_shape && !ggml_are_same_stride(ar_local, res_tensor)) {
+        return false;
+    }
+    if (ggml_is_contiguous(ar_local) != ggml_is_contiguous(add_dst)) return false;
+
+    const bool is_shard_ok = n_ranks <= HTP_OP_MAX_OUTPUTS && add_dst->data == ar_local->data;
+
+    struct htp_allreduce_kernel_params kparams;
+    if (!htp_kparams_solve_allreduce(sess, add_dst, rank, n_ranks, true, is_row_bcast, is_shard_ok, &kparams)) {
+        HEX_VERBOSE("ggml-hex: %s skip ALLREDUCE_ADD fusion: solver failed\n", htp_session_name(sess));
+        return false;
+    }
+
+    const bool scatter_ok = kparams.n_dsts > 1;
+    kparams.mode = scatter_ok ? HTP_ALLREDUCE_SHARDED_FANOUT : HTP_ALLREDUCE_FULL;
+
+    for (uint32_t r = 0; r < n_ranks; r++) {
+        const ggml_tensor * ar_src = prev.inputs[r];
+        if (!htp_tensors_overlap(add_dst, ar_src)) continue;
+        if (scatter_ok && r == rank) continue;
+        HEX_VERBOSE("ggml-hex: %s skip ALLREDUCE_ADD fusion: dst overlaps allreduce src %u\n", htp_session_name(sess), r);
+        return false;
+    }
+
+    prev.opcode = HTP_OP_ALLREDUCE_ADD;
+    prev.inputs.push_back(res_tensor);
+    if (kparams.mode != HTP_ALLREDUCE_SHARDED_FANOUT) {
+        prev.outputs = { add_dst };
+    }
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_rms_norm_mul(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_RMS_NORM) return false;
+
+    const ggml_tensor * weight = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &weight)) return false;
+
+    const ggml_tensor * src0 = prev.src0();
+    const ggml_tensor * dst  = cur.dst();
+    if (src0->type != GGML_TYPE_F32 || src0->ne[0] != weight->ne[0] || !htp_same_shape_and_contiguity(prev.dst(), dst)) return false;
+    if (htp_inputs_overlap_unsafe(prev, dst, weight)) return false;
+
+    const bool is_row_bcast  = (weight->ne[1] == 1 && weight->ne[2] == 1 && weight->ne[3] == 1);
+    const bool is_same_shape = ggml_are_same_shape(src0, weight);
+    if (!is_row_bcast && !is_same_shape) return false;
+
+    struct htp_unary_kernel_params kparams;
+    if (!htp_kparams_solve_rms_norm_mul(sess, src0, weight, dst, &kparams)) {
+        HEX_VERBOSE("ggml-hex: %s skip RMS_NORM_MUL fusion: VTCM needed (%d) > budget (%zu)\n",
+                    htp_session_name(sess), kparams.vtcm_size, htp_session_vtcm_size(sess));
+        return false;
+    }
+
+    prev.opcode  = HTP_OP_RMS_NORM_MUL;
+    prev.inputs  = { src0, weight };
+    prev.outputs = { dst };
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_rms_norm_scale(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    GGML_UNUSED(sess);
+    if (prev.opcode != HTP_OP_RMS_NORM && prev.opcode != HTP_OP_RMS_NORM_MUL) return false;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false)) return false;
+
+    const ggml_tensor * src0 = prev.src0();
+    const ggml_tensor * dst  = cur.dst();
+    if (src0->type != GGML_TYPE_F32 || !htp_same_shape_and_contiguity(prev.dst(), dst)) return false;
+    if (htp_inputs_overlap_unsafe(prev, dst)) return false;
+
+    float scale = 0.0f;
+    float bias  = 0.0f;
+    memcpy(&scale, (const float *) cur.params + 0, sizeof(float));
+    memcpy(&bias,  (const float *) cur.params + 1, sizeof(float));
+    if (bias != 0.0f || scale == 0.0f) return false;
+
+    float prev_scale = 0.0f;
+    memcpy(&prev_scale, &prev.params[1], sizeof(float));
+    if (prev_scale != 0.0f) {
+        scale *= prev_scale;
+    }
+    memcpy(&prev.params[1], &scale, sizeof(float));
+
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_mul_mat_add(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL_MAT) return false;
+
+    const ggml_tensor * src2 = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &src2)) return false;
+    if (src2->type != GGML_TYPE_F32) return false;
+
+    const ggml_tensor * src0 = prev.src0();
+    const ggml_tensor * src1 = prev.src1();
+    const ggml_tensor * dst  = cur.dst();
+    if (!ggml_are_same_shape(prev.dst(), dst)) return false;
+
+    const auto * orig_kparams = prev.kparams<htp_mm_kernel_params>();
+    struct htp_mm_kernel_params kparams;
+    if (!htp_kparams_solve_matmul_add(sess, src0, src1, src2, dst, &kparams) ||
+        orig_kparams->collapse != kparams.collapse) {
+        return false;
+    }
+
+    const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3];
+    if (kparams.n_hmx == 0 && src1_nrows != 1) return false;
+
+    if (kparams.n_hmx > 0 && orig_kparams->n_hmx > 0) {
+        if (kparams.m_chunk < orig_kparams->m_chunk ||
+            kparams.n_chunk < orig_kparams->n_chunk ||
+            kparams.n_act_threads < orig_kparams->n_act_threads) {
+            HEX_VERBOSE("ggml-hex: %s skip MUL_MAT_ADD fusion: HMX efficiency reduced (m %d->%d, n %d->%d, th %d->%d)\n",
+                        htp_session_name(sess), orig_kparams->m_chunk, kparams.m_chunk,
+                        orig_kparams->n_chunk, kparams.n_chunk,
+                        orig_kparams->n_act_threads, kparams.n_act_threads);
+            return false;
+        }
+    }
+
+    prev.opcode  = HTP_OP_MUL_MAT_ADD;
+    prev.inputs  = { src0, src1, src2 };
+    prev.outputs = { dst };
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_mul_mat_nx(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL_MAT && prev.opcode != HTP_OP_MUL_MAT_NX) return false;
+
+    const uint32_t curr_n = (uint32_t) prev.outputs.size();
+    if (curr_n >= HTP_OP_MAX_OUTPUTS || curr_n + 1 >= HTP_OP_MAX_INPUTS) return false;
+
+    const ggml_tensor * w0   = prev.inputs[0];
+    const ggml_tensor * x    = prev.inputs[curr_n];
+    const ggml_tensor * w_in = cur.src0();
+    const ggml_tensor * x_in = cur.src1();
+    const ggml_tensor * d_in = cur.dst();
+
+    auto * orig_kparams = prev.kparams<htp_mm_kernel_params>();
+    const auto * cur_kparams = cur.kparams<htp_mm_kernel_params>();
+
+    if (x_in != x || w_in->type != w0->type || w_in->ne[0] != w0->ne[0] ||
+        w0->ne[2] != 1 || w0->ne[3] != 1 || w_in->ne[2] != 1 || w_in->ne[3] != 1 ||
+        htp_tensor_is_tiled(w_in) != htp_tensor_is_tiled(w0) ||
+        orig_kparams->n_hmx != cur_kparams->n_hmx) {
+        return false;
+    }
+
+    const bool collapse = orig_kparams->collapse && htp_matmul_can_collapse(w_in, x, d_in);
+    if ((x->ne[2] > 1 || x->ne[3] > 1) && !collapse) return false;
+
+    if (prev.opcode == HTP_OP_MUL_MAT_NX) {
+        orig_kparams->n_weights = curr_n + 1;
+    } else {
+        struct htp_mm_kernel_params kparams;
+        if (collapse) {
+            const ggml_tensor x_collapsed = htp_tensor_collapse_rows(x);
+            if (!htp_kparams_solve_mmnx(sess, w0, &x_collapsed, curr_n + 1, &kparams)) return false;
+            kparams.collapse = 1;
+        } else if (!htp_kparams_solve_mmnx(sess, w0, x, curr_n + 1, &kparams)) {
+            return false;
+        }
+        if (!is_supported_mul_mat_nx_kernel(w0, &kparams)) return false;
+
+        prev.set_kparams(kparams);
+        prev.opcode = HTP_OP_MUL_MAT_NX;
+    }
+
+    prev.inputs.insert(prev.inputs.begin() + curr_n, w_in);
+    prev.outputs.push_back(d_in);
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_mul_mat_id_nx(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL_MAT_ID && prev.opcode != HTP_OP_MUL_MAT_ID_NX) return false;
+
+    const uint32_t curr_n = (uint32_t) prev.outputs.size();
+    if (curr_n >= HTP_OP_MAX_OUTPUTS || curr_n + 2 >= HTP_OP_MAX_INPUTS) return false;
+
+    const ggml_tensor * w0     = prev.inputs[0];
+    const ggml_tensor * x      = prev.inputs[curr_n];
+    const ggml_tensor * ids    = prev.inputs[curr_n + 1];
+    const ggml_tensor * w_in   = cur.src0();
+    const ggml_tensor * x_in   = cur.src1();
+    const ggml_tensor * ids_in = cur.node->src[2];
+    const ggml_tensor * d_in   = cur.dst();
+
+    auto * orig_kparams = prev.kparams<htp_mm_kernel_params>();
+    const auto * cur_kparams = cur.kparams<htp_mm_kernel_params>();
+
+    if (x_in != x || ids_in != ids || w_in->type != w0->type ||
+        w_in->ne[0] != w0->ne[0] || w_in->ne[2] != w0->ne[2] ||
+        htp_tensor_is_tiled(w_in) != htp_tensor_is_tiled(w0) ||
+        orig_kparams->n_hmx != cur_kparams->n_hmx) {
+        return false;
+    }
+
+    struct htp_mm_kernel_params kparams;
+    if (d_in->nb[1] == prev.node->nb[1]) {
+        kparams = *orig_kparams;
+        kparams.n_weights = curr_n + 1;
+    } else if (!htp_kparams_solve_mmidnx(sess, w0, x, d_in, curr_n + 1, &kparams)) {
+        return false;
+    }
+    if (!is_supported_mul_mat_id_nx_kernel(w0, &kparams)) return false;
+
+    prev.opcode = HTP_OP_MUL_MAT_ID_NX;
+    prev.inputs.insert(prev.inputs.begin() + curr_n, w_in);
+    prev.outputs.push_back(d_in);
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_gdn_cpy(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    GGML_UNUSED(sess);
+    if (prev.opcode != HTP_OP_GATED_DELTA_NET || prev.outputs.size() != 1) return false;
+
+    const ggml_tensor * gdn_out  = prev.dst();
+    const ggml_tensor * cpy_node = cur.node;
+    const ggml_tensor * cpy_src  = cur.src0();
+    const ggml_tensor * cpy_dst  = cur.dst();
+
+    if (gdn_out->type != GGML_TYPE_F32 || cpy_src->type != GGML_TYPE_F32 || cpy_dst->type != GGML_TYPE_F32) return false;
+    if ((gdn_out->flags & GGML_TENSOR_FLAG_OUTPUT) || (cpy_node->flags & GGML_TENSOR_FLAG_OUTPUT)) return false;
+
+    const ggml_tensor * v = prev.node->src[2];
+    if (!v) return false;
+
+    const int64_t S_v       = v->ne[0];
+    const int64_t H         = v->ne[1];
+    const int64_t n_tokens  = v->ne[2];
+    const int64_t n_seqs    = v->ne[3];
+    const int64_t K         = ggml_get_op_params_i32(prev.node, 0);
+    const size_t  tail_off  = (size_t) S_v * H * n_tokens * n_seqs * sizeof(float);
+    const int64_t D         = S_v * S_v * H;
+    const int64_t n_written = std::min<int64_t>(n_tokens, K);
+
+    if (cpy_src->op != GGML_OP_VIEW || !htp_same_tensor_data(cpy_src->view_src, gdn_out) ||
+        cpy_src->view_offs != tail_off || !ggml_is_contiguous(cpy_src)) {
+        return false;
+    }
+
+    if (cpy_dst->ne[0] != D || cpy_dst->ne[1] != n_seqs || cpy_dst->ne[2] != n_written || !ggml_is_contiguous(cpy_dst)) {
+        return false;
+    }
+
+    prev.outputs.push_back(cpy_dst);
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_ssm_conv_silu(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    GGML_UNUSED(sess);
+    if (prev.opcode != HTP_OP_SSM_CONV) return false;
+
+    auto * kparams = prev.kparams<htp_ssm_conv_kernel_params>();
+    if (kparams->silu) return false;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false)) return false;
+
+    const ggml_tensor * conv_dst = prev.dst();
+    const ggml_tensor * dst      = cur.dst();
+    if (!ggml_are_same_layout(conv_dst, dst)) return false;
+    if (htp_tensors_overlap(dst, prev.src0()) || htp_tensors_overlap(dst, prev.src1())) return false;
+
+    kparams->silu   = 1;
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_repeat_get_rows(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_REPEAT) return false;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false)) return false;
+
+    const ggml_tensor * repeat_out  = prev.dst();
+    const ggml_tensor * repeat_src0 = prev.src0();
+    const ggml_tensor * gr_src1     = cur.src1();
+    const ggml_tensor * gr_dst      = cur.dst();
+
+    if (repeat_src0->ne[0] != repeat_out->ne[0] || repeat_src0->ne[1] != repeat_out->ne[1]) return false;
+    if (repeat_src0->type != repeat_out->type) return false;
+
+    struct htp_get_rows_kernel_params kparams;
+    if (!htp_kparams_solve_repeat_get_rows(sess, repeat_src0, gr_src1, gr_dst, &kparams) || kparams.vtcm_size == 0) return false;
+
+    prev.opcode  = HTP_OP_GET_ROWS;
+    prev.inputs  = { repeat_src0, gr_src1 };
+    prev.outputs = { gr_dst };
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_unary_mul_glu(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    htp_op_code glu_op = HTP_OP_INVALID;
+    switch (prev.opcode) {
+        case HTP_OP_UNARY_SILU:       glu_op = HTP_OP_GLU_SWIGLU;      break;
+        case HTP_OP_UNARY_GELU:       glu_op = HTP_OP_GLU_GEGLU;       break;
+        case HTP_OP_UNARY_GELU_QUICK: glu_op = HTP_OP_GLU_GEGLU_QUICK; break;
+        case HTP_OP_UNARY_GELU_ERF:   glu_op = HTP_OP_GLU_GEGLU_ERF;   break;
+        case HTP_OP_UNARY_RELU:       glu_op = HTP_OP_GLU_REGLU;       break;
+        case HTP_OP_UNARY_SIGMOID:    glu_op = HTP_OP_GLU_SIGMOID;     break;
+        default: return false;
+    }
+
+    const ggml_tensor * up = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &up)) return false;
+
+    const ggml_tensor * gate = prev.src0();
+    const ggml_tensor * dst  = cur.dst();
+    if (htp_inputs_overlap_unsafe(prev, dst, up)) return false;
+
+    struct htp_glu_kernel_params kparams;
+    if (!htp_kparams_solve_unary_mul_glu(sess, glu_op, gate, up, dst, &kparams)) return false;
+
+    prev.opcode  = glu_op;
+    prev.inputs  = { gate, up };
+    prev.outputs = { dst };
+    prev.fused.push_back(cur.node);
+    memset(prev.params, 0, sizeof(prev.params));
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_mul_mul_scalar(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL || prev.inputs.size() != 2) return false;
+
+    const auto * k0 = prev.kparams<htp_binary_kernel_params>();
+    const auto * k1 = cur.kparams<htp_binary_kernel_params>();
+    if (k0->kernel_type != HTP_BINARY_KERNEL_SCALAR || k1->kernel_type != HTP_BINARY_KERNEL_SCALAR) return false;
+    if (k0->src1_strided || k1->src1_strided) return false;
+
+    const ggml_tensor * src2 = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false, &src2)) return false;
+
+    const ggml_tensor * src1 = prev.src1();
+    const ggml_tensor * dst  = cur.dst();
+    if (!ggml_are_same_shape(src1, src2) || !htp_same_shape_and_contiguity(prev.dst(), dst)) return false;
+    if (htp_inputs_overlap_unsafe(prev, dst, src2)) return false;
+
+    const size_t vtcm_size = htp_session_vtcm_size(sess);
+    struct htp_binary_kernel_params kparams = *k0;
+    kparams.src2_size       = kparams.src1_size;
+    kparams.rows_per_buffer = 0;
+    if (kparams.src1_size > 65536 || kparams.src1_size + kparams.src2_size > vtcm_size) return false;
+
+    struct htp_binary_vtcm_layout L;
+    htp_binary_vtcm_layout_build(&L, &kparams, vtcm_size);
+    if (L.rows_per_buffer == 0 || L.total_bytes > vtcm_size) return false;
+
+    kparams.rows_per_buffer = L.rows_per_buffer;
+    kparams.vtcm_size       = L.total_bytes;
+
+    prev.inputs.push_back(src2);
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_binary_nx(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != cur.opcode) return false;
+
+    const uint32_t curr_n = (uint32_t) prev.inputs.size();
+    if (curr_n < 2 || curr_n >= HTP_OP_MAX_INPUTS) return false;
+
+    const auto * k0 = prev.kparams<htp_binary_kernel_params>();
+    const auto * k1 = cur.kparams<htp_binary_kernel_params>();
+    const bool k0_ok = (k0->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) ||
+                       (k0->kernel_type == HTP_BINARY_KERNEL_CHUNKED && !k0->is_scalar);
+    const bool k1_ok = (k1->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) ||
+                       (k1->kernel_type == HTP_BINARY_KERNEL_CHUNKED && !k1->is_scalar);
+    if (!k0_ok || !k1_ok) return false;
+
+    const bool is_comm = (cur.opcode == HTP_OP_ADD || cur.opcode == HTP_OP_MUL);
+    const ggml_tensor * new_src = nullptr;
+    if (!htp_match_consumer(prev, cur, is_comm, &new_src)) return false;
+
+    const ggml_tensor * src0 = prev.inputs[0];
+    const ggml_tensor * dst  = cur.dst();
+    if (!ggml_are_same_shape(src0, prev.inputs[1]) ||
+        !ggml_are_same_shape(src0, new_src) ||
+        !htp_same_shape_and_contiguity(prev.dst(), dst)) {
+        return false;
+    }
+    if (htp_inputs_overlap_unsafe(prev, dst, new_src)) return false;
+
+    const size_t vtcm_size = htp_session_vtcm_size(sess);
+    struct htp_binary_kernel_params kparams = *k0;
+    struct htp_binary_vtcm_layout L;
+    memset(&L, 0, sizeof(L));
+
+    const uint32_t n_srcs = curr_n + 1;
+    if (k0->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE || k1->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) {
+        if (k0->kernel_type != HTP_BINARY_KERNEL_SAME_SHAPE) {
+            kparams = *k1;
+        }
+        kparams.n_srcs          = (uint8_t) n_srcs;
+        kparams.rows_per_buffer = 0;
+
+        auto check_split = [&](const ggml_tensor * t) {
+            if (src0->ne[2] > 1 && t->nb[2] != t->ne[1] * t->nb[1]) kparams.split_at_ne01 = 1;
+            if (src0->ne[3] > 1 && t->nb[3] != t->ne[2] * t->nb[2]) kparams.split_at_ne02 = 1;
+        };
+        for (const auto * in : prev.inputs) check_split(in);
+        check_split(new_src);
+        check_split(dst);
+
+        htp_binary_vtcm_layout_build(&L, &kparams, vtcm_size);
+    }
+
+    if (L.rows_per_buffer == 0 || L.total_bytes == 0 || L.total_bytes > vtcm_size) {
+        bool all_contig = ggml_is_contiguous(new_src) && ggml_is_contiguous(dst);
+        for (const auto * in : prev.inputs) {
+            all_contig = all_contig && ggml_is_contiguous(in);
+        }
+        if (!all_contig) return false;
+
+        const uint32_t total_elems = (uint32_t) ggml_nelements(src0);
+        const uint32_t elem_size   = (uint32_t) ggml_type_size(src0->type);
+        const uint32_t max_threads = htp_session_n_threads(sess);
+        if (!htp_binary_solve_chunked(&L, &kparams, total_elems, elem_size, max_threads, n_srcs, false, vtcm_size)) {
+            return false;
+        }
+    }
+
+    kparams.rows_per_buffer = L.rows_per_buffer;
+    kparams.vtcm_size       = (uint32_t) L.total_bytes;
+
+    prev.inputs.push_back(new_src);
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+bool htp_fusion_can_stack(const struct ggml_tensor * a) {
+    if (!a || !a->src[0] || !a->src[1]) return false;
+    switch (a->op) {
+        case GGML_OP_MUL_MAT:
+            if (!htp_fusion_enabled(HTP_FUSE_MUL_MAT_NX)) return false;
+            if (htp_is_hadamard(a)) return false;
+            if (a->src[0]->ne[2] != 1 || a->src[0]->ne[3] != 1) return false;
+            return ggml_is_quantized(a->src[0]->type);
+        case GGML_OP_MUL_MAT_ID:
+            if (!htp_fusion_enabled(HTP_FUSE_MUL_MAT_ID_NX)) return false;
+            return a->src[2] != nullptr && ggml_is_quantized(a->src[0]->type);
+        default:
+            return false;
+    }
+}
+
+bool htp_fusion_can_stack(const struct ggml_tensor * a, const struct ggml_tensor * b) {
+    if (!htp_fusion_can_stack(b) || a->op != b->op || a->src[1] != b->src[1]) {
+        return false;
+    }
+    if (a->op == GGML_OP_MUL_MAT_ID && a->src[2] != b->src[2]) {
+        return false;
+    }
+    const ggml_tensor * w0 = a->src[0];
+    const ggml_tensor * w1 = b->src[0];
+    return w0->type == w1->type && w0->ne[0] == w1->ne[0] && w0->ne[2] == w1->ne[2];
+}
+
+int htp_fusion_max(const struct ggml_cgraph * graph, int idx) {
+    const int n = graph->n_nodes;
+    const ggml_op op0 = graph->nodes[idx]->op;
+    if (op0 != GGML_OP_ADD && op0 != GGML_OP_NORM && op0 != GGML_OP_RMS_NORM) {
+        return 1;
+    }
+
+    enum ggml_op ops[HTP_OP_MAX_FUSED];
+    ops[0] = op0;
+
+    int f = idx + 1;
+    while (f < n && f < idx + HTP_OP_MAX_FUSED) {
+        const ggml_op op = graph->nodes[f]->op;
+        if (op != GGML_OP_ADD && op != GGML_OP_MUL && op != GGML_OP_SCALE &&
+            op != GGML_OP_NORM && op != GGML_OP_RMS_NORM) {
+            break;
+        }
+        ops[f - idx] = op;
+        f++;
+    }
+
+    f -= idx;
+    for (; f > 1; f--) {
+        if (ggml_can_fuse(graph, idx, ops, f)) {
+            break;
+        }
+    }
+    return f;
+}
+
+void htp_fusion_tag_graph(struct ggml_cgraph * graph) {
+    for (int i = 0; i < graph->n_nodes; i++) {
+        auto * extra = (htp_tensor_extra *) graph->nodes[i]->extra;
+        if (!extra) continue;
+
+        extra->flags &= ~HTP_TENSOR_FUSEABLE;
+
+        const ggml_op op = graph->nodes[i]->op;
+        if ((op == GGML_OP_RMS_NORM || op == GGML_OP_SCALE || op == GGML_OP_REPEAT ||
+             op == GGML_OP_UNARY    || op == GGML_OP_ADD   || op == GGML_OP_SUB    ||
+             op == GGML_OP_MUL      || op == GGML_OP_DIV   || op == GGML_OP_MUL_MAT_ID ||
+             op == GGML_OP_SSM_CONV ||
+             (op == GGML_OP_MUL_MAT && !htp_is_hadamard(graph->nodes[i]))) &&
+            ggml_node_has_n_uses(graph, i, 1)) {
+            extra->flags |= HTP_TENSOR_FUSEABLE;
+        }
+    }
+}
+
+bool htp_fusion_try_fuse(const struct ggml_hexagon_session * sess, struct htp_opnode & prev, const struct htp_opnode & cur) {
+    if (!opt_opfusion) return false;
+
+    switch (cur.opcode) {
+        case HTP_OP_ADD:
+            return (htp_fusion_enabled(HTP_FUSE_ALLREDUCE_ADD) && fuse_allreduce_add(sess, prev, cur)) ||
+                   (htp_fusion_enabled(HTP_FUSE_MUL_MAT_ADD)   && fuse_mul_mat_add(sess, prev, cur))   ||
+                   (htp_fusion_enabled(HTP_FUSE_BINARY_NX)     && fuse_binary_nx(sess, prev, cur));
+        case HTP_OP_MUL:
+            return (htp_fusion_enabled(HTP_FUSE_RMS_NORM_MUL)   && fuse_rms_norm_mul(sess, prev, cur))   ||
+                   (htp_fusion_enabled(HTP_FUSE_UNARY_MUL_GLU)  && fuse_unary_mul_glu(sess, prev, cur))  ||
+                   (htp_fusion_enabled(HTP_FUSE_MUL_MUL_SCALAR) && fuse_mul_mul_scalar(sess, prev, cur)) ||
+                   (htp_fusion_enabled(HTP_FUSE_BINARY_NX)      && fuse_binary_nx(sess, prev, cur));
+        case HTP_OP_SUB:
+        case HTP_OP_DIV:
+            return htp_fusion_enabled(HTP_FUSE_BINARY_NX) && fuse_binary_nx(sess, prev, cur);
+        case HTP_OP_SCALE:
+            return htp_fusion_enabled(HTP_FUSE_RMS_NORM_SCALE) && fuse_rms_norm_scale(sess, prev, cur);
+        case HTP_OP_UNARY_SILU:
+            return htp_fusion_enabled(HTP_FUSE_SSM_CONV_SILU) && fuse_ssm_conv_silu(sess, prev, cur);
+        case HTP_OP_GET_ROWS:
+            return htp_fusion_enabled(HTP_FUSE_REPEAT_GET_ROWS) && fuse_repeat_get_rows(sess, prev, cur);
+        case HTP_OP_CPY:
+            return htp_fusion_enabled(HTP_FUSE_GDN_CPY) && fuse_gdn_cpy(sess, prev, cur);
+        case HTP_OP_MUL_MAT:
+            return htp_fusion_enabled(HTP_FUSE_MUL_MAT_NX) && fuse_mul_mat_nx(sess, prev, cur);
+        case HTP_OP_MUL_MAT_ID:
+            return htp_fusion_enabled(HTP_FUSE_MUL_MAT_ID_NX) && fuse_mul_mat_id_nx(sess, prev, cur);
+        default:
+            return false;
+    }
+}
@@ -0,0 +1,33 @@
+#ifndef HTP_FUSION_H
+#define HTP_FUSION_H
+
+#include "htp-common.h"
+
+enum htp_fusion_flags {
+    HTP_FUSE_ALLREDUCE_ADD   = (1 << 1),  // 2
+    HTP_FUSE_RMS_NORM_MUL    = (1 << 2),  // 4
+    HTP_FUSE_MUL_MAT_ADD     = (1 << 3),  // 8
+    HTP_FUSE_MUL_MAT_NX      = (1 << 4),  // 16
+    HTP_FUSE_MUL_MAT_ID_NX   = (1 << 5),  // 32
+    HTP_FUSE_GDN_CPY         = (1 << 6),  // 64
+    HTP_FUSE_REPEAT_GET_ROWS = (1 << 7),  // 128
+    HTP_FUSE_UNARY_MUL_GLU   = (1 << 8),  // 256
+    HTP_FUSE_MUL_MUL_SCALAR  = (1 << 9),  // 512
+    HTP_FUSE_BINARY_NX       = (1 << 10), // 1024
+    HTP_FUSE_RMS_NORM_SCALE  = (1 << 11), // 2048
+    HTP_FUSE_SSM_CONV_SILU   = (1 << 12), // 4096
+};
+
+static inline bool htp_fusion_enabled(int flag) {
+    if (opt_opfusion <= 0) return false;
+    if (opt_opfusion == 1) return true; // 1 enables all
+    return (opt_opfusion & flag) != 0;
+}
+
+bool htp_fusion_can_stack(const struct ggml_tensor * a);
+bool htp_fusion_can_stack(const struct ggml_tensor * a, const struct ggml_tensor * b);
+int  htp_fusion_max(const struct ggml_cgraph * graph, int idx);
+void htp_fusion_tag_graph(struct ggml_cgraph * graph);
+bool htp_fusion_try_fuse(const struct ggml_hexagon_session * sess, struct htp_opnode & prev, const struct htp_opnode & cur);
+
+#endif // HTP_FUSION_H
@@ -0,0 +1,123 @@
+#ifndef HTP_KPARAMS_H
+#define HTP_KPARAMS_H
+
+#include "htp-common.h"
+
+#include "htp/allreduce-ops.h"
+#include "htp/argsort-ops.h"
+#include "htp/binary-ops.h"
+#include "htp/concat-ops.h"
+#include "htp/cpy-ops.h"
+#include "htp/cumsum-ops.h"
+#include "htp/diag-ops.h"
+#include "htp/fill-ops.h"
+#include "htp/flash-attn-ops.h"
+#include "htp/gated-delta-net-ops.h"
+#include "htp/get-rows-ops.h"
+#include "htp/glu-ops.h"
+#include "htp/im2col-ops.h"
+#include "htp/matmul-ops.h"
+#include "htp/pad-ops.h"
+#include "htp/pool-ops.h"
+#include "htp/repeat-ops.h"
+#include "htp/roll-ops.h"
+#include "htp/rope-ops.h"
+#include "htp/set-rows-ops.h"
+#include "htp/softmax-ops.h"
+#include "htp/solve-tri-ops.h"
+#include "htp/ssm-conv.h"
+#include "htp/sum-rows-ops.h"
+#include "htp/unary-ops.h"
+
+#include <type_traits>
+#include <utility>
+
+bool        htp_rows_stride(const int64_t * ne, const size_t * nb, size_t * stride);
+bool        htp_matmul_can_collapse(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst);
+ggml_tensor htp_tensor_collapse_rows(const struct ggml_tensor * t);
+
+template <typename... Args>
+auto htp_kparams_type(bool (*)(Args...)) -> std::remove_pointer_t<std::remove_reference_t<decltype((std::declval<Args>(), ...))>>;
+
+template <auto SolveFn, typename... Args>
+static inline bool htp_kparams_check(Args &&... args) {
+    decltype(htp_kparams_type(SolveFn)) kparams;
+    return SolveFn(std::forward<Args>(args)..., &kparams);
+}
+
+bool htp_kparams_solve(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    void * kparams = nullptr
+);
+
+static inline bool htp_kparams_check(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op
+) {
+    return htp_kparams_solve(sess, op, nullptr);
+}
+
+bool htp_kparams_solve_allreduce(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * dst,
+    uint32_t rank,
+    uint32_t n_ranks,
+    bool has_add,
+    bool is_row_bcast,
+    bool is_shard_ok,
+    struct htp_allreduce_kernel_params * kparams
+);
+
+bool htp_kparams_solve_rms_norm_mul(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * dst,
+    struct htp_unary_kernel_params * kparams
+);
+
+bool htp_kparams_solve_matmul_add(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * src2,
+    const struct ggml_tensor * dst,
+    struct htp_mm_kernel_params * kparams
+);
+
+bool htp_kparams_solve_mmnx(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * act,
+    int32_t n_weights,
+    struct htp_mm_kernel_params * kparams
+);
+
+bool htp_kparams_solve_mmidnx(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * act,
+    const struct ggml_tensor * dst,
+    int32_t n_weights,
+    struct htp_mm_kernel_params * kparams
+);
+
+bool htp_kparams_solve_repeat_get_rows(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * dst,
+    struct htp_get_rows_kernel_params * kparams
+);
+
+bool htp_kparams_solve_unary_mul_glu(
+    const struct ggml_hexagon_session * sess,
+    uint32_t op,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * dst,
+    struct htp_glu_kernel_params * kparams
+);
+
+#endif // HTP_KPARAMS_H
@@ -0,0 +1,260 @@
+#include "htp-opnode.h"
+
+#include "htp/allreduce-ops.h"
+#include "htp/argsort-ops.h"
+#include "htp/binary-ops.h"
+#include "htp/cumsum-ops.h"
+#include "htp/diag-ops.h"
+#include "htp/fill-ops.h"
+#include "htp/flash-attn-ops.h"
+#include "htp/gated-delta-net-ops.h"
+#include "htp/get-rows-ops.h"
+#include "htp/glu-ops.h"
+#include "htp/matmul-ops.h"
+#include "htp/pad-ops.h"
+#include "htp/repeat-ops.h"
+#include "htp/roll-ops.h"
+#include "htp/softmax-ops.h"
+#include "htp/solve-tri-ops.h"
+#include "htp/ssm-conv.h"
+#include "htp/sum-rows-ops.h"
+#include "htp/unary-ops.h"
+
+#include <stdio.h>
+
+std::string htp_opnode::name() const {
+    switch (opcode) {
+        case HTP_OP_MDEV_GROUP:    return "MDEV_GROUP";
+        case HTP_OP_CPY_FENCE:     return "CPY+FENCE";
+        case HTP_OP_FENCE:         return (node && node->op_params[1] == 0) ? "FENCE_WAIT" : "FENCE_SIGNAL";
+        case HTP_OP_ALLREDUCE:     return "ALLREDUCE";
+        case HTP_OP_ALLREDUCE_ADD: return "ALLREDUCE+ADD";
+        case HTP_OP_HADAMARD:      return "MUL_MAT_HADAMARD";
+        case HTP_OP_MUL_MAT_NX:    return "MUL_MAT_NX";
+        case HTP_OP_MUL_MAT_ID_NX: return "MUL_MAT_ID_NX";
+        case HTP_OP_ADD:
+        case HTP_OP_SUB:
+        case HTP_OP_MUL:
+        case HTP_OP_DIV:
+            if (!fused.empty()) {
+                const auto * kp = kparams<htp_binary_kernel_params>();
+                if (kp->kernel_type != HTP_BINARY_KERNEL_SCALAR) {
+                    switch (opcode) {
+                        case HTP_OP_ADD: return "ADD_NX";
+                        case HTP_OP_SUB: return "SUB_NX";
+                        case HTP_OP_MUL: return "MUL_NX";
+                        case HTP_OP_DIV: return "DIV_NX";
+                        default: break;
+                    }
+                }
+            }
+            break;
+        default:
+            break;
+    }
+
+    if (!node) {
+        return "";
+    }
+    if (htp_is_hadamard(node)) {
+        return "MUL_MAT_HADAMARD";
+    }
+
+    std::string res = ggml_op_desc(node);
+    for (const auto * t : fused) {
+        res += "+";
+        res += ggml_op_desc(t);
+    }
+    return res;
+}
+
+static const char * format_tensor_name(const struct ggml_tensor * t) {
+    return t ? t->name : "NONE";
+}
+
+static std::string format_tensor_dims(const struct ggml_tensor * t) {
+    if (!t) {
+        return "NONE";
+    }
+    char buf[64];
+    if (t->ne[2] == 1 && t->ne[3] == 1) {
+        snprintf(buf, sizeof(buf), "%d:%d", (int) t->ne[0], (int) t->ne[1]);
+    } else {
+        snprintf(buf, sizeof(buf), "%d:%d:%d:%d", (int) t->ne[0], (int) t->ne[1], (int) t->ne[2], (int) t->ne[3]);
+    }
+    return buf;
+}
+
+static const char * format_tensor_type(const struct ggml_tensor * t) {
+    return t ? ggml_type_name(t->type) : "NONE";
+}
+
+static std::string format_tensor_strides(const struct ggml_tensor * t) {
+    if (!t) {
+        return "NONE";
+    }
+    char buf[64];
+    const char * c = ggml_is_contiguous(t) ? "" : "!";
+    if (t->ne[2] == 1 && t->ne[3] == 1) {
+        snprintf(buf, sizeof(buf), "%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], c);
+    } else {
+        snprintf(buf, sizeof(buf), "%zu:%zu:%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], (size_t) t->nb[2], (size_t) t->nb[3], c);
+    }
+    return buf;
+}
+
+static const char * format_tensor_buff(const struct ggml_tensor * t) {
+    if (t && t->buffer) {
+        return ggml_backend_buffer_name(t->buffer);
+    }
+    return "NONE";
+}
+
+template <typename Fn>
+static std::string format_op_tensors(const htp_opnode & node, Fn fn) {
+    std::string res;
+    const auto & inputs  = node.get_inputs();
+    const auto & outputs = node.get_outputs();
+
+    if (!inputs.empty()) {
+        for (size_t i = 0; i < inputs.size(); i++) {
+            if (i > 0) res += " x ";
+            res += fn(inputs[i]);
+        }
+        res += " -> ";
+    }
+
+    if (outputs.empty()) {
+        res += fn(node.dst());
+    } else {
+        for (size_t i = 0; i < outputs.size(); i++) {
+            if (i > 0) res += " x ";
+            res += fn(outputs[i]);
+        }
+    }
+    return res;
+}
+
+static std::string format_kernel_params(const htp_opnode & node) {
+    char str[128];
+    const size_t max_size = sizeof(str);
+
+    if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID ||
+        node.opcode == HTP_OP_MUL_MAT_NX || node.opcode == HTP_OP_MUL_MAT_ID_NX ||
+        node.opcode == HTP_OP_MUL_MAT_ADD) {
+        const auto * kparams = node.kparams<htp_mm_kernel_params>();
+        const char * path = "unknown";
+        int32_t type = kparams->kernel_type;
+        if (type == HTP_MM_KERNEL_HMX_2D || type == HTP_MM_KERNEL_HMX_BATCHED) {
+            path = (kparams->tile_size > 0) ? "hmx-tiled" : "hmx-flat";
+        } else if (type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || type == HTP_MM_KERNEL_HVX_F32_F32_VTCM ||
+                   type == HTP_MM_KERNEL_HVX_BF16_F32_VTCM) {
+            path = "hvx-flat";
+        } else if (type == HTP_MM_KERNEL_HVX_QUANT_ROW || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
+            path = "hvx-tiled";
+        }
+        snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) {
+        const auto * kparams = node.kparams<htp_fa_kernel_params>();
+        const char * path = "unknown";
+        int32_t type = kparams->kernel_type;
+        if (type == HTP_FA_KERNEL_HMX) {
+            path = kparams->u.hmx.pipeline ? "hmx-pipe" : "hmx-seq";
+        } else if (type == HTP_FA_KERNEL_HVX) {
+            path = "hvx";
+        }
+        snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
+    } else if (htp_op_is_unary(node.opcode)) {
+        const auto * kparams = node.kparams<htp_unary_kernel_params>();
+        snprintf(str, max_size, "%s vtcm %d", kparams->col_tile ? "wide-row" : "row-block", (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_MDEV_GROUP && node.node) {
+        snprintf(str, max_size, "idx %d count %d", (int) node.node->op_params[0], (int) node.dst()->ne[1]);
+    } else if ((node.opcode == HTP_OP_FENCE || node.opcode == HTP_OP_CPY_FENCE) && node.node) {
+        snprintf(str, max_size, "seq 0x%x", (uint32_t) node.node->op_params[0]);
+    } else if (node.opcode == HTP_OP_ALLREDUCE && node.node) {
+        snprintf(str, max_size, "seq 0x%x -> 0x%x", (uint32_t) node.node->op_params[0], (uint32_t) node.node->op_params[1]);
+    } else if (node.opcode == HTP_OP_SSM_CONV) {
+        const auto * kparams = node.kparams<htp_ssm_conv_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 = node.kparams<htp_softmax_kernel_params>();
+        snprintf(str, max_size, "k%d nth %d vtcm %d", (int) kparams->kernel_type, (int) kparams->n_threads, (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_GATED_DELTA_NET) {
+        const auto * kparams = node.kparams<htp_gdn_kernel_params>();
+        const char * path = (kparams->kernel_type == HTP_GDN_KERNEL_HMX_CHUNKED) ? "hmx-chunked" : "hvx-recurrent";
+        snprintf(str, max_size, "%s-%s vtcm %u", path, kparams->kda ? "kda" : "scalar",
+                 (unsigned int) kparams->vtcm_size);
+    } 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 = node.kparams<htp_binary_kernel_params>();
+        snprintf(str, max_size, "vtcm %u", (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_ARGSORT || node.opcode == HTP_OP_TOP_K) {
+        const auto * kparams = node.kparams<htp_sort_kernel_params>();
+        snprintf(str, max_size, "%s nth %d nchk %d chk %d vtcm %d",
+                 node.opcode == HTP_OP_TOP_K ? "top_k" : "argsort",
+                 (int) kparams->n_threads, (int) kparams->n_chunks,
+                 (int) kparams->chunk_elems, (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_GET_ROWS) {
+        const auto * kparams = node.kparams<htp_get_rows_kernel_params>();
+        const char * ktype_str = "unknown";
+        switch (kparams->kernel_type) {
+            case HTP_GET_ROWS_KERNEL_SAMETYPE: ktype_str = "sametype"; break;
+            case HTP_GET_ROWS_KERNEL_TILED:    ktype_str = "tiled"; break;
+            case HTP_GET_ROWS_KERNEL_FLAT:     ktype_str = "flat"; break;
+        }
+        snprintf(str, max_size, "%s%s vtcm %u", ktype_str, kparams->n_threads > 1 ? "-multi" : "", (unsigned int) kparams->vtcm_size);
+    } else if (htp_op_is_glu(node.opcode)) {
+        const auto * kparams = node.kparams<htp_glu_kernel_params>();
+        snprintf(str, max_size, "nth %d vtcm %u", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_PAD) {
+        const auto * kparams = node.kparams<htp_pad_kernel_params>();
+        snprintf(str, max_size, "k%d nth %d vtcm %u", (int) kparams->kernel_type, (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_CUMSUM) {
+        const auto * kparams = node.kparams<htp_cumsum_kernel_params>();
+        snprintf(str, max_size, "%s nth %d vtcm %u", kparams->use_dma ? "dma" : "hvx", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_FILL) {
+        const auto * kparams = node.kparams<htp_fill_kernel_params>();
+        snprintf(str, max_size, "%s nth %d", kparams->kernel_type == HTP_FILL_KERNEL_CONTIGUOUS ? "contig" : "strided", (int) kparams->n_threads);
+    } else if (node.opcode == HTP_OP_ROLL) {
+        const auto * kparams = node.kparams<htp_roll_kernel_params>();
+        snprintf(str, max_size, "k%d nth %d", (int) kparams->kernel_type, (int) kparams->n_threads);
+    } else if (node.opcode == HTP_OP_DIAG) {
+        const auto * kparams = node.kparams<htp_diag_kernel_params>();
+        snprintf(str, max_size, "%s nth %d vtcm %u", kparams->use_dma ? "dma" : "hvx", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_SOLVE_TRI) {
+        const auto * kparams = node.kparams<htp_solve_tri_kernel_params>();
+        snprintf(str, max_size, "%s nth %d", kparams->batched ? "batched" : "chunked", (int) kparams->n_threads);
+    } else {
+        snprintf(str, max_size, "----");
+    }
+    return str;
+}
+
+static void append_section(std::string & res, const std::string & sec) {
+    if (!res.empty()) {
+        res += "|";
+    }
+    res += sec;
+}
+
+std::string htp_format_tensor(const ggml_tensor * t, uint32_t flags) {
+    std::string res;
+    if (flags & HTP_OPFMT_NAMES)   append_section(res, format_tensor_name(t));
+    if (flags & HTP_OPFMT_DIMS)    append_section(res, format_tensor_dims(t));
+    if (flags & HTP_OPFMT_TYPES)   append_section(res, format_tensor_type(t));
+    if (flags & HTP_OPFMT_STRIDES) append_section(res, format_tensor_strides(t));
+    if (flags & HTP_OPFMT_BUFFS)   append_section(res, format_tensor_buff(t));
+    return res;
+}
+
+std::string htp_format_op(const htp_opnode & node, uint32_t flags) {
+    std::string res;
+    if (flags & HTP_OPFMT_OP)      append_section(res, node.name());
+    if (flags & HTP_OPFMT_NAMES)   append_section(res, format_op_tensors(node, format_tensor_name));
+    if (flags & HTP_OPFMT_DIMS)    append_section(res, format_op_tensors(node, format_tensor_dims));
+    if (flags & HTP_OPFMT_TYPES)   append_section(res, format_op_tensors(node, format_tensor_type));
+    if (flags & HTP_OPFMT_STRIDES) append_section(res, format_op_tensors(node, format_tensor_strides));
+    if (flags & HTP_OPFMT_BUFFS)   append_section(res, format_op_tensors(node, format_tensor_buff));
+    if (flags & HTP_OPFMT_KPARAMS) append_section(res, format_kernel_params(node));
+    return res;
+}
@@ -1,38 +1,74 @@
 #ifndef HTP_OPNODE_H
 #define HTP_OPNODE_H

-#define GGML_COMMON_IMPL_CPP
 #include "ggml-backend-impl.h"
-#include "ggml-common.h"
+#include "ggml-impl.h"
+#include "htp/htp-ops.h"

-#include <algorithm>
+#include <cstring>
+#include <initializer_list>
+#include <memory>
 #include <string>
 #include <vector>
-#include <memory>
-#include <stdio.h>
-#include "htp-ops.h"
-#include "htp/matmul-ops.h"
-#include "htp/flash-attn-ops.h"
-#include "htp/unary-ops.h"
-#include "htp/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"
-#include "htp/argsort-ops.h"
-#include "htp/get-rows-ops.h"
+
+#define HTP_OP_MAX_FUSED 16
+
+// local std::inplace_vector replacement
+template <typename T, size_t MAX_N>
+struct htp_opvec {
+    T      data[MAX_N] {};
+    size_t n = 0;
+
+    htp_opvec() = default;
+    htp_opvec(std::initializer_list<T> init) {
+        GGML_ASSERT(init.size() <= MAX_N);
+        for (const auto & v : init) {
+            data[n++] = v;
+        }
+    }
+
+    size_t    size()  const { return n; }
+    bool      empty() const { return n == 0; }
+    void      clear()       { n = 0; }
+    T *       begin()       { return data; }
+    T *       end()         { return data + n; }
+    const T * begin() const { return data; }
+    const T * end()   const { return data + n; }
+    T &       back()        { return data[n - 1]; }
+    const T & back()  const { return data[n - 1]; }
+    T &       operator[](size_t i)       { return data[i]; }
+    const T & operator[](size_t i) const { return data[i]; }
+
+    void push_back(T val) {
+        GGML_ASSERT(n < MAX_N);
+        data[n++] = val;
+    }
+
+    void insert(T * pos, T val) {
+        GGML_ASSERT(n < MAX_N);
+        for (T * p = data + n; p > pos; --p) {
+            *p = *(p - 1);
+        }
+        *pos = val;
+        n++;
+    }
+};
+
+static inline bool htp_is_hadamard(const struct ggml_tensor * t) {
+    return t && t->op == GGML_OP_MUL_MAT && ggml_get_op_params_i32(t, 1) == GGML_HINT_SRC0_IS_HADAMARD;
+}

 struct htp_opnode {
-    ggml_tensor * node   { nullptr };
-    htp_op_code   opcode { HTP_OP_INVALID };
-    int32_t       kernel_params[HTP_OP_MAX_KERN_PARAMS] {0};
+    const ggml_tensor * node   { nullptr };
+    htp_op_code         opcode { HTP_OP_INVALID };
+    int32_t             params[HTP_OP_MAX_PARAMS] {0};
+    int32_t             kernel_params[HTP_OP_MAX_KERN_PARAMS] {0};

-    std::vector<ggml_tensor *>                fused;
-    std::vector<std::shared_ptr<ggml_tensor>> dummy;
+    std::vector<std::shared_ptr<ggml_tensor>>          dummy;

-    std::vector<const ggml_tensor *> inputs;
-    std::vector<const ggml_tensor *> outputs;
-    std::string                      name;
+    htp_opvec<const ggml_tensor *, HTP_OP_MAX_FUSED>   fused;
+    htp_opvec<const ggml_tensor *, HTP_OP_MAX_INPUTS>  inputs;
+    htp_opvec<const ggml_tensor *, HTP_OP_MAX_OUTPUTS> outputs;

    int n_active_src(const ggml_tensor * t) const {
        if (!t) return 0;
@@ -44,34 +80,36 @@ struct htp_opnode {
        return 0;
    }

-    void init(ggml_tensor * node) {
+    void init(const ggml_tensor * node) {
+        static_assert(sizeof(this->params) == sizeof(this->node->op_params), "params size mismatch");
+
        this->node = node;
-        if (this->node) {
-            if (this->node->op == GGML_OP_MUL_MAT &&
-                ggml_get_op_params_i32(this->node, 1) == GGML_HINT_SRC0_IS_HADAMARD) {
-                this->name = "MUL_MAT_HADAMARD";
-                this->inputs.push_back(this->node->src[1]);
-            } else {
-                this->name = ggml_op_desc(this->node);
+        if (!node) return;

-                // Build inputs (preserving optional nullptrs)
-                int n_inputs = n_active_src(this->node);
-                this->inputs.resize(n_inputs, nullptr);
-                for (int i = 0; i < n_inputs; i++) {
-                    this->inputs[i] = this->node->src[i];
-                }
+        memcpy(this->params, this->node->op_params, sizeof(this->params));
+
+        // Build inputs (preserving optional nullptrs)
+        if (htp_is_hadamard(this->node)) {
+            this->inputs.push_back(this->node->src[1]);
+        } else {
+            int n_inputs = n_active_src(this->node);
+            for (int i = 0; i < n_inputs; i++) {
+                this->inputs.push_back(this->node->src[i]);
            }
-
-            // Build outputs
-            this->outputs.push_back(this->dst());
        }
+
+        // Build outputs
+        this->outputs.push_back(this->dst());
    }

-    htp_opnode(htp_op_code opcode = HTP_OP_INVALID, ggml_tensor * node = nullptr) : opcode(opcode) {
+    htp_opnode(htp_op_code opcode = HTP_OP_INVALID, const ggml_tensor * node = nullptr) : opcode(opcode) {
+        init(node);
+    }
+
+    htp_opnode(const ggml_tensor * node) : opcode(HTP_OP_INVALID) {
        init(node);
    }

-    ggml_op             op()   const { return node->op; }
    const ggml_tensor * src0() const { return node->src[0]; }
    const ggml_tensor * src1() const { return node->src[1]; }
    const ggml_tensor * dst()  const { return outputs.empty() ? node : outputs.back(); }
@@ -81,343 +119,54 @@ struct htp_opnode {
        return dummy.back().get();
    }

-    void add_fused(ggml_tensor * t, bool extra_dst = false) {
-        fused.push_back(t);
-
-        name += "+";
-        name += ggml_op_desc(t);
-
-        if (extra_dst) {
-            outputs.push_back(t);
-        } else {
-            outputs.clear();
-            outputs.push_back(t);
-        }
-
-        // Remove the newly fused intermediate output tensor t from inputs (if it was there)
-        inputs.erase(std::remove(inputs.begin(), inputs.end(), t), inputs.end());
-
-        // Append new inputs from t, preserving middle nullptrs
-        int n_inputs = n_active_src(t);
-        for (int i = 0; i < n_inputs; i++) {
-            const auto * src = t->src[i];
-            if (!src) {
-                inputs.push_back(nullptr);
-            } else if (src != node &&
-                       std::find(fused.begin(), fused.end(), src) == fused.end() &&
-                       std::find(inputs.begin(), inputs.end(), src) == inputs.end()) {
-                inputs.push_back(src);
-            }
-        }
-    }
-
-    const std::vector<const ggml_tensor *> & get_inputs() const {
+    const htp_opvec<const ggml_tensor *, HTP_OP_MAX_INPUTS> & get_inputs() const {
        return inputs;
    }

-    const std::vector<const ggml_tensor *> & get_outputs() const {
+    const htp_opvec<const ggml_tensor *, HTP_OP_MAX_OUTPUTS> & get_outputs() const {
        return outputs;
    }

-    std::string op_name() const {
-        return name;
+    template <typename T>
+    T * kparams() {
+        static_assert(sizeof(T) <= sizeof(kernel_params), "kparams size exceeds HTP_OP_MAX_KERN_PARAMS");
+        static_assert(alignof(T) <= alignof(int32_t), "kparams alignment exceeds int32_t");
+        return reinterpret_cast<T *>(kernel_params);
    }

-    bool is_empty() const {
-        return ggml_op_is_empty(node->op);
+    template <typename T>
+    const T * kparams() const {
+        static_assert(sizeof(T) <= sizeof(kernel_params), "kparams size exceeds HTP_OP_MAX_KERN_PARAMS");
+        static_assert(alignof(T) <= alignof(int32_t), "kparams alignment exceeds int32_t");
+        return reinterpret_cast<const T *>(kernel_params);
    }

-    bool stackable() const {
-        switch (this->op()) {
-            case GGML_OP_MUL_MAT:
-                if (ggml_get_op_params_i32(this->node, 1) == GGML_HINT_SRC0_IS_HADAMARD) {
-                    return false;
-                }
-                return ggml_is_quantized(this->src0()->type);
-            case GGML_OP_MUL_MAT_ID:
-                return ggml_is_quantized(this->src0()->type);
-            default:
-                return false;
-        }
+    template <typename T>
+    void set_kparams(const T & kp) {
+        static_assert(sizeof(T) <= sizeof(kernel_params), "kparams size exceeds HTP_OP_MAX_KERN_PARAMS");
+        static_assert(alignof(T) <= alignof(int32_t), "kparams alignment exceeds int32_t");
+        memcpy(kernel_params, &kp, sizeof(T));
    }

-    bool same_input(const htp_opnode& n) const {
-        return n.src1() == this->src1();
-    }
+    std::string name() const;
 };

-struct htp_opformat {
-    char strides[64 * GGML_MAX_SRC];
-    char dims[64 * GGML_MAX_SRC];
-    char types[16 * GGML_MAX_SRC];
-    char buffs[64 * GGML_MAX_SRC];
-    char names[64 * GGML_MAX_SRC];
-    char kparams[128];
+enum htp_opformat_flags {
+    HTP_OPFMT_OP      = (1 << 0),
+    HTP_OPFMT_NAMES   = (1 << 1),
+    HTP_OPFMT_DIMS    = (1 << 2),
+    HTP_OPFMT_TYPES   = (1 << 3),
+    HTP_OPFMT_STRIDES = (1 << 4),
+    HTP_OPFMT_BUFFS   = (1 << 5),
+    HTP_OPFMT_KPARAMS = (1 << 6),

-    int format_tensor_dims(char * str, size_t max_size, const struct ggml_tensor * t) {
-        if (!t) {
-            return snprintf(str, max_size, "NONE");
-        }
-        if (t->ne[2] == 1 && t->ne[3] == 1) {
-            return snprintf(str, max_size, "%d:%d", (int) t->ne[0], (int) t->ne[1]);
-        } else {
-            return snprintf(str, max_size, "%d:%d:%d:%d", (int) t->ne[0], (int) t->ne[1], (int) t->ne[2], (int) t->ne[3]);
-        }
-    }
-
-    void format_op_dims(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            p += std::min((size_t)format_tensor_dims(p, p_end - p, inputs[0]), (size_t)(p_end - p));
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)format_tensor_dims(p, p_end - p, inputs[i]), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        char self[64];
-        format_tensor_dims(self, sizeof(self), node.dst());
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", self), (size_t)(p_end - p));
-        }
-    }
-
-    int format_tensor_strides(char * str, size_t max_size, const struct ggml_tensor * t) {
-        if (!t) {
-            return snprintf(str, max_size, "NONE");
-        }
-        const char * c = ggml_is_contiguous(t) ? "" : "!";
-
-        if (t->ne[2] == 1 && t->ne[3] == 1) {
-            return snprintf(str, max_size, "%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], c);
-        } else {
-            return snprintf(str, max_size, "%zu:%zu:%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], (size_t) t->nb[2], (size_t) t->nb[3], c);
-        }
-    }
-
-    void format_op_strides(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            p += std::min((size_t)format_tensor_strides(p, p_end - p, inputs[0]), (size_t)(p_end - p));
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)format_tensor_strides(p, p_end - p, inputs[i]), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        char self[64];
-        format_tensor_strides(self, sizeof(self), node.dst());
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", self), (size_t)(p_end - p));
-        }
-    }
-
-    void format_op_types(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[0] ? ggml_type_name(inputs[0]->type) : "NONE"), (size_t)(p_end - p));
-            }
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[i] ? ggml_type_name(inputs[i]->type) : "NONE"), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", ggml_type_name(node.dst()->type)), (size_t)(p_end - p));
-        }
-    }
-
-    const char * tensor_buff_name(const struct ggml_tensor * t) {
-        if (t && t->buffer) {
-            return ggml_backend_buffer_name(t->buffer);
-        }
-        return "NONE";
-    }
-
-    void format_op_buffs(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, "%s", tensor_buff_name(inputs[0])), (size_t)(p_end - p));
-            }
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, "%s", tensor_buff_name(inputs[i])), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", tensor_buff_name(node.dst())), (size_t)(p_end - p));
-        }
-    }
-
-    void format_op_names(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[0] ? inputs[0]->name : "NONE"), (size_t)(p_end - p));
-            }
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[i] ? inputs[i]->name : "NONE"), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", node.dst()->name), (size_t)(p_end - p));
-        }
-    }
-    void format_kernel_params(char * str, size_t max_size, const htp_opnode & node) {
-        if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID ||
-            node.opcode == HTP_OP_MUL_MAT_NX || node.opcode == HTP_OP_MUL_MAT_ID_NX ||
-            node.opcode == HTP_OP_MUL_MAT_ADD) {
-            const auto * kparams = (const struct htp_mm_kernel_params *) node.kernel_params;
-            const char * path = "unknown";
-            int32_t type = kparams->kernel_type;
-            if (type == HTP_MM_KERNEL_HMX_2D || type == HTP_MM_KERNEL_HMX_BATCHED) {
-                path = (kparams->tile_size > 0) ? "hmx-tiled" : "hmx-flat";
-            } else if (type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || type == HTP_MM_KERNEL_HVX_F32_F32_VTCM ||
-                       type == HTP_MM_KERNEL_HVX_BF16_F32_VTCM) {
-                path = "hvx-flat";
-            } else if (type == HTP_MM_KERNEL_HVX_QUANT_ROW || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
-                path = "hvx-tiled";
-            }
-            snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) {
-            const auto * kparams = (const struct htp_fa_kernel_params *) node.kernel_params;
-            const char * path = "unknown";
-            int32_t type = kparams->kernel_type;
-            if (type == HTP_FA_KERNEL_HMX) {
-                path = kparams->u.hmx.pipeline ? "hmx-pipe" : "hmx-seq";
-            } else if (type == HTP_FA_KERNEL_HVX) {
-                path = "hvx";
-            }
-            snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
-        } else if (htp_op_is_unary(node.opcode)) {
-            const auto * kparams = (const struct htp_unary_kernel_params *) node.kernel_params;
-            snprintf(str, max_size, "%s vtcm %d", kparams->col_tile ? "wide-row" : "row-block", (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_MDEV_GROUP && node.node) {
-            snprintf(str, max_size, "idx %d count %d", (int) node.node->op_params[0], (int) node.dst()->ne[1]);
-        } else if ((node.opcode == HTP_OP_FENCE || node.opcode == HTP_OP_CPY_FENCE) && node.node) {
-            snprintf(str, max_size, "seq 0x%x", (uint32_t) node.node->op_params[0]);
-        } else if (node.opcode == HTP_OP_ALLREDUCE && node.node) {
-            snprintf(str, max_size, "seq 0x%x -> 0x%x", (uint32_t) node.node->op_params[0], (uint32_t) node.node->op_params[1]);
-        } else 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;
-            const char * path = (kparams->kernel_type == HTP_GDN_KERNEL_HMX_CHUNKED) ? "hmx-chunked" : "hvx-recurrent";
-            snprintf(str, max_size, "%s-%s vtcm %u", path, 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 if (node.opcode == HTP_OP_ARGSORT || node.opcode == HTP_OP_TOP_K) {
-            const auto * kparams = (const struct htp_sort_kernel_params *) node.kernel_params;
-            snprintf(str, max_size, "%s nth %d nchk %d chk %d vtcm %d",
-                     node.opcode == HTP_OP_TOP_K ? "top_k" : "argsort",
-                     (int) kparams->n_threads, (int) kparams->n_chunks,
-                     (int) kparams->chunk_elems, (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_GET_ROWS) {
-            const auto * kparams = (const struct htp_get_rows_kernel_params *) node.kernel_params;
-            const char * ktype_str = "unknown";
-            switch (kparams->kernel_type) {
-                case HTP_GET_ROWS_KERNEL_SAMETYPE: ktype_str = "sametype"; break;
-                case HTP_GET_ROWS_KERNEL_TILED:    ktype_str = "tiled"; break;
-                case HTP_GET_ROWS_KERNEL_FLAT:     ktype_str = "flat"; break;
-            }
-            snprintf(str, max_size, "%s%s vtcm %u", ktype_str, kparams->n_threads > 1 ? "-multi" : "", (unsigned int) kparams->vtcm_size);
-        } else {
-            snprintf(str, max_size, "----");
-        }
-
-    }
-
-    void format(const htp_opnode & node) {
-        format_op_dims(dims, sizeof(dims), node);
-        format_op_strides(strides, sizeof(strides), node);
-        format_op_types(types, sizeof(types), node);
-        format_op_buffs(buffs, sizeof(buffs), node);
-        format_op_names(names, sizeof(names), node);
-        format_kernel_params(kparams, sizeof(kparams), node);
-    }
-
-    htp_opformat() {
-        strides[0] = '\0';
-        dims[0]    = '\0';
-        types[0]   = '\0';
-        buffs[0]   = '\0';
-        names[0]   = '\0';
-        kparams[0] = '\0';
-    }
-    htp_opformat(const htp_opnode & node) { format(node); }
+    HTP_OPFMT_EXEC    = HTP_OPFMT_OP | HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES | HTP_OPFMT_BUFFS | HTP_OPFMT_KPARAMS,
+    HTP_OPFMT_SUPP    = HTP_OPFMT_OP | HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES | HTP_OPFMT_BUFFS,
+    HTP_OPFMT_PROF    = HTP_OPFMT_OP | HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES | HTP_OPFMT_KPARAMS,
+    HTP_OPFMT_TENSOR  = HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES,
 };

+std::string htp_format_op(const htp_opnode & node, uint32_t flags = HTP_OPFMT_EXEC);
+std::string htp_format_tensor(const ggml_tensor * t, uint32_t flags = HTP_OPFMT_TENSOR);
+
 #endif // HTP_OPNODE_H
@@ -9,6 +9,11 @@
 #include <cmath>
 #include <cstring>
 #include <initializer_list>
+#include <thread>
+
+#if defined(__ARM_NEON) || defined(__aarch64__)
+#    include <arm_neon.h>
+#endif

 // --- 4-bit block-32 nibble helpers (shared by Q4_0, IQ4_NL, Q4_1, MXFP4) ---

@@ -43,8 +48,309 @@ static inline void transpose_tile_to_nibbles(uint8_t tile_quants[32][32], const
    }
 }

+#if defined(__ARM_NEON) || defined(__aarch64__)
+static void repack_neon_transpose8x8(const uint8x8_t rows[8], uint8x8_t columns[8]) {
+    const uint8x8x2_t t0 = vtrn_u8(rows[0], rows[1]);
+    const uint8x8x2_t t1 = vtrn_u8(rows[2], rows[3]);
+    const uint8x8x2_t t2 = vtrn_u8(rows[4], rows[5]);
+    const uint8x8x2_t t3 = vtrn_u8(rows[6], rows[7]);
+    const uint16x4x2_t t4 = vtrn_u16(vreinterpret_u16_u8(t0.val[0]), vreinterpret_u16_u8(t1.val[0]));
+    const uint16x4x2_t t5 = vtrn_u16(vreinterpret_u16_u8(t0.val[1]), vreinterpret_u16_u8(t1.val[1]));
+    const uint16x4x2_t t6 = vtrn_u16(vreinterpret_u16_u8(t2.val[0]), vreinterpret_u16_u8(t3.val[0]));
+    const uint16x4x2_t t7 = vtrn_u16(vreinterpret_u16_u8(t2.val[1]), vreinterpret_u16_u8(t3.val[1]));
+    const uint32x2x2_t t8 = vtrn_u32(vreinterpret_u32_u16(t4.val[0]), vreinterpret_u32_u16(t6.val[0]));
+    const uint32x2x2_t t9 = vtrn_u32(vreinterpret_u32_u16(t4.val[1]), vreinterpret_u32_u16(t6.val[1]));
+    const uint32x2x2_t ta = vtrn_u32(vreinterpret_u32_u16(t5.val[0]), vreinterpret_u32_u16(t7.val[0]));
+    const uint32x2x2_t tb = vtrn_u32(vreinterpret_u32_u16(t5.val[1]), vreinterpret_u32_u16(t7.val[1]));
+    columns[0] = vreinterpret_u8_u32(t8.val[0]);
+    columns[1] = vreinterpret_u8_u32(ta.val[0]);
+    columns[2] = vreinterpret_u8_u32(t9.val[0]);
+    columns[3] = vreinterpret_u8_u32(tb.val[0]);
+    columns[4] = vreinterpret_u8_u32(t8.val[1]);
+    columns[5] = vreinterpret_u8_u32(ta.val[1]);
+    columns[6] = vreinterpret_u8_u32(t9.val[1]);
+    columns[7] = vreinterpret_u8_u32(tb.val[1]);
+}
+
+static uint8x8_t repack_neon_pack_4bit_pairs(uint8x8_t first, uint8x8_t second, bool high) {
+    const uint8x8x2_t separated = vuzp_u8(first, second);
+    const uint8x8_t mask = vdup_n_u8(0x0F);
+    uint8x8_t even = separated.val[0];
+    uint8x8_t odd  = separated.val[1];
+
+    if (high) {
+        even = vshr_n_u8(even, 4);
+        odd  = vshr_n_u8(odd, 4);
+    }
+
+    even = vand_u8(even, mask);
+    odd  = vand_u8(odd, mask);
+    return vorr_u8(vshl_n_u8(odd, 4), even);
+}
+
+static void repack_neon_4bit_rows(const uint8_t * qs[8], uint8_t * tile_dst, int row0) {
+    uint8x8_t low_rows[8];
+    uint8x8_t high_rows[8];
+    uint8x8_t columns[8];
+
+    for (int rr = 0; rr < 8; ++rr) {
+        const uint8x16_t packed = vld1q_u8(qs[rr]);
+        const uint8x8_t first = vget_low_u8(packed);
+        const uint8x8_t second = vget_high_u8(packed);
+        low_rows[rr]  = repack_neon_pack_4bit_pairs(first, second, false);
+        high_rows[rr] = repack_neon_pack_4bit_pairs(first, second, true);
+    }
+
+    repack_neon_transpose8x8(low_rows, columns);
+    for (int cp = 0; cp < 8; ++cp) {
+        vst1_u8(tile_dst + cp * 32 + row0, columns[cp]);
+    }
+
+    repack_neon_transpose8x8(high_rows, columns);
+    for (int cp = 0; cp < 8; ++cp) {
+        vst1_u8(tile_dst + (cp + 8) * 32 + row0, columns[cp]);
+    }
+}
+
+static void repack_neon_k4_rows(const uint8_t * qs[8], uint8_t * tile_dst, int row0, bool high) {
+    uint8x8_t first_rows[8];
+    uint8x8_t second_rows[8];
+    uint8x8_t columns[8];
+
+    for (int rr = 0; rr < 8; ++rr) {
+        first_rows[rr] = repack_neon_pack_4bit_pairs(vld1_u8(qs[rr]), vld1_u8(qs[rr] + 8), high);
+        second_rows[rr] = repack_neon_pack_4bit_pairs(vld1_u8(qs[rr] + 16), vld1_u8(qs[rr] + 24), high);
+    }
+
+    repack_neon_transpose8x8(first_rows, columns);
+    for (int cp = 0; cp < 8; ++cp) {
+        vst1_u8(tile_dst + cp * 32 + row0, columns[cp]);
+    }
+
+    repack_neon_transpose8x8(second_rows, columns);
+    for (int cp = 0; cp < 8; ++cp) {
+        vst1_u8(tile_dst + (cp + 8) * 32 + row0, columns[cp]);
+    }
+}
+
+static void repack_neon_q8_rows(const uint8_t * qs[8], uint8_t * tile_dst, int row0) {
+    uint8x8_t rows[8];
+    uint8x8_t columns[8];
+
+    for (int cp = 0; cp < 16; cp += 4) {
+        for (int rr = 0; rr < 8; ++rr) {
+            rows[rr] = vld1_u8(qs[rr] + 2 * cp);
+        }
+        repack_neon_transpose8x8(rows, columns);
+
+        for (int pair = 0; pair < 4; ++pair) {
+            const uint8x8x2_t interleaved = vzip_u8(columns[2 * pair], columns[2 * pair + 1]);
+            uint8_t * dst = tile_dst + (cp + pair) * 64 + 2 * row0;
+            vst1_u8(dst, interleaved.val[0]);
+            vst1_u8(dst + 8, interleaved.val[1]);
+        }
+    }
+}
+
+static inline void repack_store_u32(uint8_t * dst, uint32_t v) {
+    memcpy(dst, &v, sizeof(v));
+}
+
+static inline uint32x4_t repack_neon_transpose_2bit_4x4(uint32x4_t w) {
+    uint32x4_t t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 6)), vdupq_n_u32(0x00CC00CCu));
+    w = veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 6)));
+    t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 12)), vdupq_n_u32(0x0000F0F0u));
+    return veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 12)));
+}
+
+static inline void repack_neon_q2_0_block(const uint8_t * qs, uint8_t * tile, int row) {
+    const uint32x4_t w = repack_neon_transpose_2bit_4x4(vreinterpretq_u32_u8(vld1q_u8(qs)));
+    repack_store_u32(tile + 0 * 128 + 4 * row, vgetq_lane_u32(w, 0));
+    repack_store_u32(tile + 1 * 128 + 4 * row, vgetq_lane_u32(w, 1));
+    repack_store_u32(tile + 2 * 128 + 4 * row, vgetq_lane_u32(w, 2));
+    repack_store_u32(tile + 3 * 128 + 4 * row, vgetq_lane_u32(w, 3));
+}
+
+static inline void repack_neon_q1_0_block(const uint8_t * qs, uint8_t * tile, int row) {
+    uint32x4_t w = vreinterpretq_u32_u8(vld1q_u8(qs));
+    uint32x4_t t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 1)), vdupq_n_u32(0x22222222u));
+    w = veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 1)));
+    t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 3)), vdupq_n_u32(0x0A0A0A0Au));
+    w = veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 3)));
+    w = repack_neon_transpose_2bit_4x4(w);
+    repack_store_u32(tile + 0 * 128 + 4 * row, vgetq_lane_u32(w, 0));
+    repack_store_u32(tile + 1 * 128 + 4 * row, vgetq_lane_u32(w, 1));
+    repack_store_u32(tile + 2 * 128 + 4 * row, vgetq_lane_u32(w, 2));
+    repack_store_u32(tile + 3 * 128 + 4 * row, vgetq_lane_u32(w, 3));
+}
+
+static inline void repack_neon_pack_2bit_plane(uint8x16_t v0, uint8x16_t v1, uint8_t * dst, int row) {
+    const uint32x4x2_t z   = vuzpq_u32(vreinterpretq_u32_u8(v0), vreinterpretq_u32_u8(v1));
+    const uint32x4_t   p01 = vorrq_u32(z.val[0], vshlq_n_u32(z.val[1], 2));
+    const uint32x2x2_t u   = vuzp_u32(vget_low_u32(p01), vget_high_u32(p01));
+    const uint32x2_t   out = vorr_u32(u.val[0], vshl_n_u32(u.val[1], 4));
+    repack_store_u32(dst + 0 * 128 + row * 4, vget_lane_u32(out, 0));
+    repack_store_u32(dst + 1 * 128 + row * 4, vget_lane_u32(out, 1));
+}
+
+static inline void repack_neon_q2_3_k_low2(const uint8_t * qs, int kt_local, uint8_t * lo_pl, int row) {
+    const uint8_t *  qs_sub = qs + (kt_local / 4) * 32;
+    const int8x16_t  shift  = vdupq_n_s8(-2 * (kt_local % 4));
+    const uint8x16_t mask   = vdupq_n_u8(0x03);
+    const uint8x16_t v0     = vandq_u8(vshlq_u8(vld1q_u8(qs_sub +  0), shift), mask);
+    const uint8x16_t v1     = vandq_u8(vshlq_u8(vld1q_u8(qs_sub + 16), shift), mask);
+    repack_neon_pack_2bit_plane(v0, v1, lo_pl, row);
+}
+
+static inline void repack_neon_q3_k_hmask(const uint8_t * hmask, int kt_local, uint8_t * neg_pl, int row) {
+    static const uint32_t mask_words[4] = { 0x01010101u, 0x02020202u, 0x04040404u, 0x08080808u };
+    const uint8x16_t m0   = vreinterpretq_u8_u32(vld1q_u32(mask_words));
+    const uint8x16_t m1   = vshlq_n_u8(m0, 4);
+    const uint8x16_t hbit = vdupq_n_u8((uint8_t) (1u << kt_local));
+    const uint8x16_t b0   = vbicq_u8(m0, vtstq_u8(vld1q_u8(hmask +  0), hbit));
+    const uint8x16_t b1   = vbicq_u8(m1, vtstq_u8(vld1q_u8(hmask + 16), hbit));
+    const uint32x4_t b01  = vreinterpretq_u32_u8(vorrq_u8(b0, b1));
+    const uint32x2_t f0   = vorr_u32(vget_low_u32(b01), vget_high_u32(b01));
+    const uint32x2_t f1   = vpadd_u32(f0, f0);
+    repack_store_u32(neg_pl + row * 4, vget_lane_u32(f1, 0));
+}
+
+static inline void repack_neon_q6_k(const block_q6_K * b, int kt_local, uint8_t * lo_pl, uint8_t * hi_pl, int row) {
+    const int        c      = kt_local / 4;
+    const int        g_sb   = kt_local % 4;
+    const uint8_t *  ql     = b->ql + c * 64 + (g_sb & 1) * 32;
+    const uint8_t *  qh     = b->qh + c * 32;
+    const int8x16_t  lshift = vdupq_n_s8((g_sb >= 2) ? -4 : 0);
+    const int8x16_t  hshift = vdupq_n_s8(-2 * g_sb);
+    const uint8x16_t mask4  = vdupq_n_u8(0x0F);
+    const uint8x16_t mask2  = vdupq_n_u8(0x03);
+
+    const uint8x16_t lo0  = vandq_u8(vshlq_u8(vld1q_u8(ql +  0), lshift), mask4);
+    const uint8x16_t lo1  = vandq_u8(vshlq_u8(vld1q_u8(ql + 16), lshift), mask4);
+    const uint32x4x2_t lz = vuzpq_u32(vreinterpretq_u32_u8(lo0), vreinterpretq_u32_u8(lo1));
+    const uint32x4_t lout = vorrq_u32(lz.val[0], vshlq_n_u32(lz.val[1], 4));
+    repack_store_u32(lo_pl + 0 * 128 + row * 4, vgetq_lane_u32(lout, 0));
+    repack_store_u32(lo_pl + 1 * 128 + row * 4, vgetq_lane_u32(lout, 1));
+    repack_store_u32(lo_pl + 2 * 128 + row * 4, vgetq_lane_u32(lout, 2));
+    repack_store_u32(lo_pl + 3 * 128 + row * 4, vgetq_lane_u32(lout, 3));
+
+    const uint8x16_t hi0 = vandq_u8(vshlq_u8(vld1q_u8(qh +  0), hshift), mask2);
+    const uint8x16_t hi1 = vandq_u8(vshlq_u8(vld1q_u8(qh + 16), hshift), mask2);
+    repack_neon_pack_2bit_plane(hi0, hi1, hi_pl, row);
+}
+
+static inline void repack_neon_q5_k_qh(const uint8_t * qh, uint8_t hbit, uint8_t * plane, int row) {
+    static const uint8_t mask_bytes[16] = { 1, 2, 1, 2, 1, 2, 1, 2, 4, 8, 4, 8, 4, 8, 4, 8 };
+    const uint8x16_t m0   = vld1q_u8(mask_bytes);
+    const uint8x16_t m1   = vshlq_n_u8(m0, 4);
+    const uint8x16_t vbit = vdupq_n_u8(hbit);
+    const uint8x16_t b0   = vandq_u8(vtstq_u8(vld1q_u8(qh +  0), vbit), m0);
+    const uint8x16_t b1   = vandq_u8(vtstq_u8(vld1q_u8(qh + 16), vbit), m1);
+    const uint8x16_t b01  = vorrq_u8(b0, b1);
+    const uint8x8_t  fold = vorr_u8(vget_low_u8(b01), vget_high_u8(b01));
+    const uint8x8_t  out  = vpadd_u8(fold, fold);
+    plane[0 * 32 + row] = vget_lane_u8(out, 0);
+    plane[1 * 32 + row] = vget_lane_u8(out, 1);
+    plane[2 * 32 + row] = vget_lane_u8(out, 2);
+    plane[3 * 32 + row] = vget_lane_u8(out, 3);
+}
+
+static inline uint16x8_t repack_neon_load_f16x8(const uint16_t * src) {
+    return vld1q_u16(src);
+}
+
+#if defined(__aarch64__) || defined(__ARM_FP16_FORMAT_IEEE)
+static inline uint16x8_t repack_neon_load_f16x8(const float * src) {
+    const float16x4_t lo = vcvt_f16_f32(vld1q_f32(src));
+    const float16x4_t hi = vcvt_f16_f32(vld1q_f32(src + 4));
+    return vreinterpretq_u16_f16(vcombine_f16(lo, hi));
+}
+
+static inline uint16x8_t repack_neon_load_f16x8(const ggml_bf16_t * src) {
+    const uint16x8_t  bf   = vld1q_u16((const uint16_t *) src);
+    const float32x4_t f_lo = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(bf), 16));
+    const float32x4_t f_hi = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(bf), 16));
+    const float16x4_t lo   = vcvt_f16_f32(f_lo);
+    const float16x4_t hi   = vcvt_f16_f32(f_hi);
+    return vreinterpretq_u16_f16(vcombine_f16(lo, hi));
+}
+#else
+static inline uint16x8_t repack_neon_load_f16x8(const float * src) {
+    uint16_t tmp[8];
+    for (int i = 0; i < 8; ++i) {
+        tmp[i] = GGML_FP32_TO_FP16(src[i]);
+    }
+    return vld1q_u16(tmp);
+}
+
+static inline uint16x8_t repack_neon_load_f16x8(const ggml_bf16_t * src) {
+    uint16_t tmp[8];
+    for (int i = 0; i < 8; ++i) {
+        tmp[i] = GGML_FP32_TO_FP16(GGML_BF16_TO_FP32(src[i]));
+    }
+    return vld1q_u16(tmp);
+}
+#endif
+
+template <typename SrcT>
+static void repack_neon_f16_tile(uint16_t * tile_dst, const SrcT * src_slice, int64_t row_stride, int ct, int kt) {
+    for (int row0 = 0; row0 < 32; row0 += 4) {
+        const SrcT * r0 = src_slice + (int64_t) (ct * 32 + row0 + 0) * row_stride + kt * 32;
+        const SrcT * r1 = src_slice + (int64_t) (ct * 32 + row0 + 1) * row_stride + kt * 32;
+        const SrcT * r2 = src_slice + (int64_t) (ct * 32 + row0 + 2) * row_stride + kt * 32;
+        const SrcT * r3 = src_slice + (int64_t) (ct * 32 + row0 + 3) * row_stride + kt * 32;
+        for (int cp = 0; cp < 16; cp += 4) {
+            const uint32x4_t v0 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r0 + 2 * cp));
+            const uint32x4_t v1 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r1 + 2 * cp));
+            const uint32x4_t v2 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r2 + 2 * cp));
+            const uint32x4_t v3 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r3 + 2 * cp));
+
+            const uint32x2x2_t t0 = vtrn_u32(vget_low_u32(v0),  vget_low_u32(v1));
+            const uint32x2x2_t t1 = vtrn_u32(vget_high_u32(v0), vget_high_u32(v1));
+            const uint32x2x2_t t2 = vtrn_u32(vget_low_u32(v2),  vget_low_u32(v3));
+            const uint32x2x2_t t3 = vtrn_u32(vget_high_u32(v2), vget_high_u32(v3));
+
+            vst1q_u16(tile_dst + (cp + 0) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t0.val[0], t2.val[0])));
+            vst1q_u16(tile_dst + (cp + 1) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t0.val[1], t2.val[1])));
+            vst1q_u16(tile_dst + (cp + 2) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t1.val[0], t3.val[0])));
+            vst1q_u16(tile_dst + (cp + 3) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t1.val[1], t3.val[1])));
+        }
+    }
+}
+#endif
+
 // --- Generic 32x32 tile grid walkers for block-32 and float types ---

+template <typename Fn>
+static void repack_parallel_for_ct(size_t slice_bytes, int n_col_tiles, Fn && fn) {
+    int n_threads = 1;
+    if (slice_bytes > 128 * 1024) {
+        for (int nt : { 4, 3, 2 }) {
+            if (n_col_tiles % nt == 0) {
+                n_threads = nt;
+                break;
+            }
+        }
+    }
+
+    if (n_threads <= 1) {
+        fn(0, n_col_tiles);
+        return;
+    }
+
+    const int chunk = n_col_tiles / n_threads;
+    std::thread workers[3];
+    for (int i = 0; i < n_threads - 1; ++i) {
+        workers[i] = std::thread([&fn, i, chunk]() {
+            fn(i * chunk, (i + 1) * chunk);
+        });
+    }
+    fn((n_threads - 1) * chunk, n_col_tiles);
+    for (int i = 0; i < n_threads - 1; ++i) {
+        workers[i].join();
+    }
+}
+
 template <typename SrcElemT, int ELEMS_PER_KTILE, typename PackTileFn>
 static void repack_to_tiled_grid_32(ggml_tensor * t, const void * data, size_t offset, size_t size,
                                    size_t tile_size, PackTileFn && pack_tile) {
@@ -70,12 +376,14 @@ static void repack_to_tiled_grid_32(ggml_tensor * t, const void * data, size_t o
        const SrcElemT * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * row_stride);
        uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size;

-        for (int ct = 0; ct < n_col_tiles; ct++) {
-            for (int kt = 0; kt < n_k_tiles; kt++) {
-                uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
-                pack_tile(tile_dst, src_slice, ne0, ne1, row_stride, ct, kt);
+        repack_parallel_for_ct(matrix_size, n_col_tiles, [&](int ct0, int ct1) {
+            for (int ct = ct0; ct < ct1; ct++) {
+                for (int kt = 0; kt < n_k_tiles; kt++) {
+                    uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
+                    pack_tile(tile_dst, src_slice, ne0, ne1, row_stride, ct, kt);
+                }
            }
-        }
+        });
    }
 }

@@ -130,8 +438,24 @@ static void unpack_from_tiled_grid_32(void * data, const ggml_tensor * t, size_t
 static void repack_q4_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
    repack_to_tiled_grid_32<block_q4_0, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_0,
        [](uint8_t * tile_dst, const block_q4_0 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
-            uint8_t tile_quants[32][32];
            ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_4bit_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    scale_dst[row] = src_slice[(ct * 32 + row) * row_stride + kt].d;
+                }
+                return;
+            }
+#endif
+            uint8_t tile_quants[32][32];
            for (int row = 0; row < 32; row++) {
                const int64_t r = ct * 32 + row;
                if (r < ne1 && kt < ne0 / 32) {
@@ -167,8 +491,26 @@ static void repack_tiled_q4_0(void * data, const ggml_tensor * t, size_t offset,
 static void repack_q4_1_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
    repack_to_tiled_grid_32<block_q4_1, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_1,
        [](uint8_t * tile_dst, const block_q4_1 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
-            uint8_t tile_quants[32][32];
            ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_4bit_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    const block_q4_1 & src_block = src_slice[(ct * 32 + row) * row_stride + kt];
+                    scale_dst[2 * row + 0] = src_block.d;
+                    scale_dst[2 * row + 1] = src_block.m;
+                }
+                return;
+            }
+#endif
+            uint8_t tile_quants[32][32];
            for (int row = 0; row < 32; row++) {
                const int64_t r = ct * 32 + row;
                if (r < ne1 && kt < ne0 / 32) {
@@ -207,8 +549,24 @@ static void repack_tiled_q4_1(void * data, const ggml_tensor * t, size_t offset,
 static void repack_mxfp4_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
    repack_to_tiled_grid_32<block_mxfp4, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_MXFP4,
        [](uint8_t * tile_dst, const block_mxfp4 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
-            uint8_t tile_quants[32][32];
            uint8_t * scale_dst = tile_dst + 512;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_4bit_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    scale_dst[row] = src_slice[(ct * 32 + row) * row_stride + kt].e;
+                }
+                return;
+            }
+#endif
+            uint8_t tile_quants[32][32];
            for (int row = 0; row < 32; row++) {
                const int64_t r = ct * 32 + row;
                if (r < ne1 && kt < ne0 / 32) {
@@ -244,6 +602,23 @@ static void repack_tiled_mxfp4(void * data, const ggml_tensor * t, size_t offset
 static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
    repack_to_tiled_grid_32<block_q8_0, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q8_0,
        [](uint8_t * tile_dst, const block_q8_0 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
+            ggml_half * scale_dst = (ggml_half *) (tile_dst + 1024);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = (const uint8_t *) src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_q8_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    scale_dst[row] = src_slice[(ct * 32 + row) * row_stride + kt].d;
+                }
+                return;
+            }
+#endif
            for (int cp = 0; cp < 16; cp++) {
                const int col0 = cp * 2;
                const int col1 = col0 + 1;
@@ -254,7 +629,6 @@ static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t offset,
                    tile_dst[cp * 64 + 2 * row + 1] = b ? b->qs[col1] : 0;
                }
            }
-            ggml_half * scale_dst = (ggml_half *) (tile_dst + 1024);
            for (int row = 0; row < 32; row++) {
                const int64_t r = ct * 32 + row;
                scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_slice[r * row_stride + kt].d : 0;
@@ -303,24 +677,33 @@ static void repack_q2_0_tiled(ggml_tensor * t, const void * data, size_t offset,
    for (int64_t slice = 0; slice < t->ne[2] * t->ne[3]; ++slice) {
        const block_q2_0 * src_slice = src + slice * t->ne[1] * blocks_per_row;
        uint8_t * dst_slice = (uint8_t *) t->data + slice * matrix_size;
-        memset(dst_slice, 0, matrix_size);
+        const size_t ct_bytes = (size_t) n_weight_tiles * HTP_TENSOR_TILE_SIZE_Q2_0;

-        for (int64_t row_index = 0; row_index < t->ne[1]; ++row_index) {
-            const int64_t ct = row_index / 32;
-            const int row = (int) (row_index % 32);
-            for (int64_t kt = 0; kt < n_weight_tiles; ++kt) {
-                const block_q2_0 * block = src_slice + row_index * blocks_per_row + kt;
-                uint8_t * tile = dst_slice + (ct * n_weight_tiles + kt) * HTP_TENSOR_TILE_SIZE_Q2_0;
-                for (int group = 0; group < 16; ++group) {
-                    for (int byte = 0; byte < 4; ++byte) {
-                        const int k = 4 * group + byte;
-                        const uint8_t code = (block->qs[k / 4] >> (2 * (k % 4))) & 3;
-                        tile[(group >> 2) * 128 + 4 * row + byte] |= code << (2 * (group & 3));
+        repack_parallel_for_ct(matrix_size, (int) n_col_tiles, [&](int ct0, int ct1) {
+            memset(dst_slice + (size_t) ct0 * ct_bytes, 0, (size_t) (ct1 - ct0) * ct_bytes);
+            const int64_t r0 = (int64_t) ct0 * 32;
+            const int64_t r1 = (std::min)((int64_t) ct1 * 32, t->ne[1]);
+            for (int64_t row_index = r0; row_index < r1; ++row_index) {
+                const int64_t ct = row_index / 32;
+                const int row = (int) (row_index % 32);
+                for (int64_t kt = 0; kt < n_weight_tiles; ++kt) {
+                    const block_q2_0 * block = src_slice + row_index * blocks_per_row + kt;
+                    uint8_t * tile = dst_slice + (ct * n_weight_tiles + kt) * HTP_TENSOR_TILE_SIZE_Q2_0;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+                    repack_neon_q2_0_block(block->qs, tile, row);
+#else
+                    for (int group = 0; group < 16; ++group) {
+                        for (int byte = 0; byte < 4; ++byte) {
+                            const int k = 4 * group + byte;
+                            const uint8_t code = (block->qs[k / 4] >> (2 * (k % 4))) & 3;
+                            tile[(group >> 2) * 128 + 4 * row + byte] |= code << (2 * (group & 3));
+                        }
                    }
+#endif
+                    memcpy(tile + 512 + 2 * row, &block->d, sizeof(block->d));
                }
-                memcpy(tile + 512 + 2 * row, &block->d, sizeof(block->d));
            }
-        }
+        });
    }
 }

@@ -368,6 +751,12 @@ static void repack_float_to_f16_tiled(ggml_tensor * t, const void * data, size_t
    repack_to_tiled_grid_32<SrcT, 32>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_F16,
        [&](uint8_t * tile_bytes, const SrcT * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
            uint16_t * tile_dst = (uint16_t *) tile_bytes;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt * 32 + 32 <= ne0) {
+                repack_neon_f16_tile(tile_dst, src_slice, row_stride, ct, kt);
+                return;
+            }
+#endif
            for (int cp = 0; cp < 16; cp++) {
                const int64_t c0 = kt * 32 + cp * 2;
                const int64_t c1 = c0 + 1;
@@ -424,20 +813,25 @@ static void repack_kquant_to_tiled_grid(ggml_tensor * t, const void * data, size
    for (int64_t slice_idx = 0; slice_idx < ne2 * ne3; slice_idx++) {
        const BlockT * src_slice = src_matrix + slice_idx * (ne1 * sb_per_row);
        uint8_t * matrix_dst     = (uint8_t *) t->data + slice_idx * matrix_size;
+        const size_t ct_bytes    = (size_t) n_k_tiles * tile_size;

-        // Padding rows and OR-ed bit planes require zeroed tiles
-        memset(matrix_dst, 0, matrix_size);
+        repack_parallel_for_ct(matrix_size, n_col_tiles, [&](int ct0, int ct1) {
+            // Padding rows and OR-ed bit planes require zeroed tiles
+            memset(matrix_dst + (size_t) ct0 * ct_bytes, 0, (size_t) (ct1 - ct0) * ct_bytes);

-        for (int64_t r = 0; r < ne1; r++) {
-            const int ct  = (int) (r / 32);
-            const int row = (int) (r % 32);
-            const BlockT * src_row = src_slice + r * sb_per_row;
+            const int64_t r0 = (int64_t) ct0 * 32;
+            const int64_t r1 = (std::min)((int64_t) ct1 * 32, ne1);
+            for (int64_t r = r0; r < r1; r++) {
+                const int ct  = (int) (r / 32);
+                const int row = (int) (r % 32);
+                const BlockT * src_row = src_slice + r * sb_per_row;

-            for (int kt = 0; kt < n_k_tiles; kt++) {
-                uint8_t * tile = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
-                pack_kt(tile, &src_row[kt / 8], row, kt % 8);
+                for (int kt = 0; kt < n_k_tiles; kt++) {
+                    uint8_t * tile = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
+                    pack_kt(tile, &src_row[kt / 8], row, kt % 8);
+                }
            }
-        }
+        });
    }
 }

@@ -493,6 +887,7 @@ static void unpack_kquant_from_tiled_grid(void * data, const ggml_tensor * t, si

 // --- K-Quant Helpers & Per-Type Implementations (Q2_K, Q3_K, Q4_K, Q5_K, Q6_K) ---

+#if !defined(__ARM_NEON) && !defined(__aarch64__)
 // unsigned 6-bit value (0..63) of element e of a Q6_K block, same bit layout as dequantize_row_q6_K
 static inline uint8_t q6_K_get_quant(const block_q6_K * b, int e) {
    const int c = e / 128;
@@ -510,6 +905,7 @@ static inline uint8_t q6_K_get_quant(const block_q6_K * b, int e) {
    }
    return (uint8_t) (lo | (hi << 4));
 }
+#endif

 static void repack_q6_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
    repack_kquant_to_tiled_grid<block_q6_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q6_K,
@@ -519,6 +915,9 @@ static void repack_q6_K_tiled(ggml_tensor * t, const void * data, size_t offset,
            uint8_t * hi_pl   = tile + 512;
            ggml_half * sc_pl = (ggml_half *) (tile + 768);

+#if defined(__ARM_NEON) || defined(__aarch64__)
+            repack_neon_q6_k(b, kt_local, lo_pl, hi_pl, row);
+#else
            for (int lk = 0; lk < 32; lk++) {
                const uint8_t q6 = q6_K_get_quant(b, kt_local * 32 + lk);
                const int g   = lk >> 2;
@@ -526,6 +925,7 @@ static void repack_q6_K_tiled(ggml_tensor * t, const void * data, size_t offset,
                lo_pl[(g >> 1) * 128 + pos] |= (uint8_t) ((q6 & 0xF) << ((g & 1) * 4));
                hi_pl[(g >> 2) * 128 + pos] |= (uint8_t) ((q6 >> 4) << ((g & 3) * 2));
            }
+#endif
            for (int sub = 0; sub < 2; sub++) {
                sc_pl[sub * 32 + row] = GGML_FP32_TO_FP16(d * (float) b->scales[kt_local * 2 + sub]);
            }
@@ -585,6 +985,7 @@ static void repack_tiled_q6_K(void * data, const ggml_tensor * t, size_t offset,
        });
 }

+#if !defined(__ARM_NEON) && !defined(__aarch64__)
 // low 2 bits (0..3) of element e of a Q2_K or Q3_K block, same bit layout as dequantize_row_q2_K / q3_K
 static inline uint8_t q2_3_K_get_low2(const uint8_t * qs, int e) {
    const int c = e / 128;
@@ -600,6 +1001,7 @@ static inline bool q3_K_get_hbit(const block_q3_K * b, int e) {
    const int l = e % 32;
    return (b->hmask[l] >> (c * 4 + j)) & 1;
 }
+#endif

 // signed 6-bit scale j (-32..31) of a Q3_K block, same packing as quantize_row_q3_K_ref
 static inline int q3_K_get_scale(const uint8_t * scales, int j) {
@@ -670,6 +1072,10 @@ static void repack_q3_K_tiled(ggml_tensor * t, const void * data, size_t offset,
            uint8_t * neg_pl  = tile + 256;
            ggml_half * sc_pl = (ggml_half *) (tile + 384);

+#if defined(__ARM_NEON) || defined(__aarch64__)
+            repack_neon_q2_3_k_low2(b->qs, kt_local, lo_pl, row);
+            repack_neon_q3_k_hmask(b->hmask, kt_local, neg_pl, row);
+#else
            for (int lk = 0; lk < 32; lk++) {
                const int e   = kt_local * 32 + lk;
                const int g   = lk >> 2;
@@ -679,6 +1085,7 @@ static void repack_q3_K_tiled(ggml_tensor * t, const void * data, size_t offset,
                    neg_pl[pos] |= (uint8_t) (1 << g);
                }
            }
+#endif
            for (int sub = 0; sub < 2; sub++) {
                sc_pl[sub * 32 + row] = GGML_FP32_TO_FP16(d * (float) q3_K_get_scale(b->scales, kt_local * 2 + sub));
            }
@@ -750,11 +1157,15 @@ static void repack_q2_K_tiled(ggml_tensor * t, const void * data, size_t offset,
            uint8_t * sc_pl   = tile + 256;
            ggml_half * dd_pl = (ggml_half *) (tile + 320);

+#if defined(__ARM_NEON) || defined(__aarch64__)
+            repack_neon_q2_3_k_low2(b->qs, kt_local, lo_pl, row);
+#else
            for (int lk = 0; lk < 32; lk++) {
                const int g   = lk >> 2;
                const int pos = row * 4 + (lk & 3);
                lo_pl[(g >> 2) * 128 + pos] |= (uint8_t) (q2_3_K_get_low2(b->qs, kt_local * 32 + lk) << ((g & 3) * 2));
            }
+#endif
            for (int sub = 0; sub < 2; sub++) {
                sc_pl[sub * 32 + row] = b->scales[kt_local * 2 + sub];
            }
@@ -846,23 +1257,46 @@ static inline void pack_k4_scales_mins(uint8_t * scales, ggml_half * d, ggml_hal
    }
 }

+template <typename BlockT>
+static inline void repack_k4_quants(uint8_t * tile_dst, const BlockT * b, int row, int kt_local,
+                                    int64_t sb_per_row, bool use_neon) {
+#if defined(__ARM_NEON) || defined(__aarch64__)
+    if (use_neon) {
+        if ((row & 7) == 0) {
+            const uint8_t * qs[8];
+            for (int rr = 0; rr < 8; ++rr) {
+                qs[rr] = b[rr * sb_per_row].qs + (kt_local / 2) * 32;
+            }
+            repack_neon_k4_rows(qs, tile_dst, row, (kt_local & 1) != 0);
+        }
+        return;
+    }
+#else
+    GGML_UNUSED(sb_per_row);
+    GGML_UNUSED(use_neon);
+#endif
+    const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32;
+    const int shift = (kt_local & 1) ? 4 : 0;
+
+    for (int cp = 0; cp < 16; cp++) {
+        const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F;
+        const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F;
+        tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0);
+    }
+}
+
 static void repack_q4_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
+    const int64_t sb_per_row = t->ne[0] / QK_K;
+    const bool    use_neon   = (t->ne[1] % 8 == 0);
    repack_kquant_to_tiled_grid<block_q4_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_K,
-        [](uint8_t * tile_dst, const block_q4_K * b, int row, int kt_local) {
+        [sb_per_row, use_neon](uint8_t * tile_dst, const block_q4_K * b, int row, int kt_local) {
            const float d    = GGML_FP16_TO_FP32(b->d);
            const float dmin = GGML_FP16_TO_FP32(b->dmin);

            uint8_t sc, m;
            get_scale_min_k4(kt_local, b->scales, &sc, &m);

-            const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32;
-            const int shift = (kt_local & 1) ? 4 : 0;
-
-            for (int cp = 0; cp < 16; cp++) {
-                const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F;
-                const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F;
-                tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0);
-            }
+            repack_k4_quants(tile_dst, b, row, kt_local, sb_per_row, use_neon);

            ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
            scale_dst[2 * row + 0] = GGML_FP32_TO_FP16(d * (float) sc);
@@ -900,8 +1334,10 @@ static void repack_tiled_q4_K(void * data, const ggml_tensor * t, size_t offset,
 }

 static void repack_q5_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
+    const int64_t sb_per_row = t->ne[0] / QK_K;
+    const bool    use_neon   = (t->ne[1] % 8 == 0);
    repack_kquant_to_tiled_grid<block_q5_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q5_K,
-        [](uint8_t * tile_dst, const block_q5_K * b, int row, int kt_local) {
+        [sb_per_row, use_neon](uint8_t * tile_dst, const block_q5_K * b, int row, int kt_local) {
            const float d     = GGML_FP16_TO_FP32(b->d);
            const float dmin  = GGML_FP16_TO_FP32(b->dmin);
            uint8_t * plane   = tile_dst + 640;
@@ -909,15 +1345,13 @@ static void repack_q5_K_tiled(ggml_tensor * t, const void * data, size_t offset,
            uint8_t sc, m;
            get_scale_min_k4(kt_local, b->scales, &sc, &m);

-            const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32;
-            const int shift    = (kt_local & 1) ? 4 : 0;
+            repack_k4_quants(tile_dst, b, row, kt_local, sb_per_row, use_neon);
+
            const uint8_t hbit = (uint8_t) (1 << kt_local);
-
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            repack_neon_q5_k_qh(b->qh, hbit, plane, row);
+#else
            for (int cp = 0; cp < 16; cp++) {
-                const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F;
-                const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F;
-                tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0);
-
                const int i    = cp / 4;
                const int lane = (cp % 4) * 32 + row;
                if (b->qh[2 * cp + 0] & hbit) {
@@ -927,6 +1361,7 @@ static void repack_q5_K_tiled(ggml_tensor * t, const void * data, size_t offset,
                    plane[lane] |= (uint8_t) (1 << (2 * i + 1));
                }
            }
+#endif

            ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
            scale_dst[2 * row + 0] = GGML_FP32_TO_FP16(d * (float) sc);
@@ -1002,31 +1437,37 @@ static void repack_q1_0_tiled(ggml_tensor * t, const void * data, size_t offset,
        const block_q1_0 * src_slice = src_matrix + (slice_idx - start_slice) * ne1 * n_k_tiles;
        uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size;

-        for (int ct = 0; ct < n_col_tiles; ct++) {
-            for (int kt = 0; kt < n_k_tiles; kt++) {
-                uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
-                ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
-                memset(tile_dst, 0, tile_size);
+        repack_parallel_for_ct(matrix_size, n_col_tiles, [&](int ct0, int ct1) {
+            for (int ct = ct0; ct < ct1; ct++) {
+                for (int kt = 0; kt < n_k_tiles; kt++) {
+                    uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
+                    ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
+                    memset(tile_dst, 0, tile_size);

-                for (int row = 0; row < 32; row++) {
-                    const int64_t r = ct * 32 + row;
-                    if (r >= ne1) {
-                        continue;
-                    }
-
-                    const block_q1_0 & block = src_slice[r * n_k_tiles + kt];
-                    scale_dst[row] = block.d;
-                    for (int k = 0; k < QK1_0; k++) {
-                        int byte_offset;
-                        int bit;
-                        q1_0_tiled_bit_location(k, row, byte_offset, bit);
-                        if (block.qs[k / 8] & (1u << (k % 8))) {
-                            tile_dst[byte_offset] |= (uint8_t) (1u << bit);
+                    for (int row = 0; row < 32; row++) {
+                        const int64_t r = ct * 32 + row;
+                        if (r >= ne1) {
+                            continue;
                        }
+
+                        const block_q1_0 & block = src_slice[r * n_k_tiles + kt];
+                        scale_dst[row] = block.d;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+                        repack_neon_q1_0_block(block.qs, tile_dst, row);
+#else
+                        for (int k = 0; k < QK1_0; k++) {
+                            int byte_offset;
+                            int bit;
+                            q1_0_tiled_bit_location(k, row, byte_offset, bit);
+                            if (block.qs[k / 8] & (1u << (k % 8))) {
+                                tile_dst[byte_offset] |= (uint8_t) (1u << bit);
+                            }
+                        }
+#endif
                    }
                }
            }
-        }
+        });
    }
 }

@@ -1127,11 +1568,11 @@ void htp_repack_tensor_tiled(ggml_tensor * tensor, const void * data, size_t siz
            break;
        case GGML_TYPE_F32:
            repack_float_to_f16_tiled<float>(tensor, data, 0, size,
-                [](float v) { return (uint16_t) ggml_fp32_to_fp16(v); });
+                [](float v) { return (uint16_t) GGML_FP32_TO_FP16(v); });
            break;
        case GGML_TYPE_BF16:
            repack_float_to_f16_tiled<ggml_bf16_t>(tensor, data, 0, size,
-                [](ggml_bf16_t v) { return (uint16_t) ggml_fp32_to_fp16(ggml_bf16_to_fp32(v)); });
+                [](ggml_bf16_t v) { return (uint16_t) GGML_FP32_TO_FP16(GGML_BF16_TO_FP32(v)); });
            break;
        default:
            break;
@@ -28,7 +28,7 @@ add_library(${HTP_LIB} SHARED
    unary-ops.c
    sum-rows-ops.c
    softmax-ops.c
-    act-ops.c
+    glu-ops.c
    rope-ops.c
    set-rows-ops.c
    get-rows-ops.c
@@ -53,9 +53,17 @@ target_compile_definitions(${HTP_LIB} PRIVATE
    $<IF:$<BOOL:${HEXAGON_HTP_DEBUG}>,HTP_DEBUG=1,NDEBUG=1>
    $<IF:$<BOOL:${HEXAGON_HTP_DEBUG}>,FARF_HIGH=1,>)

-
 build_idl(htp_iface.idl ${HTP_LIB})

 set_target_properties(${HTP_LIB} PROPERTIES EXPORT_COMPILE_COMMANDS ON)

+if (NOT Python3_EXECUTABLE)
+    find_package(Python3 REQUIRED COMPONENTS Interpreter)
+endif()
+
+add_custom_command(TARGET ${HTP_LIB} POST_BUILD
+    COMMAND ${CMAKE_COMMAND} -E echo "Adding SWIV section $<TARGET_FILE_NAME:${HTP_LIB}>"
+    COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/htp-swiv.py -i $<TARGET_FILE:${HTP_LIB}> -o $<TARGET_FILE:${HTP_LIB}>
+    VERBATIM)
+
 install(TARGETS ${HTP_LIB})
@@ -32,6 +32,7 @@ static inline size_t htp_allreduce_vtcm_buffer_count(
 }

 struct htp_allreduce_kernel_params {
+    uint8_t kernel_type;
    int32_t rank;
    int32_t n_ranks;
    int32_t n_threads;
@@ -39,7 +40,6 @@ struct htp_allreduce_kernel_params {
    int32_t elems_per_thread;     // 1D: nelem_per_thread, 2D: nrows_per_thread
    int32_t vtcm_size_per_thread;
    int32_t vtcm_size;
-    int32_t kernel_type;
    int32_t ne0;
    int32_t ne1;
    int32_t row_size_aligned;
@@ -50,6 +50,12 @@ struct htp_allreduce_kernel_params {
    int32_t mode;                 // enum htp_allreduce_mode: FULL or SHARDED_FANOUT
 };

+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_allreduce_kernel_params) <= 128, "htp_allreduce_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_allreduce_kernel_params) <= 128, "htp_allreduce_kernel_params is too large for kernel_params blob");
+#endif
+
 #ifdef __cplusplus
 }
 #endif
@@ -453,7 +453,10 @@ static void merge_runs(
 struct htp_sort_chunk_ctx {
    struct htp_ops_context *              octx;
    const struct htp_sort_kernel_params * kparams;
-    uint8_t *                             vtcm_base;
+    uint8_t *                             phase1_base;
+    size_t                                phase1_slot_size;
+    float *                               merge_values;
+    int32_t *                             merge_indices;
    uint32_t                              real_count[HTP_MAX_NTHREADS];
 };

@@ -462,26 +465,27 @@ static void htp_sort_chunk_job(unsigned int n, unsigned int i, void * data) {
    struct htp_ops_context * octx = cctx->octx;
    const struct htp_sort_kernel_params * kparams = cctx->kparams;
    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * dst  = octx->dst;

    uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
-    uint32_t ne00        = (uint32_t) kparams->ne00;
+    uint32_t ne00        = src0->ne[0];
    uint32_t chunk_base  = i * chunk_elems;

    uint32_t real_count = (chunk_base < ne00) ? MIN(chunk_elems, ne00 - chunk_base) : 0;
    cctx->real_count[i] = real_count;

-    uint8_t * spad = cctx->vtcm_base + (size_t) kparams->phase1_slot_size * i;
+    uint8_t * spad = cctx->phase1_base + cctx->phase1_slot_size * i;
    float *   values_buf  = (float *) spad;
    int32_t * indices_buf = (int32_t *) (spad + chunk_elems * sizeof(float));

-    float *   merge_values  = (float *)   (cctx->vtcm_base + kparams->merge_values_off);
-    int32_t * merge_indices = (int32_t *) (cctx->vtcm_base + kparams->merge_indices_off);
+    float *   merge_values  = cctx->merge_values;
+    int32_t * merge_indices = cctx->merge_indices;

    struct htp_thread_trace * tr = &octx->ctx->trace[i];
    const dma_addr_t src_addr = src0->data + (size_t) chunk_base * sizeof(float);

    if (kparams->is_top_k) {
-        uint32_t k = (uint32_t) kparams->k;
+        uint32_t k = dst->ne[0];
        uint32_t local_k = MIN(k, chunk_elems);

        if (real_count == 0) {
@@ -559,17 +563,17 @@ static void htp_argsort_merge_l1_job(unsigned int n, unsigned int i, void * data
    uint32_t count0 = cctx->real_count[c0];
    uint32_t count1 = cctx->real_count[c1];

-    size_t slot_size = (size_t) kparams->phase1_slot_size;
+    size_t slot_size = cctx->phase1_slot_size;
    uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;

-    const float *   val0 = (const float *)   (cctx->vtcm_base + slot_size * c0);
-    const int32_t * idx0 = (const int32_t *) (cctx->vtcm_base + slot_size * c0 + chunk_elems * sizeof(float));
-    const float *   val1 = (const float *)   (cctx->vtcm_base + slot_size * c1);
-    const int32_t * idx1 = (const int32_t *) (cctx->vtcm_base + slot_size * c1 + chunk_elems * sizeof(float));
+    const float *   val0 = (const float *)   (cctx->phase1_base + slot_size * c0);
+    const int32_t * idx0 = (const int32_t *) (cctx->phase1_base + slot_size * c0 + chunk_elems * sizeof(float));
+    const float *   val1 = (const float *)   (cctx->phase1_base + slot_size * c1);
+    const int32_t * idx1 = (const int32_t *) (cctx->phase1_base + slot_size * c1 + chunk_elems * sizeof(float));

    uint32_t out_offset = (2 * i) * chunk_elems;
-    float *   out_val = (float *)   (cctx->vtcm_base + kparams->merge_values_off)  + out_offset;
-    int32_t * out_idx = (int32_t *) (cctx->vtcm_base + kparams->merge_indices_off) + out_offset;
+    float *   out_val = cctx->merge_values  + out_offset;
+    int32_t * out_idx = cctx->merge_indices + out_offset;

    struct htp_thread_trace * tr = &cctx->octx->ctx->trace[i];
    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
@@ -587,21 +591,21 @@ static void htp_argsort_merge_l2_job(unsigned int n, unsigned int i, void * data
    struct htp_sort_chunk_ctx * cctx = mctx->cctx;
    const struct htp_sort_kernel_params * kparams = cctx->kparams;

-    size_t slot_size = (size_t) kparams->phase1_slot_size;
+    size_t slot_size = cctx->phase1_slot_size;
    uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;

    uint32_t off0 = (4 * i) * chunk_elems;
    uint32_t off1 = (4 * i + 2) * chunk_elems;

-    const float *   val0 = (const float *)   (cctx->vtcm_base + kparams->merge_values_off)  + off0;
-    const int32_t * idx0 = (const int32_t *) (cctx->vtcm_base + kparams->merge_indices_off) + off0;
-    const float *   val1 = (const float *)   (cctx->vtcm_base + kparams->merge_values_off)  + off1;
-    const int32_t * idx1 = (const int32_t *) (cctx->vtcm_base + kparams->merge_indices_off) + off1;
+    const float *   val0 = cctx->merge_values  + off0;
+    const int32_t * idx0 = cctx->merge_indices + off0;
+    const float *   val1 = cctx->merge_values  + off1;
+    const int32_t * idx1 = cctx->merge_indices + off1;

    uint32_t count0 = cctx->real_count[4*i + 0] + cctx->real_count[4*i + 1];
    uint32_t count1 = cctx->real_count[4*i + 2] + cctx->real_count[4*i + 3];

-    uint8_t * slot_base = cctx->vtcm_base + (4 * i) * slot_size;
+    uint8_t * slot_base = cctx->phase1_base + (4 * i) * slot_size;
    float *   out_val   = (float *) slot_base;
    int32_t * out_idx   = (int32_t *) (slot_base + 4 * chunk_elems * sizeof(float));

@@ -614,7 +618,10 @@ static void htp_argsort_merge_l2_job(unsigned int n, unsigned int i, void * data
 struct htp_sort_multi_row_ctx {
    struct htp_ops_context *              octx;
    const struct htp_sort_kernel_params * kparams;
-    uint8_t *                             vtcm_base;
+    uint8_t *                             phase1_base;
+    size_t                                phase1_slot_size;
+    uint32_t                              row_start;
+    uint32_t                              row_end;
    uint32_t                              nrows_per_thread;
 };

@@ -624,12 +631,18 @@ struct htp_tensor_row_coord {
    uint32_t i03;
 };

-static inline struct htp_tensor_row_coord init_row_coord(uint32_t r, uint32_t ne01, uint32_t ne02) {
+static inline struct htp_tensor_row_coord init_row_coord(
+    uint32_t r,
+    uint32_t ne01,
+    uint32_t ne02,
+    const struct fastdiv_values * div_ne01,
+    const struct fastdiv_values * div_ne01_ne02
+) {
    uint32_t ne01_ne02 = ne01 * ne02;
-    uint32_t i03 = (ne01_ne02 > 0) ? (r / ne01_ne02) : 0;
-    uint32_t rem = (ne01_ne02 > 0) ? (r % ne01_ne02) : 0;
-    uint32_t i02 = (ne01 > 0) ? (rem / ne01) : 0;
-    uint32_t i01 = (ne01 > 0) ? (rem % ne01) : 0;
+    uint32_t i03 = (ne01_ne02 > 0) ? fastdiv(r, div_ne01_ne02) : 0;
+    uint32_t rem = (ne01_ne02 > 0) ? (r - i03 * ne01_ne02) : 0;
+    uint32_t i02 = (ne01 > 0) ? fastdiv(rem, div_ne01) : 0;
+    uint32_t i01 = (ne01 > 0) ? (rem - i02 * ne01) : 0;
    return (struct htp_tensor_row_coord){ i01, i02, i03 };
 }

@@ -701,7 +714,7 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
    const struct htp_tensor * src0 = octx->src[0];
    const struct htp_tensor * dst  = octx->dst;

-    uint32_t ne00 = (uint32_t) kparams->ne00;
+    uint32_t ne00 = src0->ne[0];
    uint32_t ne01 = src0->ne[1];
    uint32_t ne02 = src0->ne[2];

@@ -713,16 +726,16 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
    uint32_t nb2 = dst->nb[2];
    uint32_t nb3 = dst->nb[3];

-    uint32_t row_start = (uint32_t) kparams->row_start;
-    uint32_t row_end   = (uint32_t) kparams->row_end;
+    uint32_t row_start = mctx->row_start;
+    uint32_t row_end   = mctx->row_end;
    uint32_t r_start   = row_start + ith * mctx->nrows_per_thread;
    uint32_t r_end     = MIN(r_start + mctx->nrows_per_thread, row_end);

    if (r_start >= r_end) return;

-    size_t spad_offset = ith * (size_t) (kparams->phase1_slot_size * kparams->n_slots);
-    uint8_t * thread_spad = mctx->vtcm_base + spad_offset;
-    size_t slot_size = (size_t) kparams->phase1_slot_size;
+    size_t slot_size = mctx->phase1_slot_size;
+    size_t spad_offset = ith * slot_size * (size_t) kparams->n_slots;
+    uint8_t * thread_spad = mctx->phase1_base + spad_offset;

    uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
    float *   values_buf  = (float *) thread_spad;
@@ -732,14 +745,14 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];

    bool is_top_k = kparams->is_top_k != 0;
-    uint32_t k = (uint32_t) kparams->k;
+    uint32_t k = dst->ne[0];
    enum ggml_sort_order order = (enum ggml_sort_order) kparams->order;

    const uint32_t src_bytes = ne00 * sizeof(float);
    const uint32_t dst_bytes = (is_top_k ? k : ne00) * sizeof(int32_t);

    if (kparams->n_slots == 2) {
-        struct htp_tensor_row_coord prime_coord = init_row_coord(r_start, ne01, ne02);
+        struct htp_tensor_row_coord prime_coord = init_row_coord(r_start, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);
        for (uint32_t r = r_start, spad_idx = 0; r < r_end && spad_idx < 2; r++, spad_idx++) {
            uint8_t * cur_spad = thread_spad + spad_idx * slot_size;
            float *   cur_values  = (float *) cur_spad;
@@ -759,8 +772,8 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
            step_row_coord(&prime_coord, ne01, ne02);
        }

-        struct htp_tensor_row_coord dst_coord = init_row_coord(r_start, ne01, ne02);
-        struct htp_tensor_row_coord next_src_coord = init_row_coord(r_start + 2, ne01, ne02);
+        struct htp_tensor_row_coord dst_coord = init_row_coord(r_start, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);
+        struct htp_tensor_row_coord next_src_coord = init_row_coord(r_start + 2, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);

        for (uint32_t r = r_start; r < r_end; r++) {
            uint32_t cur_slot = (r - r_start) & 1;
@@ -796,7 +809,7 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data

        dma_queue_flush(dma_q);
    } else {
-        struct htp_tensor_row_coord coord = init_row_coord(r_start, ne01, ne02);
+        struct htp_tensor_row_coord coord = init_row_coord(r_start, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);
        for (uint32_t r = r_start; r < r_end; r++) {
            dma_addr_t src_addr = src0->data + coord_offset(coord, nb01, nb02, nb03);
            dma_addr_t dst_addr = dst->data + coord_offset(coord, nb1, nb2, nb3);
@@ -847,52 +860,39 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
    uint32_t ne00 = src0->ne[0];
    uint32_t k    = dst->ne[0];

-    struct htp_sort_kernel_params kparams_local;
    const struct htp_sort_kernel_params * kparams = (const struct htp_sort_kernel_params *) octx->kernel_params;

-    if (kparams->n_threads == 0) {
-        struct htp_sort_vtcm_layout layout;
-        if (!htp_sort_solve_layout(&layout, ne00, total_rows, k, octx->n_threads, octx->ctx->vtcm_size, is_top_k)) {
-            return HTP_STATUS_VTCM_TOO_SMALL;
-        }
-        memset(&kparams_local, 0, sizeof(kparams_local));
-        kparams_local.n_threads         = (int32_t) layout.n_threads;
-        kparams_local.total_rows        = (int32_t) total_rows;
-        kparams_local.row_start         = (int32_t) row_start;
-        kparams_local.row_end           = (int32_t) row_end;
-        kparams_local.ne00              = (int32_t) ne00;
-        kparams_local.k                 = (int32_t) k;
-        kparams_local.order             = is_top_k ? GGML_SORT_ORDER_DESC : (int32_t) octx->op_params[0];
-        kparams_local.is_top_k          = is_top_k ? 1 : 0;
-        kparams_local.use_dma           = 1;
-        kparams_local.chunk_elems       = (int32_t) layout.chunk_elems;
-        kparams_local.n_chunks          = (int32_t) layout.n_chunks;
-        kparams_local.vtcm_size         = (int32_t) layout.total_bytes;
-        kparams_local.phase1_slot_size  = (int32_t) layout.phase1_slot_size;
-        kparams_local.merge_values_off  = (int32_t) layout.merge_values_off;
-        kparams_local.merge_indices_off = (int32_t) layout.merge_indices_off;
-        kparams_local.merge_elems       = (int32_t) layout.merge_elems;
-        kparams_local.n_slots           = (int32_t) layout.n_slots;
-        kparams = &kparams_local;
-    } else if (octx->ctx->mdev.count > 1) {
-        memcpy(&kparams_local, kparams, sizeof(kparams_local));
-        kparams_local.row_start = (int32_t) row_start;
-        kparams_local.row_end   = (int32_t) row_end;
-        kparams = &kparams_local;
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
    }

-    if (octx->ctx->vtcm_size < (size_t) kparams->vtcm_size) {
-        FARF(ERROR, "sort: VTCM size too small. Needed %d, have %zu", kparams->vtcm_size, octx->ctx->vtcm_size);
+    struct htp_sort_vtcm_layout layout;
+    htp_sort_vtcm_layout_build(&layout,
+                               total_rows,
+                               (uint32_t) kparams->chunk_elems,
+                               (uint32_t) kparams->n_chunks,
+                               (uint32_t) kparams->merge_elems,
+                               (uint32_t) kparams->n_threads,
+                               (uint32_t) kparams->n_slots);
+
+    if (octx->ctx->vtcm_size < layout.total_bytes) {
+        FARF(ERROR, "sort: VTCM size too small. Needed %zu, have %zu", layout.total_bytes, octx->ctx->vtcm_size);
        return HTP_STATUS_VTCM_TOO_SMALL;
    }

-    uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
+    uint8_t * vtcm_base     = (uint8_t *) octx->ctx->vtcm_base;
+    uint8_t * phase1_base   = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_phase1);
+    float *   merge_values  = VTCM_LAYOUT_PTR(float,   vtcm_base, layout.off_merge_values);
+    int32_t * merge_indices = VTCM_LAYOUT_PTR(int32_t, vtcm_base, layout.off_merge_indices);

    if (total_rows == 1 && kparams->n_chunks > 1) {
        struct htp_sort_chunk_ctx cctx;
-        cctx.octx      = octx;
-        cctx.kparams   = kparams;
-        cctx.vtcm_base = vtcm_base;
+        cctx.octx             = octx;
+        cctx.kparams          = kparams;
+        cctx.phase1_base      = phase1_base;
+        cctx.phase1_slot_size = layout.phase1_slot_size;
+        cctx.merge_values     = merge_values;
+        cctx.merge_indices    = merge_indices;
        memset(cctx.real_count, 0, sizeof(cctx.real_count));

        work_queue_run(octx->ctx->work_queue, htp_sort_chunk_job, &cctx, (uint32_t) kparams->n_chunks);
@@ -902,8 +902,8 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
            uint32_t local_k     = MIN(k, chunk_elems);

            if (k == 1 && ne00 >= 128*1024) {
-                float *   cand_vals = (float *)   (vtcm_base + kparams->merge_values_off);
-                int32_t * cand_idxs = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+                float *   cand_vals = merge_values;
+                int32_t * cand_idxs = merge_indices;
                float best_val = cand_vals[0];
                int32_t best_idx = cand_idxs[0];
                for (int32_t c = 1; c < kparams->n_chunks; c++) {
@@ -918,8 +918,8 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
            }

            uint32_t total_candidates = (uint32_t) kparams->n_chunks * local_k;
-            float *   cand_vals = (float *)   (vtcm_base + kparams->merge_values_off);
-            int32_t * cand_idxs = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+            float *   cand_vals = merge_values;
+            int32_t * cand_idxs = merge_indices;
            uint32_t merge_elems = (uint32_t) kparams->merge_elems;

            if (merge_elems > total_candidates) {
@@ -943,18 +943,18 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
                struct htp_argsort_merge_l2_ctx l2ctx = { &cctx, asc };
                work_queue_run(octx->ctx->work_queue, htp_argsort_merge_l2_job, &l2ctx, 2);

-                size_t slot_size = (size_t) kparams->phase1_slot_size;
+                size_t slot_size = layout.phase1_slot_size;
                uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;

-                const float *   val0 = (const float *) vtcm_base;
-                const int32_t * idx0 = (const int32_t *) (vtcm_base + 4 * chunk_elems * sizeof(float));
-                const float *   val1 = (const float *) (vtcm_base + 4 * slot_size);
-                const int32_t * idx1 = (const int32_t *) (vtcm_base + 4 * slot_size + 4 * chunk_elems * sizeof(float));
+                const float *   val0 = (const float *) phase1_base;
+                const int32_t * idx0 = (const int32_t *) (phase1_base + 4 * chunk_elems * sizeof(float));
+                const float *   val1 = (const float *) (phase1_base + 4 * slot_size);
+                const int32_t * idx1 = (const int32_t *) (phase1_base + 4 * slot_size + 4 * chunk_elems * sizeof(float));

                uint32_t n01 = cctx.real_count[0] + cctx.real_count[1] + cctx.real_count[2] + cctx.real_count[3];
                uint32_t n23 = cctx.real_count[4] + cctx.real_count[5] + cctx.real_count[6] + cctx.real_count[7];

-                int32_t * final_idx = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+                int32_t * final_idx = merge_indices;
                struct htp_thread_trace * tr = &octx->ctx->trace[0];
                htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                merge_runs(val0, idx0, n01, val1, idx1, n23, NULL, final_idx, asc);
@@ -968,27 +968,27 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
                uint32_t n01 = cctx.real_count[0] + cctx.real_count[1];
                uint32_t n23 = cctx.real_count[2] + cctx.real_count[3];

-                const float *   val01 = (const float *)   (vtcm_base + kparams->merge_values_off);
-                const int32_t * idx01 = (const int32_t *) (vtcm_base + kparams->merge_indices_off);
-                const float *   val23 = (const float *)   (vtcm_base + kparams->merge_values_off)  + 2 * chunk_elems;
-                const int32_t * idx23 = (const int32_t *) (vtcm_base + kparams->merge_indices_off) + 2 * chunk_elems;
+                const float *   val01 = merge_values;
+                const int32_t * idx01 = merge_indices;
+                const float *   val23 = merge_values  + 2 * chunk_elems;
+                const int32_t * idx23 = merge_indices + 2 * chunk_elems;

-                int32_t * final_idx = (int32_t *) vtcm_base;
+                int32_t * final_idx = (int32_t *) phase1_base;
                struct htp_thread_trace * tr = &octx->ctx->trace[0];
                htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                merge_runs(val01, idx01, n01, val23, idx23, n23, NULL, final_idx, asc);
                htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                dma_sync_write(octx->ctx->dma[0], dst->data, final_idx, ne00 * sizeof(int32_t));
            } else if (kparams->n_chunks == 2) {
-                size_t slot_size = (size_t) kparams->phase1_slot_size;
+                size_t slot_size = layout.phase1_slot_size;
                uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;

-                const float *   val0 = (const float *)   (vtcm_base);
-                const int32_t * idx0 = (const int32_t *) (vtcm_base + chunk_elems * sizeof(float));
-                const float *   val1 = (const float *)   (vtcm_base + slot_size);
-                const int32_t * idx1 = (const int32_t *) (vtcm_base + slot_size + chunk_elems * sizeof(float));
+                const float *   val0 = (const float *)   (phase1_base);
+                const int32_t * idx0 = (const int32_t *) (phase1_base + chunk_elems * sizeof(float));
+                const float *   val1 = (const float *)   (phase1_base + slot_size);
+                const int32_t * idx1 = (const int32_t *) (phase1_base + slot_size + chunk_elems * sizeof(float));

-                int32_t * final_idx = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+                int32_t * final_idx = merge_indices;
                struct htp_thread_trace * tr = &octx->ctx->trace[0];
                htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                merge_runs(val0, idx0, cctx.real_count[0], val1, idx1, cctx.real_count[1], NULL, final_idx, asc);
@@ -996,7 +996,7 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
                dma_sync_write(octx->ctx->dma[0], dst->data, final_idx, ne00 * sizeof(int32_t));
            } else {
                uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
-                int32_t * final_idx = (int32_t *) (vtcm_base + chunk_elems * sizeof(float));
+                int32_t * final_idx = (int32_t *) (phase1_base + chunk_elems * sizeof(float));
                dma_sync_write(octx->ctx->dma[0], dst->data, final_idx, ne00 * sizeof(int32_t));
            }
        }
@@ -1006,7 +1006,10 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
    struct htp_sort_multi_row_ctx mctx;
    mctx.octx             = octx;
    mctx.kparams          = kparams;
-    mctx.vtcm_base        = vtcm_base;
+    mctx.phase1_base      = phase1_base;
+    mctx.phase1_slot_size = layout.phase1_slot_size;
+    mctx.row_start        = row_start;
+    mctx.row_end          = row_end;

    uint32_t dst_row_bytes = (uint32_t) dst_row_size;
    uint32_t rows_per_chunk = 0;
@@ -1014,13 +1017,13 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
    if (rows_per_chunk == 0) {
        rows_per_chunk = 1;
    }
-    uint32_t dr = (nrows + (uint32_t) kparams->n_threads - 1) / (uint32_t) kparams->n_threads;
+    uint32_t dr = fastdiv(nrows + octx->n_threads - 1, &octx->n_threads_div);
    if (rows_per_chunk > 1) {
        dr = hex_round_up(dr, rows_per_chunk);
    }
    mctx.nrows_per_thread = dr;

-    work_queue_run(octx->ctx->work_queue, htp_sort_multi_row_job, &mctx, (uint32_t) kparams->n_threads);
+    work_queue_run(octx->ctx->work_queue, htp_sort_multi_row_job, &mctx, octx->n_threads);

    return HTP_STATUS_OK;
 }
@@ -7,33 +7,27 @@
 #include <string.h>

 #include "hex-fastdiv.h"
+#include "htp-vtcm.h"

 struct htp_sort_kernel_params {
-    int32_t  n_threads;
-    int32_t  total_rows;
-    int32_t  row_start;
-    int32_t  row_end;
-    int32_t  ne00;
-    int32_t  k;
-    int32_t  order;             // GGML_SORT_ORDER_ASC (0) or GGML_SORT_ORDER_DESC (1)
-    int32_t  is_top_k;          // 1 if TOP_K, 0 if ARGSORT
-    int32_t  use_dma;           // 1 if DMA enabled
-    int32_t  chunk_elems;
-    int32_t  n_chunks;
-    int32_t  vtcm_size;
-    int32_t  phase1_slot_size;
-    int32_t  merge_values_off;
-    int32_t  merge_indices_off;
-    int32_t  merge_elems;
-    int32_t  n_slots;
-    int32_t  pad[15];
+    uint32_t n_threads;
+    uint32_t order;             // GGML_SORT_ORDER_ASC (0) or GGML_SORT_ORDER_DESC (1)
+    uint32_t is_top_k;          // 1 if TOP_K, 0 if ARGSORT
+    uint32_t chunk_elems;
+    uint32_t n_chunks;
+    uint32_t vtcm_size;
+    uint32_t merge_elems;
+    uint32_t n_slots;
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne01_ne02;
 };

 struct htp_sort_vtcm_layout {
    size_t   total_bytes;
    size_t   phase1_slot_size;
-    size_t   merge_values_off;
-    size_t   merge_indices_off;
+    size_t   off_phase1;
+    size_t   off_merge_values;
+    size_t   off_merge_indices;
    uint32_t chunk_elems;
    uint32_t n_chunks;
    uint32_t merge_elems;
@@ -41,6 +35,39 @@ struct htp_sort_vtcm_layout {
    uint32_t n_slots;
 };

+static inline void htp_sort_vtcm_layout_build(
+    struct htp_sort_vtcm_layout * L,
+    uint32_t total_rows,
+    uint32_t chunk_elems,
+    uint32_t n_chunks,
+    uint32_t merge_elems,
+    uint32_t n_threads,
+    uint32_t n_slots) {
+
+    memset(L, 0, sizeof(*L));
+    L->chunk_elems = chunk_elems;
+    L->n_chunks    = n_chunks;
+    L->merge_elems = merge_elems;
+    L->n_threads   = n_threads;
+    L->n_slots     = n_slots;
+
+    const size_t values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
+    const size_t indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
+    L->phase1_slot_size       = ((values_size + indices_size) + 255) & ~255;
+
+    size_t off = 0;
+    if (total_rows > 1) {
+        VTCM_LAYOUT_ALLOC(off, off_phase1, L->phase1_slot_size * n_slots * n_threads);
+    } else {
+        VTCM_LAYOUT_ALLOC(off, off_phase1, L->phase1_slot_size * n_chunks * n_slots);
+        const size_t merge_values_size  = ((merge_elems * sizeof(float)) + 127) & ~127;
+        const size_t merge_indices_size = ((merge_elems * sizeof(int32_t)) + 127) & ~127;
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_merge_values,  merge_values_size,  n_chunks > 1);
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_merge_indices, merge_indices_size, n_chunks > 1);
+    }
+    L->total_bytes = off;
+}
+
 static inline bool htp_sort_solve_layout(
    struct htp_sort_vtcm_layout * layout,
    uint32_t ne00,
@@ -75,18 +102,8 @@ static inline bool htp_sort_solve_layout(
            n_threads--;
        }

-        size_t total_bytes = spad_per_thread * n_threads;
-        if (total_bytes > vtcm_budget) {
-            return false;
-        }
-
-        layout->total_bytes      = total_bytes;
-        layout->phase1_slot_size = spad_per_slot;
-        layout->chunk_elems      = ne00_padded;
-        layout->n_chunks         = 1;
-        layout->n_threads        = n_threads;
-        layout->n_slots          = n_slots;
-        return true;
+        htp_sort_vtcm_layout_build(layout, total_rows, ne00_padded, 1, 0, n_threads, n_slots);
+        return layout->total_bytes <= vtcm_budget;
    }

    uint32_t n_vec = (ne00 + 31) / 32;
@@ -102,17 +119,7 @@ static inline bool htp_sort_solve_layout(

    uint32_t chunk_n_vec = n_vec_pow2 / n_chunks;
    uint32_t chunk_elems = chunk_n_vec * 32;
-
-    size_t phase1_values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
-    size_t phase1_indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
-    size_t phase1_slot_size    = ((phase1_values_size + phase1_indices_size) + 255) & ~255;
-    size_t phase1_total_size   = phase1_slot_size * n_chunks;
-
-    size_t merge_values_size  = 0;
-    size_t merge_indices_size = 0;
-    size_t merge_values_off   = phase1_total_size;
-    size_t merge_indices_off  = 0;
-    uint32_t merge_elems      = 0;
+    uint32_t merge_elems = 0;

    if (n_chunks > 1) {
        if (is_top_k) {
@@ -125,47 +132,25 @@ static inline bool htp_sort_solve_layout(
        } else {
            merge_elems = n_vec_pow2 * 32;
        }
-        merge_values_size  = ((merge_elems * sizeof(float)) + 127) & ~127;
-        merge_indices_size = ((merge_elems * sizeof(int32_t)) + 127) & ~127;
-        merge_indices_off  = merge_values_off + merge_values_size;
    }

-    size_t total_bytes = phase1_total_size + merge_values_size + merge_indices_size;
+    htp_sort_vtcm_layout_build(layout, 1, chunk_elems, n_chunks, merge_elems, n_chunks, 1);

-    if (total_bytes > vtcm_budget && n_chunks > 1) {
-        n_chunks = 1;
+    if (layout->total_bytes > vtcm_budget && n_chunks > 1) {
+        n_chunks    = 1;
        chunk_elems = n_vec_pow2 * 32;
-        phase1_values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
-        phase1_indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
-        phase1_slot_size    = ((phase1_values_size + phase1_indices_size) + 255) & ~255;
-        phase1_total_size   = phase1_slot_size;
-        merge_values_size   = 0;
-        merge_indices_size  = 0;
-        merge_values_off    = phase1_total_size;
-        merge_indices_off   = 0;
-        merge_elems         = 0;
-        total_bytes         = phase1_total_size;
+        merge_elems = 0;
+        htp_sort_vtcm_layout_build(layout, 1, chunk_elems, n_chunks, merge_elems, 1, 1);
    }

-    if (total_bytes > vtcm_budget) {
+    if (layout->total_bytes > vtcm_budget) {
        return false;
    }

-    uint32_t n_slots = 1;
-    if (n_chunks == 1 && phase1_slot_size * 2 <= vtcm_budget) {
-        n_slots = 2;
-        total_bytes = phase1_slot_size * 2;
+    if (n_chunks == 1 && layout->phase1_slot_size * 2 <= vtcm_budget) {
+        htp_sort_vtcm_layout_build(layout, 1, chunk_elems, 1, 0, 1, 2);
    }

-    layout->total_bytes       = total_bytes;
-    layout->phase1_slot_size  = phase1_slot_size;
-    layout->merge_values_off  = merge_values_off;
-    layout->merge_indices_off = merge_indices_off;
-    layout->chunk_elems       = chunk_elems;
-    layout->n_chunks          = n_chunks;
-    layout->merge_elems       = merge_elems;
-    layout->n_threads         = n_chunks;
-    layout->n_slots           = n_slots;
    return true;
 }

@@ -42,6 +42,7 @@ struct htp_binary_context {
    uint32_t nrows_per_thread;
    uint32_t total_rows;
    uint32_t row_start;
+    uint32_t n_srcs;
    size_t   src0_row_size_aligned;
    size_t   src1_row_size_aligned;
    size_t   dst_row_size_aligned;
@@ -380,6 +381,7 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
    size_t dst_spad_half  = layout->dst_spad_half_size;

    dma_queue * dma_q = octx->ctx->dma[ith];
+    const uint32_t n_srcs = bctx->n_srcs;
    uint32_t ir_prefetch = start_row;
    int spad_idx = 0;

@@ -391,12 +393,7 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
        i02 = fastdiv(rem, &bctx->src0_dim1_div);
        i01 = rem - i02 * ne01;

-        uint32_t i13 = (ne13 == 1) ? 0 : i03;
-        uint32_t i12 = (ne12 == 1) ? 0 : i02;
-        uint32_t i11 = (ne11 == 1) ? 0 : i01;
-
        dma_addr_t src0_curr = src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01;
-        dma_addr_t src1_curr = src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11;
        dma_addr_t dst_curr  = dst->data  + i03 * nb3  + i02 * nb2  + i01 * nb1;

        uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half;
@@ -405,7 +402,15 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,

        dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0);
        dma_queue_push(dma_q, dma_make_data(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size);
-        dma_queue_push(dma_q, dma_make_data(s1_spad, src1_curr), bctx->src1_row_size_aligned, nb11, row_size_bytes, current_block_size);
+        for (uint32_t s = 1; s < n_srcs; s++) {
+            const struct htp_tensor * src_s = octx->src[s];
+            uint32_t is3 = (src_s->ne[3] == 1) ? 0 : i03;
+            uint32_t is2 = (src_s->ne[2] == 1) ? 0 : i02;
+            uint32_t is1 = (src_s->ne[1] == 1) ? 0 : i01;
+            dma_addr_t s_curr = src_s->data + is3 * src_s->nb[3] + is2 * src_s->nb[2] + is1 * src_s->nb[1];
+            uint8_t * s_spad = s1_spad + (s - 1) * 2 * src1_spad_half;
+            dma_queue_push(dma_q, dma_make_data(s_spad, s_curr), bctx->src1_row_size_aligned, src_s->nb[1], row_size_bytes, current_block_size);
+        }
        ir_prefetch += current_block_size;
        spad_idx ^= 1;
    }
@@ -422,6 +427,11 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
        compute(d_spad, s0_spad, s1_spad, current_block_size,
                bctx->dst_row_size_aligned, bctx->src0_row_size_aligned, bctx->src1_row_size_aligned, ne00);
+        for (uint32_t s = 2; s < n_srcs; s++) {
+            uint8_t * ss_spad = (uint8_t *) dma_queue_pop(dma_q).dst;
+            compute(d_spad, d_spad, ss_spad, current_block_size,
+                    bctx->dst_row_size_aligned, bctx->dst_row_size_aligned, bctx->src1_row_size_aligned, ne00);
+        }
        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);

        uint32_t i03, i02, i01, rem;
@@ -440,15 +450,18 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
            p02 = fastdiv(prem, &bctx->src0_dim1_div);
            p01 = prem - p02 * ne01;

-            uint32_t p13 = (ne13 == 1) ? 0 : p03;
-            uint32_t p12 = (ne12 == 1) ? 0 : p02;
-            uint32_t p11 = (ne11 == 1) ? 0 : p01;
-
            dma_addr_t s0_next = src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01;
-            dma_addr_t s1_next = src1->data + p13 * nb13 + p12 * nb12 + p11 * nb11;
-
            dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size);
-            dma_queue_push(dma_q, dma_make_data(s1_spad, s1_next), bctx->src1_row_size_aligned, nb11, row_size_bytes, next_block_size);
+
+            for (uint32_t s = 1; s < n_srcs; s++) {
+                const struct htp_tensor * src_s = octx->src[s];
+                uint32_t ps3 = (src_s->ne[3] == 1) ? 0 : p03;
+                uint32_t ps2 = (src_s->ne[2] == 1) ? 0 : p02;
+                uint32_t ps1 = (src_s->ne[1] == 1) ? 0 : p01;
+                dma_addr_t s_next = src_s->data + ps3 * src_s->nb[3] + ps2 * src_s->nb[2] + ps1 * src_s->nb[1];
+                uint8_t * s_spad = s1_spad + (s - 1) * 2 * src1_spad_half;
+                dma_queue_push(dma_q, dma_make_data(s_spad, s_next), bctx->src1_row_size_aligned, src_s->nb[1], row_size_bytes, next_block_size);
+            }

            ir_prefetch += next_block_size;
        }
@@ -859,6 +872,7 @@ struct binary_chunked_context {
    uint32_t                 chunk_size;
    uint32_t                 chunks_per_thread;
    uint32_t                 total_chunks;
+    uint32_t                 n_srcs;
    bool                     is_scalar;
    float                    scalar_f32;
    _Float16                 scalar_f16;
@@ -872,7 +886,6 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
    struct binary_chunked_context * ctx = (struct binary_chunked_context *) data;
    struct htp_ops_context * octx = ctx->octx;
    const struct htp_tensor * src0 = octx->src[0];
-    const struct htp_tensor * src1 = octx->src[1];
    const struct htp_tensor * dst  = octx->dst;

    const uint32_t start_chunk = ctx->chunks_per_thread * ith;
@@ -895,6 +908,7 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);

    dma_queue * dma_q = octx->ctx->dma[ith];
+    const uint32_t n_srcs = ctx->n_srcs;
    uint32_t prefetch_chunk = start_chunk;
    int spad_idx = 0;

@@ -913,9 +927,12 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
        dma_queue_push(dma_q, dma_make_data(d_curr, d_spad), chunk_bytes, chunk_bytes, cur_bytes, 0);
        dma_queue_push(dma_q, dma_make_data(s0_spad, s0_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
        if (!ctx->is_scalar) {
-            dma_addr_t s1_curr = src1->data + (size_t) c_start * elem_size;
            uint8_t * s1_spad = src1_spad_base + spad_idx * chunk_bytes;
-            dma_queue_push(dma_q, dma_make_data(s1_spad, s1_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
+            for (uint32_t s = 1; s < n_srcs; s++) {
+                dma_addr_t s_curr = octx->src[s]->data + (size_t) c_start * elem_size;
+                uint8_t * s_spad  = s1_spad + (s - 1) * 2 * chunk_bytes;
+                dma_queue_push(dma_q, dma_make_data(s_spad, s_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
+            }
        }

        prefetch_chunk++;
@@ -943,6 +960,10 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
            }
        } else {
            ctx->compute(d_spad, s0_spad, s1_spad, cur_elems);
+            for (uint32_t s = 2; s < n_srcs; s++) {
+                uint8_t * s_spad = (uint8_t *) dma_queue_pop(dma_q).dst;
+                ctx->compute(d_spad, d_spad, s_spad, cur_elems);
+            }
        }
        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) c);

@@ -958,8 +979,11 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
            dma_addr_t s0_next = src0->data + (size_t) pc_start * elem_size;
            dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), chunk_bytes, chunk_bytes, p_bytes, 1);
            if (!ctx->is_scalar) {
-                dma_addr_t s1_next = src1->data + (size_t) pc_start * elem_size;
-                dma_queue_push(dma_q, dma_make_data(s1_spad, s1_next), chunk_bytes, chunk_bytes, p_bytes, 1);
+                for (uint32_t s = 1; s < n_srcs; s++) {
+                    dma_addr_t s_next = octx->src[s]->data + (size_t) pc_start * elem_size;
+                    uint8_t * s_spad  = s1_spad + (s - 1) * 2 * chunk_bytes;
+                    dma_queue_push(dma_q, dma_make_data(s_spad, s_next), chunk_bytes, chunk_bytes, p_bytes, 1);
+                }
            }

            prefetch_chunk++;
@@ -974,6 +998,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
    const struct htp_tensor * src1 = octx->src[1];
    const struct htp_tensor * dst  = octx->dst;
    const struct htp_binary_kernel_params * kparams = (const struct htp_binary_kernel_params *) octx->kernel_params;
+    const uint32_t n_srcs = kparams->n_srcs > 2 ? kparams->n_srcs : 2;

    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];

@@ -985,7 +1010,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
    const size_t dst_row_size  = dst->ne[0]  * elem_size;

    if (kparams->kernel_type == HTP_BINARY_KERNEL_CHUNKED) {
-        const bool is_scalar = kparams->is_scalar || (src1->ne[0] == 1 && src1->ne[1] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1);
+        const bool is_scalar = kparams->is_scalar;

        compute_binary_chunked_t compute = NULL;
        compute_binary_scalar_chunked_f32_t compute_scalar_f32 = NULL;
@@ -1045,12 +1070,16 @@ static int execute_op_binary(struct htp_ops_context * octx) {
        const uint32_t elems_per_line = (src0_type == HTP_TYPE_F32) ? 32 : 64;

        if (octx->ctx->mdev.count > 1) {
-            const bool can_split = htp_tensor_mdev_data_aligned(dst) &&
-                                   htp_tensor_mdev_data_aligned(src0) &&
-                                   (is_scalar || htp_tensor_mdev_data_aligned(src1)) &&
-                                   htp_tensor_is_contiguous(dst, elem_size) &&
-                                   htp_tensor_is_contiguous(src0, elem_size) &&
-                                   (is_scalar || htp_tensor_is_contiguous(src1, elem_size));
+            bool can_split = htp_tensor_mdev_data_aligned(dst) &&
+                             htp_tensor_mdev_data_aligned(src0) &&
+                             (is_scalar || htp_tensor_mdev_data_aligned(src1)) &&
+                             htp_tensor_is_contiguous(dst, elem_size) &&
+                             htp_tensor_is_contiguous(src0, elem_size) &&
+                             (is_scalar || htp_tensor_is_contiguous(src1, elem_size));
+            for (uint32_t s = 2; can_split && s < n_srcs; s++) {
+                can_split = htp_tensor_mdev_data_aligned(octx->src[s]) &&
+                            htp_tensor_is_contiguous(octx->src[s], elem_size);
+            }
            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);
@@ -1072,7 +1101,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
            return HTP_STATUS_VTCM_TOO_SMALL;
        }

-        const uint32_t chunk_size = kparams->chunk_size > 0 ? kparams->chunk_size : (32768 / elem_size);
+        const uint32_t chunk_size = kparams->chunk_size;
        const uint32_t total_chunks = (nelem + chunk_size - 1) / chunk_size;
        const uint32_t n_threads = (total_chunks >= 2) ? MIN(octx->n_threads, total_chunks) : 1;
        const uint32_t chunks_per_thread = (total_chunks + n_threads - 1) / n_threads;
@@ -1101,6 +1130,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
            .chunk_size          = chunk_size,
            .chunks_per_thread   = chunks_per_thread,
            .total_chunks        = total_chunks,
+            .n_srcs              = n_srcs,
            .is_scalar           = is_scalar,
            .scalar_f32          = scalar_f32,
            .scalar_f16          = scalar_f16,
@@ -1146,65 +1176,6 @@ static int execute_op_binary(struct htp_ops_context * octx) {
        return HTP_STATUS_VTCM_TOO_SMALL;
    }

-    dma_queue * dma_q = octx->ctx->dma[0];
-    uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src1);
-    if (kparams->kernel_type == HTP_BINARY_KERNEL_ROW_BCAST) {
-        dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1->ne[0] * elem_size, 1);
-    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR ||
-               kparams->kernel_type == HTP_BINARY_KERNEL_REPEAT ||
-               kparams->kernel_type == HTP_BINARY_KERNEL_COMPLEX) {
-        if (!kparams->src1_strided) {
-            const size_t src1_bytes = (size_t) src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3] * elem_size;
-            dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1_bytes, 1);
-        } else {
-            const uint32_t ne11 = src1->ne[1];
-            const uint32_t ne12 = src1->ne[2];
-            const uint32_t ne13 = src1->ne[3];
-            const size_t row_bytes = src1->ne[0] * elem_size;
-            for (uint32_t i13 = 0; i13 < ne13; i13++) {
-                for (uint32_t i12 = 0; i12 < ne12; i12++) {
-                    dma_addr_t src_addr = src1->data + i13 * src1->nb[3] + i12 * src1->nb[2];
-                    uint8_t * dst_ptr = vtcm_src1 + (i13 * ne12 * ne11 + i12 * ne11) * row_bytes;
-                    while (!dma_queue_push(dma_q, dma_make_data(dst_ptr, src_addr),
-                                           row_bytes, src1->nb[1], row_bytes, ne11)) {
-                        dma_queue_pop(dma_q);
-                    }
-                }
-            }
-        }
-    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) {
-        const struct htp_tensor * src2 = octx->src[2];
-        uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
-        dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data),
-                       kparams->src1_row_size_aligned, src1->nb[1],
-                       src1->ne[0] * elem_size, src1->ne[1]);
-        const size_t src2_row_bytes = src2->ne[0] * sizeof(int32_t);
-        dma_queue_push(dma_q, dma_make_data(vtcm_src2, src2->data),
-                       src2_row_bytes, src2->nb[1],
-                       src2_row_bytes, src2->ne[1]);
-    }
-
-    bctx.octx                  = octx;
-    bctx.nrows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
-    bctx.total_rows            = nrows;
-    bctx.row_start             = row_start;
-    bctx.block_max             = bctx.vtcm_layout.rows_per_buffer;
-    bctx.src0_row_size_aligned = src0_row_size_aligned;
-    bctx.src1_row_size_aligned = src1_row_size_aligned;
-    bctx.dst_row_size_aligned  = dst_row_size_aligned;
-    bctx.row_size_bytes        = src0_row_size;
-
-    bctx.src0_dim1_div  = init_fastdiv_values(src0->ne[1]);
-    bctx.src0_dim2_div  = init_fastdiv_values(src0->ne[2]);
-    bctx.src0_dim12_div = init_fastdiv_values(src0->ne[1] * src0->ne[2]);
-
-    bctx.src1_dim1_div  = init_fastdiv_values(src1->ne[1]);
-    bctx.src1_dim2_div  = init_fastdiv_values(src1->ne[2]);
-    bctx.src1_dim3_div  = init_fastdiv_values(src1->ne[3]);
-
-    bctx.split_at_ne01  = kparams->split_at_ne01;
-    bctx.split_at_ne02  = kparams->split_at_ne02;
-
    worker_callback_t worker_func = NULL;
    void * compute_func = NULL;

@@ -1322,7 +1293,70 @@ static int execute_op_binary(struct htp_ops_context * octx) {
        return HTP_STATUS_NO_SUPPORT;
    }

-    bctx.compute = compute_func;
+    dma_queue * dma_q = octx->ctx->dma[0];
+    uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src1);
+    if (kparams->kernel_type == HTP_BINARY_KERNEL_ROW_BCAST) {
+        dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1->ne[0] * elem_size, 1);
+    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR ||
+               kparams->kernel_type == HTP_BINARY_KERNEL_REPEAT ||
+               kparams->kernel_type == HTP_BINARY_KERNEL_COMPLEX) {
+        if (!kparams->src1_strided) {
+            const size_t src1_bytes = (size_t) src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3] * elem_size;
+            dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1_bytes, 1);
+            if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR && octx->src[2] != NULL) {
+                uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
+                dma_queue_push(dma_q, dma_make_data(vtcm_src2, octx->src[2]->data), bctx.vtcm_layout.src2_size, 0, src1_bytes, 1);
+            }
+        } else {
+            const uint32_t ne11 = src1->ne[1];
+            const uint32_t ne12 = src1->ne[2];
+            const uint32_t ne13 = src1->ne[3];
+            const size_t row_bytes = src1->ne[0] * elem_size;
+            for (uint32_t i13 = 0; i13 < ne13; i13++) {
+                for (uint32_t i12 = 0; i12 < ne12; i12++) {
+                    dma_addr_t src_addr = src1->data + i13 * src1->nb[3] + i12 * src1->nb[2];
+                    uint8_t * dst_ptr = vtcm_src1 + (i13 * ne12 * ne11 + i12 * ne11) * row_bytes;
+                    while (!dma_queue_push(dma_q, dma_make_data(dst_ptr, src_addr),
+                                           row_bytes, src1->nb[1], row_bytes, ne11)) {
+                        dma_queue_pop(dma_q);
+                    }
+                }
+            }
+        }
+    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) {
+        const struct htp_tensor * src2 = octx->src[2];
+        uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
+        dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data),
+                       kparams->src1_row_size_aligned, src1->nb[1],
+                       src1->ne[0] * elem_size, src1->ne[1]);
+        const size_t src2_row_bytes = src2->ne[0] * sizeof(int32_t);
+        dma_queue_push(dma_q, dma_make_data(vtcm_src2, src2->data),
+                       src2_row_bytes, src2->nb[1],
+                       src2_row_bytes, src2->ne[1]);
+    }
+
+    bctx.octx                  = octx;
+    bctx.nrows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    bctx.total_rows            = nrows;
+    bctx.row_start             = row_start;
+    bctx.n_srcs                = n_srcs;
+    bctx.block_max             = bctx.vtcm_layout.rows_per_buffer;
+    bctx.src0_row_size_aligned = src0_row_size_aligned;
+    bctx.src1_row_size_aligned = src1_row_size_aligned;
+    bctx.dst_row_size_aligned  = dst_row_size_aligned;
+    bctx.row_size_bytes        = src0_row_size;
+
+    bctx.src0_dim1_div  = kparams->div_ne01;
+    bctx.src0_dim2_div  = kparams->div_ne02;
+    bctx.src0_dim12_div = kparams->div_ne012;
+
+    bctx.src1_dim1_div  = kparams->div_ne11;
+    bctx.src1_dim2_div  = kparams->div_ne12;
+    bctx.src1_dim3_div  = kparams->div_ne13;
+
+    bctx.split_at_ne01  = kparams->split_at_ne01;
+    bctx.split_at_ne02  = kparams->split_at_ne02;
+    bctx.compute        = compute_func;

    if (kparams->kernel_type == HTP_BINARY_KERNEL_ROW_BCAST ||
        kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR ||
@@ -1330,6 +1364,15 @@ static int execute_op_binary(struct htp_ops_context * octx) {
        kparams->kernel_type == HTP_BINARY_KERNEL_COMPLEX ||
        kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) {
        dma_queue_flush(dma_q);
+        if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR && octx->src[2] != NULL && octx->op == HTP_OP_MUL) {
+            uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
+            const uint32_t n_scalars = src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3];
+            if (src0_type == HTP_TYPE_F32) {
+                hvx_mul_f32_aaa(vtcm_src1, vtcm_src1, vtcm_src2, n_scalars);
+            } else if (src0_type == HTP_TYPE_F16) {
+                hvx_mul_f16_aaa(vtcm_src1, vtcm_src1, vtcm_src2, n_scalars);
+            }
+        }
    }

    work_queue_run(octx->ctx->work_queue, worker_func, &bctx, n_threads);
@@ -6,6 +6,7 @@
 #include <string.h>

 #include "hex-common.h"
+#include "hex-fastdiv.h"
 #include "htp-ops.h"
 #include "htp-vtcm.h"

@@ -20,7 +21,15 @@ enum htp_binary_kernel_type {
 };

 struct htp_binary_kernel_params {
-    uint32_t kernel_type;
+    uint8_t  kernel_type;
+    uint8_t  split_at_ne01;
+    uint8_t  split_at_ne02;
+    uint8_t  src1_strided;
+
+    uint8_t  n_srcs;
+    uint8_t  is_scalar;
+    uint8_t  pad[2];
+
    uint32_t n_threads;
    uint32_t rows_per_buffer;

@@ -34,12 +43,14 @@ struct htp_binary_kernel_params {

    uint32_t chunk_size;
    uint32_t chunk_bytes;
-    uint32_t is_scalar;

-    uint8_t  split_at_ne01;
-    uint8_t  split_at_ne02;
-    uint8_t  src1_strided;
-    uint8_t  pad;
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne02;
+    struct fastdiv_values div_ne012;
+
+    struct fastdiv_values div_ne11;
+    struct fastdiv_values div_ne12;
+    struct fastdiv_values div_ne13;
 };

 #if defined(__cplusplus)
@@ -100,6 +111,8 @@ static inline void htp_binary_vtcm_layout_build(
        return;
    }

+    const uint32_t n_srcs = kparams->n_srcs > 2 ? kparams->n_srcs : 2;
+
    if (kparams->kernel_type == HTP_BINARY_KERNEL_CHUNKED) {
        const size_t chunk_bytes = kparams->chunk_bytes;
        if (chunk_bytes == 0) {
@@ -107,7 +120,7 @@ static inline void htp_binary_vtcm_layout_build(
        }

        L->src0_bytes_per_thread = 2 * chunk_bytes;
-        L->src1_bytes_per_thread = kparams->is_scalar ? 0 : (2 * chunk_bytes);
+        L->src1_bytes_per_thread = kparams->is_scalar ? 0 : (2 * (n_srcs - 1) * chunk_bytes);
        L->dst_bytes_per_thread  = 2 * chunk_bytes;

        L->src0_spad_half_size = chunk_bytes;
@@ -135,7 +148,7 @@ static inline void htp_binary_vtcm_layout_build(
    }

    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 + (n_srcs - 1) * kparams->src1_row_size_aligned + kparams->dst_row_size_aligned)
        : 2 * (kparams->src0_row_size_aligned + kparams->dst_row_size_aligned);

    const size_t static_vtcm_size = kparams->src1_size + kparams->src2_size;
@@ -143,7 +156,9 @@ static inline void htp_binary_vtcm_layout_build(
        return;
    }

-    const size_t rows_per_buffer = (vtcm_size - static_vtcm_size) / (n_threads * spad_row_total);
+    const size_t rows_per_buffer = (kparams->rows_per_buffer > 0)
+        ? kparams->rows_per_buffer
+        : (vtcm_size - static_vtcm_size) / (n_threads * spad_row_total);
    if (rows_per_buffer == 0) {
        return;
    }
@@ -155,11 +170,13 @@ static inline void htp_binary_vtcm_layout_build(
    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
+        ? rows_per_buffer * 2 * (n_srcs - 1) * 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->src1_spad_half_size = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE)
+        ? rows_per_buffer * kparams->src1_row_size_aligned
+        : 0;
    L->dst_spad_half_size  = L->dst_bytes_per_thread / 2;

    const size_t src0_total = n_threads * L->src0_bytes_per_thread;
@@ -178,4 +195,53 @@ static inline void htp_binary_vtcm_layout_build(
    L->total_bytes = off;
 }

+static inline bool htp_binary_solve_chunked(
+    struct htp_binary_vtcm_layout * L,
+    struct htp_binary_kernel_params * kparams,
+    uint32_t total_elems,
+    uint32_t elem_size,
+    uint32_t max_threads,
+    uint32_t n_srcs,
+    bool is_scalar,
+    size_t vtcm_size
+) {
+    const size_t n_bufs = is_scalar ? 2 : (size_t) n_srcs + 1;
+    const size_t scalar_bytes = is_scalar ? 128 : 0;
+    const size_t budget_per_buf = (max_threads > 0 && vtcm_size > scalar_bytes)
+        ? hex_align_down((vtcm_size - scalar_bytes) / (2 * max_threads * n_bufs), 128)
+        : 0;
+    const uint32_t max_chunk_bytes = (budget_per_buf > 0 && budget_per_buf < 32768) ? (uint32_t) budget_per_buf : 32768u;
+    const uint32_t max_chunk_elems = max_chunk_bytes / elem_size;
+    if (max_chunk_elems == 0) {
+        return false;
+    }
+
+    const uint32_t min_chunk_elems = 256 < max_chunk_elems ? 256 : max_chunk_elems;
+    const uint32_t target_chunk_elems = hex_round_up((total_elems + (2 * max_threads) - 1) / (2 * max_threads), 32);
+    const uint32_t chunk_size  = (target_chunk_elems < min_chunk_elems) ? min_chunk_elems
+                               : (target_chunk_elems > max_chunk_elems) ? max_chunk_elems
+                               : target_chunk_elems;
+    const uint32_t chunk_bytes  = hex_round_up(chunk_size * elem_size, 128);
+    const uint32_t total_chunks = (total_elems + chunk_size - 1) / chunk_size;
+    const uint32_t n_threads    = (total_chunks >= 2) ? (max_threads < total_chunks ? max_threads : total_chunks) : 1;
+
+    kparams->kernel_type           = HTP_BINARY_KERNEL_CHUNKED;
+    kparams->n_threads             = n_threads;
+    kparams->n_srcs                = (uint8_t) n_srcs;
+    kparams->rows_per_buffer       = 1;
+    kparams->src1_row_size_aligned = is_scalar ? 0 : kparams->src0_row_size_aligned;
+    kparams->src1_size             = 0;
+    kparams->chunk_size            = chunk_size;
+    kparams->chunk_bytes           = chunk_bytes;
+    kparams->is_scalar             = is_scalar ? 1 : 0;
+
+    htp_binary_vtcm_layout_build(L, kparams, vtcm_size);
+    if (L->total_bytes == 0 || L->total_bytes > vtcm_size) {
+        return false;
+    }
+
+    kparams->vtcm_size = (uint32_t) L->total_bytes;
+    return true;
+}
+
 #endif
@@ -294,22 +294,21 @@ static int concat_transposed(struct htp_ops_context * octx, const struct htp_con

    const uint32_t n_threads = octx->n_threads;

-    // layout precomputed on host; kept for reference:
-    // struct htp_concat_transposed_vtcm_layout layout;
-    // htp_concat_transposed_vtcm_layout_build(&layout, octx->src[0]->ne[0], octx->src[1]->ne[0], type_size, n_threads);
+    struct htp_concat_transposed_vtcm_layout layout;
+    htp_concat_transposed_vtcm_layout_build(&layout, octx->src[0]->ne[0], octx->src[1]->ne[0], type_size, n_threads);

    uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;

    struct htp_concat_context cctx;
    cctx.octx                  = octx;
-    cctx.spad0_base            = vtcm_base;
-    cctx.spad1_base            = vtcm_base + n_threads * kparams->u.transposed.spad0_size_per_thread;
-    cctx.spad0_size_per_thread = kparams->u.transposed.spad0_size_per_thread;
-    cctx.spad1_size_per_thread = kparams->u.transposed.spad1_size_per_thread;
+    cctx.spad0_base            = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_spad0);
+    cctx.spad1_base            = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_spad1);
+    cctx.spad0_size_per_thread = layout.src0_spad_size_per_thread;
+    cctx.spad1_size_per_thread = layout.src1_spad_size_per_thread;
    cctx.row_start             = row_start;
    cctx.nrows                 = nrows;
    cctx.nplanes               = dst->ne[2] * dst->ne[3];
-    cctx.div_ne2               = init_fastdiv_values(dst->ne[2]);
+    cctx.div_ne2               = kparams->u.transposed.div_ne2;
    cctx.nrows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);

    work_queue_func_t worker_func = (type_size == 4) ? concat_2d_f32_transposed : concat_2d_f16_transposed;
@@ -2,6 +2,8 @@
 #define HTP_CONCAT_OPS_H

 #include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
 #include "gather-rows.h"
 #include <stdint.h>

@@ -13,8 +15,7 @@ enum htp_concat_kernel_type {
 };

 struct htp_concat_transposed_params {
-    uint32_t spad0_size_per_thread;
-    uint32_t spad1_size_per_thread;
+    struct fastdiv_values div_ne2;
 };

 struct htp_concat_kernel_params {
@@ -40,11 +41,13 @@ _Static_assert(sizeof(struct htp_concat_kernel_params) <= 128, "htp_concat_kerne
 struct htp_concat_transposed_vtcm_layout {
    uint32_t src0_spad_size_per_thread;
    uint32_t src1_spad_size_per_thread;
+    size_t   off_spad0;
+    size_t   off_spad1;
    uint32_t total_bytes;
 };

 static inline void htp_concat_transposed_vtcm_layout_build(
-    struct htp_concat_transposed_vtcm_layout * layout,
+    struct htp_concat_transposed_vtcm_layout * L,
    uint32_t src0_ne0,
    uint32_t src1_ne0,
    uint32_t type_size,
@@ -55,9 +58,13 @@ static inline void htp_concat_transposed_vtcm_layout_build(
    uint32_t src1_ne0_padded = hex_round_up(src1_ne0, block_i);
    uint32_t spad0_row_bytes = hex_round_up(src0_ne0 * type_size, 128) + src1_ne0_padded * type_size;

-    layout->src0_spad_size_per_thread = block_i * spad0_row_bytes;
-    layout->src1_spad_size_per_thread = src1_ne0_padded * spad1_stride;
-    layout->total_bytes = n_threads * (layout->src0_spad_size_per_thread + layout->src1_spad_size_per_thread);
+    L->src0_spad_size_per_thread = block_i * spad0_row_bytes;
+    L->src1_spad_size_per_thread = src1_ne0_padded * spad1_stride;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_spad0, (size_t) n_threads * L->src0_spad_size_per_thread);
+    VTCM_LAYOUT_ALLOC(off, off_spad1, (size_t) n_threads * L->src1_spad_size_per_thread);
+    L->total_bytes = (uint32_t) off;
 }

 #endif // HTP_CONCAT_OPS_H
@@ -32,6 +32,7 @@ struct htp_copy_context {
    uint32_t nelem;
    uint32_t elem_per_thread;

+    struct htp_copy_convert_vtcm_layout layout;
    uint8_t * vtcm_src0;
    uint8_t * vtcm_dst;
 };
@@ -176,13 +177,14 @@ static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, vo
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                   \
                                                                                             \
    const struct htp_copy_convert_params * cvt = &ct->kparams->u.convert;                    \
-    const uint32_t src0_buf_size = cvt->src0_buf_size;                                       \
-    const uint32_t dst_buf_size  = cvt->dst_buf_size;                                        \
-    const uint32_t src0_row_stride = cvt->src0_row_stride;                                   \
-    const uint32_t dst_row_stride  = cvt->dst_row_stride;                                    \
-    const uint32_t blk_rows        = cvt->blk_rows;                                          \
-    uint8_t * vtcm_src0_base = ct->vtcm_src0 + ith * cvt->spad0_size_per_thread;             \
-    uint8_t * vtcm_dst_base  = ct->vtcm_dst  + ith * cvt->spad1_size_per_thread;             \
+    const struct htp_copy_convert_vtcm_layout * L = &ct->layout;                             \
+    const uint32_t src0_buf_size   = L->src0_buf_size;                                       \
+    const uint32_t dst_buf_size    = L->dst_buf_size;                                        \
+    const uint32_t src0_row_stride = L->src0_row_stride;                                     \
+    const uint32_t dst_row_stride  = L->dst_row_stride;                                      \
+    const uint32_t blk_rows        = L->blk_rows;                                            \
+    uint8_t * vtcm_src0_base = ct->vtcm_src0 + ith * L->spad0_size_per_thread;               \
+    uint8_t * vtcm_dst_base  = ct->vtcm_dst  + ith * L->spad1_size_per_thread;               \
    const uint32_t src0_row_size = ne00 * ct->kparams->src0_type_size;                       \
    const uint32_t dst_row_size  = ne00 * ct->kparams->dst_type_size;                        \
                                                                                             \
@@ -254,17 +256,16 @@ static int cpy_scalar(struct htp_ops_context * octx, const struct htp_copy_kerne
    }

    dma_queue * dma_q = octx->ctx->dma[0];
-    dma_addr_t s_vtcm = (dma_addr_t)(uintptr_t) octx->ctx->vtcm_base;
-    dma_addr_t d_vtcm = s_vtcm + VLEN;
+    uint8_t * s_ptr = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, 0);
+    uint8_t * d_ptr = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, VLEN);
+    dma_addr_t s_vtcm = (dma_addr_t)(uintptr_t) s_ptr;
+    dma_addr_t d_vtcm = (dma_addr_t)(uintptr_t) d_ptr;
    const uint32_t s_size = kparams->src0_type_size;
    const uint32_t d_size = kparams->dst_type_size;

    dma_queue_push(dma_q, dma_make_data(s_vtcm, src0->data), s_size, s_size, s_size, 1);
    dma_queue_pop(dma_q);

-    uint8_t * s_ptr = (uint8_t *) octx->ctx->vtcm_base;
-    uint8_t * d_ptr = s_ptr + VLEN;
-
    const HVX_Vector v_src = hvx_vmem(s_ptr);
    HVX_Vector v_dst;

@@ -366,6 +367,8 @@ static int cpy_sameshape_convert(struct htp_ops_context * octx, const struct htp
    const struct htp_copy_convert_params * cvt = &kparams->u.convert;

    struct htp_copy_context ct;
+    htp_copy_convert_vtcm_layout_build(&ct.layout, src0->ne[0], kparams->src0_type_size, kparams->dst_type_size, cvt->blk_rows, n_threads);
+
    ct.octx                  = octx;
    ct.kparams               = kparams;
    ct.row_start             = row_start;
@@ -373,8 +376,8 @@ static int cpy_sameshape_convert(struct htp_ops_context * octx, const struct htp
    ct.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);

    uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
-    ct.vtcm_src0 = vtcm_base;
-    ct.vtcm_dst  = vtcm_base + (size_t) n_threads * cvt->spad0_size_per_thread;
+    ct.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, ct.layout.off_src0);
+    ct.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, ct.layout.off_dst);

    work_queue_func_t copy_fun = NULL;
    if (dst->type == HTP_TYPE_F16 && src0->type == HTP_TYPE_F32) {
@@ -3,6 +3,7 @@

 #include "hex-common.h"
 #include "hex-fastdiv.h"
+#include "htp-vtcm.h"
 #include "gather-rows.h"
 #include <stdint.h>

@@ -17,12 +18,6 @@ enum htp_copy_kernel_type {
 };

 struct htp_copy_convert_params {
-    uint32_t              src0_buf_size;
-    uint32_t              dst_buf_size;
-    uint32_t              spad0_size_per_thread;
-    uint32_t              spad1_size_per_thread;
-    uint32_t              src0_row_stride;  // row pitch in the VTCM buffers
-    uint32_t              dst_row_stride;
    uint32_t              blk_rows;         // rows moved per DMA descriptor
    struct fastdiv_values div_ne01;
    struct fastdiv_values div_ne02_ne01;
@@ -60,6 +55,9 @@ struct htp_copy_kernel_params {
 #define HTP_COPY_VTCM_BLK_BYTES (32 * 1024)

 struct htp_copy_convert_vtcm_layout {
+    size_t   total_bytes;
+    size_t   off_src0;
+    size_t   off_dst;
    uint32_t src0_row_stride;
    uint32_t dst_row_stride;
    uint32_t blk_rows;
@@ -67,11 +65,35 @@ struct htp_copy_convert_vtcm_layout {
    uint32_t dst_buf_size;
    uint32_t spad0_size_per_thread;
    uint32_t spad1_size_per_thread;
-    uint32_t total_bytes;
 };

-// ne01 bounds a block: rows of one block share the dim-1 stride.
 static inline void htp_copy_convert_vtcm_layout_build(
+    struct htp_copy_convert_vtcm_layout * L,
+    uint32_t ne00,
+    uint32_t src_type_size,
+    uint32_t dst_type_size,
+    uint32_t blk_rows,
+    uint32_t n_threads) {
+
+    const uint32_t src0_row_stride = hex_round_up(ne00 * src_type_size, 256);
+    const uint32_t dst_row_stride  = hex_round_up(ne00 * dst_type_size, 256);
+
+    L->src0_row_stride       = src0_row_stride;
+    L->dst_row_stride        = dst_row_stride;
+    L->blk_rows              = blk_rows;
+    L->src0_buf_size         = blk_rows * src0_row_stride;
+    L->dst_buf_size          = blk_rows * dst_row_stride;
+    L->spad0_size_per_thread = 2 * L->src0_buf_size;
+    L->spad1_size_per_thread = 2 * L->dst_buf_size;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, (size_t) n_threads * L->spad0_size_per_thread);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  (size_t) n_threads * L->spad1_size_per_thread);
+    L->total_bytes = off;
+}
+
+// ne01 bounds a block: rows of one block share the dim-1 stride.
+static inline void htp_copy_convert_solve_layout(
    struct htp_copy_convert_vtcm_layout * layout,
    uint32_t ne00,
    uint32_t ne01,
@@ -79,10 +101,10 @@ static inline void htp_copy_convert_vtcm_layout_build(
    uint32_t dst_type_size,
    uint32_t n_threads) {

-    layout->src0_row_stride = hex_round_up(ne00 * src_type_size, 256);
-    layout->dst_row_stride  = hex_round_up(ne00 * dst_type_size, 256);
+    const uint32_t src0_row_stride = hex_round_up(ne00 * src_type_size, 256);
+    const uint32_t dst_row_stride  = hex_round_up(ne00 * dst_type_size, 256);

-    const uint32_t max_row = layout->src0_row_stride > layout->dst_row_stride ? layout->src0_row_stride : layout->dst_row_stride;
+    const uint32_t max_row = src0_row_stride > dst_row_stride ? src0_row_stride : dst_row_stride;
    uint32_t blk_rows = HTP_COPY_VTCM_BLK_BYTES / max_row;
    if (blk_rows > ne01) {
        blk_rows = ne01;
@@ -90,13 +112,8 @@ static inline void htp_copy_convert_vtcm_layout_build(
    if (blk_rows == 0) {
        blk_rows = 1;
    }
-    layout->blk_rows = blk_rows;

-    layout->src0_buf_size = blk_rows * layout->src0_row_stride;
-    layout->dst_buf_size  = blk_rows * layout->dst_row_stride;
-    layout->spad0_size_per_thread = 2 * layout->src0_buf_size;
-    layout->spad1_size_per_thread = 2 * layout->dst_buf_size;
-    layout->total_bytes = n_threads * (layout->spad0_size_per_thread + layout->spad1_size_per_thread);
+    htp_copy_convert_vtcm_layout_build(layout, ne00, src_type_size, dst_type_size, blk_rows, n_threads);
 }

 #if defined(__cplusplus)
@@ -12,9 +12,11 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "htp-vtcm.h"
 #include "hvx-types.h"
 #include "hvx-utils.h"
 #include "dma-queue.h"
+#include "cumsum-ops.h"

 #define htp_cumsum_tensors_preamble                         \
    const struct htp_tensor * restrict src0 = octx->src[0]; \
@@ -89,8 +91,8 @@ static inline HVX_Vector hvx_splat_last_f32(HVX_Vector v) {
 }

 static inline void hvx_cumsum_row_f32(const float * restrict src, float * restrict dst, uint32_t n) {
-    const uint32_t nvec = n / VLEN_FP32;
-    const uint32_t nloe = n % VLEN_FP32;
+    const uint32_t nvec = n >> 5;
+    const uint32_t nloe = n & (VLEN_FP32 - 1);

    HVX_Vector carry = Q6_V_vsplat_R(0);

@@ -210,16 +212,19 @@ static void cumsum_thread_f32(unsigned int nth, unsigned int ith, void * data) {
 int op_cumsum_f32(struct htp_ops_context * octx) {
    const struct htp_tensor * src0 = octx->src[0];
    const struct htp_tensor * dst  = octx->dst;
+    const struct htp_cumsum_kernel_params * kparams =
+        (const struct htp_cumsum_kernel_params *) octx->kernel_params;

-    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);
+    const uint32_t total_rows   = kparams->total_rows;
+    const size_t   src_row_size = src0->nb[1];
+    const size_t   dst_row_size = dst->nb[1];

    uint32_t row_start = 0;
    uint32_t nrows     = total_rows;

    if (octx->ctx->mdev.count > 1) {
        uint32_t rows_per_chunk = 0;
-        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_data_row_size, &rows_per_chunk);
+        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;
@@ -229,43 +234,45 @@ int op_cumsum_f32(struct htp_ops_context * octx) {
        return HTP_STATUS_OK;
    }

+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
    const uint32_t n_threads = octx->n_threads;

-    const size_t src_row_size         = src0->nb[1];
-    const size_t dst_row_size         = dst->nb[1];
-    const size_t src_row_size_aligned = hex_round_up(src_row_size, VLEN);
-    const size_t dst_row_size_aligned = hex_round_up(dst_row_size, VLEN);
+    if (!kparams->use_dma && (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) {
+        return HTP_STATUS_NO_SUPPORT;
+    }

-    // 2 ping-pong buffers per thread for src and dst
-    const size_t spad_per_thread = 2 * (src_row_size_aligned + dst_row_size_aligned);
+    struct htp_cumsum_vtcm_layout layout = { 0 };
+    if (kparams->use_dma) {
+        htp_cumsum_vtcm_layout_build(&layout, src_row_size, dst_row_size, n_threads);

-    octx->src0_spad.size_per_thread = src_row_size_aligned * 2;
-    octx->dst_spad.size_per_thread  = dst_row_size_aligned * 2;
-
-    octx->src0_spad.size  = n_threads * octx->src0_spad.size_per_thread;
-    octx->dst_spad.size   = n_threads * octx->dst_spad.size_per_thread;
-
-    octx->src0_spad.data  = octx->ctx->vtcm_base;                        octx->src0_spad.src = NULL;
-    octx->dst_spad.data   = octx->src0_spad.data + octx->src0_spad.size; octx->dst_spad.src  = NULL;
+        octx->src0_spad.size_per_thread = layout.src0_bytes_per_thread;
+        octx->dst_spad.size_per_thread  = layout.dst_bytes_per_thread;
+        octx->src0_spad.size = n_threads * layout.src0_bytes_per_thread;
+        octx->dst_spad.size  = n_threads * layout.dst_bytes_per_thread;
+        octx->src0_spad.data = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_src0);
+        octx->dst_spad.data  = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_dst);
+        octx->src0_spad.src  = NULL;
+        octx->dst_spad.src   = NULL;
+    }

    struct htp_cumsum_context cctx = {
        .octx                 = octx,
        .src_row_size         = src_row_size,
        .dst_row_size         = dst_row_size,
-        .src_row_size_aligned = src_row_size_aligned,
-        .dst_row_size_aligned = dst_row_size_aligned,
+        .src_row_size_aligned = layout.src_row_size_aligned,
+        .dst_row_size_aligned = layout.dst_row_size_aligned,
        .rows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
        .total_rows           = nrows,
        .row_start            = row_start,
    };

-    if (octx->ctx->vtcm_size < spad_per_thread * n_threads) {
-        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 {
+    if (kparams->use_dma) {
        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32_dma, &cctx, n_threads);
+    } else {
+        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32, &cctx, n_threads);
    }

    return HTP_STATUS_OK;
@@ -0,0 +1,49 @@
+#ifndef HTP_CUMSUM_OPS_H
+#define HTP_CUMSUM_OPS_H
+
+#include "hex-common.h"
+#include "htp-vtcm.h"
+#include <stdint.h>
+
+struct htp_cumsum_vtcm_layout {
+    size_t total_bytes;
+    size_t off_src0;
+    size_t off_dst;
+    size_t src_row_size_aligned;
+    size_t dst_row_size_aligned;
+    size_t src0_bytes_per_thread;
+    size_t dst_bytes_per_thread;
+};
+
+static inline void htp_cumsum_vtcm_layout_build(
+    struct htp_cumsum_vtcm_layout * L,
+    size_t                          src_row_size,
+    size_t                          dst_row_size,
+    uint32_t                        n_threads
+) {
+    L->src_row_size_aligned  = hex_round_up(src_row_size, 128);
+    L->dst_row_size_aligned  = hex_round_up(dst_row_size, 128);
+    L->src0_bytes_per_thread = 2 * L->src_row_size_aligned;
+    L->dst_bytes_per_thread  = 2 * L->dst_row_size_aligned;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread * n_threads);
+
+    L->total_bytes = off;
+}
+
+struct htp_cumsum_kernel_params {
+    uint32_t n_threads;
+    uint32_t total_rows;
+    uint32_t use_dma;
+    uint32_t vtcm_size;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_cumsum_kernel_params) <= 128, "htp_cumsum_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_cumsum_kernel_params) <= 128, "htp_cumsum_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_CUMSUM_OPS_H
@@ -10,10 +10,12 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "htp-vtcm.h"
 #include "hvx-types.h"
 #include "hex-utils.h"
 #include "hvx-copy.h"
 #include "dma-queue.h"
+#include "diag-ops.h"

 #define htp_diag_tensors_preamble                           \
    const struct htp_tensor * restrict src0 = octx->src[0]; \
@@ -33,6 +35,7 @@

 struct htp_diag_context {
    struct htp_ops_context * octx;
+    const struct htp_diag_kernel_params * kparams;
    size_t          src_batch_size;
    size_t          dst_row_size;
    size_t          src_batch_size_aligned;
@@ -83,8 +86,8 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];

    for (uint32_t ib = ib0; ib < ib1; ib++) {
-        const uint32_t i3 = ib / ne02;
-        const uint32_t i2 = ib % ne02;
+        const uint32_t i3 = fastdiv(ib, &dctx->kparams->div_ne02);
+        const uint32_t i2 = ib - i3 * ne02;

        const dma_addr_t src_batch = src_data + i3 * nb03 + i2 * nb02;

@@ -134,8 +137,8 @@ static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) {
    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ib0);

    for (uint32_t ib = ib0; ib < ib1; ib++) {
-        const uint32_t i3 = ib / ne02;
-        const uint32_t i2 = ib % ne02;
+        const uint32_t i3 = fastdiv(ib, &dctx->kparams->div_ne02);
+        const uint32_t i2 = ib - i3 * ne02;

        const float * restrict src_batch = (const float *)(src_data + i3 * nb03 + i2 * nb02);

@@ -155,29 +158,17 @@ static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) {
 int op_diag_f32(struct htp_ops_context * octx) {
    const struct htp_tensor * src0 = octx->src[0];
    const struct htp_tensor * dst  = octx->dst;
+    const struct htp_diag_kernel_params * kparams =
+        (const struct htp_diag_kernel_params *) octx->kernel_params;

-    const uint32_t total_batches = src0->ne[2] * src0->ne[3];
-    const size_t dst_batch_size  = dst->ne[1] * dst->nb[1];
+    const uint32_t total_batches = kparams->total_batches;

    uint32_t batch_start = 0;
    uint32_t nbatches    = total_batches;

    if (octx->ctx->mdev.count > 1) {
-        bool can_split = htp_tensor_mdev_data_aligned(dst) && (dst->ne[0] == 1 || dst->nb[0] == sizeof(float)) && !htp_tensor_is_permuted(dst);
-        uint32_t batches_per_chunk = 1;
-        if (can_split) {
-            if (dst->ne[2] > 1 && (dst->nb[2] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0 &&
-                (dst->ne[3] <= 1 || (dst->nb[3] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0)) {
-                batches_per_chunk = 1;
-            } else if (dst->nb[2] == dst_batch_size &&
-                       (dst->ne[3] <= 1 || dst->nb[3] == dst->nb[2] * dst->ne[2])) {
-                batches_per_chunk = (dst_batch_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(dst_batch_size, HEX_L2_LINE_SIZE)) : 1;
-            } else {
-                can_split = false;
-            }
-        }
-
-        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_batches, can_split ? batches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        const bool can_split = (kparams->can_split != 0) && htp_tensor_mdev_data_aligned(dst);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_batches, can_split ? kparams->batches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
        batch_start = range.start;
        nbatches    = range.count;
    }
@@ -186,43 +177,49 @@ int op_diag_f32(struct htp_ops_context * octx) {
        return HTP_STATUS_OK;
    }

-    const uint32_t n_threads     = octx->n_threads;
+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nbatches))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }

-    const size_t src_batch_size         = src0->ne[0] * sizeof(float);
-    const size_t dst_row_size           = dst->ne[0] * sizeof(float);
-    const size_t src_batch_size_aligned = hex_round_up(src_batch_size, VLEN);
-    const size_t dst_row_size_aligned   = hex_round_up(dst_row_size, VLEN);
+    const uint32_t n_threads = octx->n_threads;

-    // 1 src buffer + 1 dst row buffer per thread
-    const size_t spad_per_thread = src_batch_size_aligned + dst_row_size_aligned;
+    if (!kparams->use_dma && (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) {
+        return HTP_STATUS_NO_SUPPORT;
+    }

-    octx->src0_spad.size_per_thread = src_batch_size_aligned;
-    octx->dst_spad.size_per_thread  = dst_row_size_aligned;
+    const size_t src_batch_size = src0->ne[0] * sizeof(float);
+    const size_t dst_row_size   = dst->ne[0] * sizeof(float);

-    octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread;
-    octx->dst_spad.size  = n_threads * octx->dst_spad.size_per_thread;
+    struct htp_diag_vtcm_layout layout = { 0 };
+    if (kparams->use_dma) {
+        htp_diag_vtcm_layout_build(&layout, src_batch_size, dst_row_size, n_threads);

-    octx->src0_spad.data = octx->ctx->vtcm_base;                        octx->src0_spad.src = NULL;
-    octx->dst_spad.data  = octx->src0_spad.data + octx->src0_spad.size; octx->dst_spad.src  = NULL;
+        octx->src0_spad.size_per_thread = layout.src0_bytes_per_thread;
+        octx->dst_spad.size_per_thread  = layout.dst_bytes_per_thread;
+        octx->src0_spad.size = n_threads * layout.src0_bytes_per_thread;
+        octx->dst_spad.size  = n_threads * layout.dst_bytes_per_thread;
+        octx->src0_spad.data = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_src0);
+        octx->dst_spad.data  = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_dst);
+        octx->src0_spad.src  = NULL;
+        octx->dst_spad.src   = NULL;
+    }

    struct htp_diag_context dctx = {
        .octx                   = octx,
+        .kparams                = kparams,
        .src_batch_size         = src_batch_size,
        .dst_row_size           = dst_row_size,
-        .src_batch_size_aligned = src_batch_size_aligned,
-        .dst_row_size_aligned   = dst_row_size_aligned,
+        .src_batch_size_aligned = layout.src_batch_size_aligned,
+        .dst_row_size_aligned   = layout.dst_row_size_aligned,
        .batches_per_thread     = fastdiv(nbatches + n_threads - 1, &octx->n_threads_div),
        .total_batches          = nbatches,
        .batch_start            = batch_start,
    };

-    if (octx->ctx->vtcm_size < spad_per_thread * n_threads) {
-        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 {
+    if (kparams->use_dma) {
        work_queue_run(octx->ctx->work_queue, diag_thread_f32_dma, &dctx, n_threads);
+    } else {
+        work_queue_run(octx->ctx->work_queue, diag_thread_f32, &dctx, n_threads);
    }

    return HTP_STATUS_OK;
@@ -0,0 +1,54 @@
+#ifndef HTP_DIAG_OPS_H
+#define HTP_DIAG_OPS_H
+
+#include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
+#include <stdint.h>
+
+struct htp_diag_vtcm_layout {
+    size_t total_bytes;
+    size_t off_src0;
+    size_t off_dst;
+    size_t src_batch_size_aligned;
+    size_t dst_row_size_aligned;
+    size_t src0_bytes_per_thread;
+    size_t dst_bytes_per_thread;
+};
+
+static inline void htp_diag_vtcm_layout_build(
+    struct htp_diag_vtcm_layout * L,
+    size_t                        src_batch_size,
+    size_t                        dst_row_size,
+    uint32_t                      n_threads
+) {
+    L->src_batch_size_aligned = hex_round_up(src_batch_size, 128);
+    L->dst_row_size_aligned   = hex_round_up(dst_row_size, 128);
+    L->src0_bytes_per_thread  = L->src_batch_size_aligned;
+    L->dst_bytes_per_thread   = L->dst_row_size_aligned;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread * n_threads);
+
+    L->total_bytes = off;
+}
+
+struct htp_diag_kernel_params {
+    uint32_t n_threads;
+    uint32_t total_batches;
+    uint32_t use_dma;
+    uint32_t can_split;
+    uint32_t batches_per_chunk;
+    uint32_t vtcm_size;
+
+    struct fastdiv_values div_ne02;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_diag_kernel_params) <= 128, "htp_diag_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_diag_kernel_params) <= 128, "htp_diag_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_DIAG_OPS_H
@@ -436,10 +436,10 @@ static inline void dma_sync_write(dma_queue * dma_q, dma_addr_t dst, const void
    }
 }

-#define DMA_CACHE_MAX_SIZE 256U
+#define DMA_CACHE_MAX_SIZE 128U

 // Fully assoc LRU cache
-typedef struct {
+typedef struct __attribute__((aligned(128))) {
    uint8_t *base;
    uint32_t line_size;
    uint32_t capacity;
@@ -485,7 +485,7 @@ static inline bool dma_cache_fa_push(dma_queue *q, dma_cache_fa *c, dma_addr_t s
 }

 // Direct mapped cache
-typedef struct {
+typedef struct __attribute__((aligned(128))) {
    uint8_t *base;
    uint32_t line_size;
    uint32_t capacity;
@@ -16,6 +16,7 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "fill-ops.h"

 // ggml op_params layout for FILL:
 //   op_params[0] (as float) - the scalar fill value
@@ -30,21 +31,19 @@
                                               \
    const uint32_t nb1 = dst->nb[1];           \
    const uint32_t nb2 = dst->nb[2];           \
-    const uint32_t nb3 = dst->nb[3];           \
-                                               \
-    const uint32_t nr = ne1 * ne2 * ne3;
+    const uint32_t nb3 = dst->nb[3];

 struct htp_fill_context {
    struct htp_ops_context * octx;
+    const struct htp_fill_kernel_params * kparams;
    uint32_t nrows_per_thread;
    uint32_t total_rows;  // ne1 * ne2 * ne3
    uint32_t row_start;
-    bool     opt_path;
    HVX_Vector splat_vec;
    uint32_t   elem_size;
 };

-static void fill_thread(unsigned int nth, unsigned int ith, void * data) {
+static void fill_thread_contiguous(unsigned int nth, unsigned int ith, void * data) {
    const struct htp_fill_context * fctx = (const struct htp_fill_context *) data;
    struct htp_ops_context        * octx = fctx->octx;
    fill_preamble;
@@ -60,45 +59,59 @@ static void fill_thread(unsigned int nth, unsigned int ith, void * data) {
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);

-    if (fctx->opt_path) {
-        // Opt path: tensor is fully contiguous, treat as flat array
-        const uint32_t elem_start = ir0 * ne0;
-        const uint32_t elem_end = ir1 * ne0;
-        uint8_t * dst_ptr = (uint8_t *) dst->data + elem_start * fctx->elem_size;
-        hvx_splat_u(dst_ptr, fctx->splat_vec, elem_end - elem_start, fctx->elem_size);
-    } else {
-        // Non-contiguous path: must respect strides
-        for (uint32_t ir = ir0; ir < ir1; ++ir) {
-            const uint32_t i1 = ir % ne1;
-            const uint32_t i2 = (ir / ne1) % ne2;
-            const uint32_t i3 = ir / (ne1 * ne2);
-            uint8_t * dst_ptr = (uint8_t *) dst->data + i1*nb1 + i2*nb2 + i3*nb3;
-            hvx_splat_u(dst_ptr, fctx->splat_vec, ne0, fctx->elem_size);
-        }
+    const uint32_t elem_start = ir0 * ne0;
+    const uint32_t elem_end   = ir1 * ne0;
+    uint8_t * dst_ptr = (uint8_t *) dst->data + elem_start * fctx->elem_size;
+    hvx_splat_u(dst_ptr, fctx->splat_vec, elem_end - elem_start, fctx->elem_size);
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir1);
+    FARF(HIGH, "fill-contig %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
+}
+
+static void fill_thread_strided(unsigned int nth, unsigned int ith, void * data) {
+    const struct htp_fill_context * fctx = (const struct htp_fill_context *) data;
+    struct htp_ops_context        * octx = fctx->octx;
+    const struct htp_fill_kernel_params * kparams = fctx->kparams;
+    fill_preamble;
+
+    // Parallelise over the flat row index spanning ne1*ne2*ne3
+    const uint32_t ir0 = fctx->row_start + fctx->nrows_per_thread * ith;
+    const uint32_t ir1 = MIN(ir0 + fctx->nrows_per_thread, fctx->row_start + fctx->total_rows);
+
+    if (ir0 >= ir1) {
+        return;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
+
+    for (uint32_t ir = ir0; ir < ir1; ++ir) {
+        const uint32_t i3  = fastdiv(ir, &kparams->div_ne1_ne2);
+        const uint32_t rem = ir - i3 * (ne1 * ne2);
+        const uint32_t i2  = fastdiv(rem, &kparams->div_ne1);
+        const uint32_t i1  = rem - i2 * ne1;
+        uint8_t * dst_ptr = (uint8_t *) dst->data + i1*nb1 + i2*nb2 + i3*nb3;
+        hvx_splat_u(dst_ptr, fctx->splat_vec, ne0, fctx->elem_size);
    }

    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir1);
-    FARF(HIGH, "fill %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
+    FARF(HIGH, "fill-strided %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
 }

 int op_fill(struct htp_ops_context * octx) {
-    fill_preamble;
-
-    if (dst->type != HTP_TYPE_F32 && dst->type != HTP_TYPE_F16) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
+    const struct htp_tensor * dst = octx->dst;
+    const struct htp_fill_kernel_params * kparams =
+        (const struct htp_fill_kernel_params *) octx->kernel_params;

    if (htp_tensor_is_extended(dst)) {
        return HTP_STATUS_NO_SUPPORT;
    }

    uint32_t row_start = 0;
-    uint32_t nrows     = nr;
+    uint32_t nrows     = kparams->total_rows;

    if (octx->ctx->mdev.count > 1) {
-        const uint32_t row_size = nb1;
-        const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
-        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(nr, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(kparams->total_rows, htp_tensor_mdev_data_aligned(dst) ? kparams->rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
        row_start = range.start;
        nrows     = range.count;
    }
@@ -107,40 +120,43 @@ int op_fill(struct htp_ops_context * octx) {
        return HTP_STATUS_OK;
    }

-    // nr = ne1*ne2*ne3 (flat row count across all outer dims); parallelise over it.
+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
    const uint32_t n_threads = octx->n_threads;

-    // Optimize if fully contiguous: skip stride arithmetic, treat as flat array
-    const bool opt_path = (nb2 == nb1 * ne1) && (nb3 == nb2 * ne2);
-
-    FARF(HIGH, "fill: (%ux%ux%ux%u) type=%u opt=%d\n",
-         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], dst->type, (int) opt_path);
+    FARF(HIGH, "fill: (%ux%ux%ux%u) type=%u ktype=%u\n",
+         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], dst->type, kparams->kernel_type);

    float val_f32 = 0.f;
    memcpy(&val_f32, &octx->op_params[0], sizeof(float));

    struct htp_fill_context fctx = {
        .octx             = octx,
+        .kparams          = kparams,
        .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
        .total_rows       = nrows,
        .row_start        = row_start,
-        .opt_path         = opt_path,
+        .elem_size        = kparams->elem_size,
    };

    switch (dst->type) {
    case HTP_TYPE_F32:
        fctx.splat_vec = hvx_vec_splat_f32(val_f32);
-        fctx.elem_size = sizeof(float);
        break;
    case HTP_TYPE_F16:
        fctx.splat_vec = hvx_vec_splat_f16((_Float16) val_f32);
-        fctx.elem_size = sizeof(_Float16);
        break;
    default:
        return HTP_STATUS_NO_SUPPORT;
    }

-    work_queue_run(octx->ctx->work_queue, fill_thread, &fctx, n_threads);
+    work_queue_func_t worker = (kparams->kernel_type == HTP_FILL_KERNEL_CONTIGUOUS)
+        ? fill_thread_contiguous
+        : fill_thread_strided;
+
+    work_queue_run(octx->ctx->work_queue, worker, &fctx, n_threads);

    return HTP_STATUS_OK;
 }
@@ -0,0 +1,29 @@
+#ifndef HTP_FILL_OPS_H
+#define HTP_FILL_OPS_H
+
+#include "hex-fastdiv.h"
+#include <stdint.h>
+
+enum htp_fill_kernel_type {
+    HTP_FILL_KERNEL_CONTIGUOUS = 0,
+    HTP_FILL_KERNEL_STRIDED    = 1,
+};
+
+struct htp_fill_kernel_params {
+    uint8_t  kernel_type;
+    uint32_t n_threads;
+    uint32_t total_rows;
+    uint32_t elem_size;
+    uint32_t rows_per_chunk;
+
+    struct fastdiv_values div_ne1;
+    struct fastdiv_values div_ne1_ne2;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_fill_kernel_params) <= 128, "htp_fill_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_fill_kernel_params) <= 128, "htp_fill_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_FILL_OPS_H
@@ -39,6 +39,26 @@
 // Must be multiple of 32
 #define FLASH_ATTN_BLOCK_SIZE (32 * 2)

+// Large scratch lives in ctx->ddr_spad to keep stack usage low.
+// Ops can run on the dspqueue callback thread, which has a small stack.
+struct hvx_fa_ddr_spad {
+    __fp16   slopes[512];
+    uint16_t active_blocks[FA_MAX_KV_BLOCKS];
+    uint16_t thread_active_blocks[HTP_MAX_NTHREADS][FA_MAX_KV_BLOCKS] __attribute__((aligned(128)));
+};
+
+struct hmx_fa_ddr_spad {
+    uint16_t active_blocks[FA_MAX_KV_BLOCKS];
+    __fp16   slopes[] __attribute__((aligned(128))); // [g_br]
+};
+
+static inline void * fa_get_ddr_spad(const struct htp_ops_context * octx, size_t size) {
+    if (!octx->ctx->ddr_spad_base || octx->ctx->ddr_spad_size < size) {
+        return NULL;
+    }
+    return octx->ctx->ddr_spad_base;
+}
+
 struct htp_fa_context {
    const struct htp_ops_context * octx;

@@ -61,7 +81,6 @@ struct htp_fa_context {
    uint32_t n_head_log2;
    float m0;
    float m1;
-    __fp16 slopes[512];

    uint32_t n_blocks;
    uint32_t n_active_blocks;
@@ -90,7 +109,9 @@ struct htp_fa_context {
    uint8_t *  spad_m;
    uint8_t *  spad_a;
    float *    spad_sinks;
-    uint16_t   active_blocks[FA_MAX_KV_BLOCKS];
+
+    struct hvx_fa_ddr_spad * ddr;
+    dma_cache_dm *           m_cache; // [n_threads]

    const struct htp_tensor * k;
    const struct htp_tensor * v;
@@ -156,6 +177,7 @@ struct hmx_fa_context {
    __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
+    __fp16 *     ddr_slopes;           // ALiBi slopes staging buffer in DDR [g_br]
    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;
@@ -164,7 +186,7 @@ struct hmx_fa_context {
    size_t       col_vec_bytes;
    size_t       d_tile_bytes;
    bool         mask_broadcast;       // true when mask->ne[2] == 1 (head-independent, single 2D DMA)
-    dma_cache_fa m_cache;
+    dma_cache_fa * m_cache;
 };

 // Scan mask rows [q_start, q_start + n_rows_q) and build compact list of KV blocks with any entry != -inf.
@@ -308,11 +330,11 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
    uint8_t * spad_m = factx->spad_m + (mask ? factx->size_m_block * HVX_FA_DMA_CACHE_SIZE : 0) * ith;
    uint8_t * spad_a = factx->spad_a + factx->size_vkq_acc * ith;

-    dma_cache_dm m_cache;
-    dma_cache_dm_init(&m_cache, spad_m, factx->size_m_block, HVX_FA_DMA_CACHE_SIZE);
+    dma_cache_dm * m_cache = &factx->m_cache[ith];
+    dma_cache_dm_init(m_cache, spad_m, factx->size_m_block, HVX_FA_DMA_CACHE_SIZE);

-    uint16_t thread_active_blocks[FA_MAX_KV_BLOCKS];
-    const uint16_t * active_blocks = factx->shared_active_blocks ? factx->active_blocks : thread_active_blocks;
+    uint16_t * thread_active_blocks = factx->ddr->thread_active_blocks[ith];
+    const uint16_t * active_blocks = factx->shared_active_blocks ? factx->ddr->active_blocks : thread_active_blocks;
    uint32_t n_active_blocks       = factx->n_active_blocks;
    dma_addr_t prev_mp_base        = 0;

@@ -385,7 +407,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
            const uint32_t r_iq1 = (r - r_iq3*neq2*neq1 - r_iq2 * neq1);

            heads[g] = r_iq2;
-            const __fp16 slope = factx->slopes[r_iq2];
+            const __fp16 slope = factx->ddr->slopes[r_iq2];
            slope_vecs[g] = hvx_vec_splat_f16(slope);

            S_vec[g] = hvx_vec_splat_f32(0.0f);
@@ -422,7 +444,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
            // Mask
            if (mask) {
                const dma_addr_t m_src = mp_base + ic_start * sizeof(__fp16);
-                dma_cache_dm_push(dma_q, &m_cache, ib, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
+                dma_cache_dm_push(dma_q, m_cache, ib, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
            }
        }

@@ -682,7 +704,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                    // Mask
                    if (mask) {
                        const dma_addr_t m_src = mp_base + next_ic_start * sizeof(__fp16);
-                        dma_cache_dm_push(dma_q, &m_cache, next_ib, m_src, next_block_size * 2, next_block_size * 2, next_block_size * 2, 1);
+                        dma_cache_dm_push(dma_q, m_cache, next_ib, m_src, next_block_size * 2, next_block_size * 2, next_block_size * 2, 1);
                    }
                }
            } // end for g
@@ -1887,7 +1909,7 @@ static __attribute__((noinline)) void fa_compute_slopes(
    const float    m0          = factx->m0;
    const float    m1          = factx->m1;

-    __fp16 temp_slopes[512] __attribute__((aligned(128)));
+    __fp16 * temp_slopes = factx->ddr_slopes;
    if (G <= 32) {
        // Fast path: Compute G unique slope values in vector registers
        HVX_Vector v_val = hvx_alibi_slopes(kv_head, G, n_head_log2, m0, m1);
@@ -1905,14 +1927,13 @@ static __attribute__((noinline)) void fa_compute_slopes(
        }
    }

-    // Allocate stack buffer to avoid scalar writes to VTCM (which generates L2 misses)
-    __fp16 local_slopes[n_rows_g] __attribute__((aligned(128)));
-    for (size_t r = 0; r < n_rows_g; ++r) {
-        local_slopes[r] = temp_slopes[fastmodulo(r, G, &factx->div_G)];
+    // Replicate the G values in DDR to avoid scalar writes to VTCM (which generates L2 misses)
+    for (size_t r = G; r < n_rows_g; ++r) {
+        temp_slopes[r] = temp_slopes[r - G];
    }

    // Copy to VTCM slopes using HVX block copy (both are aligned to 128 bytes)
-    hvx_copy_f16_aa((uint8_t *)slopes, (const uint8_t *)local_slopes, n_rows_g);
+    hvx_copy_f16_aa((uint8_t *)slopes, (const uint8_t *)temp_slopes, n_rows_g);
 }

 static void fa_push_mask_dma_gqa(
@@ -1958,7 +1979,7 @@ static inline void fa_prefetch_block(dma_queue * dma_q, const struct htp_tensor
    if (mask) {
        if (__builtin_expect(factx->mask_broadcast, true)) {
            const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + prefetch_start * sizeof(__fp16);
-            dma_cache_fa_push(dma_q, &factx->m_cache, ms_src, m_line_bytes, mask->nb[1], prefetch_rows * sizeof(__fp16), n_rows_q);
+            dma_cache_fa_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_q, mask, q_start, im3, prefetch_start, kv_head, G, m_line_bytes, prefetch_rows, n_rows_q, factx);
        }
@@ -2151,7 +2172,14 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {

    const size_t m_line_bytes = L.m_line_bytes;  // used by the mask DMAs in the KV loop

-    dma_cache_fa_init(&factx.m_cache, (uint8_t *) factx.vtcm_mask_buf, L.m_buf_slot_bytes, HMX_FA_DMA_CACHE_SIZE);
+    struct hmx_fa_ddr_spad * ddr = (struct hmx_fa_ddr_spad *) fa_get_ddr_spad(octx, sizeof(struct hmx_fa_ddr_spad) + factx.g_br * sizeof(__fp16));
+    if (!ddr) {
+        return HTP_STATUS_INTERNAL_ERR;
+    }
+
+    factx.ddr_slopes = ddr->slopes;
+    factx.m_cache    = &octx->dma_cache_fa[0];
+    dma_cache_fa_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
@@ -2210,7 +2238,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
            const size_t   g_br_actual = hex_align_up(n_rows_g, HMX_FP16_TILE_N_ROWS);
            const size_t   n_row_tiles = g_br_actual / HMX_FP16_TILE_N_ROWS;

-            uint16_t active_blocks[FA_MAX_KV_BLOCKS];
+            uint16_t * active_blocks = ddr->active_blocks;
            uint32_t n_active_blocks = factx.n_kv_blocks;
            if (mask && factx.mask_broadcast && factx.n_kv_blocks > 1) {
                const size_t scratch_bytes = factx.k_tile_bytes * (factx.pipeline ? 2 : 1);
@@ -2252,7 +2280,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                        if (factx.pipeline && mask) {
                            if (__builtin_expect(factx.mask_broadcast, true)) {
                                const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + b0_start * sizeof(__fp16);
-                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
+                                dma_cache_fa_push(dma_q, factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
                            } else {
                                fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, b0_start, kv_head, G, m_line_bytes, b0_rows, n_rows_q, &factx);
                            }
@@ -2453,7 +2481,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                        if (mask) {
                            if (__builtin_expect(factx.mask_broadcast, true)) {
                                const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
-                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
+                                dma_cache_fa_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_q, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx);
                            }
@@ -2577,7 +2605,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                        if (factx.pipeline && mask) {
                            if (__builtin_expect(factx.mask_broadcast, true)) {
                                const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + b0_start * sizeof(__fp16);
-                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
+                                dma_cache_fa_push(dma_q, factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
                            } else {
                                fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, b0_start, next_kv_head, G, m_line_bytes, b0_rows, n_rows_q, &factx);
                            }
@@ -2680,8 +2708,15 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
    if (n_head > 512) {
        return HTP_STATUS_NO_SUPPORT;
    }
+
+    factx.ddr = (struct hvx_fa_ddr_spad *) fa_get_ddr_spad(octx, sizeof(struct hvx_fa_ddr_spad));
+    if (!factx.ddr) {
+        return HTP_STATUS_INTERNAL_ERR;
+    }
+    factx.m_cache = octx->dma_cache_dm;
+
    for (uint32_t h = 0; h < n_head; ++h) {
-        factx.slopes[h] = (__fp16) ((kparams->max_bias > 0.0f) ? alibi_slope(h, factx.n_head_log2, factx.m0, factx.m1) : 1.0f);
+        factx.ddr->slopes[h] = (__fp16) ((kparams->max_bias > 0.0f) ? alibi_slope(h, factx.n_head_log2, factx.m0, factx.m1) : 1.0f);
    }

    // total rows in q
@@ -2754,7 +2789,7 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
    if (has_mask && neq1 == 1 && neq3 == 1 && mask->ne[2] == 1 && factx.n_blocks > 1 && factx.n_blocks <= FA_MAX_KV_BLOCKS) {
        factx.n_active_blocks = fa_scan_active_blocks(
            octx->ctx->dma[0], mask, 0, 1, 0, 0, k->ne[1], FLASH_ATTN_BLOCK_SIZE,
-            factx.n_blocks, factx.spad_k, factx.size_k_block * 2, factx.active_blocks);
+            factx.n_blocks, factx.spad_k, factx.size_k_block * 2, factx.ddr->active_blocks);
        factx.shared_active_blocks = true;
    }

@@ -35,7 +35,6 @@ enum htp_fa_kernel_type {
 struct htp_fa_kernel_params {
    uint8_t  kernel_type;        // enum htp_fa_kernel_type
    uint8_t  head_split;         // 1 = partition by KV heads in multicore, 0 = token partition
-    uint8_t  flags;              // reserved
    uint8_t  n_threads;          // Number of threads to run

    // Common parameters
@@ -49,9 +48,6 @@ struct htp_fa_kernel_params {
    float    logit_softcap;
    uint32_t vtcm_size;

-    uint32_t qrows;
-    uint32_t qrows_per_thread;
-    uint32_t qrow_start;
    float    m0;
    float    m1;
    uint32_t n_head_log2;
@@ -67,8 +63,6 @@ struct htp_fa_kernel_params {
    union {
        struct {
            uint32_t g_br;
-            uint32_t row_buf_stride;
-            uint32_t mask_buf_row_stride;
            int32_t  mask_broadcast;
            int32_t  pipeline;
            struct fastdiv_values div_G;
@@ -85,6 +79,8 @@ struct htp_fa_kernel_params {

 #if defined(__cplusplus)
 static_assert(sizeof(struct htp_fa_kernel_params) <= 128, "htp_fa_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_fa_kernel_params) <= 128, "htp_fa_kernel_params is too large for kernel_params blob");
 #endif

 // VTCM region layout for the HMX flash-attention kernel.
@@ -2270,46 +2270,9 @@ static int gated_delta_net_f32_hmx_chunked(
 }

 int op_gated_delta_net(struct htp_ops_context * octx) {
-    const struct htp_tensor * q     = octx->src[0];
-    const struct htp_tensor * k     = octx->src[1];
-    const struct htp_tensor * v     = octx->src[2];
-    const struct htp_tensor * g     = octx->src[3];
-    const struct htp_tensor * beta  = octx->src[4];
-    const struct htp_tensor * state = octx->src[5];
-    const struct htp_tensor * dst   = octx->dst;
+    const struct htp_tensor * dst = octx->dst;

-    if (q->type != HTP_TYPE_F32 || k->type != HTP_TYPE_F32 || v->type != HTP_TYPE_F32 ||
-        g->type != HTP_TYPE_F32 || beta->type != HTP_TYPE_F32 || state->type != HTP_TYPE_F32 ||
-        dst->type != HTP_TYPE_F32) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    const uint32_t S_v      = v->ne[0];
-    const uint32_t H        = v->ne[1];
-    const uint32_t n_tokens = v->ne[2];
-    const uint32_t n_seqs   = v->ne[3];
-    const uint32_t K        = octx->op_params[0];
-
-    if (S_v == 0 || S_v > HTP_GDN_MAX_SV || H == 0 || n_tokens == 0 || n_seqs == 0) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    if ((g->ne[0] != 1 && g->ne[0] != S_v) || beta->ne[0] != 1) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    if (q->ne[0] != S_v || k->ne[0] != S_v || q->ne[1] == 0 || k->ne[1] == 0 ||
-        q->ne[2] != n_tokens || k->ne[2] != n_tokens || q->ne[3] == 0 || k->ne[3] == 0 ||
-        (n_seqs % q->ne[3]) != 0 || (n_seqs % k->ne[3]) != 0) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    // state holds s0 only: [S_v, S_v, H, n_seqs]
-    if (state->ne[0] != S_v || state->ne[1] != S_v || state->ne[2] != H || state->ne[3] != n_seqs) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    if (dst->ne[0] != S_v * H || dst->ne[1] != n_tokens * n_seqs + S_v * n_seqs * K) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    for (int i = 0; i < 5; i++) {
+    for (int i = 0; i < 6; i++) {
        if (htp_tensor_is_extended(octx->src[i])) {
            return HTP_STATUS_NO_SUPPORT;
        }
@@ -2325,68 +2288,12 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
    }

    const struct htp_gdn_kernel_params * kparams = (const struct htp_gdn_kernel_params *) octx->kernel_params;
-    struct htp_gdn_kernel_params kparams_local;
-    if (!kparams || kparams->S_v == 0) {
-        const uint32_t rq3 = n_seqs / q->ne[3];
-        const uint32_t rk3 = n_seqs / k->ne[3];
-        const uint32_t total_rows = H * n_seqs;
-        uint32_t n_threads = (total_rows < octx->n_threads) ? total_rows : octx->n_threads;
-        if (n_threads == 0) {
-            n_threads = 1;
-        }
-
-        memset(&kparams_local, 0, sizeof(kparams_local));
-        kparams_local.n_threads           = n_threads;
-        kparams_local.S_v                 = S_v;
-        kparams_local.H                   = H;
-        kparams_local.n_tokens            = n_tokens;
-        kparams_local.n_seqs              = n_seqs;
-        kparams_local.K                   = K;
-        kparams_local.total_rows          = total_rows;
-        kparams_local.rows_per_thread     = (total_rows + n_threads - 1) / n_threads;
-        const bool can_use_hmx = (octx->ctx->hmx_enabled) &&
-                                 (S_v % 64 == 0) &&
-                                 (n_tokens >= HTP_GDN_MIN_TOKENS) &&
-                                 (g->ne[0] == 1) &&
-                                 (K == 1);
-
-        struct htp_gdn_hmx_vtcm_layout hmx_layout_local;
-        struct htp_gdn_vtcm_layout hvx_layout_local;
-        uint32_t n_heads_batch = 1;
-
-        if (can_use_hmx && htp_gdn_hmx_solve_layout(&hmx_layout_local, S_v, HTP_GDN_CHUNK_SIZE, total_rows, octx->ctx->vtcm_size, n_threads, true, &n_heads_batch)) {
-            kparams_local.kernel_type     = HTP_GDN_KERNEL_HMX_CHUNKED;
-            kparams_local.pipeline        = hmx_layout_local.pipeline ? 1 : 0;
-            kparams_local.chunk_size      = HTP_GDN_CHUNK_SIZE;
-            kparams_local.n_chunks        = (n_tokens + HTP_GDN_CHUNK_SIZE - 1) / HTP_GDN_CHUNK_SIZE;
-            kparams_local.n_heads_batch   = (uint16_t) n_heads_batch;
-            kparams_local.vtcm_size       = (uint32_t) hmx_layout_local.total_bytes;
-            kparams_local.state_aligned   = (uint32_t) hmx_layout_local.state_f32_bytes;
-            kparams_local.vtcm_per_thread = (uint32_t) (hmx_layout_local.total_bytes / (n_threads > 0 ? n_threads : 1));
-        } else {
-            htp_gdn_vtcm_layout_build(&hvx_layout_local, S_v, n_threads);
-            kparams_local.kernel_type     = HTP_GDN_KERNEL_HVX_RECURRENT;
-            kparams_local.pipeline        = 0;
-            kparams_local.n_heads_batch   = 1;
-            kparams_local.state_aligned   = (uint32_t) hvx_layout_local.state_aligned;
-            kparams_local.vtcm_per_thread = (uint32_t) hvx_layout_local.bytes_per_thread;
-            kparams_local.vtcm_size       = (uint32_t) hvx_layout_local.total_bytes;
-        }
-        kparams_local.kda                 = (g->ne[0] == S_v) ? 1 : 0;
-        kparams_local.scale               = 1.0f / sqrtf((float) S_v);
-        kparams_local.state_seq_stride    = (uint32_t) (state->nb[3] / sizeof(float));
-        kparams_local.state_size_per_snap = S_v * S_v * H * n_seqs;
-
-        kparams_local.div_H         = init_fastdiv_values(H);
-        kparams_local.div_q1        = init_fastdiv_values(q->ne[1]);
-        kparams_local.div_k1        = init_fastdiv_values(k->ne[1]);
-        kparams_local.div_rq3       = init_fastdiv_values(rq3);
-        kparams_local.div_rk3       = init_fastdiv_values(rk3);
-        kparams_local.div_n_threads = init_fastdiv_values(n_threads);
-
-        kparams = &kparams_local;
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
    }

+    const uint32_t S_v        = kparams->S_v;
+    const uint32_t n_tokens   = kparams->n_tokens;
    const uint32_t total_rows = kparams->total_rows;
    uint32_t row_start = 0;
    uint32_t nrows     = total_rows;
@@ -2416,16 +2323,16 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
        return gated_delta_net_f32_hmx_chunked(octx, kparams, row_start, nrows);
    }

-    const uint32_t n_threads = (nrows < kparams->n_threads) ? nrows : kparams->n_threads;
+    const uint32_t n_threads = (nrows < octx->n_threads) ? nrows : octx->n_threads;

    struct htp_gdn_context gctx;
    gctx.octx      = octx;
    gctx.kparams   = kparams;
    gctx.row_start = row_start;
    gctx.nrows     = nrows;
-    gctx.vtcm_base = octx->ctx->vtcm_base;

    htp_gdn_vtcm_layout_build(&gctx.layout, S_v, n_threads);
+    gctx.vtcm_base = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, gctx.layout.off_state);

    if (gctx.layout.total_bytes > octx->ctx->vtcm_size) {
        return HTP_STATUS_VTCM_TOO_SMALL;
@@ -2433,10 +2340,10 @@ int op_gated_delta_net(struct htp_ops_context * octx) {

    FARF(HIGH, "gated-delta-net-f32: q(%ux%ux%ux%u) k(%ux%ux%ux%u) v(%ux%ux%ux%u) state(%ux%ux%ux%u) -> (%ux%ux%ux%u) : "
         "vtcm-size %zu n_threads %u\n",
-         q->ne[0], q->ne[1], q->ne[2], q->ne[3],
-         k->ne[0], k->ne[1], k->ne[2], k->ne[3],
-         v->ne[0], v->ne[1], v->ne[2], v->ne[3],
-         state->ne[0], state->ne[1], state->ne[2], state->ne[3],
+         octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3],
+         octx->src[1]->ne[0], octx->src[1]->ne[1], octx->src[1]->ne[2], octx->src[1]->ne[3],
+         octx->src[2]->ne[0], octx->src[2]->ne[1], octx->src[2]->ne[2], octx->src[2]->ne[3],
+         octx->src[5]->ne[0], octx->src[5]->ne[1], octx->src[5]->ne[2], octx->src[5]->ne[3],
         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
         gctx.layout.total_bytes, n_threads);

@@ -37,13 +37,7 @@ struct htp_gdn_kernel_params {
    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;
@@ -55,7 +49,6 @@ struct htp_gdn_kernel_params {
    struct fastdiv_values div_k1;
    struct fastdiv_values div_rq3;
    struct fastdiv_values div_rk3;
-    struct fastdiv_values div_n_threads;
 };

 #if defined(__cplusplus)
@@ -67,18 +60,21 @@ _Static_assert(sizeof(struct htp_gdn_kernel_params) <= 128, "htp_gdn_kernel_para
 struct htp_gdn_vtcm_layout {
    size_t state_aligned;
    size_t bytes_per_thread;
+    size_t off_state;
    size_t total_bytes;
 };

 static inline void htp_gdn_vtcm_layout_build(
-    struct htp_gdn_vtcm_layout * layout,
+    struct htp_gdn_vtcm_layout * L,
    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;
+    L->state_aligned = hex_round_up(state_bytes, 128);
+    L->bytes_per_thread = 2 * L->state_aligned;
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_state, L->bytes_per_thread * n_threads);
+    L->total_bytes = off;
 }

 struct htp_gdn_hmx_vtcm_layout {
@@ -206,14 +206,22 @@ static void gather_rows_thread_f32(unsigned int nth, unsigned int ith, void * da

    struct htp_thread_trace * tr = &task->octx->ctx->trace[ith];

-    uint8_t * base   = task->octx->ctx->vtcm_base + ith * kparams->slice;
-    HVX_Vector * tab = (HVX_Vector *) base;
-    HVX_Vector * inc = tab + kparams->vg;
-    HVX_Vector * cur = inc + kparams->vg;
-    uint8_t * region_base = (uint8_t *) (cur + kparams->vg);
+    struct htp_gather_rows_vtcm_layout layout;
+    htp_gather_rows_vtcm_layout_build(&layout, kparams->vg, kparams->region_size, kparams->nr_max, kparams->ne, 4);

-    uint8_t * region[2] = { region_base, region_base + kparams->buf_stride };
-    uint8_t * out[2]    = { region[0] + kparams->out_off, region[1] + kparams->out_off };
+    uint8_t * base   = VTCM_LAYOUT_PTR(uint8_t, task->octx->ctx->vtcm_base, (size_t) ith * kparams->slice);
+    HVX_Vector * tab = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_tab);
+    HVX_Vector * inc = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_inc);
+    HVX_Vector * cur = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_cur);
+
+    uint8_t * region[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region1),
+    };
+    uint8_t * out[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out1),
+    };

    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, (uint16_t) r_beg);
    gather_rows_init_tables_f32(tab, inc, kparams, a, b);
@@ -279,14 +287,22 @@ static void gather_rows_thread_f16(unsigned int nth, unsigned int ith, void * da

    struct htp_thread_trace * tr = &task->octx->ctx->trace[ith];

-    uint8_t * base   = task->octx->ctx->vtcm_base + ith * kparams->slice;
-    HVX_Vector * tab = (HVX_Vector *) base;
-    HVX_Vector * inc = tab + kparams->vg;
-    HVX_Vector * cur = inc + kparams->vg;
-    uint8_t * region_base = (uint8_t *) (cur + kparams->vg);
+    struct htp_gather_rows_vtcm_layout layout;
+    htp_gather_rows_vtcm_layout_build(&layout, kparams->vg, kparams->region_size, kparams->nr_max, kparams->ne, 2);

-    uint8_t * region[2] = { region_base, region_base + kparams->buf_stride };
-    uint8_t * out[2]    = { region[0] + kparams->out_off, region[1] + kparams->out_off };
+    uint8_t * base   = VTCM_LAYOUT_PTR(uint8_t, task->octx->ctx->vtcm_base, (size_t) ith * kparams->slice);
+    HVX_Vector * tab = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_tab);
+    HVX_Vector * inc = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_inc);
+    HVX_Vector * cur = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_cur);
+
+    uint8_t * region[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region1),
+    };
+    uint8_t * out[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out1),
+    };

    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, (uint16_t) r_beg);
    gather_rows_init_tables_f16(tab, inc, kparams, a, b);
@@ -2,6 +2,7 @@
 #define HTP_GATHER_ROWS_H

 #include "hex-common.h"
+#include "htp-vtcm.h"
 #include <stdint.h>
 #include <stdbool.h>

@@ -19,8 +20,6 @@ struct htp_gather_rows_params {
    uint32_t v1;
    uint32_t span1;
    uint32_t region_size;
-    uint32_t out_off;
-    uint32_t buf_stride;
    uint32_t slice;
    uint32_t total_rows;
    uint32_t rows_per_thread;
@@ -28,6 +27,40 @@ struct htp_gather_rows_params {
    uint8_t  pad[3];
 };

+struct htp_gather_rows_vtcm_layout {
+    size_t off_tab;
+    size_t off_inc;
+    size_t off_cur;
+    size_t off_region0;
+    size_t off_out0;
+    size_t off_region1;
+    size_t off_out1;
+    size_t total_per_thread;
+};
+
+static inline void htp_gather_rows_vtcm_layout_build(
+    struct htp_gather_rows_vtcm_layout * L,
+    uint32_t vg,
+    uint32_t region_size,
+    uint32_t nr_max,
+    uint32_t ne,
+    uint32_t elem_size) {
+
+    const size_t vec_bytes      = (size_t) vg * 128;
+    const size_t region_aligned = hex_round_up(region_size, 128);
+    const size_t out_aligned    = hex_round_up(region_aligned + (size_t) nr_max * ne * elem_size, 128) - region_aligned;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_tab,     vec_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_inc,     vec_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_cur,     vec_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_region0, region_aligned);
+    VTCM_LAYOUT_ALLOC(off, off_out0,    out_aligned);
+    VTCM_LAYOUT_ALLOC(off, off_region1, region_aligned);
+    VTCM_LAYOUT_ALLOC(off, off_out1,    out_aligned);
+    L->total_per_thread = off;
+}
+
 static inline bool htp_gather_rows_solve_layout(
    struct htp_gather_rows_params * kparams,
    uint32_t na, uint32_t a_nb1,
@@ -80,8 +113,6 @@ static inline bool htp_gather_rows_solve_layout(
    uint32_t v1          = 0;
    uint32_t span1       = 0;
    uint32_t region_size = 0;
-    uint32_t out_off     = 0;
-    uint32_t buf_stride  = 0;

    while (nr_max > 0) {
        v1 = hex_round_up((nr_max - 1) * a_nb1 + na * elem_size, 128);
@@ -98,10 +129,10 @@ static inline bool htp_gather_rows_solve_layout(
            region_size = v1;
        }

-        out_off    = hex_round_up(region_size, 128);
-        buf_stride = hex_round_up(out_off + nr_max * ne * elem_size, 128);
+        struct htp_gather_rows_vtcm_layout layout;
+        htp_gather_rows_vtcm_layout_build(&layout, vg, region_size, nr_max, ne, elem_size);

-        if ((elem_size != 2 || region_size <= 32768) && fixed + 2 * buf_stride <= slice) {
+        if ((elem_size != 2 || region_size <= 32768) && layout.total_per_thread <= slice) {
            break;
        }

@@ -121,8 +152,6 @@ static inline bool htp_gather_rows_solve_layout(
    kparams->v1              = v1;
    kparams->span1           = span1;
    kparams->region_size     = region_size;
-    kparams->out_off         = out_off;
-    kparams->buf_stride      = buf_stride;
    kparams->slice           = slice;
    kparams->total_rows      = total_rows;
    kparams->rows_per_thread = rows_per_thread;
@@ -404,7 +404,7 @@ int op_get_rows(struct htp_ops_context * octx) {

    if (octx->ctx->mdev.count > 1) {
        uint32_t tasks_per_chunk = 1;
-        htp_tensor_mdev_rows_per_chunk(dst, dst_row_size / dst->ne[0], dst_row_size, &tasks_per_chunk);
+        htp_tensor_mdev_rows_per_chunk(dst, (uint32_t) dst->nb[0], dst_row_size, &tasks_per_chunk);
        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_tasks, tasks_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
        task_start = range.start;
        tasks      = range.count;
@@ -414,7 +414,7 @@ int op_get_rows(struct htp_ops_context * octx) {
        return HTP_STATUS_OK;
    }

-    if (!htp_ops_context_set_n_threads(octx, (uint32_t) kparams->n_threads)) {
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
        return HTP_STATUS_INVAL_PARAMS;
    }

@@ -428,7 +428,7 @@ int op_get_rows(struct htp_ops_context * octx) {
    grctx.compute_tiled = NULL;
    grctx.task_start = task_start;
    grctx.tasks = tasks;
-    grctx.tasks_per_thread = octx->ctx->mdev.count == 1 ? kparams->tasks_per_thread : fastdiv(tasks + n_threads - 1, &octx->n_threads_div);
+    grctx.tasks_per_thread = fastdiv(tasks + n_threads - 1, &octx->n_threads_div);
    grctx.tile_size   = 0;
    grctx.tile_stride = 0;
    grctx.src0_row_size = src0_row_size;
@@ -15,11 +15,10 @@ enum htp_get_rows_kernel_type {
 };

 struct htp_get_rows_kernel_params {
-    int32_t  n_threads;
-    int32_t  kernel_type;
-    int32_t  total_tasks;
-    int32_t  tasks_per_thread;
-    int32_t  vtcm_size;
+    uint8_t  kernel_type;
+    uint32_t n_threads;
+    uint32_t total_tasks;
+    uint32_t vtcm_size;

    // Fastdiv helpers
    struct fastdiv_values div_ne10;
@@ -17,11 +17,12 @@
 #include "hex-common.h"
 #include "htp-tensor.h"
 #include "htp-vtcm.h"
+#include "glu-ops.h"

-#define htp_act_preamble                                 \
-    const struct htp_tensor * src0 = actx->octx->src[0]; \
-    const struct htp_tensor * src1 = actx->octx->src[1]; \
-    const struct htp_tensor * dst  = actx->octx->dst;    \
+#define htp_glu_preamble                                 \
+    const struct htp_tensor * src0 = gctx->octx->src[0]; \
+    const struct htp_tensor * src1 = gctx->octx->src[1]; \
+    const struct htp_tensor * dst  = gctx->octx->dst;    \
                                                         \
    const uint32_t ne00 = src0->ne[0];                   \
    const uint32_t ne01 = src0->ne[1];                   \
@@ -53,21 +54,21 @@
    const uint32_t nb2 = dst->nb[2];                     \
    const uint32_t nb3 = dst->nb[3];

-struct htp_act_context;
+struct htp_glu_context;

 typedef void (*glu_compute_f32_fn_t)(const float * restrict src0,
                                     const float * restrict src1,
                                     float * restrict dst,
                                     const uint32_t num_rows,
-                                     const struct htp_act_context * actx);
+                                     const struct htp_glu_context * gctx);

 typedef void (*glu_compute_f16_fn_t)(const _Float16 * restrict src0,
                                     const _Float16 * restrict src1,
                                     _Float16 * restrict dst,
                                     const uint32_t num_rows,
-                                     const struct htp_act_context * actx);
+                                     const struct htp_glu_context * gctx);

-struct htp_act_context {
+struct htp_glu_context {
    struct htp_ops_context * octx;

    void *                   compute;
@@ -108,53 +109,18 @@ struct htp_act_context {
    size_t                   vtcm_dst_size_per_thread;
 };

-struct htp_act_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;
-
-    uint32_t vtcm_row_per_thread;
-};
-
-static inline void htp_act_vtcm_layout_build(struct htp_act_vtcm_layout * L,
-                                             size_t                       src0_row_size_aligned,
-                                             size_t                       src1_row_size_aligned,
-                                             size_t                       dst_row_size_aligned,
-                                             uint32_t                     n_threads,
-                                             size_t                       vtcm_size) {
-    const size_t   spad_size_per_row   = src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned;
-    const uint32_t vtcm_row_per_thread = (uint32_t) (vtcm_size / (n_threads * spad_size_per_row));
-
-    L->vtcm_row_per_thread = vtcm_row_per_thread;
-
-    L->src0_bytes_per_thread = src0_row_size_aligned * vtcm_row_per_thread;
-    L->src1_bytes_per_thread = src1_row_size_aligned * vtcm_row_per_thread;
-    L->dst_bytes_per_thread  = dst_row_size_aligned * vtcm_row_per_thread;
-
-    L->off_src0 = 0;
-    L->off_src1 = L->off_src0 + L->src0_bytes_per_thread * n_threads;
-    L->off_dst  = L->off_src1 + L->src1_bytes_per_thread * n_threads;
-
-    L->total_bytes = L->off_dst + L->dst_bytes_per_thread * n_threads;
-}
-
 #define htp_glu_op_preamble                                            \
-    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;   \
-    const int    nc                    = actx->nc;
+    const size_t src0_row_size_aligned = gctx->src0_row_size_aligned;  \
+    const size_t src1_row_size_aligned = gctx->src1_row_size_aligned;  \
+    const size_t dst_row_size_aligned  = gctx->dst_row_size_aligned;   \
+    const int    nc                    = gctx->nc;

 // 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 struct htp_glu_context * gctx) {
    htp_glu_op_preamble;

    for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -172,10 +138,10 @@ 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 struct htp_glu_context * gctx) {
    htp_glu_op_preamble;
-    const float alpha = ((const float *) (actx->octx->op_params))[2];
-    const float limit = ((const float *) (actx->octx->op_params))[3];
+    const float alpha = ((const float *) (gctx->octx->op_params))[2];
+    const float limit = ((const float *) (gctx->octx->op_params))[3];

    for (uint32_t ib = 0; ib < num_rows; ib++) {
        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
@@ -201,9 +167,9 @@ 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 struct htp_glu_context * gctx) {
    htp_glu_op_preamble;
-    const float limit = ((const float *) (actx->octx->op_params))[3];
+    const float limit = ((const float *) (gctx->octx->op_params))[3];

    for (uint32_t ib = 0; ib < num_rows; ib++) {
        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
@@ -455,7 +421,7 @@ 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 struct htp_glu_context * gctx) {
    htp_glu_op_preamble;

    for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -472,7 +438,7 @@ 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) {
+                            const struct htp_glu_context * gctx) {
    htp_glu_op_preamble;

    for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -489,7 +455,7 @@ static void geglu_erf_f32(const float * restrict src0,
                          const float * restrict src1,
                          float * restrict dst,
                          const uint32_t num_rows,
-                          const struct htp_act_context * actx) {
+                          const struct htp_glu_context * gctx) {
    htp_glu_op_preamble;

    for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -501,12 +467,124 @@ static void geglu_erf_f32(const float * restrict src0,
    }
 }

+static inline void hvx_reglu_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_zero = Q6_V_vzero();
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        HVX_Vector x = Q6_Vsf_vmax_VsfVsf(vsrc0[i], v_zero);
+        vdst[i]      = hvx_vec_mul_f32_f32(x, vsrc1[i]);
+    }
+
+    if (nloe) {
+        HVX_Vector x   = Q6_Vsf_vmax_VsfVsf(vsrc0[i], v_zero);
+        HVX_Vector res = hvx_vec_mul_f32_f32(x, vsrc1[i]);
+        hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(float), res);
+    }
+}
+
+// reglu(x, g) = relu(x) * g
+static void reglu_f32(const float * restrict src0,
+                      const float * restrict src1,
+                      float * restrict dst,
+                      const uint32_t num_rows,
+                      const struct htp_glu_context * gctx) {
+    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_reglu_f32_aa(dst_ptr, src0_ptr, src1_ptr, nc);
+    }
+}
+
+static __attribute__((noinline)) void hvx_sigglu_f32_vec(HVX_Vector * restrict vdst,
+                                                         const HVX_Vector * restrict vsrc0,
+                                                         const HVX_Vector * restrict vsrc1,
+                                                         uint32_t nvec) {
+    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);
+
+    for (uint32_t i = 0; i < nvec; i++) {
+        HVX_Vector sig = hvx_vec_fast_sigmoid_f32_guard(vsrc0[i], v_one, v_max_exp, v_min_exp);
+        vdst[i]        = hvx_vec_mul_f32_f32(sig, vsrc1[i]);
+    }
+}
+
+static __attribute__((noinline)) void hvx_sigglu_f32_tail(HVX_Vector * restrict vdst,
+                                                          const HVX_Vector * restrict vsrc0,
+                                                          const HVX_Vector * restrict vsrc1,
+                                                          uint32_t nloe) {
+    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);
+
+    HVX_Vector sig = hvx_vec_fast_sigmoid_f32_guard(vsrc0[0], v_one, v_max_exp, v_min_exp);
+    HVX_Vector res = hvx_vec_mul_f32_f32(sig, vsrc1[0]);
+    hvx_vec_store_a((void *) vdst, nloe * sizeof(float), res);
+}
+
+static inline void hvx_sigglu_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;
+
+    if (nvec) {
+        hvx_sigglu_f32_vec(vdst, vsrc0, vsrc1, nvec);
+    }
+    if (nloe) {
+        hvx_sigglu_f32_tail(vdst + nvec, vsrc0 + nvec, vsrc1 + nvec, nloe);
+    }
+}
+
+// sigglu(x, g) = sigmoid(x) * g
+static void sigglu_f32(const float * restrict src0,
+                       const float * restrict src1,
+                       float * restrict dst,
+                       const uint32_t num_rows,
+                       const struct htp_glu_context * gctx) {
+    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_sigglu_f32_aa(dst_ptr, src0_ptr, src1_ptr, nc);
+    }
+}
+
 // swiglu(x) = x1 * sigmoid(x0)
 static void swiglu_f16(const _Float16 * restrict src0,
                       const _Float16 * restrict src1,
                       _Float16 * restrict dst,
                       const uint32_t num_rows,
-                       const struct htp_act_context * actx) {
+                       const struct htp_glu_context * gctx) {
    htp_glu_op_preamble;

    for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -524,10 +602,10 @@ static void swiglu_oai_f16(const _Float16 * restrict src0,
                           const _Float16 * restrict src1,
                           _Float16 * restrict dst,
                           const uint32_t num_rows,
-                           const struct htp_act_context * actx) {
+                           const struct htp_glu_context * gctx) {
    htp_glu_op_preamble;
-    const _Float16 alpha = (_Float16) ((const float *) (actx->octx->op_params))[2];
-    const _Float16 limit = (_Float16) ((const float *) (actx->octx->op_params))[3];
+    const _Float16 alpha = (_Float16) ((const float *) (gctx->octx->op_params))[2];
+    const _Float16 limit = (_Float16) ((const float *) (gctx->octx->op_params))[3];

    for (uint32_t ib = 0; ib < num_rows; ib++) {
        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
@@ -553,9 +631,9 @@ static void swiglu_clamp_f16(const _Float16 * restrict src0,
                             const _Float16 * restrict src1,
                             _Float16 * restrict dst,
                             const uint32_t                 num_rows,
-                             const struct htp_act_context * actx) {
+                             const struct htp_glu_context * gctx) {
    htp_glu_op_preamble;
-    const _Float16 limit = (_Float16) ((const float *) (actx->octx->op_params))[3];
+    const _Float16 limit = (_Float16) ((const float *) (gctx->octx->op_params))[3];

    for (uint32_t ib = 0; ib < num_rows; ib++) {
        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
@@ -574,7 +652,7 @@ static void geglu_f16(const _Float16 * restrict src0,
                      const _Float16 * restrict src1,
                      _Float16 * restrict dst,
                      const uint32_t num_rows,
-                      const struct htp_act_context * actx) {
+                      const struct htp_glu_context * gctx) {
    htp_glu_op_preamble;

    for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -586,56 +664,155 @@ static void geglu_f16(const _Float16 * restrict src0,
    }
 }

+static inline void hvx_reglu_f16_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(_Float16);
+    const uint32_t nvec = n / epv;
+    const uint32_t nloe = n % epv;
+
+    const HVX_Vector v_zero = Q6_V_vzero();
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        HVX_Vector x = Q6_Vhf_vmax_VhfVhf(vsrc0[i], v_zero);
+        vdst[i]      = hvx_vec_mul_f16_f16(x, vsrc1[i]);
+    }
+
+    if (nloe) {
+        HVX_Vector x   = Q6_Vhf_vmax_VhfVhf(vsrc0[i], v_zero);
+        HVX_Vector res = hvx_vec_mul_f16_f16(x, vsrc1[i]);
+        hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(_Float16), res);
+    }
+}
+
+// reglu(x, g) = relu(x) * g
+static void reglu_f16(const _Float16 * restrict src0,
+                      const _Float16 * restrict src1,
+                      _Float16 * restrict dst,
+                      const uint32_t num_rows,
+                      const struct htp_glu_context * gctx) {
+    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_reglu_f16_aa(dst_ptr, src0_ptr, src1_ptr, nc);
+    }
+}
+
+static __attribute__((noinline)) void hvx_sigglu_f16_vec(HVX_Vector * restrict vdst,
+                                                         const HVX_Vector * restrict vsrc0,
+                                                         const HVX_Vector * restrict vsrc1,
+                                                         uint32_t nvec) {
+    for (uint32_t i = 0; i < nvec; i++) {
+        HVX_Vector sig = hvx_vec_fast_sigmoid_f16(vsrc0[i]);
+        vdst[i]        = hvx_vec_mul_f16_f16(sig, vsrc1[i]);
+    }
+}
+
+static __attribute__((noinline)) void hvx_sigglu_f16_tail(HVX_Vector * restrict vdst,
+                                                          const HVX_Vector * restrict vsrc0,
+                                                          const HVX_Vector * restrict vsrc1,
+                                                          uint32_t nloe) {
+    HVX_Vector sig = hvx_vec_fast_sigmoid_f16(vsrc0[0]);
+    HVX_Vector res = hvx_vec_mul_f16_f16(sig, vsrc1[0]);
+    hvx_vec_store_a((void *) vdst, nloe * sizeof(_Float16), res);
+}
+
+static inline void hvx_sigglu_f16_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(_Float16);
+    const uint32_t nvec = n / epv;
+    const uint32_t nloe = n % epv;
+
+    if (nvec) {
+        hvx_sigglu_f16_vec(vdst, vsrc0, vsrc1, nvec);
+    }
+    if (nloe) {
+        hvx_sigglu_f16_tail(vdst + nvec, vsrc0 + nvec, vsrc1 + nvec, nloe);
+    }
+}
+
+// sigglu(x, g) = sigmoid(x) * g
+static void sigglu_f16(const _Float16 * restrict src0,
+                       const _Float16 * restrict src1,
+                       _Float16 * restrict dst,
+                       const uint32_t num_rows,
+                       const struct htp_glu_context * gctx) {
+    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_sigglu_f16_aa(dst_ptr, src0_ptr, src1_ptr, nc);
+    }
+}
+
 #define DEFINE_GLU_PER_THREAD(TYPE, SUFFIX)                                                                            \
    static void glu_##SUFFIX##_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_glu_context * gctx = (struct htp_glu_context *) data;                                               \
+        htp_glu_preamble;                                                                                              \
                                                                                                                       \
-        struct htp_thread_trace * tr = actx->octx->ctx ? &actx->octx->ctx->trace[ith] : NULL;                          \
+        struct htp_thread_trace * tr = gctx->octx->ctx ? &gctx->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_size = gctx->src0_row_size;                                                                    \
+        size_t src1_row_size = gctx->src1_row_size;                                                                    \
+        size_t dst_row_size  = gctx->dst_row_size;                                                                     \
                                                                                                                       \
-        size_t src0_row_stride = actx->src0_row_stride;                                                                \
-        size_t src1_row_stride = actx->src1_row_stride;                                                                \
+        size_t src0_row_stride = gctx->src0_row_stride;                                                                \
+        size_t src1_row_stride = gctx->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_nrows            = gctx->src0_nrows;                                                       \
+        const uint32_t src0_nrows_per_thread = gctx->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);     \
+        const uint32_t src0_start_row = gctx->row_start + src0_nrows_per_thread * ith;                                 \
+        const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, gctx->row_start + src0_nrows);     \
                                                                                                                       \
        /* no work for this thread */                                                                                  \
        if (src0_start_row >= src0_end_row) {                                                                          \
            return;                                                                                                    \
        }                                                                                                              \
                                                                                                                       \
-        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 dma_addr_t data_src0 = gctx->data_src0;                                                                  \
+        const dma_addr_t data_src1 = gctx->data_src1;                                                                  \
+        const dma_addr_t data_dst  = gctx->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;                                               \
+        const size_t src0_row_size_aligned = gctx->src0_row_size_aligned;                                              \
+        const size_t src1_row_size_aligned = gctx->src1_row_size_aligned;                                              \
+        const size_t dst_row_size_aligned  = gctx->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);                  \
+        uint8_t * restrict src0_spad_data = gctx->vtcm_src0 + (ith * gctx->vtcm_src0_size_per_thread);                 \
+        uint8_t * restrict src1_spad_data = gctx->vtcm_src1 + (ith * gctx->vtcm_src1_size_per_thread);                 \
+        uint8_t * restrict dst_spad_data  = gctx->vtcm_dst  + (ith * gctx->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;                                                         \
+        size_t src0_spad_half_size = gctx->src0_spad_half_size;                                                        \
+        size_t src1_spad_half_size = gctx->src1_spad_half_size;                                                        \
+        size_t dst_spad_half_size  = gctx->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;                                                                                                    \
-        }                                                                                                              \
+        const int BLOCK = gctx->block;                                                                                 \
                                                                                                                       \
-        dma_queue * dma_q = actx->octx->ctx->dma[ith];                                                                 \
-        glu_compute_##SUFFIX##_fn_t compute = (glu_compute_##SUFFIX##_fn_t) actx->compute;                             \
+        dma_queue * dma_q = gctx->octx->ctx->dma[ith];                                                                 \
+        glu_compute_##SUFFIX##_fn_t compute = (glu_compute_##SUFFIX##_fn_t) gctx->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);                                                 \
@@ -662,7 +839,7 @@ static void geglu_f16(const _Float16 * restrict src0,
            TYPE * src1_spad = (TYPE *) 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);                                                 \
+            compute(src0_spad, src1_spad, dst_spad, block_size, gctx);                                                 \
            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),                             \
@@ -685,16 +862,11 @@ static void geglu_f16(const _Float16 * restrict src0,
 DEFINE_GLU_PER_THREAD(float, f32)
 DEFINE_GLU_PER_THREAD(_Float16, f16)

-static int execute_op_activations_f32(struct htp_ops_context * octx) {
+static int execute_op_glu_f32(struct htp_ops_context * octx) {
    const struct htp_tensor * src0 = octx->src[0];
    const struct htp_tensor * src1 = octx->src[1];
    const struct htp_tensor * dst  = octx->dst;

-    if ((dst->ne[0] * SIZEOF_FP32) != dst->nb[1]) {
-        FARF(ERROR, "Non-contiguous dst is not supported at this time \n");
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
    glu_compute_f32_fn_t compute_fn = NULL;
    const char *         op_type    = NULL;

@@ -729,11 +901,27 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
            op_type    = "geglu-erf-f32";
            break;

+        case HTP_OP_GLU_REGLU:
+            compute_fn = reglu_f32;
+            op_type    = "reglu-f32";
+            break;
+
+        case HTP_OP_GLU_SIGMOID:
+            compute_fn = sigglu_f32;
+            op_type    = "sigglu-f32";
+            break;
+
        default:
-            FARF(ERROR, "Unsupported activations Op %u\n", octx->op);
+            FARF(ERROR, "Unsupported GLU Op %u\n", octx->op);
            return HTP_STATUS_NO_SUPPORT;
    }

+    const struct htp_glu_kernel_params * kparams = (const struct htp_glu_kernel_params *) octx->kernel_params;
+
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
    const size_t dst_row_size = dst->ne[0] * SIZEOF_FP32;

@@ -762,17 +950,16 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
    const size_t src0_row_stride = src0->nb[1];
    const size_t src1_row_stride = src1 ? src1->nb[1] : src0->nb[1];

-    const size_t src0_row_size_aligned = hex_round_up(src0_row_size, VLEN);
-    const size_t src1_row_size_aligned = hex_round_up(src1_row_size, VLEN);
-    const size_t dst_row_size_aligned  = hex_round_up(dst_row_size, VLEN);
+    struct htp_glu_vtcm_layout L;
+    htp_glu_vtcm_layout_build(&L, dst_row_size, n_threads, kparams->vtcm_row_per_thread);

-    struct htp_act_vtcm_layout L;
-    htp_act_vtcm_layout_build(&L, src0_row_size_aligned, src1_row_size_aligned, dst_row_size_aligned, n_threads,
-                              octx->ctx->vtcm_size);
+    const size_t src0_row_size_aligned = L.row_size_aligned;
+    const size_t src1_row_size_aligned = L.row_size_aligned;
+    const size_t dst_row_size_aligned  = L.row_size_aligned;

    // Make sure the reserved vtcm size is sufficient
-    if (L.vtcm_row_per_thread == 0) {
-        FARF(ERROR, "act-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
+    if (kparams->vtcm_row_per_thread == 0 || octx->ctx->vtcm_size < kparams->vtcm_size) {
+        FARF(ERROR, "glu-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
             (src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned) * n_threads);
        return HTP_STATUS_VTCM_TOO_SMALL;
    }
@@ -789,58 +976,53 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
    }

    // Prepare context
-    struct htp_act_context actx;
-    actx.octx    = octx;
-    actx.compute = (void *) compute_fn;
-    actx.op_str  = op_type;
+    struct htp_glu_context gctx;
+    gctx.octx    = octx;
+    gctx.compute = (void *) compute_fn;
+    gctx.op_str  = op_type;

-    actx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    gctx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);

-    actx.src0_row_size = src0_row_size;
-    actx.src1_row_size = src1_row_size;
-    actx.dst_row_size  = dst_row_size;
+    gctx.src0_row_size = src0_row_size;
+    gctx.src1_row_size = src1_row_size;
+    gctx.dst_row_size  = dst_row_size;

-    actx.src0_row_size_aligned = src0_row_size_aligned;
-    actx.src1_row_size_aligned = src1_row_size_aligned;
-    actx.dst_row_size_aligned  = dst_row_size_aligned;
+    gctx.src0_row_size_aligned = src0_row_size_aligned;
+    gctx.src1_row_size_aligned = src1_row_size_aligned;
+    gctx.dst_row_size_aligned  = dst_row_size_aligned;

-    actx.src0_row_stride = src0_row_stride;
-    actx.src1_row_stride = src1_row_stride;
+    gctx.src0_row_stride = src0_row_stride;
+    gctx.src1_row_stride = src1_row_stride;

    uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
-    actx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
-    actx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
-    actx.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
+    gctx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+    gctx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+    gctx.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);

-    actx.vtcm_src0_size_per_thread = L.src0_bytes_per_thread;
-    actx.vtcm_src1_size_per_thread = L.src1_bytes_per_thread;
-    actx.vtcm_dst_size_per_thread  = L.dst_bytes_per_thread;
+    gctx.vtcm_src0_size_per_thread = (uint32_t) L.src0_bytes_per_thread;
+    gctx.vtcm_src1_size_per_thread = (uint32_t) L.src1_bytes_per_thread;
+    gctx.vtcm_dst_size_per_thread  = (uint32_t) L.dst_bytes_per_thread;

-    actx.src0_spad_half_size = L.src0_bytes_per_thread / 2;
-    actx.src1_spad_half_size = L.src1_bytes_per_thread / 2;
-    actx.dst_spad_half_size  = L.dst_bytes_per_thread / 2;
+    gctx.src0_spad_half_size = gctx.vtcm_src0_size_per_thread >> 1;
+    gctx.src1_spad_half_size = gctx.vtcm_src1_size_per_thread >> 1;
+    gctx.dst_spad_half_size  = gctx.vtcm_dst_size_per_thread  >> 1;

-    actx.block = actx.src0_spad_half_size / actx.src0_row_size_aligned;
-    actx.src0_nrows = nrows;
-    actx.row_start  = row_start;
+    gctx.block      = L.block;
+    gctx.src0_nrows = nrows;
+    gctx.row_start  = row_start;

-    actx.nc = dst->ne[0];
+    gctx.nc = dst->ne[0];

    // 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 ||
-                  octx->op == HTP_OP_GLU_GEGLU_QUICK ||
-                  octx->op == HTP_OP_GLU_GEGLU_ERF)) {
+    if (!src1) {
         const int32_t swapped = octx->op_params[1];
         data_src1 = data_src0;
-         actx.src1_row_size = actx.src0_row_size;
+         gctx.src1_row_size = gctx.src0_row_size;

-         size_t nc_in_bytes = actx.nc * SIZEOF_FP32;
+         size_t nc_in_bytes = gctx.nc * SIZEOF_FP32;
         if (swapped) {
             data_src0 += nc_in_bytes;
         } else {
@@ -848,24 +1030,19 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
         }
    }

-    actx.data_src0 = data_src0;
-    actx.data_src1 = data_src1;
-    actx.data_dst  = dst->data;
+    gctx.data_src0 = data_src0;
+    gctx.data_src1 = data_src1;
+    gctx.data_dst  = dst->data;

-    work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f32_per_thread, &actx, n_threads);
+    work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f32_per_thread, &gctx, n_threads);
    return HTP_STATUS_OK;
 }

-static int execute_op_activations_f16(struct htp_ops_context * octx) {
+static int execute_op_glu_f16(struct htp_ops_context * octx) {
    const struct htp_tensor * src0 = octx->src[0];
    const struct htp_tensor * src1 = octx->src[1];
    const struct htp_tensor * dst  = octx->dst;

-    if ((dst->ne[0] * SIZEOF_FP16) != dst->nb[1]) {
-        FARF(ERROR, "Non-contiguous dst is not supported at this time \n");
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
    glu_compute_f16_fn_t compute_fn = NULL;
    const char *         op_type    = NULL;

@@ -890,11 +1067,27 @@ static int execute_op_activations_f16(struct htp_ops_context * octx) {
            op_type    = "geglu-f16";
            break;

+        case HTP_OP_GLU_REGLU:
+            compute_fn = reglu_f16;
+            op_type    = "reglu-f16";
+            break;
+
+        case HTP_OP_GLU_SIGMOID:
+            compute_fn = sigglu_f16;
+            op_type    = "sigglu-f16";
+            break;
+
        default:
-            FARF(ERROR, "Unsupported activations Op %u\n", octx->op);
+            FARF(ERROR, "Unsupported GLU Op %u\n", octx->op);
            return HTP_STATUS_NO_SUPPORT;
    }

+    const struct htp_glu_kernel_params * kparams = (const struct htp_glu_kernel_params *) octx->kernel_params;
+
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
    const size_t dst_row_size = dst->ne[0] * SIZEOF_FP16;

@@ -923,17 +1116,16 @@ static int execute_op_activations_f16(struct htp_ops_context * octx) {
    const size_t src0_row_stride = src0->nb[1];
    const size_t src1_row_stride = src1 ? src1->nb[1] : src0->nb[1];

-    const size_t src0_row_size_aligned = hex_round_up(src0_row_size, VLEN);
-    const size_t src1_row_size_aligned = hex_round_up(src1_row_size, VLEN);
-    const size_t dst_row_size_aligned  = hex_round_up(dst_row_size, VLEN);
+    struct htp_glu_vtcm_layout L;
+    htp_glu_vtcm_layout_build(&L, dst_row_size, n_threads, kparams->vtcm_row_per_thread);

-    struct htp_act_vtcm_layout L;
-    htp_act_vtcm_layout_build(&L, src0_row_size_aligned, src1_row_size_aligned, dst_row_size_aligned, n_threads,
-                              octx->ctx->vtcm_size);
+    const size_t src0_row_size_aligned = L.row_size_aligned;
+    const size_t src1_row_size_aligned = L.row_size_aligned;
+    const size_t dst_row_size_aligned  = L.row_size_aligned;

    // Make sure the reserved vtcm size is sufficient
-    if (L.vtcm_row_per_thread == 0) {
-        FARF(ERROR, "act-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
+    if (kparams->vtcm_row_per_thread == 0 || octx->ctx->vtcm_size < kparams->vtcm_size) {
+        FARF(ERROR, "glu-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
             (src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned) * n_threads);
        return HTP_STATUS_VTCM_TOO_SMALL;
    }
@@ -950,56 +1142,53 @@ static int execute_op_activations_f16(struct htp_ops_context * octx) {
    }

    // Prepare context
-    struct htp_act_context actx;
-    actx.octx    = octx;
-    actx.compute = (void *) compute_fn;
-    actx.op_str  = op_type;
+    struct htp_glu_context gctx;
+    gctx.octx    = octx;
+    gctx.compute = (void *) compute_fn;
+    gctx.op_str  = op_type;

-    actx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    gctx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);

-    actx.src0_row_size = src0_row_size;
-    actx.src1_row_size = src1_row_size;
-    actx.dst_row_size  = dst_row_size;
+    gctx.src0_row_size = src0_row_size;
+    gctx.src1_row_size = src1_row_size;
+    gctx.dst_row_size  = dst_row_size;

-    actx.src0_row_size_aligned = src0_row_size_aligned;
-    actx.src1_row_size_aligned = src1_row_size_aligned;
-    actx.dst_row_size_aligned  = dst_row_size_aligned;
+    gctx.src0_row_size_aligned = src0_row_size_aligned;
+    gctx.src1_row_size_aligned = src1_row_size_aligned;
+    gctx.dst_row_size_aligned  = dst_row_size_aligned;

-    actx.src0_row_stride = src0_row_stride;
-    actx.src1_row_stride = src1_row_stride;
+    gctx.src0_row_stride = src0_row_stride;
+    gctx.src1_row_stride = src1_row_stride;

    uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
-    actx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
-    actx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
-    actx.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
+    gctx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+    gctx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+    gctx.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);

-    actx.vtcm_src0_size_per_thread = L.src0_bytes_per_thread;
-    actx.vtcm_src1_size_per_thread = L.src1_bytes_per_thread;
-    actx.vtcm_dst_size_per_thread  = L.dst_bytes_per_thread;
+    gctx.vtcm_src0_size_per_thread = (uint32_t) L.src0_bytes_per_thread;
+    gctx.vtcm_src1_size_per_thread = (uint32_t) L.src1_bytes_per_thread;
+    gctx.vtcm_dst_size_per_thread  = (uint32_t) L.dst_bytes_per_thread;

-    actx.src0_spad_half_size = L.src0_bytes_per_thread / 2;
-    actx.src1_spad_half_size = L.src1_bytes_per_thread / 2;
-    actx.dst_spad_half_size  = L.dst_bytes_per_thread / 2;
+    gctx.src0_spad_half_size = gctx.vtcm_src0_size_per_thread >> 1;
+    gctx.src1_spad_half_size = gctx.vtcm_src1_size_per_thread >> 1;
+    gctx.dst_spad_half_size  = gctx.vtcm_dst_size_per_thread  >> 1;

-    actx.block = actx.src0_spad_half_size / actx.src0_row_size_aligned;
-    actx.src0_nrows = nrows;
-    actx.row_start  = row_start;
+    gctx.block      = L.block;
+    gctx.src0_nrows = nrows;
+    gctx.row_start  = row_start;

-    actx.nc = dst->ne[0];
+    gctx.nc = dst->ne[0];

    // 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) {
         const int32_t swapped = octx->op_params[1];
         data_src1 = data_src0;
-         actx.src1_row_size = actx.src0_row_size;
+         gctx.src1_row_size = gctx.src0_row_size;

-         size_t nc_in_bytes = actx.nc * SIZEOF_FP16;
+         size_t nc_in_bytes = gctx.nc * SIZEOF_FP16;
         if (swapped) {
             data_src0 += nc_in_bytes;
         } else {
@@ -1007,22 +1196,22 @@ static int execute_op_activations_f16(struct htp_ops_context * octx) {
         }
    }

-    actx.data_src0 = data_src0;
-    actx.data_src1 = data_src1;
-    actx.data_dst  = dst->data;
+    gctx.data_src0 = data_src0;
+    gctx.data_src1 = data_src1;
+    gctx.data_dst  = dst->data;

-    work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f16_per_thread, &actx, n_threads);
+    work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f16_per_thread, &gctx, n_threads);
    return HTP_STATUS_OK;
 }


-int op_activations(struct htp_ops_context * octx) {
+int op_glu(struct htp_ops_context * octx) {
    switch (octx->src[0]->type) {
        case HTP_TYPE_F32:
-            return execute_op_activations_f32(octx);
+            return execute_op_glu_f32(octx);

        case HTP_TYPE_F16:
-            return execute_op_activations_f16(octx);
+            return execute_op_glu_f16(octx);

        default:
            return HTP_STATUS_NO_SUPPORT;
@@ -0,0 +1,76 @@
+#ifndef HTP_GLU_OPS_H
+#define HTP_GLU_OPS_H
+
+#include "hex-common.h"
+#include "htp-ops.h"
+#include "htp-vtcm.h"
+#include <stdint.h>
+
+struct htp_glu_kernel_params {
+    uint32_t n_threads;
+    uint32_t vtcm_row_per_thread;
+    uint32_t vtcm_size;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_glu_kernel_params) <= 128, "htp_glu_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_glu_kernel_params) <= 128, "htp_glu_kernel_params is too large for kernel_params blob");
+#endif
+
+static inline bool htp_op_is_glu(uint32_t opcode) {
+    switch (opcode) {
+        case HTP_OP_GLU_SWIGLU:
+        case HTP_OP_GLU_SWIGLU_OAI:
+        case HTP_OP_GLU_SWIGLU_CLAMP:
+        case HTP_OP_GLU_GEGLU:
+        case HTP_OP_GLU_GEGLU_QUICK:
+        case HTP_OP_GLU_GEGLU_ERF:
+        case HTP_OP_GLU_REGLU:
+        case HTP_OP_GLU_SIGMOID:
+            return true;
+        default:
+            return false;
+    }
+}
+
+struct htp_glu_vtcm_layout {
+    size_t total_bytes;
+    size_t off_src0;
+    size_t off_src1;
+    size_t off_dst;
+
+    size_t row_size_aligned;
+    size_t src0_bytes_per_thread;
+    size_t src1_bytes_per_thread;
+    size_t dst_bytes_per_thread;
+
+    uint32_t vtcm_row_per_thread;
+    uint32_t block;
+};
+
+static inline void htp_glu_vtcm_layout_build(
+    struct htp_glu_vtcm_layout * L,
+    size_t                       dst_row_size,
+    uint32_t                     n_threads,
+    uint32_t                     vtcm_row_per_thread
+) {
+    const size_t row_size_aligned = hex_round_up(dst_row_size, 128);
+
+    L->row_size_aligned      = row_size_aligned;
+    L->vtcm_row_per_thread   = vtcm_row_per_thread;
+    L->block                 = vtcm_row_per_thread >> 1;
+
+    L->src0_bytes_per_thread = row_size_aligned * vtcm_row_per_thread;
+    L->src1_bytes_per_thread = row_size_aligned * vtcm_row_per_thread;
+    L->dst_bytes_per_thread  = row_size_aligned * vtcm_row_per_thread;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_src1, L->src1_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread * n_threads);
+
+    L->total_bytes = off;
+}
+
+#endif // HTP_GLU_OPS_H
@@ -17,6 +17,10 @@
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
 #endif

+#ifndef HEX_L2_LINE_SIZE
+#define HEX_L2_LINE_SIZE 128
+#endif
+
 static inline uint32_t hex_ceil_pow2(uint32_t x) {
    if (x <= 1) { return 1; }
    int p = 2;
@@ -44,12 +44,14 @@ struct htp_thread_trace {
 };

 static inline void htp_trace_event(struct htp_thread_trace * tr, uint16_t id, uint16_t info, uint32_t type) {
-    if (tr->count < tr->max_events) {
-        uint32_t i = tr->count;
+    if (!tr->max_events) {
+        return;
+    }
+    uint32_t i = tr->count++;
+    if (i < tr->max_events) {
        tr->events[i].id     = id;
        tr->events[i].info   = info | (type == HTP_TRACE_EVT_STOP ? 0x8000 : 0);
        tr->events[i].cycles = (uint32_t) hex_get_cycles();
-        tr->count++;
    }
 }

@@ -12,8 +12,10 @@
 #include <assert.h>
 #include <dspqueue.h>
 #include <stdatomic.h>
+#include <stddef.h>
 #include <stdint.h>
 #include <stdbool.h>
+#include <string.h>

 #ifndef HTP_MAX_NTHREADS
 #define HTP_MAX_NTHREADS 10
@@ -66,14 +68,18 @@ struct htp_ops_context {
    // TODO convert these to an array
    struct htp_spad src0_spad;
    struct htp_spad src1_spad;
-    struct htp_spad src2_spad;
-    struct htp_spad src3_spad;
    struct htp_spad dst_spad;

    uint32_t              flags;
    uint32_t              n_threads;
    struct fastdiv_values n_threads_div;
    int                   status;
+
+    // Must stay last: not cleared by htp_ops_context_reset(), ops init the caches they use
+    union {
+        dma_cache_dm dma_cache_dm[HTP_MAX_NTHREADS];
+        dma_cache_fa dma_cache_fa[HTP_MAX_NTHREADS];
+    } __attribute__((aligned(128)));
 };

 // Main context for htp DSP backend
@@ -118,6 +124,13 @@ struct htp_context {
    size_t                 footprint;
 };

+static inline void htp_ops_context_reset(struct htp_ops_context * octx, struct htp_context * ctx) {
+    memset(octx, 0, offsetof(struct htp_ops_context, dma_cache_dm));
+    octx->ctx           = ctx;
+    octx->n_threads     = ctx->n_threads;
+    octx->n_threads_div = ctx->n_threads_div;
+}
+
 static inline bool htp_ops_context_set_n_threads(struct htp_ops_context * octx, uint32_t n_threads) {
    if (n_threads == 0 || n_threads > octx->ctx->n_threads) {
        return false;
@@ -148,7 +161,7 @@ int op_unary(struct htp_ops_context * octx);
 int op_sum(struct htp_ops_context * octx);
 int op_sum_rows(struct htp_ops_context * octx);
 int op_argmax(struct htp_ops_context * octx);
-int op_activations(struct htp_ops_context * octx);
+int op_glu(struct htp_ops_context * octx);
 int op_softmax(struct htp_ops_context * octx);
 int op_add_id(struct htp_ops_context * octx);
 int op_rope(struct htp_ops_context * octx);
@@ -38,11 +38,6 @@ enum htp_data_type {
    HTP_TYPE_INVALID
 };

-// Constants for internal types
-#define QK_Q8_0_TILED  128  // 32x32 Q8_0 tiled layout
-
-
-
 // Do not reorder first 4 (used as an index)
 enum htp_op_code {
    HTP_OP_MUL = 0,
@@ -115,6 +110,8 @@ enum htp_op_code {
    HTP_OP_UNARY_GELU_QUICK,
    HTP_OP_IM2COL_3D,
    HTP_OP_HADAMARD,
+    HTP_OP_GLU_REGLU,
+    HTP_OP_GLU_SIGMOID,

    HTP_OP_INVALID
 };
@@ -125,8 +122,8 @@ enum htp_op_code {
 #define HTP_OP_MAX_PARAMS  16   // aka GGML_MAX_OP_PARAMS
 #define HTP_OP_MAX_KERN_PARAMS 32

-#define HTP_OP_MAX_BUFS    16
-#define HTP_OP_MAX_TENSORS 8192 // must stay under 64K (uint16)
+#define HTP_OP_MAX_BUFS    64
+#define HTP_OP_MAX_TENSORS 65534 // must stay under 64K (uint16)

 #define HTP_FENCE_TIMEOUT  (1000000000ULL)
 #define HTP_FENCE_SLOT_SIZE 128
@@ -136,9 +133,10 @@ enum htp_op_code {
 #define HTP_MMAP_MAX_VMEM  (2147483648u)

 enum htp_tensor_flags {
-    HTP_TENSOR_WEIGHT  = (1U << 0), // Tensor buffer model weight data (not compute)
-    HTP_TENSOR_TILED   = (1U << 1), // Tensor is in tiled format
-    HTP_TENSOR_FENCE   = (1U << 2)  // Tensor is synchronization fence (explicitly managed)
+    HTP_TENSOR_WEIGHT   = (1U << 0), // Tensor buffer model weight data (not compute)
+    HTP_TENSOR_TILED    = (1U << 1), // Tensor is in tiled format
+    HTP_TENSOR_FENCE    = (1U << 2), // Tensor is synchronization fence (explicitly managed)
+    HTP_TENSOR_FUSEABLE = (1U << 3), // Tensor can be fused (host-only)
 };

 enum htp_buf_flags {
--- a/Show More
+++ b/Show More