engine: separate engine and server

model : Granite Switch Decision (#30676 )
* feat(convert): Add decision.chat to support decision models that can also chat Branch: GraniteSwitchDecision AI-usage: full (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart <ghart@us.ibm.com> * feat(convert): Add conversion for Granite Switch as decision model The template follows from the usage sample: https://huggingface.co/ibm-granite/granite-switch-4.1-8b#guardian-and-decision-side-by-side Branch: GraniteSwitchDecision AI-usage: draft (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart <ghart@us.ibm.com> * feat(common): Add support for chat+decision models in common Branch: GraniteSwitchDecision AI-usage: full (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart <ghart@us.ibm.com> * feat(common): Add granite switch decision model enum Branch: GraniteSwitchDecision AI-usage: full (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart <ghart@us.ibm.com> * feat(server): Plumb through decision_chat bool to support decision+chat models Branch: GraniteSwitchDecision AI-usage: full (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart <ghart@us.ibm.com> * feat(server): Add Granite Switch decision support Branch: GraniteSwitchDecision AI-usage: draft (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart <ghart@us.ibm.com> * docs: Update server README for granite switch decisions / chat+decisions Branch: GraniteSwitchDecision AI-usage: draft (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart <ghart@us.ibm.com> * Revert "feat(convert): Add decision.chat to support decision models that can also chat" This reverts commit 952ef56764. * Revert "feat(common): Add support for chat+decision models in common" This reverts commit 93ed782e3a. * Revert "feat(server): Plumb through decision_chat bool to support decision+chat models" This reverts commit 97e9af60dd. * fix: If Granite Switch supports decisions, it always supports text too Branch: GraniteSwitchDecision AI-usage: none Signed-off-by: Gabe Goodhart <ghart@us.ibm.com> * docs: Remove reference to decision.chat in server README Branch: GraniteSwitchDecision AI-usage: none Signed-off-by: Gabe Goodhart <ghart@us.ibm.com>
2026-10-17 10:17:35 -05:00 · 2026-10-16 20:17:40 +02:00 · 2026-10-16 20:00:41 +02:00 · 2026-10-16 17:01:37 +02:00 · 2026-10-16 15:17:30 +02:00 · 2026-10-16 11:52:18 +02:00
217 changed files with 18407 additions and 7429 deletions
@@ -19,7 +19,12 @@ RUN npm ci
 COPY tools/ui/ ./
 RUN LLAMA_BUILD_NUMBER="$APP_VERSION" npm run build

-FROM docker.io/intel/oneapi-toolkit:$ONEAPI_VERSION AS build
+FROM docker.io/intel/deep-learning-essentials:$ONEAPI_VERSION AS build
+
+RUN echo "Download oneDNN installation package" && \
+    wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/94c3dbac-0852-45be-a57d-21c204cada3e/intel-onednn-2026.0.2.46_offline.sh -O intel-onednn_offline.sh && \
+    echo "Install oneDNN components" && \
+    bash intel-onednn_offline.sh -s -a --silent --eula accept

 ARG GGML_SYCL_F16=ON
 ARG LEVEL_ZERO_VERSION=1.28.2
@@ -45,7 +50,7 @@ RUN if [ "${GGML_SYCL_F16}" = "ON" ]; then \
    fi && \
    echo "Building with dynamic libs" && \
    cmake -B build -DGGML_NATIVE=OFF -DGGML_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON -DLLAMA_BUILD_TESTS=OFF ${OPT_SYCL_F16} && \
-    cmake --build build --config Release -j$(nproc)
+    cmake --build build --config Release -j$((($(nproc)+1)/2))

 RUN mkdir -p /app/lib && \
    find build -name "*.so*" -exec cp -P {} /app/lib \;
@@ -102,7 +107,7 @@ RUN mkdir /tmp/neo/ && cd /tmp/neo/ \
  && dpkg --install *.deb

 RUN apt-get update \
-    && apt-get install -y libgomp1 curl ffmpeg \
+    && apt-get install -y libgomp1 curl ffmpeg openssl libssl-dev \
    && apt autoremove -y \
    && apt clean -y \
    && rm -rf /tmp/* /var/tmp/* \
@@ -46,6 +46,7 @@ examples:
                - examples/**
                - tools/**
            - all-globs-to-all-files:
+                - '!tools/engine/**'
                - '!tools/server/**'
                - '!tools/mtmd/**'
                - '!tools/ui/**'
@@ -76,6 +77,7 @@ server/ui:
 server:
    - changed-files:
        - any-glob-to-any-file:
+            - tools/engine/**
            - tools/server/**
 mtmd:
    - changed-files:
@@ -67,8 +67,15 @@ jobs:
        shell: bash
        run: |
          cd /tmp
-          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/5996e26b-f48a-42b1-8db0-b002ad0bd8d7/intel-oneapi-toolkit-2026.1.1.33_offline.sh -O intel-oneapi-toolkit_offline.sh
-          sudo bash intel-oneapi-toolkit_offline.sh -s -a --silent --eula accept
+          echo "Download oneAPI Deep Learning Essentials installation package"
+          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/b5c100a7-8052-4e88-90e2-63d9916a998b/intel-deep-learning-essentials-2026.1.4.17_offline.sh -O intel-deep-learning-essentials_offline.sh
+          echo "Install oneAPI Deep Learning Essentials components"
+          sudo bash intel-deep-learning-essentials_offline.sh -s -a --silent --eula accept
+
+          echo "Download oneDNN installation package"
+          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/94c3dbac-0852-45be-a57d-21c204cada3e/intel-onednn-2026.0.2.46_offline.sh -O intel-onednn_offline.sh
+          echo "Install oneDNN components"
+          sudo bash intel-onednn_offline.sh -s -a --silent --eula accept

      - name: Install Level Zero SDK
        shell: bash
@@ -79,6 +86,12 @@ jobs:
          wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/libze-dev_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O libze-dev.deb
          sudo apt-get install -y ./libze1.deb ./libze-dev.deb

+      - name: Dependencies
+        id: depends
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y libssl-dev
+
      - name: ccache
        uses: ggml-org/ccache-action@v1.2.24
        with:
@@ -104,7 +117,6 @@ jobs:
            -DGGML_SYCL=ON \
            -DCMAKE_C_COMPILER=icx \
            -DCMAKE_CXX_COMPILER=icpx \
-            -DLLAMA_OPENSSL=OFF \
            -DGGML_NATIVE=OFF \
            -DGGML_SYCL_F16=${{ matrix.fp16 }} \
            -DGGML_SYCL_SUPPORT_LEVEL_ZERO_API=ON \
@@ -138,8 +150,10 @@ jobs:
        shell: bash

    env:
-      WINDOWS_BASEKIT_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/0cb67a0d-67f6-410b-868b-f4a0a17ff0cf/intel-oneapi-toolkit-2026.1.1.32_offline.exe
-      WINDOWS_DPCPP_MKL: intel.oneapi.win.cpp-dpcpp-common:intel.oneapi.win.mkl.devel:intel.oneapi.win.dnnl:intel.oneapi.win.tbb.devel
+      WINDOWS_DLE_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/a51104eb-6f4a-4a0e-9dae-1b81b0c81ae1/intel-deep-learning-essentials-2026.1.4.16_offline.exe
+      WINDOWS_ONEDNN_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/34e28083-d04b-46a3-826c-8c1d00bf1e1c/intel-onednn-2026.0.2.47_offline.exe
+      WINDOWS_DPCPP_MKL: intel.oneapi.win.cpp-dpcpp-common:intel.oneapi.win.mkl.devel:intel.oneapi.win.tbb.devel
+      WINDOWS_DNNL: intel.oneapi.win.dnnl
      LEVEL_ZERO_SDK_URL: https://github.com/oneapi-src/level-zero/releases/download/v1.33.1/level-zero-win-sdk-1.33.1.zip
      ONEAPI_ROOT: "C:/Program Files (x86)/Intel/oneAPI"
      ONEAPI_INSTALLER_VERSION: "2026.1"
@@ -151,7 +165,8 @@ jobs:
      - name: Download & Install oneAPI
        shell: bash
        run: |
-          scripts/install-oneapi.bat $WINDOWS_BASEKIT_URL $WINDOWS_DPCPP_MKL
+          scripts/install-oneapi.bat $WINDOWS_DLE_URL $WINDOWS_DPCPP_MKL
+          scripts/install-oneapi.bat $WINDOWS_ONEDNN_URL $WINDOWS_DNNL

      - name: Install Level Zero SDK
        shell: pwsh
@@ -35,7 +35,7 @@ jobs:
        uses: actions/setup-python@v6
        with:
          python-version: "3.11"
-          pip-install: -r requirements/requirements-all.txt ty==0.0.84
+          pip-install: -r requirements/requirements-all.txt ty==0.0.88
      # - name: Type-check with Pyright
      #   uses: jakebailey/pyright-action@v2
      #   with:
@@ -1300,8 +1300,11 @@ jobs:
        shell: bash

    env:
-      WINDOWS_BASEKIT_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/0cb67a0d-67f6-410b-868b-f4a0a17ff0cf/intel-oneapi-toolkit-2026.1.1.32_offline.exe
-      WINDOWS_DPCPP_MKL: intel.oneapi.win.cpp-dpcpp-common:intel.oneapi.win.mkl.devel:intel.oneapi.win.dnnl:intel.oneapi.win.tbb.devel
+      WINDOWS_DLE_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/a51104eb-6f4a-4a0e-9dae-1b81b0c81ae1/intel-deep-learning-essentials-2026.1.4.16_offline.exe
+      WINDOWS_ONEDNN_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/34e28083-d04b-46a3-826c-8c1d00bf1e1c/intel-onednn-2026.0.2.47_offline.exe
+      WINDOWS_DPCPP_MKL: intel.oneapi.win.cpp-dpcpp-common:intel.oneapi.win.mkl.devel:intel.oneapi.win.tbb.devel
+      WINDOWS_DNNL: intel.oneapi.win.dnnl
+
      LEVEL_ZERO_SDK_URL: https://github.com/oneapi-src/level-zero/releases/download/v1.33.1/level-zero-win-sdk-1.33.1.zip
      ONEAPI_ROOT: "C:/Program Files (x86)/Intel/oneAPI"
      ONEAPI_INSTALLER_VERSION: "2026.1"
@@ -1314,7 +1317,8 @@ jobs:
      - name: Download & Install oneAPI
        shell: bash
        run: |
-          scripts/install-oneapi.bat $WINDOWS_BASEKIT_URL $WINDOWS_DPCPP_MKL
+          scripts/install-oneapi.bat $WINDOWS_DLE_URL $WINDOWS_DPCPP_MKL
+          scripts/install-oneapi.bat $WINDOWS_ONEDNN_URL $WINDOWS_DNNL

      - name: Install Level Zero SDK
        shell: pwsh
@@ -1425,8 +1429,15 @@ jobs:
        shell: bash
        run: |
          cd /tmp
-          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/5996e26b-f48a-42b1-8db0-b002ad0bd8d7/intel-oneapi-toolkit-2026.1.1.33_offline.sh -O intel-oneapi-toolkit_offline.sh
-          sudo bash intel-oneapi-toolkit_offline.sh -s -a --silent --eula accept
+          echo "Download oneAPI Deep Learning Essentials installation package"
+          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/b5c100a7-8052-4e88-90e2-63d9916a998b/intel-deep-learning-essentials-2026.1.4.17_offline.sh -O intel-deep-learning-essentials_offline.sh
+          echo "Install oneAPI Deep Learning Essentials components"
+          sudo bash intel-deep-learning-essentials_offline.sh -s -a --silent --eula accept
+
+          echo "Download oneDNN installation package"
+          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/94c3dbac-0852-45be-a57d-21c204cada3e/intel-onednn-2026.0.2.46_offline.sh -O intel-onednn_offline.sh
+          echo "Install oneDNN components"
+          sudo bash intel-onednn_offline.sh -s -a --silent --eula accept

      - name: Install Level Zero SDK
        shell: bash
@@ -1442,6 +1453,12 @@ jobs:
          name: llama-ui.zip
          path: tools/ui/dist

+      - name: Dependencies
+        id: depends
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y libssl-dev
+
      - name: ccache
        uses: ggml-org/ccache-action@v1.2.24
        with:
@@ -1460,7 +1477,6 @@ jobs:
            -DCMAKE_CXX_COMPILER=icpx \
            -DCMAKE_INSTALL_RPATH='$ORIGIN' \
            -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
-            -DLLAMA_OPENSSL=OFF \
            -DGGML_NATIVE=OFF \
            -DGGML_SYCL_F16=${{ matrix.fp16 }}
          time cmake --build build --config Release -j $(nproc)
@@ -6,10 +6,12 @@ General:
 - PR and commit titles format: `<module> : <title>`. Lookup recents for examples
 - Don't try to build or run the code unless you are explicitly asked to do so
 - Use the `gh` CLI tool when querying PRs, issues, or other GitHub resources
- When [MODEL] is needed, first try to get it from the `PI_MODEL_NAME` env var before asking the user
+- When a model name is needed, first try to get it from the `PI_MODEL_NAME` env var before asking the user
+- When a hardware description is needed, first try to get it from the `PI_HARDWARE` env var before asking the user
 - Never read the `AGENTS.md` file

 Coding:
+- Never write inline comments in the code unless the user tells you to
 - When in doubt, always refer to the CONTRIBUTING.md file of the project
 - In `test-backend-ops.cpp`, do not mention specific backends (e.g. Metal, CUDA) in comments
 - When referencing issues or PRs in comments, use the format:
@@ -21,13 +23,22 @@ Pull requests (PRs):
 - Before opening a pull request, ask the user to confirm the description
 - Don't explicitly wrap lines in the PR description (each paragraph and bullet is a single line)
 - When creating a pull request, look for the repository's PR template and follow it
- For the AI usage disclosure section, write "YES. pi:llama.cpp/[MODEL]"
- If `PI_MODEL_NAME` env var is not set, ask the user to tell you what model was used and write it in place of [MODEL]
+- For the AI usage disclosure section, write "YES", then these 4 bullets:
+  - Harness: `pi:llama.cpp`
+  - Hardware: `<hardware>`
+  - Model: `<model>`, linked to its Hugging Face page - the link text is the plain model name
+  - Stats: `<stats>`
+- Run `~/.pi/stats.sh "$PI_SESSION_FILE"` to get `<stats>` and write it as
+  `◷ <time> · ↑ <input tokens> · ↓ <output tokens> · ≡ <messages> · ⌨ <user messages>`,
+  e.g. `◷ 12.3m · ↑ 120k · ↓ 45k · ≡ 87 · ⌨ 6`
+- Write these symbols as literal unicode in the PR description, never `:shortcodes:`, and keep
+  commit trailers ASCII
+- If `PI_MODEL_NAME` or `PI_HARDWARE` is not set, ask the user and use their answer for `<model>` or `<hardware>`
 - Always create the pull requests in draft mode
 - Never reply to review comments or post comments on issues/PRs without explicit permission from the user

 Commits:
- On every commit that you make, include a "Assisted-by: pi:llama.cpp/[MODEL]" tag
+- On every commit that you make, include an `Assisted-by: pi:llama.cpp/<model> (<hardware>)` tag
 - Do not explicitly set the git author in commits - rely on the default git config
 - Always use `--no-gpg-sign` when committing
 - Never `git push` without explicit confirmation from the user
@@ -4,7 +4,7 @@ include(CheckIncludeFileCXX)

 ### llama.cpp version
 set(LLAMA_VERSION_MAJOR 0)
-set(LLAMA_VERSION_MINOR 6)
+set(LLAMA_VERSION_MINOR 7)
 set(LLAMA_VERSION_PATCH 0)
 set(LLAMA_VERSION_BASE "${LLAMA_VERSION_MAJOR}.${LLAMA_VERSION_MINOR}.${LLAMA_VERSION_PATCH}")

@@ -101,6 +101,7 @@
 /tools/batched-bench/                   @ggerganov
 /tools/cli/                             @ngxson
 /tools/completion/                      @ggerganov
+/tools/engine/                          @ggml-org/llama-server
 /tools/mtmd/                            @ggml-org/llama-mtmd
 /tools/perplexity/                      @ggerganov
 /tools/parser/                          @pwilkin
@@ -1110,6 +1110,12 @@ static common_chat_params_init_fn common_chat_template_detect_params_init(const
        return common_chat_params_init_gpt_oss;
    }

+    // Muse Spark shares the Muse recipient protocol and adds a commentary recipient.
+    if (src.find("to=commentary") != std::string::npos && src.find("<|eom|>") != std::string::npos) {
+        LOG_DBG("Using specialized template: Muse Spark\n");
+        return common_chat_params_init_muse_spark;
+    }
+
    // Muse Glimmer format using " to=<recipient>" recipients and <|eom|>/<|eot|> message terminators.
    if (src.find("<atem:function_calls>") != std::string::npos && src.find("<|eom|>") != std::string::npos) {
        LOG_DBG("Using specialized template: Muse Glimmer\n");
@@ -1174,6 +1174,7 @@ static const std::map<common_decision_type, std::string> COMMON_DECISION_TYPE_NA
    { COMMON_DECISION_TYPE_VLLM_D3,        "vllm-d3"        },
    { COMMON_DECISION_TYPE_LFM2_D1,        "lfm2-d1"       },
    { COMMON_DECISION_TYPE_LFM2_D1_OMNI,   "lfm2-d1-omni"  },
+    { COMMON_DECISION_TYPE_GRANITE_SWITCH, "granite-switch"},
    // for future models, ALWAYS add org name prefix to avoid name collisions, e.g. "orgname-modelname"
 };

@@ -970,6 +970,7 @@ enum common_decision_type {
    COMMON_DECISION_TYPE_VLLM_D3, // same as pplx-decider, non-causal attention
    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_UNKNOWN, // a decision model of a type that is not supported
 };

@@ -1,5 +1,7 @@
 #pragma once

+#include "log.h"
+
 #include <cpp-httplib/httplib.h>

 #ifdef _WIN32
@@ -97,6 +99,118 @@ static common_http_url common_http_parse_url(const std::string & url) {
    return parts;
 }

+// match a loopback destination: a remote proxy can never reach the client machine's own
+// loopback, so these hosts are always contacted directly, before NO_PROXY is consulted
+// (similar to Go's net/http)
+static bool common_http_is_loopback(const std::string & host) {
+    auto lower = [](std::string s) {
+        for (auto & c : s) {
+            if (c >= 'A' && c <= 'Z') {
+                c += 'a' - 'A';
+            }
+        }
+        return s;
+    };
+
+    const std::string h = lower(host);
+
+    if (h == "localhost" || h == "::1" || h == "0.0.0.0") {
+        return true;
+    }
+
+    if (h.size() > 10 && h.compare(h.size() - 10, 10, ".localhost") == 0) {
+        return true;
+    }
+
+    // IPv4 literal: exactly four all-digit octets with values 0-255, loopback is 127.0.0.0/8
+    int octets[4] = { 0, 0, 0, 0 };
+    size_t pos = 0;
+    int n = 0;
+    bool complete = false;
+    while (n < 4 && pos <= h.size()) {
+        size_t dot = h.find('.', pos);
+        if (dot == std::string::npos) {
+            dot = h.size();
+        }
+        if (dot == pos || dot - pos > 3) {
+            return false;
+        }
+        int val = 0;
+        for (size_t i = pos; i < dot; i++) {
+            if (h[i] < '0' || h[i] > '9') {
+                return false;
+            }
+            val = val * 10 + (h[i] - '0');
+        }
+        if (val > 255) {
+            return false;
+        }
+        octets[n++] = val;
+        if (dot == h.size()) {
+            complete = (n == 4);
+            break;
+        }
+        pos = dot + 1;
+    }
+
+    return complete && octets[0] == 127;
+}
+
+// match a host against a NO_PROXY list: comma separated entries, "*" bypasses the proxy for every host,
+// an entry matches the host itself or any of its subdomains, a leading dot is optional, ports are ignored
+static bool common_http_proxy_bypass(const std::string & no_proxy, const std::string & host) {
+    auto lower = [](std::string s) {
+        for (auto & c : s) {
+            if (c >= 'A' && c <= 'Z') {
+                c += 'a' - 'A';
+            }
+        }
+        return s;
+    };
+
+    const std::string h = lower(host);
+
+    for (size_t pos = 0; pos < no_proxy.size(); ) {
+        size_t end = no_proxy.find(',', pos);
+        if (end == std::string::npos) {
+            end = no_proxy.size();
+        }
+
+        std::string entry = lower(no_proxy.substr(pos, end - pos));
+        pos = end + 1;
+
+        entry.erase(0, entry.find_first_not_of(" \t"));
+        auto last = entry.find_last_not_of(" \t");
+        if (last == std::string::npos) {
+            continue;
+        }
+        entry.erase(last + 1);
+
+        if (entry == "*") {
+            return true;
+        }
+
+        if (entry.front() == '.') {
+            entry.erase(0, 1);
+        }
+
+        if (entry.empty()) {
+            continue;
+        }
+
+        if (h == entry) {
+            return true;
+        }
+
+        if (h.size() > entry.size() && h[h.size() - entry.size() - 1] == '.' &&
+            h.compare(h.size() - entry.size(), entry.size(), entry) == 0) {
+            return true;
+        }
+    }
+
+    return false;
+}
+
 static std::pair<httplib::Client, common_http_url> common_http_client(const std::string & url) {
    common_http_url parts = common_http_parse_url(url);

@@ -123,6 +237,33 @@ static std::pair<httplib::Client, common_http_url> common_http_client(const std:

    cli.set_follow_location(true);

+    // Honor HTTP_PROXY / HTTPS_PROXY environment variables (both lower- and upper-case).
+    auto getenv_s = [](const char * upper, const char * lower) -> std::string {
+        const char * val = std::getenv(upper);
+        if (!val) {
+            val = std::getenv(lower);
+        }
+        return val ? val : "";
+    };
+
+    const std::string proxy_url = (parts.scheme == "https")
+        ? getenv_s("HTTPS_PROXY", "https_proxy")
+        : getenv_s("HTTP_PROXY",  "http_proxy");
+
+    if (!proxy_url.empty() && !common_http_is_loopback(parts.host) &&
+        !common_http_proxy_bypass(getenv_s("NO_PROXY", "no_proxy"), parts.host)) {
+        try {
+            common_http_url proxy = common_http_parse_url(proxy_url);
+            cli.set_proxy(proxy.host, proxy.port);
+            if (!proxy.user.empty()) {
+                cli.set_proxy_basic_auth(proxy.user, proxy.password);
+            }
+        } catch (const std::exception & e) {
+            // fall back to a direct connection, the URL itself is not logged as it can carry credentials
+            LOG_WRN("%s: ignoring malformed proxy URL from the environment: %s\n", __func__, e.what());
+        }
+    }
+
    return { std::move(cli), std::move(parts) };
 }

@@ -600,6 +600,7 @@ value for_statement::execute_impl(context & ctx) const {
    bool noIteration = true;
    for (size_t i = 0; i < filtered_items.size(); i++) {
        JJ_DEBUG("For loop iteration %zu/%zu", i + 1, filtered_items.size());
+        context loop_scope(scope);
        value_object loop_obj = mk_val<value_object>();
        loop_obj->has_builtins = false; // loop object has no builtins
        loop_obj->insert("index", mk_val<value_int>(i + 1));
@@ -611,11 +612,11 @@ value for_statement::execute_impl(context & ctx) const {
        loop_obj->insert("length", mk_val<value_int>(filtered_items.size()));
        loop_obj->insert("previtem", i > 0 ? filtered_items[i - 1] : mk_val<value_undefined>("previtem"));
        loop_obj->insert("nextitem", i < filtered_items.size() - 1 ? filtered_items[i + 1] : mk_val<value_undefined>("nextitem"));
-        scope.set_val("loop", loop_obj);
-        scope_update_fns[i](scope);
+        loop_scope.set_val("loop", loop_obj);
+        scope_update_fns[i](loop_scope);
        try {
            for (auto & stmt : body) {
-                value val = stmt->execute(scope);
+                value val = stmt->execute(loop_scope);
                result->push_back(val);
            }
        } catch (const continue_statement::signal &) {
@@ -74,16 +74,12 @@ struct context {
    }
    ~context() = default;

-    context(const context & parent) : context() {
+    context(const context & parent) :
+        src(parent.src),
+        current_time(parent.current_time),
+        is_get_stats(parent.is_get_stats),
        // inherit variables (for example, when entering a new scope)
-        auto & pvar = parent.env->as_ordered_object();
-        for (const auto & pair : pvar) {
-            set_val(pair.first, pair.second);
-        }
-        current_time = parent.current_time;
-        is_get_stats = parent.is_get_stats;
-        src = parent.src;
-    }
+        env(mk_val<value_object>(*parent.env)) {}

    value get_val(const std::string & name) {
        value default_val = mk_val<value_undefined>(name);
@@ -244,6 +244,8 @@ static std::string json_ensure_ascii_preserving_format(const std::string & json_
    return output;
 }

+static std::string float_repr(double val);
+
 static value tojson(const func_args & args) {
    args.ensure_count(1, 5);
    value val_ascii      = args.get_kwarg_or_pos("ensure_ascii", 1);
@@ -634,6 +636,11 @@ const func_builtins & value_int_t::get_builtins() const {


 const func_builtins & value_float_t::get_builtins() const {
+    // Python's str() of a float, nan and inf stay as they are (tojson writes NaN and Infinity)
+    static const func_handler tostring = [](const func_args & args) -> value {
+        args.ensure_vals<value_float>();
+        return mk_val<value_string>(float_repr(args.get_pos(0)->as_float()));
+    };
    static const func_builtins builtins = {
        {"default", default_value},
        {"abs", [](const func_args & args) -> value {
@@ -650,8 +657,8 @@ const func_builtins & value_float_t::get_builtins() const {
            args.ensure_vals<value_float>();
            return mk_val<value_float>(args.get_pos(0)->as_float());
        }},
-        {"safe", tojson},
-        {"string", tojson},
+        {"safe", tostring},
+        {"string", tostring},
        {"tojson", tojson},
    };
    return builtins;
@@ -1522,10 +1529,10 @@ void global_from_json(context & ctx, const common_json & json_obj, bool mark_inp
    }
 }

-// a float as Python writes it (repr, json.dumps): the shortest digits that read back as the same value
-static std::string float_to_json(double val) {
+// a float as Python's repr() and str() write it: the shortest digits that read back as the same value
+static std::string float_repr(double val) {
    if (!std::isfinite(val)) {
-        return std::isnan(val) ? "nan" : val < 0 ? "-inf" : "inf"; // as before (json.dumps writes NaN and Infinity, str() writes nan and inf)
+        return std::isnan(val) ? "nan" : val < 0 ? "-inf" : "inf";
    }
    const double mag = std::fabs(val);

@@ -1551,6 +1558,7 @@ static std::string float_to_json(double val) {
        if (prec == 16 || reads_back(d)) {
            digits = d;
        } else if (is_pow2 && d.back() != '9') {
+            // a last digit 9 is not bumped (it would carry): no power of 2, from 2^-1074 to 2^1023, needs it
            d.back()++;
            if (reads_back(d)) {
                digits = d;
@@ -1581,6 +1589,17 @@ static std::string float_to_json(double val) {
    return out;
 }

+// a float as json.dumps writes it: nan and inf as NaN and Infinity, which are not valid JSON
+static std::string float_json(double val) {
+    if (std::isnan(val)) {
+        return "NaN";
+    }
+    if (std::isinf(val)) {
+        return val < 0 ? "-Infinity" : "Infinity";
+    }
+    return float_repr(val);
+}
+
 // recursively convert value to JSON string
 // TODO: avoid circular references
 static void value_to_json_internal(std::ostringstream & oss, const value & val, int curr_lvl, int indent, const std::string_view item_sep, const std::string_view key_sep) {
@@ -1598,7 +1617,7 @@ static void value_to_json_internal(std::ostringstream & oss, const value & val,
    } else if (is_val<value_int>(val)) {
        oss << val->as_int();
    } else if (is_val<value_float>(val)) {
-        oss << float_to_json(val->as_float());
+        oss << float_json(val->as_float());
    } else if (is_val<value_string>(val)) {
        oss << "\"";
        for (char c : val->as_string().str()) {
@@ -1645,7 +1664,8 @@ static void value_to_json_internal(std::ostringstream & oss, const value & val,
            size_t i = 0;
            for (const auto & pair : obj) {
                oss << indent_str() << (indent > 0 ? std::string(indent, ' ') : "");
-                value_to_json_internal(oss, mk_val<value_string>(pair.first->as_string().str()), curr_lvl + 1, indent, item_sep, key_sep);
+                const std::string key = is_val<value_float>(pair.first) ? float_json(pair.first->as_float()) : pair.first->as_string().str();
+                value_to_json_internal(oss, mk_val<value_string>(key), curr_lvl + 1, indent, item_sep, key_sep);
                oss << key_sep;
                value_to_json_internal(oss, pair.second, curr_lvl + 1, indent, item_sep, key_sep);
                if (i < obj.size() - 1) {
@@ -68,7 +68,7 @@ static llama_token try_draft(common_ngram_cache & nc_static, const common_ngram
    if (part_static_it == nc_static.end()) {
        return LLAMA_TOKEN_NULL;
    }
-    const common_ngram_cache_part part_static = part_static_it->second;
+    const common_ngram_cache_part & part_static = part_static_it->second;

    int max_count_static  = 0;
    int sum_count_static  = 0;
@@ -96,7 +96,7 @@ static llama_token try_draft(common_ngram_cache & nc_static, const common_ngram

 // Try to draft a token from primary cache (context/dynamic), validate with static cache:
 static llama_token try_draft(
-    common_ngram_cache & nc_primary, const std::vector<common_ngram> & ngrams_primary, common_ngram_cache_part & part_static,
+    common_ngram_cache & nc_primary, const std::vector<common_ngram> & ngrams_primary, const common_ngram_cache_part & part_static,
    const int * min_sample_size, const int * min_percent) {

    llama_token drafted_token = LLAMA_TOKEN_NULL;
@@ -108,7 +108,7 @@ static llama_token try_draft(
        if (part_primary_it == nc_primary.end()) {
            continue;
        }
-        const common_ngram_cache_part part_primary = part_primary_it->second;
+        const common_ngram_cache_part & part_primary = part_primary_it->second;

        int max_count_primary = 0;
        int max_count_static  = 0;
@@ -118,7 +118,7 @@ static llama_token try_draft(
        for (std::pair<llama_token, int> token_count_primary : part_primary) {
            const llama_token token = token_count_primary.first;

-            common_ngram_cache_part::iterator token_count_static_it = part_static.find(token);
+            common_ngram_cache_part::const_iterator token_count_static_it = part_static.find(token);

            const int32_t count_primary = token_count_primary.second;
            const int32_t count_static  = token_count_static_it != part_static.end() ? 100*token_count_static_it->second : 1;
@@ -162,11 +162,9 @@ void common_ngram_cache_draft(
        for (int j = ngram_start_static; j < ngram_start_static + LLAMA_NGRAM_STATIC; ++j) {
            ngram_static.tokens[j-ngram_start_static] = get_token(inp, draft, j);
        }
-        common_ngram_cache::iterator part_static_it = nc_static.find(ngram_static);
-        common_ngram_cache_part part_static;
-        if (part_static_it != nc_static.end()) {
-            part_static = part_static_it->second;
-        }
+        static const common_ngram_cache_part part_static_empty;
+        common_ngram_cache::const_iterator part_static_it = nc_static.find(ngram_static);
+        const common_ngram_cache_part & part_static = part_static_it != nc_static.end() ? part_static_it->second : part_static_empty;

        // cd = context + dynamic
        std::vector<common_ngram> ngrams_cd;
@@ -0,0 +1,146 @@
+#include "parsers.h"
+
+// An assistant turn is rendered as one or more messages, each
+// "<|start|>assistant to=<recipient><|message|>{content}{END}" where END is
+// <|eom|> (more messages follow) or <|eot|> (end of turn):
+//   - chain-of-thought:  to=self, terminated by <|eom|>
+//   - user facing aside: to=commentary, terminated by <|eom|>
+//   - final answer:      to=user, terminated by <|eot|>
+// A commentary message carries content the user sees while the assistant keeps
+// working, typically before or between tool calls.
+// The generation prompt is just "<|start|>assistant"; the model emits its own
+// " to=...<|message|>".
+common_chat_params common_chat_params_init_muse_spark(const common_chat_template &          tmpl,
+                                                               const autoparser::generation_params & inputs) {
+    common_chat_params data;
+
+    data.prompt            = common_chat_template_direct_apply_impl(tmpl, inputs);
+    data.generation_prompt = "<|start|>assistant";
+    data.format            = COMMON_CHAT_FORMAT_PEG_NATIVE;
+    data.supports_thinking = true;
+
+    data.preserved_tokens = {
+        "<|start|>", "<|message|>", "<|eom|>", "<|eot|>",
+        // ATEM tool-call markup emitted on " to=<tool>" turns.
+        "<atem:function_calls>", "<atem:invoke", "<atem:parameter", "</atem:parameter>",
+        "</atem:invoke>", "</atem:function_calls>",
+    };
+
+    data.message_delimiters = {
+        { COMMON_CHAT_ROLE_ASSISTANT, "<|start|>assistant" },
+        { COMMON_CHAT_ROLE_USER,      "<|start|>user"      },
+        { COMMON_CHAT_ROLE_SYSTEM,    "<|start|>system"    },
+        { COMMON_CHAT_ROLE_TOOL,      "<|start|>tool"      },
+    };
+
+    if (inputs.has_continuation()) {
+        const auto & msg = inputs.continue_msg;
+
+        data.generation_prompt = "<|start|>assistant to=self<|message|>" + msg.reasoning_content;
+        if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) {
+            data.generation_prompt += "<|eom|><|start|>assistant to=user<|message|>" + msg.render_content();
+        }
+
+        data.prompt.append(data.generation_prompt);
+    }
+
+    auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
+
+    auto has_tools = inputs.tools.is_array() && !inputs.tools.empty();
+    // Constrained grammar whenever tools are offered.
+    auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE;
+
+    data.parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
+        auto start = p.rule("start", p.literal("<|start|>assistant"));
+
+        if (!extract_reasoning && !include_grammar) {
+            return start + p.content(p.rest());
+        }
+
+        if (extract_reasoning) {
+            p.rule("analysis", p.literal(" to=self<|message|>") + p.reasoning(p.until("<|eom|>")) + p.literal("<|eom|>"));
+        } else {
+            p.rule("analysis", p.literal(" to=self<|message|>") + p.content(p.until("<|eom|>")) + p.literal("<|eom|>"));
+        }
+        auto analysis = p.ref("analysis");
+
+        // Commentary is user facing content emitted mid turn, so it reads as
+        // content and always ends on <|eom|>.
+        auto preamble = p.rule("preamble", p.literal(" to=commentary<|message|>") +
+                                               p.content(p.until("<|eom|>")) + p.literal("<|eom|>"));
+
+        auto aside = p.rule("aside", analysis | preamble);
+
+        auto recipient  = p.optional(p.literal(" to=user"));
+        auto final_msg  = p.rule("final", recipient + p.literal("<|message|>") +
+                                              p.content(p.until_one_of({ "<|eot|>", "<|eom|>" })));
+
+        if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
+            auto string_value = p.ac(
+                p.tool_arg_string_value(p.until("</atem:parameter>")) + p.tool_arg_close(p.literal("</atem:parameter>")),
+                "</atem:parameter>");
+
+            auto tool_choice = p.choice();
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
+                const auto &      function = tool.at("function");
+                const std::string name     = function.at("name");
+
+                std::vector<common_peg_parser> arg_rules;
+                foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) {
+                    auto value_parser = p.eps();
+                    if (prop.schema->may_be_string()) {
+                        value_parser = string_value;
+                    } else {
+                        value_parser = p.tool_arg_json_value(
+                                p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index) + "-schema", doc, *prop.schema))
+                            + p.tool_arg_close(p.literal("</atem:parameter>"));
+                    }
+
+                    arg_rules.push_back(p.tool_arg(
+                        p.tool_arg_open(p.literal("<atem:parameter name=\"") + p.tool_arg_name(p.literal(prop.name)) + p.literal("\">")) +
+                        value_parser));
+                });
+
+                auto args = p.eps();
+                if (!arg_rules.empty()) {
+                    args = p.zero_or_more(p.choice(arg_rules) + p.space());
+                }
+
+                auto tool_parser = p.tool(
+                    p.tool_open(p.literal(" to=") + p.until("<|message|>") +
+                                p.literal("<|message|><atem:function_calls>") + p.space() +
+                                p.literal("<atem:invoke name=\"") + p.tool_name(p.literal(name)) + p.literal("\">") + p.space())
+                    << p.tool_args(args)
+                    << p.tool_close(p.literal("</atem:invoke>") + p.space() + p.literal("</atem:function_calls>")));
+
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_parser);
+            });
+
+            auto tool_calls = inputs.parallel_tool_calls
+                ? p.trigger_rule("tool-call", tool_choice + p.zero_or_more(p.literal("<|eom|>") + start + tool_choice))
+                : p.trigger_rule("tool-call", tool_choice);
+
+
+            if (inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED) {
+                return p.zero_or_more(start + aside) + start + tool_calls;
+            }
+            auto trailing_calls = p.optional(p.literal("<|eom|>") + start + tool_calls);
+            return p.zero_or_more(start + aside) + start + (tool_calls | (final_msg + trailing_calls));
+        }
+
+        return p.zero_or_more(start + aside) + start + final_msg;
+    });
+
+    if (include_grammar) {
+        data.grammar_lazy = inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_REQUIRED;
+        data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
+            data.parser.build_grammar(builder, data.grammar_lazy);
+        });
+        data.grammar_triggers = {
+            { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN,
+              "<\\|start\\|>assistant( to=(?!self<\\|message\\|>)(?!commentary<\\|message\\|>)(?!user<\\|message\\|>)[^<]*?<\\|message\\|>)" },
+        };
+    }
+
+    return data;
+}
@@ -79,6 +79,7 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
 common_chat_params common_chat_params_init_ministral_3(const common_chat_template & tmpl, const autoparser::generation_params & inputs);

 common_chat_params common_chat_params_init_muse_glimmer(const common_chat_template & tmpl, const autoparser::generation_params & inputs);
+common_chat_params common_chat_params_init_muse_spark(const common_chat_template & tmpl, const autoparser::generation_params & inputs);

 common_chat_params common_chat_params_init_qwen3_coder(const common_chat_template & tmpl, const autoparser::generation_params & inputs);

@@ -19,6 +19,7 @@ set(LLAMA_CHAT_PARSERS_SOURCES
    ${CMAKE_CURRENT_LIST_DIR}/minimax-m3.cpp
    ${CMAKE_CURRENT_LIST_DIR}/ministral3.cpp
    ${CMAKE_CURRENT_LIST_DIR}/muse-glimmer.cpp
+    ${CMAKE_CURRENT_LIST_DIR}/muse-spark.cpp
    ${CMAKE_CURRENT_LIST_DIR}/qwen3-coder.cpp
    ${CMAKE_CURRENT_LIST_DIR}/translate-gemma.cpp
 )
@@ -1338,7 +1338,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
                GGML_ASSERT(lattice && "DFlash2 selector produced no lattice");

                int32_t predecessor = 0;
-                for (int32_t i = 1; i < n_block_tokens; ++i) {
+                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;
                    const int32_t top_k = is_dflash2 ? selector_top_k : lilicorr_top_k;
                    const float * scores = row + top_k + (size_t) predecessor * top_k;
@@ -199,6 +199,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "MiniMaxM3SparseForConditionalGeneration": "minimax",
    "Ministral3ForCausalLM": "mistral3",
    "Mistral3ForConditionalGeneration": "mistral3",
+    "VoxtralRealtimeForConditionalGeneration": "mistral3",
    "MistralForCausalLM": "llama",
    "MixtralForCausalLM": "llama",
    "ModernBertDecisionModel": "bert",
@@ -220,6 +221,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "OlmoeForCausalLM": "olmo",
    "MuseGlimmerAssistantModel": "muse_glimmer",
    "MuseGlimmerForConditionalGeneration": "muse_glimmer",
+    "MuseSparkForConditionalGeneration": "muse_spark",
    "OpenELMForCausalLM": "openelm",
    "OpenJevModel": "qwen",
    "OrionForCausalLM": "orion",
@@ -353,8 +355,10 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "MiniCPMV4_6ForConditionalGeneration": "minicpm",
    "MiniCPMV4_7ForConditionalGeneration": "minicpm",
    "Mistral3ForConditionalGeneration": "llava",
+    "NemotronH_Omni_Reasoning_V3": "nemotron",
    "NemotronH_Nano_VL_V2": "nemotron",
    "MuseGlimmerForConditionalGeneration": "muse_glimmer",
+    "MuseSparkForConditionalGeneration": "muse_spark",
    "PaddleOCRVisionModel": "ernie",
    "Phi4ForCausalLMV": "phi",
    "Qwen2AudioForConditionalGeneration": "ultravox",
@@ -380,6 +384,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "UltravoxModel": "ultravox",
    "UnlimitedOCRForCausalLM": "deepseek",
    "VoxtralForConditionalGeneration": "ultravox",
+    "VoxtralRealtimeForConditionalGeneration": "ultravox",
    "YoutuVLForConditionalGeneration": "youtuvl",
 }

@@ -568,7 +568,7 @@ class ModelBase:
                                tensors_to_remove.append(base_name + "_zero_point")
                else:
                    raise NotImplementedError(f"Quant format {quant_format!r} for method {quant_method!r} is not yet supported")
-            elif quant_method == "modelopt":
+            elif quant_method in ("modelopt", "modelopt_mixed"):
                dequant_fp8()
            elif quant_method is not None:
                raise NotImplementedError(f"Quant method is not yet supported: {quant_method!r}")
@@ -1869,6 +1869,9 @@ class TextModel(ModelBase):
        if chkhsh == "66b8d4e19ab16c3bfd89bce5d785fb7e0155e8648708a1f42077cb9fe002c273":
            # ref: https://huggingface.co/alvarobartt/grok-2-tokenizer
            res = "grok-2"
+        if chkhsh == "f5c965e4198da6e7525072d87d1ecf4c51d981fd486c7c4ec728ef9cec7e03fc":
+            # ref: https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B
+            res = "gpt-4o"
        if chkhsh == "b3d1dd861f1d4c5c0d2569ce36baf3f90fe8a102db3de50dd71ff860d91be3df":
            # ref: https://huggingface.co/aari1995/German_Semantic_V3
            res = "jina-v2-de"
@@ -243,7 +243,7 @@ class DeepseekV2Model(TextModel):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
-        hparams: dict = ModelBase.load_hparams(self.dir_model, is_mistral_format=False)
+        hparams: dict = ModelBase.load_hparams(self.dir_model, is_mistral_format=self.is_mistral_format)
        self.origin_hf_arch = hparams.get('architectures', [None])[0]

        # special handling for Deepseek OCR
@@ -8,7 +8,7 @@ import torch
 if TYPE_CHECKING:
    from torch import Tensor

-from .base import MmprojModel, ModelBase, gguf, logger
+from .base import MmprojModel, ModelBase, gguf, jinja_str_or_json, logger

 from .llama import LlamaModel
 from .mamba import Mamba2Model
@@ -239,8 +239,21 @@ class GraniteSwitchModel(GraniteMoeModel):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
-        # the weightless switch reserves one cache slot: one fewer block than num_hidden_layers
-        self.block_count = self.block_count - 1
+        # The weightless switch layer(s) are counted in num_hidden_layers but carry no tensors.
+        # The count is not fixed, nor is it present in hparams, so we derive it here based on tensor block names.
+        layer_ids = {
+            int(m.group(1))
+            for name in self.model_tensors
+            if (m := re.match(r"(?:model\.)?layers\.(\d+)\.", name))
+        }
+        if not layer_ids:
+            raise ValueError("graniteswitch: no 'model.layers.N' tensors found in checkpoint")
+        n_weight_layers = max(layer_ids) + 1
+        logger.info(
+            "gguf: (graniteswitch) num_hidden_layers %s counts %d weightless switch layer(s); using %d blocks",
+            self.hparams["num_hidden_layers"], self.block_count - n_weight_layers, n_weight_layers,
+        )
+        self.block_count = n_weight_layers
        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)

        self._n_adapters = int(self.hparams["num_adapters"])
@@ -271,10 +284,39 @@ class GraniteSwitchModel(GraniteMoeModel):
        self.gguf_writer.add_adapter_lora_rank(self._max_lora_rank)
        self.gguf_writer.add_adapter_token_ids_activate(self.hparams["adapter_token_ids"])
        self.gguf_writer.add_adapter_token_ids_substitute(self.hparams["adapter_substitute_token_ids"])
+        if self._has_decision_adapter():
+            self.gguf_writer.add_decision_type(gguf.DecisionType.GRANITE_SWITCH)
        router_gain = float(self.hparams.get("control_token_gain", 15.0))
        self.gguf_writer.add_adapter_router_gain(router_gain)
        logger.info("gguf: (graniteswitch) num_adapters=%s max_lora_rank=%s n_slots=%s router_gain=%s", self._n_adapters, self._max_lora_rank, self._n_slots, router_gain)

+    def _has_decision_adapter(self) -> bool:
+        return "decision" in self.hparams.get("adapter_names", [])
+
+    def set_vocab(self):
+        super().set_vocab()
+        if self._has_decision_adapter():
+            self.gguf_writer.add_chat_template([{"name": "systemone", "template": self._systemone_template()}])
+
+    def _systemone_template(self) -> str:
+        # prompt of the decision aLoRA: a state turn, then a question turn that opens with the control token
+        # (it replaces the first token of "### DECISION"), the answer is the label token that comes next
+        description = jinja_str_or_json("o.description")
+        option = (
+            "{% if type != 'noul' %}{{ letters[loop.index0] }}: {% if o.description %}" + description + "{% endif %}"
+            "{% elif o.key == 'true' %}A: {% if o.description %}" + description + "{% else %}The statement is true.{% endif %}"
+            "{% else %}B: {% if o.description %}" + description + "{% else %}The statement is false.{% endif %}{% endif %}"
+        )
+        return (
+            "{% set letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz' %}"
+            "<|start_of_role|>user<|end_of_role|>State:\n" + jinja_str_or_json("state") + "<|end_of_text|>\n"
+            "<|start_of_role|>user<|end_of_role|><|decision|> DECISION\nQuestion:\n" + jinja_str_or_json("instructions")
+            + "{% if type == 'score' %} Rate along the ordered levels below (lowest first).{% endif %}"
+            "{{ '\n\nOptions:\n' }}"
+            "{% for o in options %}" + option + "{{ '\n' }}{% endfor %}"
+            "{{ '\nReply with exactly one option label.<|end_of_text|>\n<|start_of_role|>assistant<|end_of_role|>' }}"
+        )
+
    def _lora_a(self, data: Tensor) -> Tensor:
        # on-disk A: [n_adapters, 1, max_rank, in] -> [n_adapters+1, max_rank, in]
        a = data.squeeze(1)
@@ -188,7 +188,7 @@ def _load_d1_omni_hparams(dir_model: Path) -> dict[str, Any]:
 class D1OmniModel(LFM2Model):
    model_arch = gguf.MODEL_ARCH.LFM2

-    # the server cuts the text to these lengths, see server-decision.cpp
+    # the server cuts the text to these lengths, see engine-decision.cpp
    _MAX_LENGTH = 16384
    _IMAGE_TEXT_LENGTH = 896
    _AUDIO_TEXT_LENGTH = 15360
@@ -1,6 +1,6 @@
 from __future__ import annotations

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

 import torch

@@ -14,6 +14,54 @@ from .base import ModelBase, TextModel, gguf, logger
@ModelBase.example("JetBrains/Mellum2-12B-A2.5B-Base")
 class MellumModel(TextModel):
    model_arch = gguf.MODEL_ARCH.MELLUM
+    supports_mtp_export = True
+
+    _n_main_layers: int | None = None
+
+    _experts: list[dict[str, Tensor]] | None = None
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # number of exported MTP blocks, appended after the trunk layers
+        self.n_nextn = 0 if self.no_mtp else (self.find_hparam(["num_nextn_predict_layers", "num_mtp_layers"], optional=True) or 0)
+        if self.mtp_only and self.n_nextn == 0:
+            raise ValueError("--mtp requested, but the checkpoint has no MTP layers (num_nextn_predict_layers is 0)")
+        if self.n_nextn > 0:
+            if any(t != "sparse" for t in self.hparams.get("mtp_mlp_layer_types") or []):
+                raise ValueError("Mellum MTP layers with dense MLP are not supported")
+            self.block_count += self.n_nextn
+            self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
+
+    def index_tensors(self, remote_hf_model_id: str | None = None):
+        type(self)._n_main_layers = self.hparams["num_hidden_layers"]
+        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        if (titem := super().filter_tensors(item)) is None:
+            return None
+        name, gen = titem
+
+        if name.startswith("mtp."):
+            if cls.no_mtp:
+                return None
+            assert cls._n_main_layers is not None
+            # mtp.layers.{i}.* -> model.layers.{n_main + i}.*, mtp.{fc,norm,...} -> nextn names of the first MTP block
+            remapper = {
+                "fc":                    "eh_proj",
+                "pre_fc_norm_embedding": "enorm",
+                "pre_fc_norm_hidden":    "hnorm",
+                "norm":                  "shared_head.norm",
+            }
+            parts = name.split(".", 3)
+            if len(parts) == 4 and parts[1] == "layers" and parts[2].isdecimal():
+                name = f"model.layers.{cls._n_main_layers + int(parts[2])}.{parts[3]}"
+            elif len(parts) == 3 and parts[1] in remapper:
+                name = f"model.layers.{cls._n_main_layers}.{remapper[parts[1]]}.{parts[2]}"
+        elif cls.mtp_only and name not in ("model.embed_tokens.weight", "model.norm.weight", "lm_head.weight"):
+            return None
+
+        return name, gen

    def set_gguf_parameters(self):
        super().set_gguf_parameters()
@@ -26,10 +74,31 @@ class MellumModel(TextModel):
        if (use_sliding_window is True or use_sliding_window is None) and sliding_window is not None:
            self.gguf_writer.add_sliding_window(sliding_window)
            logger.info(f"gguf: sliding window = {sliding_window}")
-            self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in self.hparams["layer_types"]])
-            logger.info(f"gguf: sliding window pattern length = {len(self.hparams['layer_types'])}")
+            layer_types = list(self.hparams["layer_types"])
+            # MTP blocks default to sliding window attention, as in transformers
+            if self.n_nextn > 0:
+                layer_types += self.hparams.get("mtp_layer_types") or ["sliding_attention"] * self.n_nextn
+            self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in layer_types])
+            logger.info(f"gguf: sliding window pattern length = {len(layer_types)}")

-    _experts: list[dict[str, Tensor]] | None = None
+        if (norm_topk_prob := self.hparams.get("norm_topk_prob")) is not None:
+            self.gguf_writer.add_expert_weights_norm(norm_topk_prob)
+
+        if self.n_nextn > 0:
+            self.gguf_writer.add_nextn_predict_layers(self.n_nextn)
+
+    def prepare_metadata(self, vocab_only: bool):
+        from_dir = self.fname_out.is_dir()
+        super().prepare_metadata(vocab_only=vocab_only)
+
+        if not self.mtp_only or not from_dir:
+            return
+
+        output_type: str = self.ftype.name.partition("_")[2]
+        fname_default: str = gguf.naming_convention(
+            self.metadata.name, self.metadata.basename, self.metadata.finetune,
+            self.metadata.version, size_label=None, output_type=output_type, model_type=None)
+        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"

    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
        if name.find("experts") != -1:
@@ -66,3 +66,13 @@ class Mistral3Model(TextModel):

    def write(self):
        self.impl.write()
+
+
+@ModelBase.register("VoxtralRealtimeForConditionalGeneration")
+@ModelBase.example("mistralai/Voxtral-Mini-4B-Realtime-2602")
+class VoxtralRealtimeModel(LlamaModel):
+    model_arch = gguf.MODEL_ARCH.MISTRAL3
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_realtime_delay_tokens(self.hparams["default_num_delay_tokens"])
@@ -94,7 +94,7 @@ class MuseGlimmerVisionModel(MmprojModel):
        assert self.hparams_vision is not None
        c = self.hparams_vision  # enriched vision_config from get_vision_config()

-        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.MUSE_GLIMMER)
+        self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.MUSE_GLIMMER)
        self.gguf_writer.add_vision_attention_layernorm_eps(float(c["layer_norm_eps"]))
        self.gguf_writer.add_vision_spatial_merge_size(int(c["merge_size"]))

@@ -0,0 +1,221 @@
+from __future__ import annotations
+
+from typing import Any, Iterable, TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import ModelBase, MmprojModel, TextModel, gguf
+from .muse_glimmer import MuseGlimmerVisionModel
+
+
+# The reference clamps the squared stream coefficient with this floor.
+RESIDUAL_GATE_EPS = 1e-3
+
+# Block norms hold a delta around one; the output norm holds an absolute gain.
+OFFSET_NORM_SUFFIXES = (
+    "input_layernorm.weight",
+    "pre_feedforward_layernorm.weight",
+    "post_feedforward_layernorm.weight",
+    "pre_expert_norm.weight",
+    "post_expert_norm.weight",
+)
+
+
+def _unpermute_for_rope(tensor: "Tensor", n_heads: int) -> "Tensor":
+    """Invert transformers' `_permute_for_rope`: HF stores Q/K in rotate_half layout,
+    llama.cpp consumes the interleaved (NORM) layout. Vision Q/K carry a bias too."""
+    if tensor.ndim == 1:
+        dim1, = tensor.shape
+        return tensor.view(n_heads, 2, dim1 // n_heads // 2).transpose(1, 2).reshape(dim1)
+    dim1, dim2 = tensor.shape
+    return tensor.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2)
+
+
+@ModelBase.register("MuseSparkForConditionalGeneration")
+@ModelBase.example("meta-models/Muse-Spark-Open-816B-A42B")
+class MuseSparkModel(TextModel):
+    model_arch = gguf.MODEL_ARCH.MUSE_SPARK
+
+    def set_vocab(self):
+        self._set_vocab_gpt2()
+
+        from transformers import AutoTokenizer
+        tok = AutoTokenizer.from_pretrained(self.dir_model)
+        eot_id = tok.convert_tokens_to_ids("<|eot|>")  # ty: ignore[unresolved-attribute]
+        if isinstance(eot_id, int) and eot_id >= 0:
+            self.gguf_writer.add_eot_token_id(eot_id)
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        hparams = self.hparams
+
+        # MetaP resolves to two scalars
+        # The output multiplier lands here and the attention scale rides on the synthesized QK-norm weights
+        self.gguf_writer.add_logit_scale(hparams["output_multiplier"])
+        self.gguf_writer.add_final_logit_softcapping(hparams["final_logit_softcapping"])
+
+        self.gguf_writer.add_post_norm_eps(hparams["post_norm_eps"])
+
+        # Only the sliding layers carry rope, so the graph takes it from the pattern
+        sliding = [t == "sliding_attention" for t in hparams["layer_types"]]
+        theta = hparams["rope_parameters"]["rope_theta"]
+        assert sliding == [bool(t) for t in hparams["layer_rope_theta"]]
+        assert all(t in (0, theta) for t in hparams["layer_rope_theta"])
+        assert all(t == "sparse" for t in hparams["mlp_layer_types"])
+
+        self.gguf_writer.add_sliding_window(hparams["sliding_window"])
+        self.gguf_writer.add_rope_freq_base_swa(theta)
+        self.gguf_writer.add_sliding_window_pattern(sliding)
+
+        # Experts live in a latent space narrower than the residual stream.
+        self.gguf_writer.add_moe_latent_size(hparams["moe_hidden_size"])
+        self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"])
+        self.gguf_writer.add_expert_weights_scale(hparams["routed_scaling_factor"])
+        self.gguf_writer.add_expert_weights_norm(True)
+        # the config carries no score function, the reference hardcodes sigmoid
+        self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name.endswith(OFFSET_NORM_SUFFIXES):
+            data_torch = data_torch + 1.0
+
+        if name.endswith(".self_attn.q_proj.weight"):
+            assert bid is not None
+            n_head = self.hparams["num_attention_heads"]
+            head_dim = self.hparams["head_dim"]
+
+            yield self.map_tensor_name(name), _unpermute_for_rope(data_torch, n_head)
+
+            # Synthesize QK-norm to carry the MetaP scale so the graph keeps the plain 1 / sqrt(head_dim)
+            ones = torch.ones(head_dim, dtype=torch.float32)
+            q_norm = ones * self.hparams["qk_scale_factor"]
+            yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q_NORM, bid), q_norm
+            yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K_NORM, bid), ones
+            return
+
+        if name.endswith(".self_attn.k_proj.weight"):
+            n_head_kv = self.hparams["num_key_value_heads"]
+            yield self.map_tensor_name(name), _unpermute_for_rope(data_torch, n_head_kv)
+            return
+
+        # gate_up_proj fuses both projections over the output dim, gate first.
+        if name.endswith(".mlp.experts.gate_up_proj"):
+            assert bid is not None
+            n_ff = data_torch.shape[1] // 2
+            yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_EXP, bid), data_torch[:, :n_ff, :]
+            yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_EXP, bid), data_torch[:, n_ff:, :]
+            return
+
+        if name.endswith(".mlp.experts.down_proj"):
+            assert bid is not None
+            yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_EXP, bid), data_torch
+            return
+
+        # The router scales its logits by the MetaP output multiplier before the sigmoid
+        # Fold it into the weight to keep a plain matmul
+        if name.endswith(".mlp.gate.weight"):
+            data_torch = data_torch.float() * self.hparams["output_multiplier"]
+            yield from super().modify_tensors(data_torch, name, bid)
+            return
+
+        # Norm preserving residual gate: alpha scales the stream, beta the branch, with alpha^2 + beta^2 = 1
+        # Both follow from the weight, so resolve them here
+        if name.endswith("_residual_gate.gate"):
+            assert bid is not None
+            attn = "post_attention" in name
+
+            w = data_torch.float() / self.hparams["residual_gate_temperature"]
+            beta = torch.sigmoid(w)
+            alpha_sq = torch.sigmoid(-w) * (1.0 + beta)
+            alpha = torch.sqrt(torch.clamp(alpha_sq, min=RESIDUAL_GATE_EPS))
+
+            t_alpha = gguf.MODEL_TENSOR.ATTN_RESID_ALPHA if attn else gguf.MODEL_TENSOR.FFN_RESID_ALPHA
+            t_beta = gguf.MODEL_TENSOR.ATTN_RESID_BETA if attn else gguf.MODEL_TENSOR.FFN_RESID_BETA
+            yield self.format_tensor_name(t_alpha, bid), alpha
+            yield self.format_tensor_name(t_beta, bid), beta
+            return
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+
+@ModelBase.register("MuseSparkForConditionalGeneration")
+@ModelBase.example("meta-models/Muse-Spark-Open-816B-A42B")
+class MuseSparkVisionModel(MuseGlimmerVisionModel):
+    has_vision_encoder = True # same vision encoder as Muse Glimmer
+    has_audio_encoder = True
+
+    def get_audio_config(self) -> dict[str, Any]:
+        c = self.global_config["audio_config"]
+        # the base class wants one width and one layer count: use the last stage width and the total number of conv blocks
+        return {
+            **c,
+            "hidden_size": c["hidden_sizes"][-1],
+            "intermediate_size": int(c["hidden_sizes"][-1] * c["mlp_ratio"]),
+            "num_attention_heads": 1,
+            "num_hidden_layers": sum(c["depths"]),
+        }
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+
+        assert self.hparams_audio
+        a = self.hparams_audio
+
+        self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.MUSE_SPARK_AUDIO)
+        # the input is the raw waveform, one channel takes the place of the mel bins
+        self.gguf_writer.add_audio_num_mel_bins(1)
+        self.gguf_writer.add_audio_attention_layernorm_eps(float(a["rms_norm_eps"]))
+        self.gguf_writer.add_audio_conv_kernel_size(int(a["conv_kernel_size"]))
+        self.gguf_writer.add_audio_hidden_sizes([int(x) for x in a["hidden_sizes"]])
+        self.gguf_writer.add_audio_patch_sizes([int(x) for x in a["patch_sizes"]])
+        self.gguf_writer.add_audio_depths([int(x) for x in a["depths"]])
+
+    @classmethod
+    def filter_tensors(cls, item):
+        name, gen = item
+        if name.startswith("model.audio_tower."):
+            return MmprojModel.filter_tensors((name, gen))
+        return super().filter_tensors((name, gen))
+
+    _AUDIO_HEAD = {
+        "model.audio_tower.norm": gguf.MODEL_TENSOR.A_PRE_NORM,
+        "model.audio_tower.proj": gguf.MODEL_TENSOR.A_MMPROJ_FC,
+        "model.audio_tower.post_proj_norm": gguf.MODEL_TENSOR.A_POST_NORM,
+    }
+
+    def _audio_tensor(self, data_torch, name, suffix):
+        assert self.hparams_audio
+
+        stem = name[: -len(suffix) - 1]
+
+        if stem in self._AUDIO_HEAD:
+            return self.format_tensor_name(self._AUDIO_HEAD[stem], suffix="." + suffix), data_torch
+
+        parts = stem.split(".")
+        stage = int(parts[3])
+
+        if parts[4] == "patch_embedder":
+            key = (gguf.MODEL_TENSOR.A_ENC_PATCH_PROJ if parts[5] == "proj"
+                   else gguf.MODEL_TENSOR.A_ENC_PATCH_NORM)
+            return self.format_tensor_name(key, bid=stage, suffix="." + suffix), data_torch
+
+        # conv blocks are numbered per stage upstream and flat here
+        depths = [int(d) for d in self.hparams_audio["depths"]]
+        bid = sum(depths[:stage]) + int(parts[5])
+        tail = ".".join(parts[6:])
+        key = {
+            "conv": gguf.MODEL_TENSOR.A_ENC_CONV1D,
+            "norm": gguf.MODEL_TENSOR.A_ENC_CONV1D_NORM,
+            "mlp.fc1": gguf.MODEL_TENSOR.A_ENC_FFN_UP,
+            "mlp.fc2": gguf.MODEL_TENSOR.A_ENC_FFN_DOWN,
+        }[tail]
+        return self.format_tensor_name(key, bid=bid, suffix="." + suffix), data_torch
+
+    def modify_tensors(self, data_torch, name, bid):
+        if name.startswith("model.audio_tower."):
+            yield self._audio_tensor(data_torch, name, name.rpartition(".")[2])
+            return
+        yield from super().modify_tensors(data_torch, name, bid)
@@ -1,6 +1,6 @@
 from __future__ import annotations

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

 import numpy as np
 import torch
@@ -89,19 +89,22 @@ class LiLiCorrDraftModel(DFlashModel):
        self.gguf_writer.add_lilicorr_vector_eps(float(dflash_config["lilicorr_vector_eps"]))
        self.gguf_writer.add_lilicorr_logit_scale(float(dflash_config["lilicorr_logit_scale"]))

+    def map_tensor_name(self, name: str, try_suffixes: Sequence[str] = (".weight", ".bias")) -> str:
+        if name.startswith("model.lilicorr.layers."):
+            parts = name.split(".", 4)
+            if not parts[4].startswith("attn.in_proj."):
+                parts[3] = str(self.block_count + int(parts[3]))
+            name = ".".join(parts)
+        elif name.startswith("model.lilicorr.feature_norm."):
+            feature_bid = self.block_count + self.lilicorr_geometry["num_layers"]
+            name = name.replace("model.lilicorr.feature_norm.", f"model.lilicorr.feature_mlp.{feature_bid}.norm.")
+        elif name.startswith("model.lilicorr.feature_mlp."):
+            feature_bid = self.block_count + self.lilicorr_geometry["num_layers"]
+            name = name.replace("model.lilicorr.feature_mlp.", f"model.lilicorr.feature_mlp.{feature_bid}.")
+        return super().map_tensor_name(name, try_suffixes)
+
    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
        if name.startswith("model.lilicorr."):
-            if name.startswith("model.lilicorr.layers."):
-                parts = name.split(".", 4)
-                if not parts[4].startswith("attn.in_proj."):
-                    parts[3] = str(self.block_count + int(parts[3]))
-                name = ".".join(parts)
-            elif name.startswith("model.lilicorr.feature_norm."):
-                feature_bid = self.block_count + self.lilicorr_geometry["num_layers"]
-                name = name.replace("model.lilicorr.feature_norm.", f"model.lilicorr.feature_mlp.{feature_bid}.norm.")
-            elif name.startswith("model.lilicorr.feature_mlp."):
-                feature_bid = self.block_count + self.lilicorr_geometry["num_layers"]
-                name = name.replace("model.lilicorr.feature_mlp.", f"model.lilicorr.feature_mlp.{feature_bid}.")
            yield self.map_tensor_name(name), data_torch
            return
        yield from super().modify_tensors(data_torch, name, bid)
@@ -251,6 +254,113 @@ class NemotronNanoV2VLModel(MmprojModel):
            yield mapped_name, tensor


+@ModelBase.register("NemotronH_Omni_Reasoning_V3")
+class NemotronOmniVLModel(NemotronNanoV2VLModel):
+    _omni_tensor_prefixes = (
+        "vision_model.embeddings.",
+        "vision_model.encoder.layer.",
+        "vision_projector.mlp1.",
+    )
+
+    def get_vision_config(self) -> dict[str, Any] | None:
+        vision_config = dict(self.global_config["vision_config"])
+        vision_config["intermediate_size"] = int(vision_config["hidden_size"] * vision_config["mlp_ratio"])
+        vision_config["image_size"] = 512  # warmup only; the full position table is resized at runtime
+        return vision_config
+
+    def set_gguf_parameters(self):
+        preprocessor_config = self.preprocessor_config
+
+        preprocessor_config["image_mean"] = preprocessor_config["norm_mean"]
+        preprocessor_config["image_std"] = preprocessor_config["norm_std"]
+
+        MmprojModel.set_gguf_parameters(self)
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_OMNI)
+        self.gguf_writer.add_vision_attention_layernorm_eps(1e-6)
+        self.gguf_writer.add_vision_use_gelu(True)
+        downsample_ratio = self.global_config.get("downsample_ratio", 0.5)
+        self.gguf_writer.add_vision_projector_scale_factor(int(1.0 / downsample_ratio))
+        patch_area = self.hparams["patch_size"] ** 2
+        self.gguf_writer.add_vision_min_pixels(preprocessor_config["min_num_patches"] * patch_area)
+        self.gguf_writer.add_vision_max_pixels(preprocessor_config["max_num_patches"] * patch_area)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+
+        # Transformers can rename RADIO and projector tensors in Omni checkpoints.
+        # Accept both checkpoint layouts.
+        if name == "vision_model.summary_idxs":
+            # This buffer is only used by RadioModel.summary. The multimodal
+            # projector consumes RadioModel.features for both images and video.
+            return None
+        if name.startswith(cls._omni_tensor_prefixes):
+            return name, gen
+
+        if "radio_model.model.patch_generator.video_embedder" in name or name.startswith("vision_projector.vision_final_layernorm."):
+            return MmprojModel.filter_tensors((name, gen))
+
+        return super().filter_tensors(item)
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if new_name in ("v.position_embd.weight", "v.class_embd"):
+            return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name.startswith(self._omni_tensor_prefixes):
+            if name == "vision_model.embeddings.video_patch_projection.weight":
+                data_torch = data_torch.reshape(
+                    self.hparams["hidden_size"],
+                    6,
+                    self.hparams["patch_size"],
+                    self.hparams["patch_size"],
+                )
+                yield gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.V_ENC_EMBD_PATCH] + ".weight.1", data_torch
+                return
+
+            if name.startswith("vision_projector.mlp1."):
+                projector_mapping = {
+                    "vision_projector.mlp1.norm.weight": "mlp1.0.weight",
+                    "vision_projector.mlp1.linear1.weight": "mlp1.1.weight",
+                    "vision_projector.mlp1.linear2.weight": "mlp1.3.weight",
+                }
+                name = projector_mapping[name]
+            elif name == "vision_model.embeddings.patch_projection.weight":
+                data_torch = data_torch.reshape(
+                    self.hparams["hidden_size"],
+                    3,
+                    self.hparams["patch_size"],
+                    self.hparams["patch_size"],
+                )
+            elif name == "vision_model.embeddings.position_embedding":
+                data_torch = data_torch.squeeze(0)
+                name += ".weight"
+            elif ".layer_scale" in name and name.endswith(".lambda1"):
+                name = name.removesuffix("lambda1") + "weight"
+
+            yield from MmprojModel.modify_tensors(self, data_torch, name, bid)
+            return
+
+        if "patch_generator.pos_embed" in name:
+            # Keep the full table for runtime interpolation.
+            yield from MmprojModel.modify_tensors(self, data_torch, name, bid)
+            return
+
+        if "patch_generator.video_embedder" in name:
+            patch_size = self.hparams["patch_size"]
+            n_embd = self.hparams["hidden_size"]
+            temporal_patch_size = self.global_config["video_temporal_patch_size"]
+            if temporal_patch_size != 2:
+                raise ValueError("Current implementation only supports a temporal patch size of 2")
+            data_torch = data_torch.reshape(n_embd, temporal_patch_size * 3, patch_size, patch_size)
+            new_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.V_ENC_EMBD_PATCH] + ".weight.1"
+            yield new_name, data_torch
+            return
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+
@ModelBase.register("NemotronForCausalLM")
@ModelBase.example("nvidia/Minitron-4B-Base")
 class NemotronModel(TextModel):
@@ -382,6 +492,8 @@ class NemotronHModel(GraniteHybridModel):
    @classmethod
    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
        name, gen = item
+        if name.startswith("language_model."):
+            name = name[len("language_model."):]
        if name.startswith("mtp."):
            # --no-mtp: drop the MTP head entirely
            if cls.no_mtp:
@@ -531,6 +643,14 @@ class NemotronHModel(GraniteHybridModel):
        "mtp.layers.1.final_layernorm.weight": "model.layers.{bid}.shared_head.norm.weight",
    }

+    def _generate_nvfp4_tensors(self):
+        if self._mtp_bid is not None:
+            for name in list(self.model_tensors):
+                if name.startswith("mtp.layers.1.mixer"):
+                    new_name = f"backbone.layers.{self._mtp_bid}.{name.split('.', 3)[3]}"
+                    self.model_tensors[new_name] = self.model_tensors.pop(name)
+        super()._generate_nvfp4_tensors()
+
    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
        #   mtp.layers.0: NextN input fusion + attention
        #   mtp.layers.1: MoE + final head norm
@@ -3,7 +3,7 @@ from __future__ import annotations
 import json

 from pathlib import Path
-from typing import Any, Callable, Iterable, TYPE_CHECKING
+from typing import Any, Iterable, TYPE_CHECKING

 import torch

@@ -35,6 +35,7 @@ def _load_pplx_decider_hparams(dir_model: Path) -> dict[str, Any]:

@ModelBase.register("PplxDeciderModel")
@ModelBase.example("perplexity-ai/pplx-decider-v1-27b")
+@ModelBase.example("FINAL-Bench/Darwin-27B-ZTC-v5")
 class PplxDeciderModel(Qwen3_5TextModel):
    model_arch = gguf.MODEL_ARCH.QWEN35
    no_mtp = True  # the checkpoint has no MTP head
@@ -77,14 +78,6 @@ class PplxDeciderModel(Qwen3_5TextModel):
        for name in ("choice", "score", "noul"):
            self.gguf_writer.add_decision_temperature(name, self.hparams["decision"]["temperature"])

-    @classmethod
-    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
-        name, gen = item
-        # the checkpoint is the bare backbone, its text tensors have no "model." prefix
-        if name.startswith("language_model."):
-            name = "model." + name
-        return super().filter_tensors((name, gen))
-
    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
        yield from super().generate_extra_tensors()
        from safetensors.torch import load_file
@@ -92,7 +85,7 @@ class PplxDeciderModel(Qwen3_5TextModel):
        # the readout has one row per option label, store it as an LM head that is zero for the other tokens
        readout = load_file(self.dir_model / "readout.safetensors")["weight"]
        token_ids = self.hparams["decision"]["token_ids"]
-        n_vocab = self.hparams["text_config"]["vocab_size"]
+        n_vocab = self.hparams["vocab_size"]
        assert readout.shape[0] == len(token_ids) == len(set(token_ids))
        lm_head = torch.zeros(n_vocab, readout.shape[1], dtype=readout.dtype)
        lm_head[token_ids] = readout
@@ -6,6 +6,7 @@ if TYPE_CHECKING:
    from torch import Tensor

 from .base import MmprojModel, ModelBase, TextModel, gguf
+from .llama import LlamaModel


@ModelBase.register("UltravoxModel")
@@ -208,3 +209,47 @@ class AudioFlamingo3WhisperEncoderModel(WhisperEncoderModel):
            # Was trained in BF16, being safe, avoiding quantizing to FP16
            return gguf.GGMLQuantizationType.F32
        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+
+@ModelBase.register("VoxtralRealtimeForConditionalGeneration")
+@ModelBase.example("mistralai/Voxtral-Mini-4B-Realtime-2602")
+class VoxtralRealtimeEncoderModel(MmprojModel):
+    has_vision_encoder = False
+    has_audio_encoder = True
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        assert self.hparams_audio is not None
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.VOXTRAL_RT)
+        self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["num_mel_bins"])
+        self.gguf_writer.add_audio_head_dim(self.hparams_audio["head_dim"])
+        self.gguf_writer.add_audio_attention_layernorm_eps(self.hparams_audio.get("rms_norm_eps", 1e-5))
+        self.gguf_writer.add_audio_window_size(self.hparams_audio["sliding_window"])
+        self.gguf_writer.add_audio_stack_factor(self.global_config["downsample_factor"])
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if ".conv" in name and ".weight" in name:
+            return gguf.GGMLQuantizationType.F16
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+
+        # prevent clash naming with vision tensors
+        if name.startswith("multi_modal_projector"):
+            name = "audio." + name
+
+        return super().filter_tensors((name, gen))
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if "conv1.bias" in name or "conv2.bias" in name:
+            data_torch = data_torch.unsqueeze(-1)
+
+        # HF stores q/k in rotate-half layout, the encoder uses interleaved (normal) RoPE
+        if ".self_attn.q_proj." in name or ".self_attn.k_proj." in name:
+            assert self.hparams_audio is not None
+            n_head = self.hparams_audio["num_attention_heads"]
+            data_torch = LlamaModel.permute(data_torch, n_head, n_head)
+
+        yield from super().modify_tensors(data_torch, name, bid)
@@ -189,6 +189,8 @@ pre_computed_hashes = [
    {"name": "qwen2",     "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/Qwen/Qwen3-Embedding-0.6B", "chkhsh": "d4540891389ea895b53b399da6ac824becc30f2fba0e9ddbb98f92e55ca0e97c"},
    {"name": "qwen35",    "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/openbmb/MiniCPM-V-4_6", "chkhsh": "1444df51289cfa8063b96f0e62b1125440111bc79a52003ea14b6eac7016fd5f"},
    {"name": "grok-2",    "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/alvarobartt/grok-2-tokenizer", "chkhsh": "66b8d4e19ab16c3bfd89bce5d785fb7e0155e8648708a1f42077cb9fe002c273"},
+    # muse-spark uses the l4_202k tokenizer, whose pre-tokenizer regex is gpt-4o's
+    {"name": "gpt-4o",    "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B", "chkhsh": "f5c965e4198da6e7525072d87d1ecf4c51d981fd486c7c4ec728ef9cec7e03fc"},
    # jina-v2-de variants
    {"name": "jina-v2-de", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/aari1995/German_Semantic_V3", "chkhsh": "b3d1dd861f1d4c5c0d2569ce36baf3f90fe8a102db3de50dd71ff860d91be3df"},
    {"name": "gpt-2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/evilfreelancer/ruGPT3XL", "chkhsh": "0fe1cf6eda062318a1af7270f3331a85c539a01778ff948e24388e949c5282f4"},
@@ -134,6 +134,18 @@ Writing high-performance operators for Hexagon requires following specific guide
 - **Distinction**: `octx->src[N]` pointers *can* be NULL by design and must be checked when optional.
  Examples include attention masks, optional bias or weights in fused kernels, and frequency factors.

+### `supports_op` Independence from Buffer and Repack State
+
+- `supports_op` (and op-specific checkers such as `supported_mul_mat()`, `supported_get_rows()`) runs during GGML graph scheduling **before** buffers are allocated and before tensor data is populated.
+- **Never check buffer usage or repacking state in `supports_op`**:
+  - `t->buffer` may be NULL or uninitialized.
+  - `t->data` may be NULL.
+  - `buffer->usage` is unknown at scheduling time.
+  - `is_tiled` or `extra->flags & GGML_HEXAGON_TENSOR_TILED` is not set yet.
+- Operator support decisions must depend strictly on static tensor properties: operand types (`t->type`), shapes (`t->ne`), strides (`t->nb`), and operation parameters.
+- Repacking (converting weights to tiled layout) is an internal optimization applied during `set_tensor()` or rt-repack. Non-weight tensors or test harness tensors with the same type remain in flat layout.
+- Therefore, any supported type (such as BF16, F16, quantized types) must be supported both when repacked into tiles and when stored flat in memory.
+
 ### Multiline Macro Formatting

 - Keep trailing backslashes in multiline `#define` macros cleanly aligned to a consistent column.
@@ -190,7 +202,7 @@ sessions.
 Common partitioning logic is factored into reusable inline helpers in
 [`htp-tensor.h`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h):

-1. [`htp_tensor_mdev_rows_per_chunk`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L71):
+1. [`htp_tensor_mdev_rows_per_chunk`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L104):
   Determines the minimum number of rows per chunk so that the chunk byte size is a multiple of 128 bytes:

   ```
@@ -201,7 +213,7 @@ Common partitioning logic is factored into reusable inline helpers in
   Returns `false` if the tensor cannot be safely row-partitioned (such as unaligned base pointer, permuted layout,
   or non-128-byte aligned outer strides).

-2. [`htp_tensor_mdev_partition`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L98):
+2. [`htp_tensor_mdev_partition`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L131):
   Calculates the per-device work range `struct htp_tensor_mdev_range { uint32_t start; uint32_t count; }` given
   `total_units`, `units_per_chunk`, `mdev_idx`, `mdev_count`, and the precomputed `mdev_count_div`.
   Handles chunk distribution across devices, assigns remainder units to the last device, and automatically triggers
@@ -240,11 +252,11 @@ if (nrows == 0) {
 For flat element-wise operations that write directly to DDR:
 - Partition total linear elements N = ne0 * ne1 * ne2 * ne3 in 128-byte cache line chunks (`elems_per_line = (elem_size == 4) ? 32 : 64`).
 - Requires strict 1D contiguity:
-  [`htp_tensor_is_contiguous(dst, elem_size)`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L32)
+  [`htp_tensor_is_contiguous(dst, elem_size)`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L65)
  and 128-byte aligned destination pointer
-  [`htp_tensor_mdev_data_aligned(dst)`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L51).
+  [`htp_tensor_mdev_data_aligned(dst)`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L84).
 - If contiguous and aligned, pass `elems_per_line` to
-  [`htp_tensor_mdev_partition`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L98);
+  [`htp_tensor_mdev_partition`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L131);
  otherwise pass 0 to trigger Device 0 fallback.

 ### Single-Device Fallback (Device 0)
@@ -312,7 +324,7 @@ Multi-device execution synchronizes worker sessions through atomic fence slots a
  (Output Slice 0)                             (Output Slice 1)
       |                                            |
  Tensor Cache Flush                           Tensor Cache Flush
-  (htp_tensor_flush_all)                       (htp_tensor_flush_all)
+  (htp_dcache_flush_dirty_tensors)             (htp_dcache_flush_dirty_tensors)
       |                                            |
  Post-Op/Batch Barrier ---------------------- Post-Op/Batch Barrier
  (htp_mdev_group_barrier)                     (htp_mdev_group_barrier)
@@ -350,8 +362,8 @@ Multi-device execution synchronizes worker sessions through atomic fence slots a
 ### Tensor Cache Flush and Pipeline Completion

 - In the kernel, ensure all pushed DMA operations have been popped in strict FIFO order to drain the queue.
- Use [`htp_tensor_flush_all()`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h) to flush specific dirty tensors back to DDR:
-  - [`htp_tensor_flush_all()`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h) flushes modified tensor address ranges, or the
+- Use [`htp_dcache_flush_dirty_tensors()`](../../../ggml/src/ggml-hexagon/htp/htp-dcache.h) to flush specific dirty tensors back to DDR:
+  - [`htp_dcache_flush_dirty_tensors()`](../../../ggml/src/ggml-hexagon/htp/htp-dcache.h) flushes modified tensor address ranges, or the
    full D-cache when their total size exceeds the flush threshold, ensuring peer devices and the host CPU observe consistent
    data in DDR.
 - Never signal completion before all DMA transfers are drained and dirty tensor flushes have completed.
@@ -147,11 +147,11 @@ The conversion is the same as above, but a new model needs its own `DecisionType

 > [!IMPORTANT]
 >
-> Most of the logic is handled in `tools/server/server-decision.cpp`, to avoid too many changes to `libllama`.
+> Most of the logic is handled in `tools/engine/engine-decision.cpp`, to avoid too many changes to `libllama`.

 Note:
 - If a new public API is needed in `libllama`, add it to `llama-ext.h`.
- Metadata with a single use case must be hard-coded in `server-decision.cpp` instead of being saved to the GGUF. This avoids bloating the conversion code.
+- Metadata with a single use case must be hard-coded in `engine-decision.cpp` instead of being saved to the GGUF. This avoids bloating the conversion code.
 - Most importantly, keep your change as small and as self-contained as possible. Reuse the existing infrastructure whenever you can.

 For more information, see [PR #29818](https://github.com/ggml-org/llama.cpp/pull/29818).
@@ -24,3 +24,267 @@ python3 llama-eval.py \
  --dataset aime2025 --n_cases 240 \
  --grader-type regex
 ```
+
+### Authentication
+
+A server started with `--api-key` (or any OpenAI-compatible endpoint behind a
+token) needs one here too. `--api-key` is sent as `Authorization: Bearer` on
+every eval request; it falls back to `$LLAMA_API_KEY`, which keeps the secret
+out of your shell history and off the process list.
+
+```bash
+# One key for every server
+python3 llama-eval.py --server http://localhost:8033 --api-key sk-... ...
+
+# From the environment
+export LLAMA_API_KEY=sk-...
+python3 llama-eval.py --server http://localhost:8033 ...
+
+# One key per server, in --server order
+python3 llama-eval.py \
+  --server http://server1:8033,http://server2:8033 \
+  --api-key sk-one,sk-two --threads 16,16 ...
+```
+
+The LLM grader authenticates too. It reuses the key of the server it runs on,
+so `--grader-type llm` without `--grader-server` needs nothing extra; a grader
+on a separate endpoint takes `--grader-api-key`.
+
+Endpoints on an opencode host additionally get an `x-opencode-session` header,
+generated fresh for every task. An agentic episode sends the same id on all of
+its turns, so opencode sees one session per task instead of one per request;
+other servers never see the header.
+
+## Agentic suite
+
+`--dataset agentic` is a different shape from the other code suites. It runs
+**SACB** (Simple Agent Coding Benchmark), hosted at
+[`ilintar/SACB`](https://huggingface.co/datasets/ilintar/SACB) -- 60 tasks over
+three hand-written repositories, in Python and TypeScript. Rather than
+answering one prompt, the model is given a small but real repository, a
+deliberately narrow tool set, and a bug report, and drives its own multi-turn
+conversation until it declares itself finished. It is then graded by running a
+test suite it was never allowed to see.
+
+```bash
+python3 llama-eval.py \
+  --server http://localhost:8033 --model my-model \
+  --dataset agentic --threads 1 --temperature 0
+```
+
+### The tools
+
+| tool | purpose |
+|---|---|
+| `list_files` | enumerate the tree |
+| `read_file` | read, with 1-based line numbers and optional line-range narrowing |
+| `search` | regular-expression search across the repository |
+| `edit_lines` | replace an inclusive line range |
+| `edit_replace` | replace an exact snippet, which must be unique |
+| `write_file` | write a whole file |
+| `lint` | run the language's linter and type checker |
+| `finish` | declare the change complete |
+
+Three constraints are deliberate:
+
+**There is no shell and no way to run the tests.** The only feedback channel is
+`lint`. A syntax or type error can be driven out mechanically; a logic error has
+to be reasoned about. A harness that lets an agent iterate against the test
+suite measures something closer to search than to understanding.
+
+**There are two edit tools**, one line-addressed and one content-addressed.
+Which one a model reaches for, and whether it keeps line numbers straight after
+its own earlier edit, is itself a signal.
+
+**The tests never touch the working tree.** They are held outside it and copied
+into a throwaway copy only after the agent has stopped, so they can be neither
+read nor edited. Every tool path is resolved and confinement-checked against the
+repository root, so a symlink planted inside the tree cannot reach outside it.
+
+### Scoring
+
+A task is **resolved** when every test that was failing now passes *and* nothing
+that was already passing broke. Both sets are derived by running the tests twice
+when the corpus is built -- once against the defective tree and once against the
+reference fix -- never hand-written, so a task whose defect no test exercises is
+rejected rather than shipped.
+
+Because many tasks contain several independent defects, the fraction of target
+tests passed is reported alongside, which separates "fixed two of the three
+faults" from "changed nothing". Process metrics are reported too: turns, tool
+calls, rejected calls, edits, and the peak conversation size.
+
+### Dependencies
+
+Nothing here is installed unless you select this suite, and then only for the
+language you actually run. `--agentic-lang python` provisions ruff and mypy;
+`--agentic-lang typescript` provisions the TypeScript compiler and uses node's
+own test runner. Neither touches the other, and no other suite touches either.
+
+### Useful flags
+
+| flag | meaning |
+|---|---|
+| `--agentic-lang` | restrict to one language, and provision only its toolchain |
+| `--agentic-max-turns` | give up after this many assistant turns (default 50) |
+| `--agentic-max-tokens` | cap generation per turn (default: no cap); bounds a model that falls into a repetition loop, but also cuts reasoning off mid-thought. Truncated turns are counted as `truncated_turns` |
+| `--agentic-context-limit` | abandon an episode once its context exceeds this many tokens |
+| `--agentic-transcripts DIR` | save each episode's full message history to `DIR/<task_id>.json`, plus a `DIR/<task_id>.html` showing every reasoning block, reply, tool call and tool result; the main report links to it |
+
+
+## Code suites (HumanEval, ClassEval)
+
+`--dataset humaneval` runs the 164 OpenAI HumanEval problems; `--dataset
+classeval` runs the 100 ClassEval class-generation tasks. Unlike the other
+suites these are not graded by matching text: the model's reply is parsed for a
+```` ```python ```` block, that code is combined with the task's own unit tests,
+and the result is **executed**.
+
+ClassEval asks for a whole class rather than one function, so it is scored the
+way upstream does - **class level** (every one of the task's tests passes, which
+is what feeds pass@k) and **method level** (what fraction of all 2196 test
+methods passed). The second number separates "generated nothing usable" from
+"got most of the class right", and both are printed at the end of a run.
+
+Give it more room to generate than HumanEval - a class with ten methods does not
+fit in a short budget:
+
+```bash
+python3 llama-eval.py \
+  --server http://localhost:8033 --model my-model \
+  --dataset classeval --threads 8 --n_predict 4096 --temperature 0
+```
+
+```bash
+# pass@1, greedy
+python3 llama-eval.py \
+  --server http://localhost:8033 --model my-model \
+  --dataset humaneval --threads 8 --n_predict 1024 --temperature 0
+
+# pass@10 -- each extra multiple of 164 is another sample per problem
+python3 llama-eval.py \
+  --server http://localhost:8033 --model my-model \
+  --dataset humaneval --n_cases 1640 --temperature 0.8 --top-p 0.95
+```
+
+`--grader-type exec` is the default for code suites, so neither it nor
+`--dataset-source` normally needs to be passed.
+
+### Execution sandbox - no Docker required
+
+Generated code is untrusted, but these suites need far less than a container to
+run safely. **The package set is derived from the dataset's own imports at
+startup**, not hardcoded and not taken from the upstream `requirements.txt` -
+which for ClassEval lists `torch`, `openai` and `transformers` that its *tests*
+never touch. So each suite installs exactly what it needs:
+
+| suite | third-party packages |
+|---|---|
+| HumanEval | **none** - every problem imports only stdlib (`typing, math, random, copy, string, collections, hashlib, re`) |
+| ClassEval | 13: `beautifulsoup4, gensim, lxml, netifaces, nltk, numpy, openpyxl, pandas, pillow, pypdf2, python-docx, reportlab, scipy` |
+
+Two of ClassEval's are *declared* rather than derived, because no import
+statement mentions them: `lxml` is selected by string
+(`BeautifulSoup(html, 'lxml')`, ClassEval_44 - without it that task silently
+drops from 23/23 to 7/23), and `scipy` is a hard gensim requirement. Note also
+that `PyPDF2` must **not** be modernised to `pypdf`: ClassEval_69's test imports
+`PdfFileReader`, a name pypdf removed.
+
+The venv is cached under `~/.cache/llama-eval/` (override with
+`LLAMA_EVAL_CACHE`) and rebuilt only when the package set, interpreter version
+or provisioning steps change; force one with `--rebuild-sandbox`. `uv` is used
+when available, which makes even the ClassEval venv build in seconds.
+
+**Interpreter selection.** A suite can prefer a different Python than the one
+running the tool, and declares it: ClassEval asks for 3.13, because gensim
+publishes no wheel past cp313 and its generated C no longer compiles on 3.14
+(it still uses `PyLongObject.ob_digit`). That interpreter is found automatically,
+including via `uv python`; `--sandbox-python` overrides it.
+
+3.14 is fully supported anyway. ClassEval uses exactly two gensim APIs
+(`utils.decode_htmlentities`, `matutils.unitvec`), both pure Python upstream, so
+on 3.14+ the suite drops gensim from the install set and provisions a small shim
+instead. Both paths score **100/100**.
+
+**Offline provisioning.** A suite that would otherwise fetch data at test time
+declares a provisioning step that runs once at venv build time, while the
+network is still up. ClassEval uses this to pre-seed exactly three nltk corpora
+(`punkt_tab`, `averaged_perceptron_tagger_eng`, `wordnet`, ~33 MB), so the tests
+themselves run with no network at all. Those are the *current* names - the
+dataset asks for `punkt` and `averaged_perceptron_tagger`, which nltk has since
+renamed, and seeding the legacy names too is not harmless: with a network
+available nltk considers them satisfied and re-fetches the modern ones anyway.
+The corpora are written *inside* the venv, not `~/nltk_data`, because `$HOME` is
+a tmpfs under the strongest isolation tier.
+
+Each candidate program runs in a throwaway directory as its own process group,
+under the strongest isolation the host offers:
+
+| tier | requires | isolates |
+|---|---|---|
+| `bwrap` | bubblewrap installed | read-only root, tmpfs `/tmp` and `$HOME`, **no network** |
+| `unshare` | `unshare --map-current-user` | **no network** |
+| `unshare-root` | any `unshare -r` | no network, but **runs as uid 0** - see below |
+| `rlimit` | always available | cpu/memory/file-size/process caps only |
+
+**Candidate code must not run as root**, which is why `unshare -r` is a
+last-resort tier that prints a warning. Root bypasses file permission bits, so a
+test asserting that writing to a `chmod 0444` file *fails* instead sees it
+succeed - ClassEval_50 scores 14/16 under uid 0 and 16/16 under uid 1000, with
+nothing in the output to suggest the sandbox caused it. Install bubblewrap.
+
+All tiers additionally cap CPU time, address space (4 GiB), file size and core
+dumps, and kill the whole process group on timeout (`--exec-timeout`, default
+15 s - raise it for ClassEval, whose tests do real work). Writes to the
+throwaway working directory are allowed, which is all the file-handling tasks in
+ClassEval need.
+
+Note that the weakest tier does not contain filesystem access. If you are
+scoring output from a model you do not trust, check the tier printed at startup.
+
+### Reproducibility
+
+Several tasks draw their test inputs from the *global, unseeded* RNG, so a
+verdict is not reproducible run to run - which adds noise exactly when you are
+comparing two models or two quantisations. `--exec-seed` (default `0`) seeds the
+RNG before each test; pass `--exec-seed -1` for the stock upstream behaviour.
+This does not change what the model is asked to compute.
+
+It matters more than it sounds. HumanEval/38, /50 and /53 are affected, and
+ClassEval_58 (duplicate mine placement) fails about 30% of the time on its own
+*reference* solution - measured here at 21/30 passes unseeded versus **30/30
+seeded**.
+
+### Dataset fidelity
+
+The canonical dataset runs **unmodified on current Python** - all 164 reference
+solutions pass on CPython 3.14 with no patches, so no forked or re-hosted copy
+is needed. `--dataset-source` accepts a local `.jsonl` or a HuggingFace repo id
+if you do want to point at a modified copy.
+
+HumanEval does contain a handful of long-standing defects where the *prompt*
+contradicts the graded test (notably HumanEval/47, /116 and /148). Those
+penalise a model that reads the docstring carefully, but correcting them makes
+the benchmark marginally easier and the resulting scores no longer comparable
+with published HumanEval numbers, so nothing is patched by default.
+
+ClassEval is different: the upstream data has defects that cap the achievable
+score below 100% for reasons unrelated to the model, so `--dataset classeval`
+defaults to a **patched fork**,
+[`ilintar/ClassEval`](https://huggingface.co/datasets/ilintar/ClassEval).
+Six edits over five tasks lift the reference solutions from 96/100 classes and
+2182/2196 test methods to **100/100 and 2196/2196**. Three repair tests nothing
+could pass:
+
+- **ClassEval_17** was a time bomb - a hard-coded 2024 date checked against
+  `datetime.now()`, unpassable since January 2024.
+- **ClassEval_48** asserted `get_hostname('0.0.0.0') == 'LAPTOP-2CS86KUM'`, the
+  dataset author's own machine name.
+- **ClassEval_31** compared floats with `assertEqual`, failing on a 4-ULP numpy
+  difference.
+
+The other three fix only the *reference* solutions (NumPy 2.0 removals in
+ClassEval_51, a 30%-flaky mine generator in ClassEval_58) and cannot affect a
+model's score. No `skeleton` was touched, so what the model is asked to write is
+byte-identical to upstream. Pass `--dataset-source FudanSELab/ClassEval` for the
+unpatched original.
@@ -0,0 +1,918 @@
+#!/usr/bin/env python3
+"""Multi-turn agentic coding eval: tools over a real working tree.
+
+The existing code suites are single-shot -- one prompt, one reply, execute,
+grade. This one is not. The model is given a small but real repository and a
+deliberately narrow tool set, and has to find its way to the defect before it
+can fix it. That shape is what puts a run in the tens-of-thousands-of-tokens
+range without any artificial padding: the context grows because the model
+chose to read things.
+
+Three deliberate constraints:
+
+* No shell, and no way to run the tests. The only feedback channel is `lint`.
+  A syntax or type error can be driven out mechanically; a logic error has to
+  be reasoned about. Suites that let an agent iterate against the test suite
+  measure something closer to search than to understanding.
+* Two edit tools, one line-addressed and one content-addressed. Which one a
+  model reaches for, and whether it keeps line numbers straight after its own
+  earlier edit, is itself a signal.
+* Tests never touch the working tree. They are held outside it and copied in
+  after the agent has stopped, so they can be neither read nor edited.
+
+Everything here is imported lazily by llama-eval.py, and every dependency it
+needs is scoped to the language actually under test, so none of it is
+installed for anyone running the other suites.
+"""
+
+import html
+import json
+import os
+import re
+import shutil
+import tempfile
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+MAX_READ_LINES = 400
+MAX_TOOL_CHARS = 12000
+MAX_SEARCH_HITS = 60
+
+
+@dataclass
+class LangSpec:
+    """Everything language-specific: how to lint, how to test, what to install.
+
+    Kept as data so adding a language is a corpus concern rather than a harness
+    change, and so a run only ever provisions the toolchain for the language it
+    is actually evaluating.
+    """
+    name: str
+    exts: Tuple[str, ...]
+    packages: Tuple[str, ...] = ()
+    provision: Tuple[str, ...] = ()
+    env: Dict[str, str] = field(default_factory=dict)
+    address_space_limit: Optional[int] = 4 * 1024 * 1024 * 1024
+
+    def source_files(self, root: Path) -> List[Path]:
+        return sorted(p for p in root.rglob("*")
+                      if p.is_file() and p.suffix in self.exts)
+
+
+PYTHON = LangSpec(
+    name="python",
+    exts=(".py",),
+    packages=("ruff==0.14.2", "mypy==1.18.2"),
+)
+
+TYPESCRIPT = LangSpec(
+    name="typescript",
+    exts=(".ts",),
+    address_space_limit=16 * 1024 * 1024 * 1024,
+    provision=(
+        "import os, subprocess, shutil, sys\n"
+        "npm = shutil.which('npm')\n"
+        "if not npm: sys.exit('npm not found on PATH; needed for the typescript suite')\n"
+        "dest = os.path.join(os.environ['SANDBOX_VENV'], 'node')\n"
+        "os.makedirs(dest, exist_ok=True)\n"
+        "r = subprocess.run([npm, 'install', '--silent', '--no-audit', '--no-fund',\n"
+        "                    '--prefix', dest, 'typescript@5.9.3'],\n"
+        "                   capture_output=True, text=True)\n"
+        "sys.exit(0 if r.returncode == 0 else (r.stderr or r.stdout)[-2000:])\n",
+    ),
+)
+
+LANGS: Dict[str, LangSpec] = {"python": PYTHON, "typescript": TYPESCRIPT}
+
+
+class PathEscape(Exception):
+    """A tool was handed a path outside the repository."""
+
+
+class Workspace:
+    """A private, writable copy of one task's repository.
+
+    Every tool path is resolved through here. Confinement is checked on the
+    fully-resolved path, so a symlink planted inside the tree cannot be used to
+    read or write outside it.
+    """
+
+    def __init__(self, root: Path):
+        self.root = root.resolve()
+        self.escape_attempts = 0
+
+    def resolve(self, rel: str) -> Path:
+        if rel is None or not str(rel).strip():
+            raise PathEscape("no path given")
+        raw = Path(str(rel))
+        base = raw if raw.is_absolute() else (self.root / raw)
+        candidate = base.resolve()
+        if candidate != self.root and self.root not in candidate.parents:
+            self.escape_attempts += 1
+            raise PathEscape(f"path {rel!r} is outside the repository")
+        return candidate
+
+    def rel(self, path: Path) -> str:
+        return str(path.relative_to(self.root))
+
+    def files(self) -> List[Path]:
+        return sorted(p for p in self.root.rglob("*")
+                      if p.is_file() and "__pycache__" not in p.parts)
+
+    def snapshot(self) -> Dict[str, str]:
+        out = {}
+        for p in self.files():
+            try:
+                out[self.rel(p)] = p.read_text()
+            except (UnicodeDecodeError, OSError):
+                pass
+        return out
+
+
+def _clip(text: str, limit: int = MAX_TOOL_CHARS) -> str:
+    if len(text) <= limit:
+        return text
+    return text[:limit] + f"\n... [truncated, {len(text) - limit} more characters]"
+
+
+def _numbered(lines: Sequence[str], start: int) -> str:
+    width = len(str(start + len(lines) - 1))
+    return "\n".join(f"{start + i:>{width}}\t{ln}" for i, ln in enumerate(lines))
+
+
+def tool_schemas() -> List[Dict[str, Any]]:
+    """OpenAI-format function schemas for the whole tool set."""
+    def fn(name, desc, props, required=()):
+        return {"type": "function", "function": {
+            "name": name, "description": desc,
+            "parameters": {"type": "object", "properties": props,
+                           "required": list(required)}}}
+
+    return [
+        fn("list_files", "List files in the repository.",
+           {"path": {"type": "string",
+                     "description": "Directory relative to the repository root. "
+                                    "Omit for the whole tree."}}),
+        fn("read_file",
+           "Read a file. Output is prefixed with 1-based line numbers. Give "
+           "start_line/end_line to read only part of a large file.",
+           {"path": {"type": "string", "description": "File path relative to the repository root."},
+            "start_line": {"type": "integer", "description": "First line to read, 1-based, inclusive."},
+            "end_line": {"type": "integer", "description": "Last line to read, 1-based, inclusive."}},
+           ["path"]),
+        fn("search",
+           "Search the repository with a regular expression. Returns matching "
+           "lines with their file and line number.",
+           {"pattern": {"type": "string", "description": "Python regular expression."},
+            "path": {"type": "string", "description": "Limit the search to this subdirectory."},
+            "glob": {"type": "string", "description": "Limit to files matching this glob, e.g. '*.py'."}},
+           ["pattern"]),
+        fn("edit_lines",
+           "Replace an inclusive 1-based line range with new text. Line numbers "
+           "refer to the file as it is now, so re-read it after any edit that "
+           "changed the number of lines.",
+           {"path": {"type": "string", "description": "File path relative to the repository root."},
+            "start_line": {"type": "integer", "description": "First line to replace, 1-based, inclusive."},
+            "end_line": {"type": "integer", "description": "Last line to replace, 1-based, inclusive."},
+            "new_text": {"type": "string", "description": "Replacement text. Use an empty string to delete the range."}},
+           ["path", "start_line", "end_line", "new_text"]),
+        fn("edit_replace",
+           "Replace an exact snippet of text. The snippet must occur exactly "
+           "once in the file unless replace_all is set.",
+           {"path": {"type": "string", "description": "File path relative to the repository root."},
+            "old_text": {"type": "string", "description": "Exact text to find, including indentation."},
+            "new_text": {"type": "string", "description": "Replacement text."},
+            "replace_all": {"type": "boolean", "description": "Replace every occurrence instead of requiring exactly one."}},
+           ["path", "old_text", "new_text"]),
+        fn("write_file",
+           "Write a whole file, creating or overwriting it. Prefer the edit "
+           "tools for changes to existing files.",
+           {"path": {"type": "string", "description": "File path relative to the repository root."},
+            "content": {"type": "string", "description": "Full new file contents."}},
+           ["path", "content"]),
+        fn("lint",
+           "Run the language's linter and type checker over the repository and "
+           "return the diagnostics. This is the only feedback available; the "
+           "test suite cannot be run.",
+           {"path": {"type": "string", "description": "Limit to this file or directory."}}),
+        fn("finish",
+           "Declare the task complete and end the session. This is the only "
+           "way to end it; call this once the change is made.",
+           {"summary": {"type": "string", "description": "Brief description of the change made."}}),
+    ]
+
+
+class ToolBox:
+    """Executes tool calls against one workspace. Returns plain text results.
+
+    Failures are reported as ordinary text rather than raised, because how a
+    model recovers from a rejected edit is part of what is being measured.
+    """
+
+    def __init__(self, ws: Workspace, linter=None):
+        self.ws = ws
+        self.linter = linter
+        self.calls = 0
+        self.errors = 0
+        self.edits = 0
+        self.finished = False
+        self.finish_summary: Optional[str] = None
+        self.per_tool: Dict[str, int] = {}
+
+    def dispatch(self, name: str, args: Dict[str, Any]) -> str:
+        self.calls += 1
+        self.per_tool[name] = self.per_tool.get(name, 0) + 1
+        handler = getattr(self, f"_t_{name}", None)
+        if handler is None:
+            result = f"Error: unknown tool {name!r}."
+        else:
+            try:
+                result = handler(args)
+            except PathEscape as e:
+                result = f"Error: {e}"
+            except KeyError as e:
+                result = f"Error: missing required argument {e}."
+            except Exception as e:
+                result = f"Error: {type(e).__name__}: {e}"
+        if result.startswith("Error:"):
+            self.errors += 1
+        return _clip(result)
+
+    def _t_list_files(self, a):
+        base = self.ws.resolve(a.get("path") or ".")
+        if not base.exists():
+            return f"Error: {a.get('path')!r} does not exist."
+        paths = [p for p in self.ws.files() if p == base or base in p.parents]
+        if not paths:
+            return "(no files)"
+        return "\n".join(f"{self.ws.rel(p)}\t{p.stat().st_size} bytes" for p in paths)
+
+    def _t_read_file(self, a):
+        path = self.ws.resolve(a["path"])
+        if not path.is_file():
+            return f"Error: {a['path']!r} is not a file."
+        lines = path.read_text().splitlines()
+        start = max(1, int(a.get("start_line") or 1))
+        end = int(a.get("end_line") or len(lines))
+        end = min(end, len(lines))
+        if start > len(lines):
+            return f"Error: {a['path']!r} has only {len(lines)} lines."
+        window = lines[start - 1:end]
+        note = ""
+        if len(window) > MAX_READ_LINES:
+            window = window[:MAX_READ_LINES]
+            note = (f"\n... [showing {MAX_READ_LINES} of {end - start + 1} requested "
+                    f"lines; re-read with start_line/end_line for the rest]")
+        header = f"{a['path']} ({len(lines)} lines total)\n"
+        return header + _numbered(window, start) + note
+
+    def _t_search(self, a):
+        try:
+            rx = re.compile(a["pattern"])
+        except re.error as e:
+            return f"Error: bad regular expression: {e}"
+        base = self.ws.resolve(a.get("path") or ".")
+        glob = a.get("glob")
+        hits, scanned = [], 0
+        for p in self.ws.files():
+            if not (p == base or base in p.parents):
+                continue
+            if glob and not p.match(glob):
+                continue
+            try:
+                text = p.read_text()
+            except (UnicodeDecodeError, OSError):
+                continue
+            scanned += 1
+            for i, line in enumerate(text.splitlines(), 1):
+                if rx.search(line):
+                    hits.append(f"{self.ws.rel(p)}:{i}:{line.strip()[:200]}")
+                    if len(hits) >= MAX_SEARCH_HITS:
+                        return ("\n".join(hits) +
+                                f"\n... [stopped at {MAX_SEARCH_HITS} matches; narrow the pattern]")
+        if not hits:
+            return f"No matches for {a['pattern']!r} in {scanned} file(s)."
+        return "\n".join(hits)
+
+    def _t_edit_lines(self, a):
+        path = self.ws.resolve(a["path"])
+        if not path.is_file():
+            return f"Error: {a['path']!r} is not a file."
+        lines = path.read_text().splitlines()
+        start, end = int(a["start_line"]), int(a["end_line"])
+        if start < 1 or start > len(lines):
+            return (f"Error: start_line {start} out of range; "
+                    f"{a['path']!r} has {len(lines)} lines.")
+        if end < start or end > len(lines):
+            return (f"Error: end_line {end} out of range; must be between "
+                    f"{start} and {len(lines)}.")
+        new = str(a["new_text"]).splitlines()
+        path.write_text("\n".join(lines[:start - 1] + new + lines[end:]) + "\n")
+        self.edits += 1
+        delta = len(new) - (end - start + 1)
+        return (f"Replaced lines {start}-{end} of {a['path']} "
+                f"({end - start + 1} -> {len(new)} lines, {delta:+d}). "
+                f"Later line numbers have shifted." if delta else
+                f"Replaced lines {start}-{end} of {a['path']}.")
+
+    def _t_edit_replace(self, a):
+        path = self.ws.resolve(a["path"])
+        if not path.is_file():
+            return f"Error: {a['path']!r} is not a file."
+        text = path.read_text()
+        old, new = str(a["old_text"]), str(a["new_text"])
+        if not old:
+            return "Error: old_text must not be empty."
+        n = text.count(old)
+        if n == 0:
+            return (f"Error: old_text not found in {a['path']}. It must match "
+                    f"exactly, including indentation.")
+        if n > 1 and not a.get("replace_all"):
+            return (f"Error: old_text occurs {n} times in {a['path']}. Include "
+                    f"more surrounding context to make it unique, or set "
+                    f"replace_all.")
+        path.write_text(text.replace(old, new))
+        self.edits += 1
+        return f"Replaced {n} occurrence(s) in {a['path']}."
+
+    def _t_write_file(self, a):
+        path = self.ws.resolve(a["path"])
+        path.parent.mkdir(parents=True, exist_ok=True)
+        existed = path.is_file()
+        path.write_text(str(a["content"]))
+        self.edits += 1
+        return f"{'Overwrote' if existed else 'Created'} {a['path']}."
+
+    def _t_lint(self, a):
+        if self.linter is None:
+            return "Error: no linter is configured for this task."
+        target = a.get("path")
+        if target:
+            self.ws.resolve(target)
+        return self.linter(target)
+
+    def _t_finish(self, a):
+        self.finished = True
+        self.finish_summary = a.get("summary")
+        return "Task marked complete."
+
+
+PY_TEST_RUNNER = r'''
+import json, os, random, sys, unittest
+sys.path.insert(0, os.getcwd())
+random.seed(0)
+
+class Collect(unittest.TestResult):
+    def __init__(self):
+        super().__init__()
+        self.outcomes = {}
+    def addSuccess(self, t): self.outcomes[t.id()] = "pass"
+    def addFailure(self, t, e): self.outcomes[t.id()] = "fail"
+    def addError(self, t, e): self.outcomes[t.id()] = "error"
+    def addSkip(self, t, r): self.outcomes[t.id()] = "skip"
+    def addExpectedFailure(self, t, e): self.outcomes[t.id()] = "pass"
+    def addUnexpectedSuccess(self, t): self.outcomes[t.id()] = "fail"
+
+def flatten(s):
+    for t in s:
+        if isinstance(t, unittest.TestSuite):
+            yield from flatten(t)
+        else:
+            yield t
+
+out = {"tests": {}}
+try:
+    suite = unittest.TestLoader().discover(os.environ.get("TEST_DIR", "tests"),
+                                           top_level_dir=".")
+    res = Collect()
+    for t in flatten(suite):
+        tid = t.id()
+        try:
+            t.run(res)
+        except Exception:
+            pass
+        res.outcomes.setdefault(tid, "error")
+    out["tests"] = res.outcomes
+except Exception as e:
+    out["error"] = f"{type(e).__name__}: {e}"
+
+print("###RESULT###")
+print(json.dumps(out))
+'''
+
+
+def _node_roots() -> List[Path]:
+    """Directories that must stay visible for node to start under bwrap."""
+    node = shutil.which("node")
+    if not node:
+        return []
+    real = Path(node).resolve()
+    return [Path(node).parent.parent, real.parent.parent]
+
+
+class Runner:
+    """Runs the linter and the hidden tests for one language, in the sandbox."""
+
+    def __init__(self, sandbox, lang: LangSpec, timeout: float = 90.0):
+        self.sandbox = sandbox
+        self.lang = lang
+        self.timeout = timeout
+
+    def _binds(self) -> List[Path]:
+        return _node_roots() if self.lang.name == "typescript" else []
+
+    def _env(self) -> Dict[str, str]:
+        if self.lang.name != "typescript":
+            return {}
+        node = shutil.which("node")
+        return {"PATH": f"{Path(node).parent}:/usr/bin:/bin"} if node else {}
+
+    def lint(self, workdir: Path, target: Optional[str] = None) -> str:
+        if self.lang.name == "python":
+            return self._lint_python(workdir, target)
+        return self._lint_typescript(workdir, target)
+
+    def _lint_python(self, workdir: Path, target: Optional[str]) -> str:
+        venv = self.sandbox.venv_dir
+        tgt = target or "."
+        chunks = []
+        r = self.sandbox.run_argv(
+            [str(venv / "bin" / "ruff"), "check", "--no-cache",
+             "--output-format", "concise", tgt],
+            workdir, timeout=self.timeout)
+        text = (r.stdout or r.stderr).strip()
+        if text and "All checks passed" not in text:
+            chunks.append(text)
+        r = self.sandbox.run_argv(
+            [str(venv / "bin" / "mypy"), "--no-color-output", "--no-error-summary",
+             "--cache-dir", "/dev/null", "--ignore-missing-imports", tgt],
+            workdir, timeout=self.timeout)
+        text = (r.stdout or r.stderr).strip()
+        if text and not text.startswith("Success:"):
+            chunks.append(text)
+        return "\n".join(chunks) if chunks else "No problems found."
+
+    def _lint_typescript(self, workdir: Path, target: Optional[str]) -> str:
+        tsc = self.sandbox.venv_dir / "node" / "node_modules" / "typescript" / "bin" / "tsc"
+        files = [str(p.relative_to(workdir)) for p in sorted(workdir.rglob("*.ts"))
+                 if "node_modules" not in p.parts]
+        if target:
+            t = (workdir / target).resolve()
+            files = [f for f in files if (workdir / f).resolve() == t
+                     or t in (workdir / f).resolve().parents]
+        if not files:
+            return "No TypeScript files to check."
+        r = self.sandbox.run_argv(
+            ["node", str(tsc), "--noEmit", "--strict", "--target", "es2022",
+             "--module", "esnext", "--moduleResolution", "bundler",
+             "--allowImportingTsExtensions", *files],
+            workdir, timeout=self.timeout, env=self._env(), ro_binds=self._binds())
+        out = (r.stdout or "") + (r.stderr or "")
+        return out.strip() or "No problems found."
+
+    def test(self, workdir: Path) -> Dict[str, Any]:
+        if self.lang.name == "python":
+            return self._test_python(workdir)
+        return self._test_typescript(workdir)
+
+    def _test_python(self, workdir: Path) -> Dict[str, Any]:
+        runner = workdir / "__run_tests.py"
+        runner.write_text(PY_TEST_RUNNER)
+        r = self.sandbox.run_argv(
+            [str(self.sandbox.python), "-I", "-B", str(runner)],
+            workdir, timeout=self.timeout, env={"TEST_DIR": "tests"})
+        return self._parse_json_result(r)
+
+    def _parse_json_result(self, r) -> Dict[str, Any]:
+        if r.status == "timeout":
+            return {"tests": {}, "error": "timeout"}
+        marker = "###RESULT###"
+        if marker not in (r.stdout or ""):
+            return {"tests": {},
+                    "error": f"no result: {(r.stderr or r.stdout or '')[-400:]}"}
+        try:
+            return json.loads(r.stdout.split(marker, 1)[1].strip())
+        except Exception as e:
+            return {"tests": {}, "error": f"unparseable result: {e}"}
+
+    def _test_typescript(self, workdir: Path) -> Dict[str, Any]:
+        tests = sorted(p for p in (workdir / "tests").rglob("*.ts")) \
+            if (workdir / "tests").is_dir() else []
+        if not tests:
+            return {"tests": {}, "error": "no test files"}
+        r = self.sandbox.run_argv(
+            ["node", "--test", "--experimental-strip-types",
+             "--test-reporter=tap",
+             *[str(p.relative_to(workdir)) for p in tests]],
+            workdir, timeout=self.timeout, env=self._env(), ro_binds=self._binds())
+        return self._parse_tap(r)
+
+    @staticmethod
+    def _parse_tap(r) -> Dict[str, Any]:
+        if r.status == "timeout":
+            return {"tests": {}, "error": "timeout"}
+        text = (r.stdout or "") + "\n" + (r.stderr or "")
+        outcomes: Dict[str, str] = {}
+        for line in text.splitlines():
+            m = re.match(r"\s*(not )?ok\s+\d+\s*-\s*(.+?)\s*$", line)
+            if not m:
+                continue
+            name = m.group(2).strip()
+            if name.endswith(" # SKIP") or " # SKIP" in name:
+                outcomes[name.split(" # ")[0]] = "skip"
+                continue
+            verdict = "fail" if m.group(1) else "pass"
+            if outcomes.get(name) != "fail":
+                outcomes[name] = verdict
+        if not outcomes:
+            return {"tests": {},
+                    "error": f"no TAP output: {text.strip()[-400:]}"}
+        return {"tests": outcomes}
+
+
+SYSTEM_PROMPT = """You are an expert {lang} engineer working in an existing repository.
+
+Resolve the user's request by editing files with the tools provided.
+
+How this session works:
+- You are working autonomously. There is no human watching; nobody will answer
+  questions, approve a plan, or read progress reports.
+- Every reply you send must contain at least one tool call until the task is
+  done. A reply that contains only text does nothing: no work happens and the
+  turn is wasted. Do not stop to describe what you are about to do; do it.
+- The session ends only when you call `finish`. Call it once the change is
+  made, and not before. Only the files as you leave them are evaluated.
+
+Important constraints:
+- You cannot run the code, the tests, or any shell command. The only feedback
+  available is the `lint` tool, which reports syntax and type problems.
+- The repository already has a test suite that will be run against your work
+  after you finish, but you cannot see it or run it. Do not create test files.
+- Read enough of the code to understand it before editing. A change that only
+  makes the symptom go away will usually fail the tests.
+- Make the smallest correct change. Do not reformat or refactor unrelated code.
+
+Keep calling tools until the change is complete, then call `finish`."""
+
+
+@dataclass
+class AgenticTask:
+    task_id: str
+    lang: str
+    category: str
+    difficulty: int
+    instruction: str
+    files: Dict[str, str]
+    tests: Dict[str, str]
+    gold: Dict[str, str] = field(default_factory=dict)
+    fail_to_pass: List[str] = field(default_factory=list)
+    pass_to_pass: List[str] = field(default_factory=list)
+
+    @classmethod
+    def from_record(cls, rec: Dict[str, Any]) -> "AgenticTask":
+        def as_map(v):
+            return v if isinstance(v, dict) else json.loads(v or "{}")
+        def as_list(v):
+            return v if isinstance(v, list) else json.loads(v or "[]")
+        return cls(
+            task_id=rec["task_id"], lang=rec["lang"], category=rec["category"],
+            difficulty=int(rec.get("difficulty", 3)),
+            instruction=rec["instruction"],
+            files=as_map(rec["files"]), tests=as_map(rec["tests"]),
+            gold=as_map(rec.get("gold")),
+            fail_to_pass=as_list(rec.get("fail_to_pass")),
+            pass_to_pass=as_list(rec.get("pass_to_pass")),
+        )
+
+    @property
+    def spec(self) -> LangSpec:
+        return LANGS[self.lang]
+
+
+def materialise(files: Dict[str, str], dest: Path) -> Path:
+    """Write a file map out as a real directory tree."""
+    dest.mkdir(parents=True, exist_ok=True)
+    for rel, content in files.items():
+        p = dest / rel
+        p.parent.mkdir(parents=True, exist_ok=True)
+        p.write_text(content)
+    return dest
+
+
+def repo_overview(files: Dict[str, str]) -> str:
+    """A compact tree, so a run is not lost merely to not finding the files."""
+    lines = []
+    for rel in sorted(files):
+        n = files[rel].count("\n") + 1
+        lines.append(f"  {rel} ({n} lines)")
+    return "\n".join(lines)
+
+
+def user_prompt(task: AgenticTask) -> str:
+    return (f"Repository files:\n{repo_overview(task.files)}\n\n"
+            f"Task:\n{task.instruction}")
+
+
+@dataclass
+class EpisodeResult:
+    stop_reason: str = "unknown"
+    turns: int = 0
+    tool_calls: int = 0
+    tool_errors: int = 0
+    edits: int = 0
+    escape_attempts: int = 0
+    finished: bool = False
+    bad_tool_json: int = 0
+    nudges: int = 0
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    peak_context: int = 0
+    truncated_turns: int = 0
+    per_tool: Dict[str, int] = field(default_factory=dict)
+    error: Optional[str] = None
+
+
+NUDGE = ("Your last reply contained no tool call, so nothing happened. There is "
+         "no human in this session to respond to text. If your change is "
+         "complete, call `finish`. Otherwise continue working by calling the "
+         "tools; do not reply with text alone.")
+
+
+def run_episode(task: AgenticTask, toolbox: ToolBox, chat,
+                max_turns: int = 50,
+                context_limit: Optional[int] = None,
+                max_nudges: int = 3) -> Tuple[EpisodeResult, List[Dict]]:
+    """Drive one multi-turn tool-calling episode.
+
+    `chat(messages, tools) -> response dict` is supplied by the caller so this
+    module stays free of any HTTP concerns.
+    """
+    tools = tool_schemas()
+    messages: List[Dict[str, Any]] = [
+        {"role": "system", "content": SYSTEM_PROMPT.format(lang=task.spec.name)},
+        {"role": "user", "content": user_prompt(task)},
+    ]
+    res = EpisodeResult()
+
+    for turn in range(max_turns):
+        res.turns = turn + 1
+        try:
+            reply = chat(messages, tools)
+        except Exception as e:
+            res.stop_reason, res.error = "request_failed", f"{type(e).__name__}: {e}"
+            break
+
+        usage = reply.get("usage") or {}
+        pt = int(usage.get("prompt_tokens") or 0)
+        ct = int(usage.get("completion_tokens") or 0)
+        res.prompt_tokens += pt
+        res.completion_tokens += ct
+        res.peak_context = max(res.peak_context, pt + ct)
+
+        choices = reply.get("choices") or []
+        if not choices:
+            res.stop_reason, res.error = "no_choices", json.dumps(reply)[:300]
+            break
+        msg = choices[0].get("message") or {}
+        if choices[0].get("finish_reason") == "length":
+            res.truncated_turns += 1
+        calls = msg.get("tool_calls") or []
+
+        reasoning = msg.get("reasoning_content")
+        messages.append({
+            "role": "assistant",
+            "content": msg.get("content") or "",
+            **({"reasoning_content": reasoning} if reasoning else {}),
+            **({"tool_calls": calls} if calls else {}),
+        })
+
+        if not calls:
+            if toolbox.finished:
+                res.stop_reason = "finished"
+                break
+            if res.nudges >= max_nudges:
+                res.stop_reason = "no_tool_calls"
+                break
+            res.nudges += 1
+            messages.append({"role": "user", "content": NUDGE})
+            continue
+
+        for call in calls:
+            fn = call.get("function") or {}
+            name = fn.get("name") or ""
+            raw = fn.get("arguments")
+            if isinstance(raw, dict):
+                args = raw
+            else:
+                try:
+                    args = json.loads(raw or "{}")
+                except (json.JSONDecodeError, TypeError):
+                    res.bad_tool_json += 1
+                    args = None
+            output = (toolbox.dispatch(name, args) if args is not None else
+                      f"Error: arguments for {name!r} were not valid JSON.")
+            messages.append({"role": "tool",
+                             "tool_call_id": call.get("id") or f"call_{res.tool_calls}",
+                             "content": output})
+
+        if toolbox.finished:
+            res.stop_reason = "finished"
+            break
+        if context_limit and res.peak_context > context_limit:
+            res.stop_reason = "context_limit"
+            break
+    else:
+        res.stop_reason = "max_turns"
+
+    res.tool_calls = toolbox.calls
+    res.tool_errors = toolbox.errors
+    res.edits = toolbox.edits
+    res.escape_attempts = toolbox.ws.escape_attempts
+    res.finished = toolbox.finished
+    res.per_tool = dict(toolbox.per_tool)
+    return res, messages
+
+
+def run_tests_against(files: Dict[str, str], task: AgenticTask, runner: Runner,
+                      workroot: Path, label: str) -> Dict[str, str]:
+    """Materialise a tree, drop the hidden tests in, and run them.
+
+    The tests are only ever written into this throwaway copy, never into the
+    tree the agent worked in, so there is no window in which they could have
+    been read or edited.
+    """
+    dest = workroot / f"grade-{label}"
+    shutil.rmtree(dest, ignore_errors=True)
+    materialise(files, dest)
+    materialise(task.tests, dest)
+    return runner.test(dest).get("tests", {})
+
+
+def score(task: AgenticTask, outcomes: Dict[str, str]) -> Dict[str, Any]:
+    """Resolved means every target test passes and nothing previously passing broke."""
+    def passing(tid):
+        return outcomes.get(tid) == "pass"
+
+    f2p_ok = [t for t in task.fail_to_pass if passing(t)]
+    p2p_ok = [t for t in task.pass_to_pass if passing(t)]
+    regressions = [t for t in task.pass_to_pass if not passing(t)]
+    resolved = (len(f2p_ok) == len(task.fail_to_pass)
+                and not regressions and bool(task.fail_to_pass))
+    return {
+        "resolved": resolved,
+        "f2p_passed": len(f2p_ok), "f2p_total": len(task.fail_to_pass),
+        "p2p_passed": len(p2p_ok), "p2p_total": len(task.pass_to_pass),
+        "regressions": regressions[:10],
+        "n_regressions": len(regressions),
+        "partial": (len(f2p_ok) / len(task.fail_to_pass)) if task.fail_to_pass else 0.0,
+    }
+
+
+def make_sandbox(lang_name: str, sandbox_cls, cache_root=None, base_python=None):
+    """A sandbox carrying only the toolchain for `lang_name`.
+
+    Keyed per language rather than per suite: running the Python tasks must
+    never provision node, and vice versa.
+    """
+    lang = LANGS[lang_name]
+    return sandbox_cls(
+        name=f"agentic-{lang_name}",
+        packages=list(lang.packages),
+        provision=list(lang.provision),
+        env=dict(lang.env),
+        address_space_limit=lang.address_space_limit,
+        cache_root=cache_root,
+        base_python=base_python,
+    )
+
+
+def run_task(task: AgenticTask, chat, sandbox, workroot: Optional[Path] = None,
+             max_turns: int = 50, context_limit: Optional[int] = None,
+             test_timeout: float = 90.0) -> Dict[str, Any]:
+    """One full attempt: fresh workspace, agent episode, then hidden tests."""
+    root = Path(tempfile.mkdtemp(prefix=f"agentic-{task.task_id}-",
+                                 dir=str(workroot) if workroot else None))
+    try:
+        ws_root = root / "repo"
+        materialise(task.files, ws_root)
+        ws = Workspace(ws_root)
+        runner = Runner(sandbox, task.spec, timeout=test_timeout)
+        toolbox = ToolBox(ws, linter=lambda t: runner.lint(ws_root, t))
+
+        episode, messages = run_episode(task, toolbox, chat,
+                                        max_turns=max_turns,
+                                        context_limit=context_limit)
+        final = ws.snapshot()
+        changed = sorted(k for k in final
+                         if final[k] != task.files.get(k)) + \
+                  sorted(k for k in task.files if k not in final)
+
+        outcomes = run_tests_against(final, task, runner, root, "final")
+        result = score(task, outcomes)
+        result.update({
+            "task_id": task.task_id, "lang": task.lang,
+            "category": task.category, "difficulty": task.difficulty,
+            "changed_files": changed,
+            "episode": {k: v for k, v in vars(episode).items()},
+            "messages": messages,
+        })
+        return result
+    finally:
+        shutil.rmtree(root, ignore_errors=True)
+
+
+TRANSCRIPT_CSS = """
+body { font-family: system-ui, sans-serif; margin: 0; padding: 16px; background: #fff; color: #222; max-width: 1100px; }
+h1 { font-size: 16px; margin: 0 0 4px; }
+.meta { font-size: 13px; color: #555; font-family: 'SF Mono', 'Menlo', 'Consolas', monospace; margin-bottom: 12px; white-space: pre-wrap; }
+.msg { border: 1px solid #e1e4e8; border-left-width: 4px; margin: 8px 0; padding: 6px 10px; background: #fafbfc; }
+.msg .role { font-size: 12px; font-weight: 600; color: #555; margin-bottom: 4px; }
+.system { border-left-color: #888; }
+.user { border-left-color: #0969da; }
+.nudge { border-left-color: #9a6700; }
+.assistant { border-left-color: #1a7f37; background: #fff; }
+.tool { border-left-color: #8250df; }
+.tool.err { border-left-color: #cf222e; }
+pre { font-family: 'SF Mono', 'Menlo', 'Consolas', monospace; font-size: 12px; background: #fff; border: 1px solid #e1e4e8; padding: 6px 8px; margin: 4px 0; white-space: pre-wrap; word-wrap: break-word; }
+.reasoning pre { color: #57606a; background: #f6f8fa; }
+.call { font-size: 12px; font-weight: 600; color: #8250df; margin-top: 6px; }
+details > summary { cursor: pointer; font-size: 12px; color: #555; }
+.empty { font-size: 12px; color: #888; font-style: italic; }
+"""
+
+
+def _pre(text: str) -> str:
+    return f"<pre>{html.escape(text)}</pre>"
+
+
+def _render_call(call: Dict[str, Any]) -> str:
+    fn = call.get("function") or {}
+    raw = fn.get("arguments")
+    try:
+        args = raw if isinstance(raw, dict) else json.loads(raw or "{}")
+        shown = json.dumps(args, indent=2, ensure_ascii=False)
+    except (json.JSONDecodeError, TypeError):
+        shown = f"(invalid JSON)\n{raw}"
+    return (f'<div class="call">&rarr; {html.escape(fn.get("name") or "?")} '
+            f'<span style="font-weight:400;color:#888">{html.escape(call.get("id") or "")}</span></div>'
+            + _pre(shown))
+
+
+def render_transcript_html(messages: List[Dict[str, Any]], title: str,
+                           meta: str = "") -> str:
+    """One page per episode: every turn, with reasoning, tool calls and results."""
+    names = {}
+    for m in messages:
+        for c in m.get("tool_calls") or []:
+            names[c.get("id")] = (c.get("function") or {}).get("name") or "?"
+
+    parts = []
+    turn = 0
+    for m in messages:
+        role = m.get("role") or "?"
+        content = m.get("content") or ""
+        if role == "assistant":
+            turn += 1
+            body = ""
+            if m.get("reasoning_content"):
+                body += ('<details class="reasoning" open><summary>reasoning '
+                         f'({len(m["reasoning_content"])} chars)</summary>'
+                         + _pre(m["reasoning_content"]) + "</details>")
+            body += _pre(content) if content.strip() else ""
+            body += "".join(_render_call(c) for c in m.get("tool_calls") or [])
+            if not body:
+                body = '<div class="empty">(empty reply)</div>'
+            parts.append(f'<div class="msg assistant"><div class="role">assistant '
+                         f'&middot; turn {turn}</div>{body}</div>')
+        elif role == "tool":
+            name = names.get(m.get("tool_call_id"), "?")
+            err = " err" if content.startswith("Error") else ""
+            parts.append(f'<div class="msg tool{err}"><div class="role">tool result '
+                         f'&middot; {html.escape(name)}</div>{_pre(content)}</div>')
+        elif role == "system":
+            parts.append('<div class="msg system"><details><summary>system prompt'
+                         f'</summary>{_pre(content)}</details></div>')
+        else:
+            cls = "nudge" if content == NUDGE else role
+            label = "user (nudge)" if content == NUDGE else role
+            parts.append(f'<div class="msg {html.escape(cls)}"><div class="role">'
+                         f'{html.escape(label)}</div>{_pre(content)}</div>')
+
+    return (f'<!DOCTYPE html>\n<html>\n<head>\n<meta charset="UTF-8">\n'
+            f'<title>{html.escape(title)}</title>\n<style>{TRANSCRIPT_CSS}</style>\n'
+            f'</head>\n<body>\n<h1>{html.escape(title)}</h1>\n'
+            f'<div class="meta">{html.escape(meta)}</div>\n'
+            + "\n".join(parts) + "\n</body>\n</html>\n")
+
+
+def write_transcript(messages: List[Dict[str, Any]], dest: Path, title: str,
+                     meta: str = "") -> Path:
+    """Write `<dest>.json` (the raw message history) and `<dest>.html`."""
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    json_path = dest.with_name(dest.name + ".json")
+    with open(json_path, "w") as f:
+        json.dump(messages, f, indent=2, ensure_ascii=False)
+    with open(dest.with_name(dest.name + ".html"), "w") as f:
+        f.write(render_transcript_html(messages, title, meta))
+    return json_path
@@ -0,0 +1,488 @@
+#!/usr/bin/env python3
+"""Lightweight sandboxed execution for code-generation evals (HumanEval, ClassEval).
+
+Replaces the usual Docker-per-suite setup with a cached venv plus whatever
+process isolation the host can give us without root:
+
+    tier 1  bubblewrap   read-only root, tmpfs /tmp and $HOME, no network
+    tier 2  unshare -rn  no network (home still visible)
+    tier 3  rlimits only last resort, still capped on cpu/memory/procs/files
+
+Every tier also applies rlimits and runs the candidate program in a throwaway
+working directory as its own process group, so a runaway solution is killed
+along with anything it spawned.
+"""
+
+import json
+import os
+import re
+import resource
+import shutil
+import signal
+import subprocess
+import sys
+import tempfile
+import time
+import venv
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional, Sequence
+
+DEFAULT_CACHE_ROOT = Path(
+    os.environ.get("LLAMA_EVAL_CACHE", Path.home() / ".cache" / "llama-eval")
+)
+
+LIMIT_AS_BYTES = 4 * 1024 * 1024 * 1024
+LIMIT_CPU_SECONDS_SLACK = 5
+LIMIT_NPROC = 256
+LIMIT_FSIZE_BYTES = 64 * 1024 * 1024
+
+
+@dataclass
+class ExecResult:
+    passed: bool
+    status: str
+    stdout: str = ""
+    stderr: str = ""
+    returncode: Optional[int] = None
+    duration: float = 0.0
+
+    def summary(self, max_chars: int = 400) -> str:
+        """Short human-readable reason, for grader logs."""
+        if self.passed:
+            return "ok"
+        if self.status == "timeout":
+            return "timeout"
+        tail = (self.stderr or self.stdout).strip()
+        if not tail:
+            return f"{self.status} (rc={self.returncode})"
+        lines = [ln for ln in tail.splitlines() if ln.strip()]
+        reason = lines[-1] if lines else tail
+        return reason[:max_chars]
+
+
+def resolve_python(spec: Optional[str]) -> Optional[str]:
+    """Turn '3.13' or a path into a usable interpreter, or None if unavailable.
+
+    Suites can need an older interpreter than the one running this tool, because
+    a dependency has no wheel (or no working source build) for the newest
+    CPython. Accepts an explicit path, a bare version, or an executable name.
+    """
+    if not spec:
+        return None
+
+    candidate = Path(spec).expanduser()
+    if candidate.exists() and candidate.is_file():
+        return str(candidate)
+
+    uv = shutil.which("uv")
+    if uv and re.fullmatch(r"\d+\.\d+(\.\d+)?", spec):
+        try:
+            r = subprocess.run([uv, "python", "find", spec],
+                               capture_output=True, text=True, timeout=120)
+            path = r.stdout.strip()
+            if r.returncode == 0 and path and Path(path).exists():
+                return path
+        except Exception:
+            pass
+
+    return shutil.which(spec) or shutil.which(f"python{spec}")
+
+
+def _detect_isolation(probe_python: Optional[str] = None,
+                      ro_binds: Sequence[Path] = ()) -> str:
+    """Pick the strongest isolation tier this host supports, once.
+
+    The probe must run the *same* interpreter the suite will, with the same
+    binds. Probing a bare sys.executable silently downgrades the tier whenever
+    this tool is itself run from a venv under $HOME -- which is the usual case,
+    since it needs `datasets`. bwrap tmpfs-mounts the home directory, the
+    interpreter disappears, the probe fails, and the sandbox falls back to a
+    tier that does not hide $HOME at all. Nothing about the results looks wrong
+    when that happens, so it has to be got right here.
+    """
+    py = probe_python or sys.executable
+    probe = "import sys; sys.exit(0)"
+
+    if shutil.which("bwrap"):
+        try:
+            r = subprocess.run(
+                _bwrap_argv(Path.cwd(), [py, "-I", "-c", probe], ro_binds=ro_binds),
+                capture_output=True, timeout=20,
+            )
+            if r.returncode == 0:
+                return "bwrap"
+        except Exception:
+            pass
+
+    if shutil.which("unshare"):
+        for tier, argv in (
+            ("unshare", ["unshare", "--user", "--map-current-user", "--net"]),
+            ("unshare-root", ["unshare", "-rn"]),
+        ):
+            try:
+                r = subprocess.run(
+                    [*argv, py, "-I", "-c", probe],
+                    capture_output=True, timeout=20,
+                )
+                if r.returncode == 0:
+                    return tier
+            except Exception:
+                continue
+
+    return "rlimit"
+
+
+UNSHARE_ARGV = {
+    "unshare": ["unshare", "--user", "--map-current-user", "--net"],
+    "unshare-root": ["unshare", "-rn"],
+}
+
+
+def _under(path: Path, parent: str) -> bool:
+    return bool(parent) and (str(path) == parent
+                             or str(path).startswith(parent.rstrip("/") + os.sep))
+
+
+def _bwrap_argv(workdir: Path, inner: Sequence[str],
+                ro_binds: Sequence[Path] = ()) -> List[str]:
+    home = os.path.expanduser("~")
+    argv = [
+        "bwrap",
+        "--ro-bind", "/", "/",
+        "--dev", "/dev",
+        "--proc", "/proc",
+        "--tmpfs", "/tmp",
+    ]
+    if home and home != "/":
+        argv += ["--tmpfs", home]
+    for path in ro_binds:
+        if _under(path, home):
+            argv += ["--ro-bind", str(path), str(path)]
+    argv += [
+        "--bind", str(workdir), str(workdir),
+        "--chdir", str(workdir),
+        "--unshare-all",
+        "--die-with-parent",
+        "--new-session",
+        "--",
+    ]
+    return argv + list(inner)
+
+
+def _nproc_ceiling() -> int:
+    """A process cap that bounds a fork bomb without tripping on existing load.
+
+    RLIMIT_NPROC is enforced per real UID across the whole system, not per
+    process tree, so it has to sit above whatever the user is already running.
+    """
+    try:
+        current = len(os.listdir("/proc"))
+    except OSError:
+        current = 512
+    return max(current * 2, 1024)
+
+
+def _apply_rlimits(timeout: float, with_nproc: bool,
+                   as_bytes: Optional[int] = LIMIT_AS_BYTES):
+    def _preexec():
+        os.setsid()
+        cpu = int(timeout) + LIMIT_CPU_SECONDS_SLACK
+        limits = [
+            (resource.RLIMIT_CPU, (cpu, cpu)),
+            (resource.RLIMIT_FSIZE, (LIMIT_FSIZE_BYTES, LIMIT_FSIZE_BYTES)),
+            (resource.RLIMIT_CORE, (0, 0)),
+        ]
+        if as_bytes:
+            limits.append((resource.RLIMIT_AS, (as_bytes, as_bytes)))
+        if with_nproc:
+            ceiling = _nproc_ceiling()
+            limits.append((resource.RLIMIT_NPROC, (ceiling, ceiling)))
+        for res, lim in limits:
+            try:
+                resource.setrlimit(res, lim)
+            except (ValueError, OSError):
+                pass
+    return _preexec
+
+
+class Sandbox:
+    """A cached venv plus an isolated runner for untrusted candidate programs."""
+
+    def __init__(
+        self,
+        name: str,
+        packages: Optional[Sequence[str]] = None,
+        cache_root: Optional[Path] = None,
+        base_python: Optional[str] = None,
+        isolation: Optional[str] = None,
+        provision: Optional[Sequence[str]] = None,
+        env: Optional[Dict[str, str]] = None,
+        address_space_limit: Optional[int] = LIMIT_AS_BYTES,
+    ):
+        self.name = name
+        self.packages = sorted(set(packages or []))
+        self.provision = list(provision or [])
+        self.env = dict(env or {})
+        # None drops the RLIMIT_AS cap entirely -- see _apply_rlimits
+        self.address_space_limit = address_space_limit
+        self.cache_root = Path(cache_root or DEFAULT_CACHE_ROOT)
+        self.base_python = base_python or sys.executable
+        self.venv_dir = self.cache_root / f"{name}-venv"
+        self._python: Optional[Path] = None
+        self._isolation = isolation
+        self._base_prefix: List[Path] = []
+
+    @property
+    def python(self) -> Path:
+        if self._python is None:
+            raise RuntimeError("Sandbox.ensure() must be called first")
+        return self._python
+
+    @property
+    def isolation(self) -> str:
+        if self._isolation is None:
+            if self._python:
+                py = str(self._python)
+                roots = [self.venv_dir, *(self._base_prefix or [])]
+            else:
+                py = sys.executable
+                roots = self._detect_base_prefix(Path(py))
+            self._isolation = _detect_isolation(py, roots)
+        return self._isolation
+
+    def _stamp_path(self) -> Path:
+        return self.venv_dir / ".llama-eval-stamp.json"
+
+    def _stamp_matches(self) -> bool:
+        try:
+            stamp = json.loads(self._stamp_path().read_text())
+        except Exception:
+            return False
+        return (
+            stamp.get("packages") == self.packages
+            and stamp.get("py") == self._base_version()
+            and stamp.get("provision") == self.provision
+        )
+
+    def _resolved_env(self) -> Dict[str, str]:
+        return {k: v.replace("{venv}", str(self.venv_dir)) for k, v in self.env.items()}
+
+    def _base_version(self) -> str:
+        r = subprocess.run(
+            [self.base_python, "-c",
+             "import sys;print('.'.join(map(str,sys.version_info[:3])))"],
+            capture_output=True, text=True,
+        )
+        return r.stdout.strip()
+
+    def ensure(self, force: bool = False, quiet: bool = False) -> Path:
+        """Create/refresh the venv. Idempotent and cheap when already valid."""
+        py = self.venv_dir / "bin" / "python"
+        if force and self.venv_dir.exists():
+            shutil.rmtree(self.venv_dir)
+
+        if py.exists() and self._stamp_matches():
+            self._python = py
+            self._base_prefix = self._detect_base_prefix(py)
+            return py
+
+        if not quiet:
+            pkgs = ", ".join(self.packages) if self.packages else "stdlib only"
+            print(f"[sandbox] provisioning venv {self.venv_dir} ({pkgs})")
+
+        self.cache_root.mkdir(parents=True, exist_ok=True)
+        if self.venv_dir.exists():
+            shutil.rmtree(self.venv_dir)
+
+        uv = shutil.which("uv")
+        if uv:
+            subprocess.run(
+                [uv, "venv", "--python", self.base_python, str(self.venv_dir)],
+                check=True, capture_output=True,
+            )
+        elif Path(self.base_python).resolve() == Path(sys.executable).resolve():
+            venv.EnvBuilder(with_pip=bool(self.packages), clear=True).create(self.venv_dir)
+        else:
+            subprocess.run(
+                [self.base_python, "-m", "venv", str(self.venv_dir)],
+                check=True, capture_output=True,
+            )
+
+        if self.packages:
+            if uv:
+                cmd = [uv, "pip", "install", "--python", str(py), *self.packages]
+            else:
+                cmd = [str(py), "-m", "pip", "install", "--quiet",
+                       "--disable-pip-version-check", *self.packages]
+            r = subprocess.run(cmd, capture_output=True, text=True)
+            if r.returncode != 0:
+                raise RuntimeError(
+                    f"failed to install {len(self.packages)} package(s) into "
+                    f"{self.venv_dir}:\n{(r.stderr or r.stdout)[-3000:]}"
+                )
+
+        for i, snippet in enumerate(self.provision, 1):
+            if not quiet:
+                print(f"[sandbox] provisioning step {i}/{len(self.provision)}")
+            prov_env = dict(os.environ)
+            prov_env["SANDBOX_VENV"] = str(self.venv_dir)
+            prov_env.update(self._resolved_env())
+            r = subprocess.run([str(py), "-c", snippet], env=prov_env,
+                               capture_output=True, text=True, timeout=1800)
+            if r.returncode != 0:
+                raise RuntimeError(
+                    f"provisioning step {i} failed:\n{(r.stderr or r.stdout)[-3000:]}"
+                )
+
+        self._stamp_path().write_text(json.dumps({
+            "packages": self.packages,
+            "py": self._base_version(),
+            "provision": self.provision,
+        }))
+        self._python = py
+        self._base_prefix = self._detect_base_prefix(py)
+        if not quiet:
+            print(f"[sandbox] ready ({self.isolation} isolation)")
+        if self.isolation == "unshare-root":
+            print("Warning: this host's unshare cannot keep the caller's uid, so "
+                  "candidate code runs as root inside the namespace. Tests that "
+                  "assert a permission denial will not see one. Install "
+                  "bubblewrap (bwrap) for correct results.")
+        return py
+
+    @staticmethod
+    def _detect_base_prefix(py: Path) -> List[Path]:
+        """Every directory that must stay visible for the venv to start.
+
+        A venv is only symlinks into its base interpreter, so the whole install
+        tree has to survive the tmpfs that hides $HOME -- and binding just the
+        resolved tree is not enough. Version-managed interpreters (uv, pyenv)
+        are commonly reached through an unversioned alias symlink
+        (cpython-3.13-... -> cpython-3.13.14-...), and startup resolves through
+        the alias path, so both spellings need to be mounted. Walk the chain
+        and keep every install root it passes through.
+        """
+        roots: List[Path] = []
+
+        def add(path: Path):
+            if path.exists() and path not in roots:
+                roots.append(path)
+
+        cur, hops = py, 0
+        while hops < 16:
+            hops += 1
+            add(cur.parent.parent)
+            if cur.is_symlink():
+                target = Path(os.readlink(cur))
+                cur = target if target.is_absolute() else (cur.parent / target)
+            else:
+                break
+
+        try:
+            r = subprocess.run([str(py), "-c", "import sys; print(sys.base_prefix)"],
+                               capture_output=True, text=True, timeout=60)
+            if r.returncode == 0 and r.stdout.strip():
+                add(Path(r.stdout.strip()))
+        except Exception:
+            pass
+
+        return [p for p in roots if str(p) not in ("/", "/usr", "/usr/local")]
+
+    def _isolate(self, argv: Sequence[str], workdir: Path,
+                 ro_binds: Sequence[Path] = ()) -> List[str]:
+        """Wrap a command in the strongest isolation tier this host supports."""
+        if self.isolation == "bwrap":
+            return _bwrap_argv(
+                workdir, argv,
+                ro_binds=[self.venv_dir, *(self._base_prefix or []), *ro_binds],
+            )
+        if self.isolation in UNSHARE_ARGV:
+            return [*UNSHARE_ARGV[self.isolation], *argv]
+        return list(argv)
+
+    def _base_env(self, workdir: Path) -> Dict[str, str]:
+        env = {
+            "PATH": "/usr/bin:/bin",
+            "HOME": str(workdir),
+            "TMPDIR": str(workdir),
+            "LC_ALL": "C.UTF-8",
+            "PYTHONHASHSEED": "0",
+            "PYTHONDONTWRITEBYTECODE": "1",
+            "OPENBLAS_NUM_THREADS": "1",
+            "OMP_NUM_THREADS": "1",
+        }
+        env.update(self._resolved_env())
+        return env
+
+    def run_argv(self, argv: Sequence[str], workdir: Path,
+                 timeout: float = 15.0, env: Optional[Dict[str, str]] = None,
+                 ro_binds: Sequence[Path] = ()) -> ExecResult:
+        """Run an arbitrary command in `workdir`, which stays writable.
+
+        Unlike run(), the caller owns the directory and it survives the call.
+        That is what a multi-step agent needs: its edits have to persist across
+        many tool invocations before anything is graded. Taking an argv rather
+        than a source string is what lets a suite drive a linter or a test
+        runner for a language other than Python.
+        """
+        run_env = self._base_env(workdir)
+        if env:
+            run_env.update(env)
+
+        started = time.monotonic()
+        try:
+            proc = subprocess.Popen(
+                self._isolate(argv, workdir, ro_binds),
+                cwd=str(workdir),
+                stdin=subprocess.DEVNULL,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                errors="replace",
+                env=run_env,
+                preexec_fn=_apply_rlimits(timeout,
+                                          with_nproc=self.isolation == "rlimit",
+                                          as_bytes=self.address_space_limit),
+            )
+        except Exception as e:
+            return ExecResult(False, "error", "", f"{type(e).__name__}: {e}")
+
+        try:
+            out, err = proc.communicate(timeout=timeout)
+        except subprocess.TimeoutExpired:
+            _kill_tree(proc)
+            out, err = proc.communicate()
+            return ExecResult(False, "timeout", out or "", err or "",
+                              None, time.monotonic() - started)
+
+        duration = time.monotonic() - started
+        if proc.returncode == 0:
+            return ExecResult(True, "ok", out, err, 0, duration)
+        return ExecResult(False, "failed", out, err, proc.returncode, duration)
+
+    def run(self, code: str, timeout: float = 15.0,
+            env: Optional[Dict[str, str]] = None) -> ExecResult:
+        """Run `code` to completion in the sandbox. Exit 0 == passed."""
+        workdir = Path(tempfile.mkdtemp(prefix=f"{self.name}-", dir="/tmp"))
+        try:
+            prog = workdir / "candidate.py"
+            prog.write_text(code)
+            return self.run_argv([str(self.python), "-I", "-B", str(prog)],
+                                 workdir, timeout=timeout, env=env)
+        except Exception as e:
+            return ExecResult(False, "error", "", f"{type(e).__name__}: {e}")
+        finally:
+            shutil.rmtree(workdir, ignore_errors=True)
+
+
+def _kill_tree(proc: subprocess.Popen):
+    """Kill the candidate and everything it spawned."""
+    for sig in (signal.SIGKILL,):
+        try:
+            os.killpg(os.getpgid(proc.pid), sig)
+        except (ProcessLookupError, PermissionError):
+            try:
+                proc.kill()
+            except ProcessLookupError:
+                pass
@@ -13,8 +13,6 @@ from typing import Dict, List, Optional
 from dataclasses import dataclass
 from pathlib import Path

-import datasets
-
 # Set cache directory for HuggingFace datasets
 cache_dir = Path.home() / ".cache" / "huggingface" / "datasets"
 cache_dir.mkdir(parents=True, exist_ok=True)
@@ -76,6 +74,8 @@ class AimeDataset:
        return question.get("problem", question.get("question", ""))

    def _load_dataset(self):
+        import datasets
+
        if self.dataset_type == "aime":
            print(f"Loading AIME dataset (split: {self.split})...")
            cache_path = Path.home() / ".cache" / "huggingface" / "datasets" / "AI-MO___aimo-validation-aime" / "default" / "0.0.0"
@@ -164,6 +164,58 @@ class AimeDataset:
            return str(normalized) if normalized is not None else answer
        return str(answer)

+class CodeDataset:
+    """Stand-in for a code model: replays canonical solutions as chat replies.
+
+    Handles both code suites: HumanEval keys off the function stub, ClassEval
+    off the class skeleton.
+    """
+
+    DEFAULT_SOURCE = {
+        "humaneval": "openai/openai_humaneval",
+        "classeval": "FudanSELab/ClassEval",
+    }
+
+    def __init__(self, dataset_type: str, source: Optional[str] = None):
+        self.dataset_type = dataset_type
+        self.source = source or self.DEFAULT_SOURCE[dataset_type]
+        self.questions: List[Dict] = []
+        self._load_dataset()
+
+    def _load_dataset(self):
+        print(f"Loading {self.dataset_type} dataset (source: {self.source})...")
+        local = Path(self.source)
+        if local.exists():
+            text = local.read_text()
+            if local.suffix == ".json":
+                rows = json.loads(text)
+            else:
+                rows = [json.loads(ln) for ln in text.splitlines() if ln.strip()]
+        else:
+            import datasets
+            rows = [dict(r) for r in datasets.load_dataset(self.source, split="test")]
+        self.questions = rows
+        print(f"{self.dataset_type} dataset loaded: {len(self.questions)} questions")
+
+    def _stub_field(self) -> str:
+        return "skeleton" if self.dataset_type == "classeval" else "prompt"
+
+    def _get_question_text(self, question: Dict) -> str:
+        return question[self._stub_field()]
+
+    def find_question(self, request_text: str) -> Optional[Dict]:
+        for question in self.questions:
+            stub = question[self._stub_field()].rstrip()
+            if stub and stub in request_text:
+                return question
+        return None
+
+    def get_answer(self, question: Dict) -> str:
+        if self.dataset_type == "classeval":
+            return question["solution_code"]
+        return question["prompt"] + question["canonical_solution"]
+
+
 class Simulator:
    def __init__(
        self,
@@ -171,12 +223,17 @@ class Simulator:
        host: str = "localhost",
        success_rate: float = 0.8,
        dataset_split: str = "train",
-        dataset_type: str = "aime"
+        dataset_type: str = "aime",
+        dataset_source: Optional[str] = None
    ):
        self.port = port
        self.host = host
        self.success_rate = success_rate
-        self.dataset = AimeDataset(dataset_split, dataset_type)
+        self.dataset_type = dataset_type
+        if dataset_type in ("humaneval", "classeval"):
+            self.dataset = CodeDataset(dataset_type, dataset_source)
+        else:
+            self.dataset = AimeDataset(dataset_split, dataset_type)
        self.eval_state = EvalState(
            id=dataset_type,
            tasks=[dataset_type],
@@ -196,6 +253,12 @@ class Simulator:
        else:
            response_text = self._generate_wrong_answer(question)

+        if self.dataset_type in ("humaneval", "classeval"):
+            response_text = (
+                f"Here is the implementation:\n\n```python\n{response_text}\n```\n\n"
+                "Let me know if you want it optimised further."
+            )
+
        comp_tokens = random.randint(10000, 60000)
        tps_gen = random.uniform(90.0, 110.0)
        t_gen_ms = comp_tokens / tps_gen * 1000
@@ -229,6 +292,12 @@ class Simulator:
    def _generate_wrong_answer(self, question: Dict) -> str:
        expected_answer = self.dataset.get_answer(question)

+        if self.dataset_type == "classeval":
+            return re.sub(r"(\n\s+def [^\n]+:\n)", r"\1        return None\n",
+                          question["skeleton"])
+        if self.dataset_type == "humaneval":
+            return question["prompt"] + "    return None\n"
+
        if expected_answer.isdigit():
            wrong_answer = str(int(expected_answer) + 1)
        else:
@@ -335,7 +404,7 @@ def main():
        "--dataset",
        type=str,
        default="aime",
-        choices=["aime", "aime2025"],
+        choices=["aime", "aime2025", "humaneval", "classeval"],
        help="Dataset type (default: aime)"
    )
    parser.add_argument(
@@ -344,6 +413,12 @@ def main():
        default="train",
        help="AIME dataset split to use (default: train)"
    )
+    parser.add_argument(
+        "--dataset-source",
+        type=str,
+        default=None,
+        help="Override the dataset location (local .jsonl path or HuggingFace repo id)"
+    )

    args = parser.parse_args()

@@ -353,7 +428,8 @@ def main():
        host=args.host,
        success_rate=args.success_rate,
        dataset_split=args.dataset_split,
-        dataset_type=args.dataset
+        dataset_type=args.dataset,
+        dataset_source=args.dataset_source
    )

    server = HTTPServer((args.host, args.port), RequestHandler)
@@ -4,15 +4,18 @@
 #  SPDX-License-Identifier: MIT

 print_usage() {
-    echo "Usage: ./build.sh [fp32|fp16] [--help]"
+    echo "Usage: ./build.sh [fp32|fp16] [ssl|nossl] [--help]"
    echo ""
    echo "Options:"
    echo "  fp32    Build with FP32 precision (default)"
    echo "  fp16    Build with FP16 precision (faster for long-prompt inference)"
+    echo "  ssl     Build with OpenSSL (default)"
+    echo "  nossl   Build without OpenSSL"
    echo "  --help  Print this help message"
 }

 PRECISION=fp32
+SSL=ON

 for arg in "$@"; do
    case "$arg" in
@@ -23,6 +26,12 @@ for arg in "$@"; do
        fp32|fp16)
            PRECISION="$arg"
            ;;
+        ssl)
+            SSL=ON
+            ;;
+        nossl)
+            SSL=OFF
+            ;;
        *)
            echo "Error: unknown option '$arg'"
            print_usage
@@ -37,10 +46,13 @@ source /opt/intel/oneapi/setvars.sh

 if [ "$PRECISION" = "fp16" ]; then
    #for FP16
-    cmake .. -DGGML_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DGGML_SYCL_F16=ON -DLLAMA_OPENSSL=ON # faster for long-prompt inference
+    echo "cmake .. -DGGML_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DGGML_SYCL_F16=ON -DLLAMA_OPENSSL=$SSL # faster for long-prompt inference
+"
+    cmake .. -DGGML_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DGGML_SYCL_F16=ON -DLLAMA_OPENSSL="$SSL" # faster for long-prompt inference
 else
    #for FP32
-    cmake .. -DGGML_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DLLAMA_OPENSSL=ON
+    echo "cmake .. -DGGML_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DLLAMA_OPENSSL=$SSL"
+    cmake .. -DGGML_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DLLAMA_OPENSSL="$SSL"
 fi

 #build example/main
@@ -50,4 +62,5 @@ fi
 #cmake --build . --config Release --target llama-bench

 #build all binary
+echo "cmake --build . --config Release -j$((($(nproc)+1)/2)) -v"
 cmake --build . --config Release -j$((($(nproc)+1)/2)) -v
@@ -3,23 +3,62 @@
 ::  Copyright (C) 2024 Intel Corporation
 ::  SPDX-License-Identifier: MIT

-IF /I "%1"=="--help" (
-    echo Usage: win-build-sycl.bat [fp32^|fp16] [--help]
-    echo.
-    echo Options:
-    echo   fp32    Build with FP32 precision ^(default^)
-    echo   fp16    Build with FP16 precision ^(faster for long-prompt inference^)
-    echo   --help  Print this help message
-    exit /B 0
-)
+SET "PRECISION=fp32"
+SET "SSL=ON"
+SET "OPENSSL_OPTION="

-SET PRECISION=%1
-IF "%PRECISION%"=="" SET PRECISION=fp32
-IF /I NOT "%PRECISION%"=="fp32" IF /I NOT "%PRECISION%"=="fp16" (
-    echo Error: invalid value '%PRECISION%'. Use 'fp32' or 'fp16'.
-    echo Usage: win-build-sycl.bat [fp32^|fp16] [--help]
-    exit /B 1
-)
+:PARSE_ARGS
+IF "%~1"=="" GOTO ARGS_DONE
+IF /I "%~1"=="--help" GOTO HELP
+IF /I "%~1"=="fp32" GOTO ARG_FP32
+IF /I "%~1"=="fp16" GOTO ARG_FP16
+IF /I "%~1"=="ssl" GOTO ARG_SSL
+IF /I "%~1"=="nossl" GOTO ARG_NOSSL
+GOTO INVALID_ARG
+
+:ARG_FP32
+SET "PRECISION=fp32"
+GOTO NEXT_ARG
+
+:ARG_FP16
+SET "PRECISION=fp16"
+GOTO NEXT_ARG
+
+:ARG_SSL
+SET "SSL=ON"
+SET "OPENSSL_OPTION="
+GOTO NEXT_ARG
+
+:ARG_NOSSL
+SET "SSL=OFF"
+SET "OPENSSL_OPTION=-DLLAMA_OPENSSL=OFF"
+GOTO NEXT_ARG
+
+:NEXT_ARG
+SHIFT
+GOTO PARSE_ARGS
+
+:HELP
+SET "ARG_ERROR=0"
+GOTO PRINT_USAGE
+
+:INVALID_ARG
+SET "ARG_ERROR=1"
+echo Error: unknown option '%~1'
+
+:PRINT_USAGE
+echo Usage: win-build-sycl.bat [fp32^|fp16] [ssl^|nossl] [--help]
+echo.
+echo Options:
+echo   fp32      Build with FP32 precision ^(default^)
+echo   fp16      Build with FP16 precision ^(faster for long-prompt inference^)
+echo   ssl       Build with BoringSSL ^(default^)
+echo   nossl     Build without BoringSSL
+echo   --help    Print this help message
+IF "%ARG_ERROR%"=="1" exit /B 1
+exit /B 0
+
+:ARGS_DONE

 IF not exist build (mkdir build)
 cd build
@@ -31,10 +70,10 @@ if %errorlevel% neq 0 goto ERROR
 IF /I "%PRECISION%"=="fp16" (
    ::  for FP16
    ::  faster for long-prompt inference
-    cmake -G "MinGW Makefiles" .. -DLLAMA_OPENSSL=OFF -DGGML_SYCL=ON -DCMAKE_CXX_COMPILER=icx -DBUILD_SHARED_LIBS=ON -DCMAKE_BUILD_TYPE=Release -DGGML_SYCL_F16=ON
+    cmake -G "MinGW Makefiles" .. -DLLAMA_BUILD_BORINGSSL=%SSL% %OPENSSL_OPTION% -DGGML_SYCL=ON -DCMAKE_CXX_COMPILER=icx -DBUILD_SHARED_LIBS=ON -DCMAKE_BUILD_TYPE=Release -DGGML_SYCL_F16=ON
 ) ELSE (
    ::  for FP32
-    cmake -G "Ninja" .. -DLLAMA_OPENSSL=OFF -DGGML_SYCL=ON -DCMAKE_C_COMPILER=cl -DCMAKE_CXX_COMPILER=icx -DBUILD_SHARED_LIBS=ON -DCMAKE_BUILD_TYPE=Release
+    cmake -G "Ninja" .. -DLLAMA_BUILD_BORINGSSL=%SSL% %OPENSSL_OPTION% -DGGML_SYCL=ON -DCMAKE_C_COMPILER=cl -DCMAKE_CXX_COMPILER=icx -DBUILD_SHARED_LIBS=ON -DCMAKE_BUILD_TYPE=Release
 )
 if %errorlevel% neq 0 goto ERROR

@@ -4,7 +4,7 @@ project("ggml" C CXX ASM)

 ### GGML Version
 set(GGML_VERSION_MAJOR 0)
-set(GGML_VERSION_MINOR 26)
+set(GGML_VERSION_MINOR 27)
 set(GGML_VERSION_PATCH 0)
 set(GGML_VERSION_BASE "${GGML_VERSION_MAJOR}.${GGML_VERSION_MINOR}.${GGML_VERSION_PATCH}")

@@ -310,7 +310,7 @@ void ggml_vec_dot_q4_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const voi
        const int8x16_t v_yl = vec_xl(0      , y[ib].qs);
        const int8x16_t v_yh = vec_xl(QK8_1/2, y[ib].qs);

-        const int32x4_t v_xy_ = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_xl, v_yl), v_xh, v_yh);
+        const int32x4_t v_xy_ = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_xl, v_yl), v_xh, v_yh);
        const float32x4_t v_xy = vec_float(v_xy_);

        const float32x4_t v_d = vec_splats(GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d));
@@ -379,8 +379,8 @@ void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo
        const int8x16_t v_y1l = vec_xl(0,       y1->qs);
        const int8x16_t v_y1h = vec_xl(QK8_0/2, y1->qs);

-        const int32x4_t v_xy0 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x0l, v_y0l), v_x0h, v_y0h);
-        const int32x4_t v_xy1 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x1l, v_y1l), v_x1h, v_y1h);
+        const int32x4_t v_xy0 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x0l, v_y0l), v_x0h, v_y0h);
+        const int32x4_t v_xy1 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x1l, v_y1l), v_x1h, v_y1h);

        const float32x4_t v_xy0f = vec_float(v_xy0);
        const float32x4_t v_xy1f = vec_float(v_xy1);
@@ -407,7 +407,7 @@ void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo
        const int8x16_t v_yl = vec_xl(0,       y0->qs);
        const int8x16_t v_yh = vec_xl(QK8_0/2, y0->qs);

-        const int32x4_t v_xy = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_xl, v_yl), v_xh, v_yh);
+        const int32x4_t v_xy = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_xl, v_yl), v_xh, v_yh);
        const float32x4_t v_xyf = vec_float(v_xy);

        const float32x4_t v_d = vec_splats(GGML_E8M0_TO_FP32_HALF(x0->e) * GGML_CPU_FP16_TO_FP32(y0->d));
@@ -502,8 +502,8 @@ void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi
        const int8x16_t v_y1l = vec_xl(0,       (const int8_t *)y1->qs);
        const int8x16_t v_y1h = vec_xl(QK8_0/2, (const int8_t *)y1->qs);

-        const int32x4_t v_xy0 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x0lf, v_y0l), v_x0hf, v_y0h);
-        const int32x4_t v_xy1 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x1lf, v_y1l), v_x1hf, v_y1h);
+        const int32x4_t v_xy0 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x0lf, v_y0l), v_x0hf, v_y0h);
+        const int32x4_t v_xy1 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x1lf, v_y1l), v_x1hf, v_y1h);

        const float32x4_t v_xy0f = vec_float(v_xy0);
        const float32x4_t v_xy1f = vec_float(v_xy1);
@@ -548,7 +548,7 @@ void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi
        const int8x16_t v_yl = vec_xl(0,       (const int8_t *)y0->qs);
        const int8x16_t v_yh = vec_xl(QK8_0/2, (const int8_t *)y0->qs);

-        const int32x4_t v_xy = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_xlf, v_yl), v_xhf, v_yh);
+        const int32x4_t v_xy = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_xlf, v_yl), v_xhf, v_yh);
        const float32x4_t v_xyf = vec_float(v_xy);

        const float32x4_t v_d = vec_splats(GGML_CPU_FP16_TO_FP32(x0->d) * GGML_CPU_FP16_TO_FP32(y0->d));
@@ -653,8 +653,8 @@ void ggml_vec_dot_q5_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const voi
        const int8x16_t v_y1l = vec_xl(0      , y1->qs);
        const int8x16_t v_y1h = vec_xl(QK8_1/2, y1->qs);

-        const int32x4_t v_xy0 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x0lf, v_y0l), v_x0hf, v_y0h);
-        const int32x4_t v_xy1 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x1lf, v_y1l), v_x1hf, v_y1h);
+        const int32x4_t v_xy0 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x0lf, v_y0l), v_x0hf, v_y0h);
+        const int32x4_t v_xy1 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x1lf, v_y1l), v_x1hf, v_y1h);

        const float32x4_t v_xy0f = vec_float(v_xy0);
        const float32x4_t v_xy1f = vec_float(v_xy1);
@@ -701,7 +701,7 @@ void ggml_vec_dot_q5_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const voi
        const int8x16_t v_yl = vec_xl(0      , y0->qs);
        const int8x16_t v_yh = vec_xl(QK8_1/2, y0->qs);

-        const int32x4_t v_xy = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_xlf, v_yl), v_xhf, v_yh);
+        const int32x4_t v_xy = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_xlf, v_yl), v_xhf, v_yh);
        const float32x4_t v_xyf = vec_float(v_xy);

        const float32x4_t v_d = vec_splats(GGML_CPU_FP16_TO_FP32(x0->d) * GGML_CPU_FP16_TO_FP32(y0->d));
@@ -751,7 +751,7 @@ void ggml_vec_dot_q8_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi
        const int8x16_t v_yl = vec_xl(0      , y[ib].qs);
        const int8x16_t v_yh = vec_xl(QK8_0/2, y[ib].qs);

-        const int32x4_t v_xy_ = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_xl, v_yl), v_xh, v_yh);
+        const int32x4_t v_xy_ = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_xl, v_yl), v_xh, v_yh);
        const float32x4_t v_xy = vec_float(v_xy_);
        const float32x4_t v_d = vec_splats(GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d));

@@ -855,10 +855,10 @@ void ggml_vec_dot_q3_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi
            q3bytes[2] = vec_sub((int8x16_t)vec_and(vec_sr(q3b[0], 2), v_3m), (int8x16_t)q3h[2]);
            q3bytes[3] = vec_sub((int8x16_t)vec_and(vec_sr(q3b[1], 2), v_3m), (int8x16_t)q3h[3]);

-            isum0 = ggml_vec_dot(v_z, q3bytes[0], q8bytes[0]);
-            isum1 = ggml_vec_dot(v_z, q3bytes[1], q8bytes[1]);
-            isum2 = ggml_vec_dot(v_z, q3bytes[2], q8bytes[2]);
-            isum3 = ggml_vec_dot(v_z, q3bytes[3], q8bytes[3]);
+            isum0 = ggml_vec_dot_s32(v_z, q3bytes[0], q8bytes[0]);
+            isum1 = ggml_vec_dot_s32(v_z, q3bytes[1], q8bytes[1]);
+            isum2 = ggml_vec_dot_s32(v_z, q3bytes[2], q8bytes[2]);
+            isum3 = ggml_vec_dot_s32(v_z, q3bytes[3], q8bytes[3]);

            isum += (isum0[0] + isum0[1] + isum0[2] + isum0[3]) * scale[0];
            isum += (isum1[0] + isum1[1] + isum1[2] + isum1[3]) * scale[1];
@@ -877,10 +877,10 @@ void ggml_vec_dot_q3_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi
            q3bytes[2] = vec_sub((int8x16_t)vec_and(vec_sr(q3b[0], 6), v_3m), (int8x16_t)q3h[2]);
            q3bytes[3] = vec_sub((int8x16_t)vec_and(vec_sr(q3b[1], 6), v_3m), (int8x16_t)q3h[3]);

-            isum0 = ggml_vec_dot(v_z, q3bytes[0], q8bytes[4]);
-            isum1 = ggml_vec_dot(v_z, q3bytes[1], q8bytes[5]);
-            isum2 = ggml_vec_dot(v_z, q3bytes[2], q8bytes[6]);
-            isum3 = ggml_vec_dot(v_z, q3bytes[3], q8bytes[7]);
+            isum0 = ggml_vec_dot_s32(v_z, q3bytes[0], q8bytes[4]);
+            isum1 = ggml_vec_dot_s32(v_z, q3bytes[1], q8bytes[5]);
+            isum2 = ggml_vec_dot_s32(v_z, q3bytes[2], q8bytes[6]);
+            isum3 = ggml_vec_dot_s32(v_z, q3bytes[3], q8bytes[7]);

            isum += vec_hsum_i32x4(isum0) * scale[0];
            isum += vec_hsum_i32x4(isum1) * scale[1];
@@ -981,7 +981,7 @@ void ggml_vec_dot_q4_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi
            v_xl[0] = (int8x16_t)vec_and(v_x[0], v_lm);
            v_xl[1] = (int8x16_t)vec_and(v_x[1], v_lm);

-            const int32x4_t p1 = ggml_vec_dot(ggml_vec_dot(v_z, v_xl[0], v_y[0]), v_xl[1], v_y[1]);
+            const int32x4_t p1 = ggml_vec_dot_s32(ggml_vec_dot_s32(v_z, v_xl[0], v_y[0]), v_xl[1], v_y[1]);
            sumi1 += vec_hsum_i32x4(p1) * scales[2*j+0];

            v_y[0] = vec_xl(0 , y0);
@@ -991,7 +991,7 @@ void ggml_vec_dot_q4_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi
            v_xl[0] = (int8x16_t)vec_sr(v_x[0], 4);
            v_xl[1] = (int8x16_t)vec_sr(v_x[1], 4);

-            const int32x4_t p2 = ggml_vec_dot(ggml_vec_dot(v_z, v_xl[0], v_y[0]), v_xl[1], v_y[1]);
+            const int32x4_t p2 = ggml_vec_dot_s32(ggml_vec_dot_s32(v_z, v_xl[0], v_y[0]), v_xl[1], v_y[1]);
            sumi2 += vec_hsum_i32x4(p2) * scales[2*j+1];
        }

@@ -1107,8 +1107,8 @@ void ggml_vec_dot_q5_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi
            q5b[2] = (int8x16_t)vec_or(vec_sr(v_xl[0], 4), q5h[2]);
            q5b[3] = (int8x16_t)vec_or(vec_sr(v_xl[1], 4), q5h[3]);

-            int32x4_t sumi0 = ggml_vec_dot(ggml_vec_dot(v_z, q5b[0], v_y[0]), q5b[1], v_y[1]);
-            int32x4_t sumi1 = ggml_vec_dot(ggml_vec_dot(v_z, q5b[2], v_y[2]), q5b[3], v_y[3]);
+            int32x4_t sumi0 = ggml_vec_dot_s32(ggml_vec_dot_s32(v_z, q5b[0], v_y[0]), q5b[1], v_y[1]);
+            int32x4_t sumi1 = ggml_vec_dot_s32(ggml_vec_dot_s32(v_z, q5b[2], v_y[2]), q5b[3], v_y[3]);

            sumi += vec_hsum_i32x4(sumi0) * *scales++;
            sumi += vec_hsum_i32x4(sumi1) * *scales++;
@@ -1217,10 +1217,10 @@ void ggml_vec_dot_q6_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi
            q6b[2] = (int8x16_t)(vec_or(vec_and(v_xl[2], v_lm), q6h[2]));
            q6b[3] = (int8x16_t)(vec_or(vec_and(v_xl[3], v_lm), q6h[3]));

-            int32x4_t summs0 = ggml_vec_dot(v_z, q6b[0], v_y[0]);
-            int32x4_t summs1 = ggml_vec_dot(v_z, q6b[1], v_y[1]);
-            int32x4_t summs2 = ggml_vec_dot(v_z, q6b[2], v_y[2]);
-            int32x4_t summs3 = ggml_vec_dot(v_z, q6b[3], v_y[3]);
+            int32x4_t summs0 = ggml_vec_dot_s32(v_z, q6b[0], v_y[0]);
+            int32x4_t summs1 = ggml_vec_dot_s32(v_z, q6b[1], v_y[1]);
+            int32x4_t summs2 = ggml_vec_dot_s32(v_z, q6b[2], v_y[2]);
+            int32x4_t summs3 = ggml_vec_dot_s32(v_z, q6b[3], v_y[3]);

            isum += vec_hsum_i32x4(summs0) * scale[0] +
                    vec_hsum_i32x4(summs1) * scale[1] +
@@ -1251,10 +1251,10 @@ void ggml_vec_dot_q6_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi
            q6b[2] = (int8x16_t)(vec_or(vec_sr(v_xl[2], 4), q6h[2]));
            q6b[3] = (int8x16_t)(vec_or(vec_sr(v_xl[3], 4), q6h[3]));

-            summs0 = ggml_vec_dot(v_z, q6b[0], v_y[0]);
-            summs1 = ggml_vec_dot(v_z, q6b[1], v_y[1]);
-            summs2 = ggml_vec_dot(v_z, q6b[2], v_y[2]);
-            summs3 = ggml_vec_dot(v_z, q6b[3], v_y[3]);
+            summs0 = ggml_vec_dot_s32(v_z, q6b[0], v_y[0]);
+            summs1 = ggml_vec_dot_s32(v_z, q6b[1], v_y[1]);
+            summs2 = ggml_vec_dot_s32(v_z, q6b[2], v_y[2]);
+            summs3 = ggml_vec_dot_s32(v_z, q6b[3], v_y[3]);

            isum += vec_hsum_i32x4(summs0) * scale[0] +
                    vec_hsum_i32x4(summs1) * scale[1] +
@@ -1367,8 +1367,8 @@ void ggml_vec_dot_q6_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi
 //            q2u2 = vec_mul(q2u2, q2s2);
 //            q2u3 = vec_mul(q2u3, q2s3);

-//            const int32x4_t p1 = ggml_vec_dot(ggml_vec_dot(vec_splat_s32(0), q2u0, q8b0), q2u1, q8b1);
-//            const int32x4_t p2 = ggml_vec_dot(ggml_vec_dot(vec_splat_s32(0), q2u2, q8b2), q2u3, q8b3);
+//            const int32x4_t p1 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splat_s32(0), q2u0, q8b0), q2u1, q8b1);
+//            const int32x4_t p2 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splat_s32(0), q2u2, q8b2), q2u3, q8b3);

 //            sumf1 += (p1[0] + p1[1] + p1[2] + p1[3]) * (0.5f + (aux32[1] >> 28));
 //            sumf2 += (p2[0] + p2[1] + p2[2] + p2[3]) * (0.5f + (aux32[3] >> 28));
@@ -1445,7 +1445,7 @@ void ggml_vec_dot_iq4_nl_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const v

        const int8x16_t v_yl = vec_xl(0      , y0->qs);
        const int8x16_t v_yh = vec_xl(QK8_0/2, y0->qs);
-        const int32x4_t v_xy = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_xl, v_yl), v_xh, v_yh);
+        const int32x4_t v_xy = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_xl, v_yl), v_xh, v_yh);

        sumf += GGML_CPU_FP16_TO_FP32(x0->d) * GGML_CPU_FP16_TO_FP32(y0->d) * vec_hsum_i32x4(v_xy);
    }
@@ -1508,8 +1508,8 @@ void ggml_vec_dot_iq4_xs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const v
            const int8x16_t v_y3 = vec_xl(48, q8);
            q8 += 64;

-            int32x4_t vsumi0 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x0l, v_y0), v_x0h, v_y1);
-            int32x4_t vsumi1 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x1l, v_y2), v_x1h, v_y3);
+            int32x4_t vsumi0 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x0l, v_y0), v_x0h, v_y1);
+            int32x4_t vsumi1 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x1l, v_y2), v_x1h, v_y3);

            int ls1 = ((x[ibl].scales_l[ib] & 0xF) | ((h << 4) & 0x30)) - 32;
            int ls2 = ((x[ibl].scales_l[ib] >>  4) | ((h << 2) & 0x30)) - 32;
@@ -70,23 +70,6 @@ void ggml_quantize_mat_q8_0_4x4(const float * GGML_RESTRICT x, void * GGML_RESTR
 #endif
 }

-#if defined(__VXE__) || defined(__VXE2__)
-static inline int16x8_t vxe_dot_acc(const int8x16_t v_x, const int8x16_t v_y, const int16x8_t v_acc) {
-    return vec_meadd(v_x, v_y, vec_moadd(v_x, v_y, v_acc));
-}
-
-static inline int8x16_t vxe_splat_granule(const int8_t * qs) {
-    uint32_t g;
-    memcpy(&g, qs, sizeof(g));
-    return (int8x16_t)vec_splats(g);
-}
-
-static inline int32x4_t vxe_fold(const int16x8_t v_sumi) {
-    const int16x8_t v_ones = vec_splats((int16_t)1);
-    return vec_add(vec_mule(v_sumi, v_ones), vec_mulo(v_sumi, v_ones));
-}
-#endif
-
 void ggml_gemv_q4_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) {
    const int qk = QK8_0;
    const int nb = n / qk;
@@ -131,21 +114,21 @@ void ggml_gemv_q4_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo

            int16x8_t v_sumi = vec_splats((int16_t)0);

-            v_sumi = vxe_dot_acc(v_x0l, vxe_splat_granule(y_lo +  0), v_sumi);
-            v_sumi = vxe_dot_acc(v_x1l, vxe_splat_granule(y_lo +  4), v_sumi);
-            v_sumi = vxe_dot_acc(v_x2l, vxe_splat_granule(y_lo +  8), v_sumi);
-            v_sumi = vxe_dot_acc(v_x3l, vxe_splat_granule(y_lo + 12), v_sumi);
+            v_sumi = ggml_vec_dot_s16(v_sumi, v_x0l, ggml_vec_xl_s8x4_splat(y_lo +  0));
+            v_sumi = ggml_vec_dot_s16(v_sumi, v_x1l, ggml_vec_xl_s8x4_splat(y_lo +  4));
+            v_sumi = ggml_vec_dot_s16(v_sumi, v_x2l, ggml_vec_xl_s8x4_splat(y_lo +  8));
+            v_sumi = ggml_vec_dot_s16(v_sumi, v_x3l, ggml_vec_xl_s8x4_splat(y_lo + 12));

-            v_sumi = vxe_dot_acc(v_x0h, vxe_splat_granule(y_hi +  0), v_sumi);
-            v_sumi = vxe_dot_acc(v_x1h, vxe_splat_granule(y_hi +  4), v_sumi);
-            v_sumi = vxe_dot_acc(v_x2h, vxe_splat_granule(y_hi +  8), v_sumi);
-            v_sumi = vxe_dot_acc(v_x3h, vxe_splat_granule(y_hi + 12), v_sumi);
+            v_sumi = ggml_vec_dot_s16(v_sumi, v_x0h, ggml_vec_xl_s8x4_splat(y_hi +  0));
+            v_sumi = ggml_vec_dot_s16(v_sumi, v_x1h, ggml_vec_xl_s8x4_splat(y_hi +  4));
+            v_sumi = ggml_vec_dot_s16(v_sumi, v_x2h, ggml_vec_xl_s8x4_splat(y_hi +  8));
+            v_sumi = ggml_vec_dot_s16(v_sumi, v_x3h, ggml_vec_xl_s8x4_splat(y_hi + 12));

            const float32x4_t v_yd = vec_splats(GGML_CPU_FP16_TO_FP32(a_ptr[l].d));
-            const float32x4_t v_xd = __lzs_f16cx4_load(b_ptr[l].d);
+            const float32x4_t v_xd = ggml_vec_xl_f16x4_finite(b_ptr[l].d);
            const float32x4_t v_d  = vec_mul(v_yd, v_xd);

-            v_sumf = vec_madd(vec_float(vxe_fold(v_sumi)), v_d, v_sumf);
+            v_sumf = vec_madd(vec_float(vec_paddl_s16(v_sumi)), v_d, v_sumf);
        }

        vec_xst(v_sumf, 0, res_ptr + x * ncols_interleaved);
@@ -192,24 +175,24 @@ void ggml_gemm_q4_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo
                    const int8_t * y_lo = a_ptr[l].qs + 16 * k;
                    const int8_t * y_hi = y_lo + qk / 2 * 4;

-                    v_sumi0 = vxe_dot_acc(v_xl, vxe_splat_granule(y_lo +  0), v_sumi0);
-                    v_sumi1 = vxe_dot_acc(v_xl, vxe_splat_granule(y_lo +  4), v_sumi1);
-                    v_sumi2 = vxe_dot_acc(v_xl, vxe_splat_granule(y_lo +  8), v_sumi2);
-                    v_sumi3 = vxe_dot_acc(v_xl, vxe_splat_granule(y_lo + 12), v_sumi3);
+                    v_sumi0 = ggml_vec_dot_s16(v_sumi0, v_xl, ggml_vec_xl_s8x4_splat(y_lo +  0));
+                    v_sumi1 = ggml_vec_dot_s16(v_sumi1, v_xl, ggml_vec_xl_s8x4_splat(y_lo +  4));
+                    v_sumi2 = ggml_vec_dot_s16(v_sumi2, v_xl, ggml_vec_xl_s8x4_splat(y_lo +  8));
+                    v_sumi3 = ggml_vec_dot_s16(v_sumi3, v_xl, ggml_vec_xl_s8x4_splat(y_lo + 12));

-                    v_sumi0 = vxe_dot_acc(v_xh, vxe_splat_granule(y_hi +  0), v_sumi0);
-                    v_sumi1 = vxe_dot_acc(v_xh, vxe_splat_granule(y_hi +  4), v_sumi1);
-                    v_sumi2 = vxe_dot_acc(v_xh, vxe_splat_granule(y_hi +  8), v_sumi2);
-                    v_sumi3 = vxe_dot_acc(v_xh, vxe_splat_granule(y_hi + 12), v_sumi3);
+                    v_sumi0 = ggml_vec_dot_s16(v_sumi0, v_xh, ggml_vec_xl_s8x4_splat(y_hi +  0));
+                    v_sumi1 = ggml_vec_dot_s16(v_sumi1, v_xh, ggml_vec_xl_s8x4_splat(y_hi +  4));
+                    v_sumi2 = ggml_vec_dot_s16(v_sumi2, v_xh, ggml_vec_xl_s8x4_splat(y_hi +  8));
+                    v_sumi3 = ggml_vec_dot_s16(v_sumi3, v_xh, ggml_vec_xl_s8x4_splat(y_hi + 12));
                }

                const float32x4_t v_yd = __lzs_f16cx4_load(a_ptr[l].d);
-                const float32x4_t v_xd = __lzs_f16cx4_load(b_ptr[l].d);
+                const float32x4_t v_xd = ggml_vec_xl_f16x4_finite(b_ptr[l].d);

-                v_sumf[0] = vec_madd(vec_float(vxe_fold(v_sumi0)), vec_mul(v_xd, vec_splat(v_yd, 0)), v_sumf[0]);
-                v_sumf[1] = vec_madd(vec_float(vxe_fold(v_sumi1)), vec_mul(v_xd, vec_splat(v_yd, 1)), v_sumf[1]);
-                v_sumf[2] = vec_madd(vec_float(vxe_fold(v_sumi2)), vec_mul(v_xd, vec_splat(v_yd, 2)), v_sumf[2]);
-                v_sumf[3] = vec_madd(vec_float(vxe_fold(v_sumi3)), vec_mul(v_xd, vec_splat(v_yd, 3)), v_sumf[3]);
+                v_sumf[0] = vec_madd(vec_float(vec_paddl_s16(v_sumi0)), vec_mul(v_xd, vec_splat(v_yd, 0)), v_sumf[0]);
+                v_sumf[1] = vec_madd(vec_float(vec_paddl_s16(v_sumi1)), vec_mul(v_xd, vec_splat(v_yd, 1)), v_sumf[1]);
+                v_sumf[2] = vec_madd(vec_float(vec_paddl_s16(v_sumi2)), vec_mul(v_xd, vec_splat(v_yd, 2)), v_sumf[2]);
+                v_sumf[3] = vec_madd(vec_float(vec_paddl_s16(v_sumi3)), vec_mul(v_xd, vec_splat(v_yd, 3)), v_sumf[3]);
            }

            for (int m = 0; m < 4; m++) {
@@ -443,6 +443,13 @@ inline static ggml_int8x16x4_t ggml_vec_xl_s8x4(const int8_t * ptr) {
    return res;
 }

+// load 4 bytes and replicate them into all 4 word lanes
+inline static int8x16_t ggml_vec_xl_s8x4_splat(const int8_t * ptr) {
+    uint32_t tmp;
+    memcpy(&tmp, ptr, sizeof(tmp));
+    return (int8x16_t)vec_splats(tmp);
+}
+
 typedef struct ggml_int16x8x2_t {
    int16x8_t val[2];
 } ggml_int16x8x2_t;
@@ -456,31 +463,17 @@ inline static ggml_int16x8x2_t ggml_vec_xl_s16x2(const int16_t * ptr) {
    return res;
 }

-/*
-    ! WARNING: Very slow. Use vec_perm if possible. Refer to iq4_xs
-    !          or iq4_nl for example implementation.
-*/
-inline static int8x16_t ggml_vec_tbl(int8x16_t a, uint8x16_t b) {
-    int8x16_t res;
+// faster fp16 -> fp32 conversion using vector intrinsics
+// WARNING: fp16 inf converts into 65536.0f. use this only for finite values.
+inline static float32x4_t ggml_vec_xl_f16x4_finite(const ggml_fp16_t * x) {
+    uint64_t tmp;
+    memcpy(&tmp, x, sizeof(tmp));

-    res[ 0] = a[b[ 0]];
-    res[ 1] = a[b[ 1]];
-    res[ 2] = a[b[ 2]];
-    res[ 3] = a[b[ 3]];
-    res[ 4] = a[b[ 4]];
-    res[ 5] = a[b[ 5]];
-    res[ 6] = a[b[ 6]];
-    res[ 7] = a[b[ 7]];
-    res[ 8] = a[b[ 8]];
-    res[ 9] = a[b[ 9]];
-    res[10] = a[b[10]];
-    res[11] = a[b[11]];
-    res[12] = a[b[12]];
-    res[13] = a[b[13]];
-    res[14] = a[b[14]];
-    res[15] = a[b[15]];
-
-    return res;
+    const uint64x2_t  v_x    = { tmp, tmp };
+    const uint32x4_t  v_xh   = vec_unpackh((uint16x8_t)v_x);
+    const uint32x4_t  v_bits = vec_sl(vec_and(v_xh, vec_splats((uint32_t)0x7FFF)), 13);
+    const float32x4_t v_abs  = vec_mul((float32x4_t)v_bits, vec_splats(0x1.0p+112f));
+    return (float32x4_t)vec_or((uint32x4_t)v_abs, vec_sl(vec_and(v_xh, vec_splats((uint32_t)0x8000)), 16));
 }

 inline static int16x8_t vec_padd_s16(int16x8_t a, int16x8_t b) {
@@ -492,6 +485,11 @@ inline static int16x8_t vec_padd_s16(int16x8_t a, int16x8_t b) {
    return v_abo + v_abe;
 }

+inline static int32x4_t vec_paddl_s16(int16x8_t a) {
+    const int16x8_t v_ones = vec_splats((int16_t)1);
+    return vec_mule(a, v_ones) + vec_mulo(a, v_ones);
+}
+
 /**
 * @see https://github.com/ggml-org/llama.cpp/pull/14037
 */
@@ -505,7 +503,11 @@ inline static int32_t vec_hsum_i32x4(int32x4_t v) {
    return v_temp[0] + v_temp[1];
 }

-inline static int32x4_t ggml_vec_dot(int32x4_t acc, int8x16_t a, int8x16_t b) {
+inline static int16x8_t ggml_vec_dot_s16(int16x8_t acc, int8x16_t a, int8x16_t b) {
+    return vec_meadd(a, b, vec_moadd(a, b, acc));
+}
+
+inline static int32x4_t ggml_vec_dot_s32(int32x4_t acc, int8x16_t a, int8x16_t b) {
    const int16x8_t p = vec_mule(a, b) + vec_mulo(a, b);
    return acc + (vec_unpackh(p) + vec_unpackl(p));
 }
@@ -3548,6 +3548,73 @@ static void ggml_compute_forward_swiglu_oai_f32(
    }
 }

+static void ggml_compute_forward_swiglu_oai_f16(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+
+    const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
+    char * src0_d = (char *) src0->data;
+    char * src1_d = (char *) (src1 ? src1->data : src0->data);
+    const size_t src0_o = src0->nb[1];
+    const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1];
+
+    GGML_ASSERT(ggml_is_contiguous_1(src0));
+    GGML_ASSERT(ggml_is_contiguous_1(dst));
+
+    if (src1) {
+        GGML_ASSERT(ggml_is_contiguous_1(src1));
+        GGML_ASSERT(src0->type == src1->type);
+    }
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2;
+    const int nr = ggml_nrows(src0);
+
+    GGML_ASSERT(dst->ne[0] == nc);
+    GGML_ASSERT(ggml_nrows(dst) == nr);
+
+    const int32_t swapped = ggml_get_op_params_i32(dst, 1);
+    const float alpha = ggml_get_op_params_f32(dst, 2);
+    const float limit = ggml_get_op_params_f32(dst, 3);
+
+    // rows per thread
+    const int dr = (nr + nth - 1)/nth;
+
+    // row range for this thread
+    const int ir0 = dr*ith;
+    const int ir1 = MIN(ir0 + dr, nr);
+
+    for (int i1 = ir0; i1 < ir1; i1++) {
+        ggml_fp16_t * src0_p = (ggml_fp16_t *) (src0_d + i1*src0_o);
+        ggml_fp16_t * src1_p = (ggml_fp16_t *) (src1_d + i1*src1_o);
+        ggml_fp16_t * dst_p  = (ggml_fp16_t *) ((char *) dst->data + i1*(dst->nb[1]));
+
+        if (!src1) {
+            src0_p += swapped ? nc : 0;
+            src1_p += swapped ? 0 : nc;
+        }
+
+        for (int k = 0; k < nc; k++) {
+            const float x = std::min(GGML_CPU_FP16_TO_FP32(src0_p[k]), limit);
+            const float y = std::clamp(GGML_CPU_FP16_TO_FP32(src1_p[k]), -limit, limit);
+            const float out_glu = x / (1.f + expf(alpha * (-x)));
+            dst_p[k] = GGML_CPU_FP32_TO_FP16(out_glu * (y + 1.f));
+        }
+
+#ifndef NDEBUG
+        for (int k = 0; k < nc; k++) {
+            const float x = GGML_CPU_FP16_TO_FP32(dst_p[k]);
+            GGML_UNUSED(x);
+            assert(!isnan(x));
+            assert(!isinf(x));
+        }
+#endif // NDEBUG
+    }
+}
+
 static void ggml_compute_forward_swiglu_oai(
        const ggml_compute_params * params,
        ggml_tensor * dst) {
@@ -3559,6 +3626,10 @@ static void ggml_compute_forward_swiglu_oai(
            {
                ggml_compute_forward_swiglu_oai_f32(params, dst);
            } break;
+        case GGML_TYPE_F16:
+            {
+                ggml_compute_forward_swiglu_oai_f16(params, dst);
+            } break;
        default:
            {
                GGML_ABORT("fatal error");
@@ -8444,40 +8515,47 @@ static void ggml_compute_forward_pad_f32(
    const int32_t lp3 = ggml_get_op_params_i32(dst, 6);
    const int32_t rp3 = ggml_get_op_params_i32(dst, 7);

-    // TODO: optimize
+    const int64_t nrows = ne1 * ne2 * ne3;
+    const bool src_row_contig = src0->nb[0] == sizeof(float);
+
+    for (int64_t row = ith; row < nrows; row += nth) {
+        const int64_t i1 = row % ne1;
+        const int64_t i2 = (row / ne1) % ne2;
+        const int64_t i3 = row / (ne1 * ne2);
+
+        float * dst_row = dst_ptr + row * ne0;
+
+        if constexpr (circular_t) {
+            // circular padding wraps around on a torus
+            const int64_t src_i1 = ggml_wrap_around(i1 - lp1, ne01);
+            const int64_t src_i2 = ggml_wrap_around(i2 - lp2, ne02);
+            const int64_t src_i3 = ggml_wrap_around(i3 - lp3, ne03);

-    for (int64_t i2 = 0; i2 < ne2; ++i2) {
-        for (int64_t i1 = ith; i1 < ne1; i1 += nth) {
            for (int64_t i0 = 0; i0 < ne0; ++i0) {
-                for (int64_t i3 = 0; i3 < ne3; ++i3) {
-                    // circular means wrap around on a torus, so x and y loop around
-                    if constexpr (circular_t) {
-                        const int64_t dst_idx = i3*(ne0*ne1*ne2) + i2*(ne0*ne1) + i1*ne0 + i0;
-                        const int64_t src_i0 = ggml_wrap_around(i0 - lp0, ne00);
-                        const int64_t src_i1 = ggml_wrap_around(i1 - lp1, ne01);
-                        const int64_t src_i2 = ggml_wrap_around(i2 - lp2, ne02);
-                        const int64_t src_i3 = ggml_wrap_around(i3 - lp3, ne03);
+                const int64_t src_i0 = ggml_wrap_around(i0 - lp0, ne00);
+                dst_row[i0] = *(const float *)((const char *) src0->data + src_i3*nb03 + src_i2*nb02 + src_i1*nb01 + src_i0*nb00);
+            }
+        } else {
+            const bool interior = (i1 >= lp1 && i1 < ne1 - rp1)
+                               && (i2 >= lp2 && i2 < ne2 - rp2)
+                               && (i3 >= lp3 && i3 < ne3 - rp3);
+            if (!interior) {
+                memset(dst_row, 0, ne0 * sizeof(float));
+                continue;
+            }

-                        const int64_t src_idx =
-                            src_i3*nb03 +
-                            src_i2*nb02 +
-                            src_i1*nb01 +
-                            src_i0*nb00;
+            const char * src_row = (const char *) src0->data + (i3 - lp3)*nb03 + (i2 - lp2)*nb02 + (i1 - lp1)*nb01;

-                        const float * src_ptr = (const float *)((char *) src0->data + src_idx);
-                        dst_ptr[dst_idx] = *src_ptr;
+            if (src_row_contig && lp0 >= 0 && rp0 >= 0) {
+                memset(dst_row, 0, lp0 * sizeof(float));
+                memcpy(dst_row + lp0, src_row, ne00 * sizeof(float));
+                memset(dst_row + lp0 + ne00, 0, rp0 * sizeof(float));
+            } else {
+                for (int64_t i0 = 0; i0 < ne0; ++i0) {
+                    if (i0 >= lp0 && i0 < ne0 - rp0) {
+                        dst_row[i0] = *(const float *)(src_row + (i0 - lp0)*nb00);
                    } else {
-                        const int64_t dst_idx = i3*(ne0*ne1*ne2) + i2*(ne0*ne1) + i1*ne0 + i0;
-                        if ((i0 >= lp0 && i0 < ne0 - rp0) \
-                            && (i1 >= lp1 && i1 < ne1 - rp1) \
-                            && (i2 >= lp2 && i2 < ne2 - rp2) \
-                            && (i3 >= lp3 && i3 < ne3 - rp3)) {
-                            const int64_t src_idx = (i3 - lp3)*nb03 + (i2 - lp2)*nb02 + (i1 - lp1)*nb01 + (i0 - lp0)*nb00;
-                            const float * src_ptr = (const float *)((char *) src0->data + src_idx);
-                            dst_ptr[dst_idx] = *src_ptr;
-                        } else {
-                            dst_ptr[dst_idx] = 0;
-                        }
+                        dst_row[i0] = 0.0f;
                    }
                }
            }
@@ -244,16 +244,6 @@ static void tiled_run_microtile_avx2(const tiled_tile_src0 & src0, const tiled_t
    const int bs_stride = (nkr == 1) ? TILED_TILE_ROWS : TILED_MICRO;
    const int bs_off = seff * TILED_MICRO;

-    // Horizontal sum helper: converts 8x32-bit int lane inside a YMM register to scalar int32
-    auto hsum256_epi32 = [](const __m256i v) -> int32_t {
-        __m128i low = _mm256_castsi256_si128(v);
-        __m128i high = _mm256_extracti128_si256(v, 1);
-        __m128i sum128 = _mm_add_epi32(low, high);
-        sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, _MM_SHUFFLE(2, 3, 0, 1)));
-        sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, _MM_SHUFFLE(1, 0, 3, 2)));
-        return _mm_cvtsi128_si32(sum128);
-    };
-
    for (int i = 0; i < TILED_MICRO; i++) {
        const int ar = i0 + i;
        const int d_off = seff * TILED_MICRO + ar;
@@ -362,15 +352,20 @@ static void tiled_run_microtile_avx2(const tiled_tile_src0 & src0, const tiled_t
                }
            }

-            // OPTIMIZATION 2: Vectorized Epilogue with FMA3 gating
-            // Reduce accumulators into a 256-bit vector across 8 columns
-            alignas(32) int32_t s1_vals[GROUP];
-            #pragma GCC unroll 8
-            for (int t = 0; t < GROUP; t++) {
-                s1_vals[t] = hsum256_epi32(acc[t]);
-            }
-            __m256i s1_vec = _mm256_load_si256((const __m256i *) s1_vals);
+            // Reduce the eight accumulators to column sums in acc[0]..acc[7] order, with additions modulo 2^32.
+            const __m256i ab = _mm256_hadd_epi32(acc[0], acc[1]);
+            const __m256i cd = _mm256_hadd_epi32(acc[2], acc[3]);
+            const __m256i ef = _mm256_hadd_epi32(acc[4], acc[5]);
+            const __m256i gh = _mm256_hadd_epi32(acc[6], acc[7]);

+            // Each 128-bit half now holds four column sums over input lanes 0..3 or 4..7.
+            const __m256i abcd = _mm256_hadd_epi32(ab, cd);
+            const __m256i efgh = _mm256_hadd_epi32(ef, gh);
+            const __m256i lo = _mm256_permute2x128_si256(abcd, efgh, 0x20);
+            const __m256i hi = _mm256_permute2x128_si256(abcd, efgh, 0x31);
+            __m256i s1_vec = _mm256_add_epi32(lo, hi);
+
+            // OPTIMIZATION 2: Vectorized Epilogue with FMA3 gating
            if constexpr (BIAS != 0 && ACTBIAS) {
                // act-bias correction: raw = desired + 128*sum_s scales[s]*w_bsum[s]; corr precomputed in repack_src0
                s1_vec = _mm256_sub_epi32(s1_vec, _mm256_set1_epi32(mins_row[0]));
@@ -1348,6 +1348,12 @@ static ggml_backend_buffer_t ggml_backend_cuda_host_buffer_type_alloc_buffer(ggm
    return buffer;
 }

+static size_t ggml_backend_cuda_host_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) {
+    GGML_UNUSED(buft);
+    // very large pinned allocations can fail (e.g. on Windows) and fall back to pageable memory, smaller ones succeed
+    return 4ull*1024*1024*1024;
+}
+
 ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type() {
    static struct ggml_backend_buffer_type ggml_backend_cuda_buffer_type_host = {
        /* .iface    = */ {
@@ -1355,7 +1361,7 @@ ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type() {
            /* .alloc_buffer        = */ ggml_backend_cuda_host_buffer_type_alloc_buffer,
            /* .alloc_buffer_n      = */ NULL,
            /* .get_alignment       = */ ggml_backend_cpu_buffer_type()->iface.get_alignment,
-            /* .get_max_size        = */ NULL, // defaults to SIZE_MAX
+            /* .get_max_size        = */ ggml_backend_cuda_host_buffer_type_get_max_size,
            /* .get_alloc_size      = */ ggml_backend_cpu_buffer_type()->iface.get_alloc_size,
            /* .get_alloc_size_n    = */ NULL,
            /* .is_host             = */ ggml_backend_cpu_buffer_type()->iface.is_host,
@@ -5474,6 +5480,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                            (ggml_get_glu_op(op) == GGML_GLU_OP_SWIGLU_OAI || ggml_get_glu_op(op) == GGML_GLU_OP_SWIGLU_CLAMP)) {
                        return false;
                    }
+                    if (op->src[0]->type == GGML_TYPE_F16 && ggml_get_glu_op(op) == GGML_GLU_OP_SWIGLU_OAI) {
+                        return false;
+                    }
                    return ggml_is_contiguous_1(op->src[0]);
                default:
                    return false;
@@ -81,7 +81,7 @@ static __global__ void group_norm_f32(const float * x, float * dst, const int gr
    }
 }

-template <int block_size, bool do_multiply = false, bool do_add = false, bool do_scale = false>
+template <int block_size, int vpt, bool do_vec, bool do_multiply = false, bool do_add = false, bool do_scale = false>
 static __global__ void rms_norm_f32(const float * x,
                                    float *       dst,
                                    const int     ncols,
@@ -142,31 +142,82 @@ static __global__ void rms_norm_f32(const float * x,
                addc = add + add_sample * add_stride_sample + add_channel * add_stride_channel + add_row * add_stride_row;
            }

-            float tmp = 0.0f; // partial sum for thread in warp
+            if constexpr (do_vec) {
+                const int      ncols4 = ncols / 4;
+                const float4 * x4     = (const float4 *) xc;
+                const float4 * mul4   = (const float4 *) mulc;
+                const float4 * add4   = (const float4 *) addc;
+                float4 *       dst4   = (float4 *) dstc;

-            for (int col = tid; col < ncols; col += block_size) {
-                const float xi = xc[col];
-                tmp += xi * xi;
-            }
+                float4 xs[vpt];
+                float  tmp = 0.0f;

-            // sum up partial sums
-            tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum);
+#pragma unroll
+                for (int k = 0; k < vpt; k++) {
+                    const int i = tid + k*block_size;
+                    if (i < ncols4) {
+                        xs[k] = x4[i];
+                        tmp += xs[k].x*xs[k].x + xs[k].y*xs[k].y + xs[k].z*xs[k].z + xs[k].w*xs[k].w;
+                    }
+                }
+                tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum);
+                const float scale = rsqrtf(tmp / ncols + eps);

-            const float mean = tmp / ncols;
-            const float scale = rsqrtf(mean + eps);
+#pragma unroll
+                for (int k = 0; k < vpt; k++) {
+                    const int i = tid + k*block_size;
+                    if (i < ncols4) {
+                        float4 r = make_float4(scale*xs[k].x, scale*xs[k].y, scale*xs[k].z, scale*xs[k].w);
+                        if constexpr (do_multiply) {
+                            const float4 m = mul4[i];
+                            r.x *= m.x;
+                            r.y *= m.y;
+                            r.z *= m.z;
+                            r.w *= m.w;
+                        }
+                        if constexpr (do_add) {
+                            const float4 a = add4[i];
+                            r.x += a.x;
+                            r.y += a.y;
+                            r.z += a.z;
+                            r.w += a.w;
+                        }
+                        if constexpr (do_scale) {
+                            r.x *= scale_out;
+                            r.y *= scale_out;
+                            r.z *= scale_out;
+                            r.w *= scale_out;
+                        }
+                        dst4[i] = r;
+                    }
+                }
+            } else {
+                float tmp = 0.0f; // partial sum for thread in warp

-            for (int col = tid; col < ncols; col += block_size) {
-                if constexpr (do_multiply && do_add) {
-                    const int mul_col = fastmodulo(col, mul_ncols_packed);
-                    const int add_col = fastmodulo(col, add_ncols_packed);
-                    dstc[col]         = scale * xc[col] * mulc[mul_col] + addc[add_col];
-                } else if constexpr (do_multiply) {
-                    const int mul_col = fastmodulo(col, mul_ncols_packed);
-                    dstc[col]         = scale * xc[col] * mulc[mul_col];
-                } else if constexpr (do_scale) {
-                    dstc[col] = scale_out * (scale * xc[col]);
-                } else {
-                    dstc[col] = scale * xc[col];
+                for (int col = tid; col < ncols; col += block_size) {
+                    const float xi = xc[col];
+                    tmp += xi * xi;
+                }
+
+                // sum up partial sums
+                tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum);
+
+                const float mean = tmp / ncols;
+                const float scale = rsqrtf(mean + eps);
+
+                for (int col = tid; col < ncols; col += block_size) {
+                    if constexpr (do_multiply && do_add) {
+                        const int mul_col = fastmodulo(col, mul_ncols_packed);
+                        const int add_col = fastmodulo(col, add_ncols_packed);
+                        dstc[col]         = scale * xc[col] * mulc[mul_col] + addc[add_col];
+                    } else if constexpr (do_multiply) {
+                        const int mul_col = fastmodulo(col, mul_ncols_packed);
+                        dstc[col]         = scale * xc[col] * mulc[mul_col];
+                    } else if constexpr (do_scale) {
+                        dstc[col] = scale_out * (scale * xc[col]);
+                    } else {
+                        dstc[col] = scale * xc[col];
+                    }
                }
            }

@@ -332,34 +383,89 @@ static void group_norm_f32_cuda(
    }
 }

+#define RMS_NORM_LAUNCH(BLOCK_SIZE, VPT, VEC) \
+    do { \
+        const dim3 block_dims(BLOCK_SIZE, 1, 1); \
+        const ggml_cuda_kernel_launch_params launch_params{ blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float) : 0, stream }; \
+        ggml_cuda_kernel_launch(rms_norm_f32<BLOCK_SIZE, VPT, VEC, do_multiply, do_add, do_scale>, launch_params, \
+            x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, \
+            mul, mul_stride_row, mul_stride_channel, mul_stride_sample, \
+            mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, \
+            add, add_stride_row, add_stride_channel, add_stride_sample, \
+            add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, scale_out); \
+    } while (0)
+
+template <bool do_multiply, bool do_add, bool do_scale>
+static void rms_norm_f32_dispatch(
+        const dim3     blocks_num,
+        const ggml_tensor * x_t,   float * dst,   const int ncols, const int nchannels, const int nsamples,
+        const int64_t  stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps,
+        const ggml_tensor * mul_t, const int64_t mul_stride_row, const int64_t mul_stride_channel, const int64_t mul_stride_sample,
+        const uint3    mul_ncols_packed, const uint3 mul_nrows_packed, const uint3 mul_nchannels_packed, const uint3 mul_nsamples_packed,
+        const ggml_tensor * add_t, const int64_t add_stride_row, const int64_t add_stride_channel, const int64_t add_stride_sample,
+        const uint3    add_ncols_packed, const uint3 add_nrows_packed, const uint3 add_nchannels_packed, const uint3 add_nsamples_packed,
+        cudaStream_t   stream, const float scale_out = 1.0f) {
+    static_assert(!do_add || do_multiply, "fusing add is not supported without multiplying");
+    static_assert(!do_scale || !do_multiply, "fusing scale is not supported with multiplying");
+
+    const float * x   = (const float *) x_t->data;
+    const float * mul = mul_t ? (const float *) mul_t->data : nullptr;
+    const float * add = add_t ? (const float *) add_t->data : nullptr;
+
+    // dst is written at row * ncols, i.e. always at a multiple of ncols, so only its base needs checking here
+    bool vec = ncols % 4 == 0 && ncols > 0 && ((uintptr_t) dst & 15) == 0;
+    vec = vec && ggml_cuda_is_aligned(x_t, 16);
+    if constexpr (do_multiply) {
+        vec = vec && mul_ncols_packed.z == (uint32_t) ncols
+            && ggml_cuda_is_aligned(mul_t, 16);
+    }
+    if constexpr (do_add) {
+        vec = vec && add_ncols_packed.z == (uint32_t) ncols
+            && ggml_cuda_is_aligned(add_t, 16);
+    }
+
+    const int n4 = ncols / 4;
+    if (vec) {
+        if (n4 <= 32) {
+            RMS_NORM_LAUNCH(32,   1, true);
+        } else if (n4 <= 128) {
+            RMS_NORM_LAUNCH(128,  1, true);
+        } else if (n4 <= 256) {
+            RMS_NORM_LAUNCH(256,  1, true);
+        } else if (n4 <= 512) {
+            RMS_NORM_LAUNCH(512,  1, true);
+        } else if (n4 <= 1024) {
+            RMS_NORM_LAUNCH(1024, 1, true);
+        } else if (n4 <= 4096) {
+            RMS_NORM_LAUNCH(1024, 4, true);
+        } else {
+            RMS_NORM_LAUNCH(1024, 1, false);
+        }
+    } else if (ncols < 1024) {
+        RMS_NORM_LAUNCH(256,  1, false);
+    } else {
+        RMS_NORM_LAUNCH(1024, 1, false);
+    }
+}
+#undef RMS_NORM_LAUNCH
+
 template <bool do_scale = false>
 static void rms_norm_f32_cuda(
-        const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples,
+        const ggml_tensor * x_t, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples,
        const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream,
        const float scale_out = 1.0f) {
    const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX));
-    if (ncols < 1024) {
-        const dim3 block_dims(256, 1, 1);
-        const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
-        ggml_cuda_kernel_launch(rms_norm_f32<256, false, false, do_scale>, launch_params,
-            x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
-        // underlying cudaLaunchKernelEx does not support default params
+    rms_norm_f32_dispatch</*do_multiply*/false, /*do_add*/false, do_scale>(
+        blocks_num, x_t, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
        nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0),
-        nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), scale_out);
-    } else {
-        const dim3 block_dims(1024, 1, 1);
-        const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
-        ggml_cuda_kernel_launch(rms_norm_f32<1024, false, false, do_scale>, launch_params, x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
-        // underlying cudaLaunchKernelEx does not support default params
        nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0),
-        nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), scale_out);
-    }
+        stream, scale_out);
 }

-static void rms_norm_mul_f32_cuda(const float *  x,
-                                  const float *  mul,
-                                  const float *  add,
-                                  float *        dst,
+static void rms_norm_mul_f32_cuda(const ggml_tensor * x_t,
+                                  const ggml_tensor * mul_t,
+                                  const ggml_tensor * add_t,
+                                  float             * dst,
                                  const int      ncols,
                                  const int      nrows,
                                  const int      nchannels,
@@ -384,59 +490,33 @@ static void rms_norm_mul_f32_cuda(const float *  x,
                                  const float    eps,
                                  cudaStream_t   stream) {
    const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX));
-    if (mul == nullptr) {
-        rms_norm_f32_cuda(x, dst, ncols, nrows, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, stream);
+    if (mul_t == nullptr) {
+        rms_norm_f32_cuda(x_t, dst, ncols, nrows, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, stream);
        return;
    }
-    if (add == nullptr) {
-        const uint3 mul_ncols_packed     = init_fastdiv_values(mul_ncols);
-        const uint3 mul_nrows_packed     = init_fastdiv_values(mul_nrows);
-        const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels);
-        const uint3 mul_nsamples_packed  = init_fastdiv_values(mul_nsamples);
-        if (ncols < 1024) {
-            const dim3 block_dims(256, 1, 1);
-            const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
-            ggml_cuda_kernel_launch(rms_norm_f32<256, true>, launch_params,
-                x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
-                mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
-                // underlying cudaLaunchKernelEx does not support default params
-            nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), 1.0f);
-        } else {
-            const dim3 block_dims(1024, 1, 1);
-            const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
-            ggml_cuda_kernel_launch(rms_norm_f32<1024, true>, launch_params,
-                x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
-                mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
-                // underlying cudaLaunchKernelEx does not support default params
-            nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), 1.0f);
-        }
-    } else {
-        const uint3 mul_ncols_packed     = init_fastdiv_values(mul_ncols);
-        const uint3 mul_nrows_packed     = init_fastdiv_values(mul_nrows);
-        const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels);
-        const uint3 mul_nsamples_packed  = init_fastdiv_values(mul_nsamples);
+    const uint3 mul_ncols_packed     = init_fastdiv_values(mul_ncols);
+    const uint3 mul_nrows_packed     = init_fastdiv_values(mul_nrows);
+    const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels);
+    const uint3 mul_nsamples_packed  = init_fastdiv_values(mul_nsamples);
+    const uint3 zero_packed          = make_uint3(0, 0, 0);

+    if (add_t == nullptr) {
+        rms_norm_f32_dispatch</*do_multiply*/true, /*do_add*/false, /*do_scale*/false>(
+            blocks_num, x_t, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
+            mul_t, mul_stride_row, mul_stride_channel, mul_stride_sample,
+            mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
+            nullptr, 0, 0, 0, zero_packed, zero_packed, zero_packed, zero_packed, stream);
+    } else {
        const uint3 add_ncols_packed     = init_fastdiv_values(add_ncols);
        const uint3 add_nrows_packed     = init_fastdiv_values(add_nrows);
        const uint3 add_nchannels_packed = init_fastdiv_values(add_nchannels);
        const uint3 add_nsamples_packed  = init_fastdiv_values(add_nsamples);
-        if (ncols < 1024) {
-            const dim3 block_dims(256, 1, 1);
-            const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims,block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
-            ggml_cuda_kernel_launch(rms_norm_f32<256, true, true>, launch_params,
-                x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
-                mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add,
-                add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed,
-                add_nchannels_packed, add_nsamples_packed, 1.0f);
-        } else {
-            const dim3 block_dims(1024, 1, 1);
-            const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
-            ggml_cuda_kernel_launch(rms_norm_f32<1024, true, true>, launch_params,
-                x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
-                mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add,
-                add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed,
-                add_nchannels_packed, add_nsamples_packed, 1.0f);
-        }
+        rms_norm_f32_dispatch</*do_multiply*/true, /*do_add*/true, /*do_scale*/false>(
+            blocks_num, x_t, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
+            mul_t, mul_stride_row, mul_stride_channel, mul_stride_sample,
+            mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
+            add_t, add_stride_row, add_stride_channel, add_stride_sample,
+            add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, stream);
    }
 }

@@ -510,7 +590,6 @@ void ggml_cuda_op_group_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst)

 void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
    const ggml_tensor * src0 = dst->src[0];
-    const float * src0_d = (const float *) src0->data;
    float * dst_d = (float *) dst->data;
    cudaStream_t stream = ctx.stream();

@@ -529,12 +608,11 @@ void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
    const int64_t s02 = nb02 / ts0;
    const int64_t s03 = nb03 / ts0;

-    rms_norm_f32_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream);
+    rms_norm_f32_cuda(src0, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream);
 }

 void ggml_cuda_op_rms_norm_scale_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * scale_tensor) {
    const ggml_tensor * src0   = dst->src[0];
-    const float *       src0_d = (const float *) src0->data;
    float *             dst_d  = (float *) scale_tensor->data;
    cudaStream_t        stream = ctx.stream();

@@ -556,7 +634,7 @@ void ggml_cuda_op_rms_norm_scale_fused(ggml_backend_cuda_context & ctx, ggml_ten
    const int64_t s02 = nb02 / ts0;
    const int64_t s03 = nb03 / ts0;

-    rms_norm_f32_cuda<true>(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream, scale);
+    rms_norm_f32_cuda<true>(src0, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream, scale);
 }

 void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor) {
@@ -565,15 +643,11 @@ void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor *

    memcpy(&eps, dst->op_params, sizeof(float));

-    const float * src0_d = (const float *) rms_norm_src->data;
-    const float * mul_d = nullptr;
    const ggml_tensor * mul_src = nullptr;

    if (mul_tensor->src[0] == dst) {
-        mul_d = (float *) mul_tensor->src[1]->data;
        mul_src = mul_tensor->src[1];
    } else if(mul_tensor->src[1] == dst) {
-        mul_d = (float *) mul_tensor->src[0]->data;
        mul_src = mul_tensor->src[0];
    } else {
        GGML_ASSERT(false);
@@ -609,7 +683,7 @@ void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor *
    const int mul_nchannels = mul_src->ne[2];
    const int mul_nsamples  = mul_src->ne[3];

-    rms_norm_mul_f32_cuda(src0_d, mul_d, nullptr, dst_d,
+    rms_norm_mul_f32_cuda(rms_norm_src, mul_src, nullptr, dst_d,
                          ne00, ne01, ne02, ne03,
                          /*s00*/ s01, s02, s03,
                          /*mul_s00*/ mul_s01, mul_s02, mul_s03,
@@ -628,28 +702,21 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx,

    memcpy(&eps, dst->op_params, sizeof(float));

-    const float *       src0_d  = (const float *) rms_norm_src->data;
-    const float *       mul_d   = nullptr;
    const ggml_tensor * mul_src = nullptr;

    if (mul_tensor->src[0] == dst) {
-        mul_d   = (float *) mul_tensor->src[1]->data;
        mul_src = mul_tensor->src[1];
    } else if (mul_tensor->src[1] == dst) {
-        mul_d   = (float *) mul_tensor->src[0]->data;
        mul_src = mul_tensor->src[0];
    } else {
        GGML_ASSERT(false);
    }

-    const float *       add_d   = nullptr;
    const ggml_tensor * add_src = nullptr;

    if (add_tensor->src[0] == mul_tensor) {
-        add_d   = (float *) add_tensor->src[1]->data;
        add_src = add_tensor->src[1];
    } else if (add_tensor->src[1] == mul_tensor) {
-        add_d   = (float *) add_tensor->src[0]->data;
        add_src = add_tensor->src[0];
    } else {
        GGML_ASSERT(false);
@@ -697,7 +764,7 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx,
    const int add_nchannels = add_src->ne[2];
    const int add_nsamples  = add_src->ne[3];

-    rms_norm_mul_f32_cuda(src0_d, mul_d,add_d,dst_d,
+    rms_norm_mul_f32_cuda(rms_norm_src, mul_src,add_src,dst_d,
                          ne00,ne01, ne02, ne03,
                          /*s00*/ s01, s02, s03,
                          /*mul_s00*/ mul_s01, mul_s02, mul_s03,
@@ -50,6 +50,8 @@ endif()
 set(TARGET_NAME ggml-hexagon)
 ggml_add_backend_library(${TARGET_NAME}
    ggml-hexagon.cpp
+    htp-repack.cpp
+    htp-repack.h
    htp-drv.cpp
    htp-drv.h
    libdl.h
@@ -326,10 +326,12 @@ struct htp_opformat {
            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_F16_BATCHED) {
-                path = "hmx-tiled";
+            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_QUANT_ROW    || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
+                       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);
@@ -0,0 +1,64 @@
+#ifndef HTP_REPACK_H
+#define HTP_REPACK_H
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include "ggml.h"
+#include "htp/htp-tensor.h"
+
+static inline bool htp_is_tiled_weight_type(enum ggml_type type) {
+    return type == GGML_TYPE_Q4_0  || type == GGML_TYPE_Q4_1   ||
+           type == GGML_TYPE_Q8_0  || type == GGML_TYPE_IQ4_NL ||
+           type == GGML_TYPE_MXFP4 || type == GGML_TYPE_Q6_K   ||
+           type == GGML_TYPE_Q4_K  || type == GGML_TYPE_Q5_K   ||
+           type == GGML_TYPE_Q3_K  || type == GGML_TYPE_Q2_K   ||
+           type == GGML_TYPE_F16   || type == GGML_TYPE_F32    ||
+           type == GGML_TYPE_BF16;
+}
+
+static inline bool htp_is_tiled_weight_tensor(const struct ggml_tensor * t) {
+    if (!ggml_is_contiguous(t) || ggml_is_view(t)) {
+        return false;
+    }
+    if (!htp_is_tiled_weight_type(t->type)) {
+        return false;
+    }
+    if (t->ne[0] < 32) {
+        return false;
+    }
+    if (!ggml_is_quantized(t->type)) {
+        return (t->ne[0] % 32 == 0) && (t->ne[1] % 32 == 0) && (t->ne[2] == 1) && (t->ne[3] == 1);
+    }
+    return t->ne[3] == 1;
+}
+
+// Size of one tiled row in the DSP tiled layout. The Q3_K..Q6_K tiles store uncompressed scales/mins,
+// so they are larger than the ggml blocks. For the other tiled types (Q2_K included) the tile has the same size.
+static inline size_t htp_tiled_row_size(enum ggml_type type, int64_t ne0) {
+    if (type == GGML_TYPE_F16 || type == GGML_TYPE_F32 || type == GGML_TYPE_BF16) {
+        return (size_t) ne0 * sizeof(ggml_fp16_t);
+    }
+    if (type == GGML_TYPE_Q6_K) {
+        return (size_t) (ne0 / 32) * (HTP_TENSOR_TILE_SIZE_Q6_K / 32);
+    }
+    if (type == GGML_TYPE_Q4_K) {
+        return (size_t) (ne0 / 32) * (HTP_TENSOR_TILE_SIZE_Q4_K / 32);
+    }
+    if (type == GGML_TYPE_Q5_K) {
+        return (size_t) (ne0 / 32) * (HTP_TENSOR_TILE_SIZE_Q5_K / 32);
+    }
+    if (type == GGML_TYPE_Q3_K) {
+        return (size_t) (ne0 / 32) * (HTP_TENSOR_TILE_SIZE_Q3_K / 32);
+    }
+    return ggml_row_size(type, ne0);
+}
+
+static inline bool htp_is_hmx_weight_type(enum ggml_type type) {
+    return htp_is_tiled_weight_type(type);
+}
+
+void htp_repack_tensor_tiled(struct ggml_tensor * t, const void * data, size_t size);
+void htp_unpack_tensor_tiled(void * data, const struct ggml_tensor * t, size_t offset, size_t size);
+
+#endif // HTP_REPACK_H
@@ -20,7 +20,7 @@ add_library(${HTP_LIB} SHARED
    work-queue.c
    dma-queue.c
    hmx-queue.c
-    htp-tensor.c
+    htp-dcache.c
    matmul-ops.c
    flash-attn-ops.c
    gated-delta-net-ops.c
@@ -361,7 +361,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir);
                for (uint32_t r = 0; r < current_block_size; ++r) {
                    __fp16 * row_k = (__fp16 *)(k_base + r * factx->size_k_row_padded);
-                    hvx_dequantize_row_q8_0_f16(row_k, row_k, DK);
+                    hvx_dequantize_row_flat_q8_0_f16(row_k, row_k, DK);
                }
                htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir);
            }
@@ -369,7 +369,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, ir);
                for (uint32_t r = 0; r < current_block_size; ++r) {
                    __fp16 * row_v = (__fp16 *)(v_base + r * factx->size_v_row_padded);
-                    hvx_dequantize_row_q8_0_f16(row_v, row_v, DV);
+                    hvx_dequantize_row_flat_q8_0_f16(row_v, row_v, DV);
                }
                htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, ir);
            }
@@ -625,7 +625,7 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data)
    if (factx->k->type == HTP_TYPE_Q8_0) {
        for (uint32_t r = start; r < end; ++r) {
            __fp16 * row_k = (__fp16 *)((char *)args->curr_k + r * args->src_stride * sizeof(__fp16));
-            hvx_dequantize_row_q8_0_f16(row_k, row_k, factx->DK);
+            hvx_dequantize_row_flat_q8_0_f16(row_k, row_k, factx->DK);
        }
    }
    hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles[args->buf_idx], (const __fp16 *) args->curr_k, total_rows, factx->DK_pad,
@@ -679,7 +679,7 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data)
    if (factx->v->type == HTP_TYPE_Q8_0) {
        for (uint32_t r = start; r < end; ++r) {
            __fp16 * row_v = (__fp16 *)((char *)args->v_src + r * args->src_stride * sizeof(__fp16));
-            hvx_dequantize_row_q8_0_f16(row_v, row_v, factx->DV);
+            hvx_dequantize_row_flat_q8_0_f16(row_v, row_v, factx->DV);
        }
    }
    hmx_interleave_cols_to_tiles(v_tiles_dst, (const __fp16 *) args->v_src, total_rows, factx->DV_pad,
@@ -755,40 +755,11 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
        const size_t m_start       = i * m_bytes_per_t;
        const size_t m_end         = hex_smin(m_start + m_bytes_per_t, col_vec_bytes);

-        if (factx->sinks) {
-            const float * sinks_data = factx->vtcm_sinks;
-            float *       m_vec      = (float *) factx->vtcm_m_vec;
-            const size_t  r_start    = l_start / sizeof(float);
-            const size_t  r_end      = l_end / sizeof(float);
-            const float   scale_factor = EXP_LOG2E_F;
-
-            const HVX_Vector v_scale = hvx_vec_splat_f32(scale_factor);
-
-            for (size_t r = r_start; r < r_end; r += 32) {
-                HVX_VectorAlias local_m;
-                for (size_t j = 0; j < 32; ++j) {
-                    size_t curr_r = r + j;
-                    if (curr_r < n_rows_g) {
-                        const size_t h_idx = fastmodulo(curr_r, G, &factx->div_G);
-                        const size_t head  = args->kv_head * G + h_idx;
-                        local_m.fp32[j] = sinks_data[head];
-                    } else {
-                        local_m.fp32[j] = HTP_FA_M_INITIAL_VAL;
-                    }
-                }
-                HVX_Vector v_scaled = HVX_OP_MUL_F32(local_m.v, v_scale);
-                *(HVX_Vector *) (m_vec + r) = v_scaled;
-            }
-            if (l_start < col_vec_bytes) {
-                hvx_splat_u8_a((char *) factx->vtcm_l_vec + l_start, 0, l_end - l_start);
-            }
-        } else {
-            if (l_start < col_vec_bytes) {
-                hvx_splat_u8_a((char *) factx->vtcm_l_vec + l_start, 0, l_end - l_start);
-            }
-            if (m_start < col_vec_bytes) {
-                hvx_splat_f32_a((char *) factx->vtcm_m_vec + m_start, HTP_FA_M_INITIAL_VAL, (m_end - m_start) / sizeof(float));
-            }
+        if (l_start < col_vec_bytes) {
+            hvx_splat_u8_a((char *) factx->vtcm_l_vec + l_start, 0, l_end - l_start);
+        }
+        if (m_start < col_vec_bytes) {
+            hvx_splat_f32_a((char *) factx->vtcm_m_vec + m_start, HTP_FA_M_INITIAL_VAL, (m_end - m_start) / sizeof(float));
        }

        // Zero the whole rescale region: vtcm_d_tiles[0], the optional vtcm_d_tiles[1]
@@ -1447,21 +1418,10 @@ static inline void fa_softmax_impl(
        HVX_Vector     v_rowsum_acc_f32_0 = Q6_V_lo_W(rowsum_acc_pair);
        HVX_Vector     v_rowsum_acc_f32_1 = Q6_V_hi_W(rowsum_acc_pair);

-        HVX_Vector v_l_curr0;
-        HVX_Vector v_l_curr1;
-        if (args->kv_start == 0 && factx->sinks != NULL) {
-            // First KV block with sinks: m_prev holds the seeded sink value (not -inf),
-            // so exp_m_diff = exp2(sink - m_curr) is the sink's contribution to the
-            // denominator. l_prev is 0 here, so add exp_m_diff directly instead of
-            // multiplying the (uninitialized) l_prev term.
-            v_l_curr0 = HVX_OP_ADD_F32(exp_m_diff0, v_rowsum_acc_f32_0);
-            v_l_curr1 = HVX_OP_ADD_F32(exp_m_diff1, v_rowsum_acc_f32_1);
-        } else {
-            HVX_Vector l_prev_v0 = factx->vtcm_l_vec[r_vec_idx * 2 + 0];
-            HVX_Vector l_prev_v1 = factx->vtcm_l_vec[r_vec_idx * 2 + 1];
-            v_l_curr0 = HVX_OP_ADD_F32(HVX_OP_MUL_F32(l_prev_v0, exp_m_diff0), v_rowsum_acc_f32_0);
-            v_l_curr1 = HVX_OP_ADD_F32(HVX_OP_MUL_F32(l_prev_v1, exp_m_diff1), v_rowsum_acc_f32_1);
-        }
+        HVX_Vector l_prev_v0 = factx->vtcm_l_vec[r_vec_idx * 2 + 0];
+        HVX_Vector l_prev_v1 = factx->vtcm_l_vec[r_vec_idx * 2 + 1];
+        HVX_Vector v_l_curr0 = HVX_OP_ADD_F32(HVX_OP_MUL_F32(l_prev_v0, exp_m_diff0), v_rowsum_acc_f32_0);
+        HVX_Vector v_l_curr1 = HVX_OP_ADD_F32(HVX_OP_MUL_F32(l_prev_v1, exp_m_diff1), v_rowsum_acc_f32_1);

        factx->vtcm_m_vec[r_vec_idx * 2 + 0] = v_m_curr0;
        factx->vtcm_m_vec[r_vec_idx * 2 + 1] = v_m_curr1;
@@ -1532,9 +1492,66 @@ static void fa_softmax_thread(unsigned int n, unsigned int i, void * data) {
    fa_softmax_impl(n, i, data, has_mask, mask_broadcast, is_g1, has_alibi, has_softcap);
 }

+static inline HVX_Vector hvx_fa_load_sinks(const float * sinks_data, size_t idx, size_t count) {
+    const HVX_Vector * v_src    = (const HVX_Vector *) sinks_data + (idx >> 5);
+    const size_t       lane_off = idx & 31;
+    const size_t       byte_off = lane_off * sizeof(float);
+    return (lane_off + count > 32)
+        ? Q6_V_valign_VVR(v_src[1], v_src[0], byte_off)
+        : Q6_V_vror_VR(v_src[0], byte_off);
+}
+
 static __attribute__((noinline)) void fa_build_d_diag_inv_l(struct hmx_fa_context * factx,
                                                            size_t                  n_row_tiles,
-                                                            size_t                  n_row_tiles_g_br) {
+                                                            size_t                  n_row_tiles_g_br,
+                                                            uint32_t                kv_head,
+                                                            size_t                  n_rows_g) {
+    (void) n_row_tiles_g_br;
+
+    // sinks: l += exp(sink - m * ln2) in f32, m and l are in the log2 domain
+    if (factx->sinks) {
+        const float *    sinks_data   = factx->vtcm_sinks;
+        const size_t     G            = factx->G;
+        const size_t     head_base    = (size_t) kv_head * G;
+        const HVX_Vector v_ln2        = Q6_V_vsplat_R(EXP_LOGN2);
+        const HVX_Vector v_head_sinks = (G <= 32) ? hvx_fa_load_sinks(sinks_data, head_base, G) : Q6_V_vzero();
+        HVX_Vector       v_sink       = Q6_V_vzero();
+        size_t           g_prev       = SIZE_MAX;
+
+        for (size_t i = 0; i < n_row_tiles; ++i) {
+            // the sink pattern of a tile only depends on (i * 32) % G, rebuild it when that changes
+            const size_t r0 = i * 32;
+            const size_t g0 = fastmodulo(r0, G, &factx->div_G);
+            if (g0 != g_prev) {
+                const size_t rem_g = G - g0;
+                if (G <= 32) {
+                    const HVX_Vector v_lo = Q6_V_vror_VR(v_head_sinks, g0 * sizeof(float));
+                    const HVX_Vector v_hi = Q6_V_vror_VR(v_head_sinks, 128 - rem_g * sizeof(float));
+                    v_sink = Q6_V_vmux_QVV(Q6_Q_vsetq2_R(rem_g * sizeof(float)), v_lo, v_hi);
+                    for (size_t len = G; len < 32; len <<= 1) {
+                        const HVX_VectorPred q_keep    = Q6_Q_vsetq_R(len * sizeof(float));
+                        const HVX_Vector     v_shifted = Q6_V_vror_VR(v_sink, 128 - len * sizeof(float));
+                        v_sink = Q6_V_vmux_QVV(q_keep, v_sink, v_shifted);
+                    }
+                } else if (rem_g >= 32) {
+                    v_sink = hvx_fa_load_sinks(sinks_data, head_base + g0, 32);
+                } else {
+                    const HVX_Vector v_lo = hvx_fa_load_sinks(sinks_data, head_base + g0, rem_g);
+                    const HVX_Vector v_hi = hvx_fa_load_sinks(sinks_data, head_base, 32 - rem_g);
+                    v_sink = Q6_V_vmux_QVV(Q6_Q_vsetq_R(rem_g * sizeof(float)), v_lo, Q6_V_vror_VR(v_hi, 128 - rem_g * sizeof(float)));
+                }
+                g_prev = g0;
+            }
+
+            HVX_Vector v_d = HVX_OP_SUB_F32(v_sink, HVX_OP_MUL_F32(factx->vtcm_m_vec[i], v_ln2));
+            if (r0 + 32 > n_rows_g) {
+                const HVX_VectorPred q_valid = Q6_Q_vsetq_R((n_rows_g - r0) * sizeof(float));
+                v_d = Q6_V_vmux_QVV(q_valid, v_d, Q6_V_vzero());
+            }
+            factx->vtcm_l_vec[i] = HVX_OP_ADD_F32(factx->vtcm_l_vec[i], hvx_vec_exp_f32(v_d));
+        }
+    }
+
    const HVX_Vector     v_offsets = *(const HVX_Vector *) d_tile_scatter_offsets;
    const HVX_VectorPred q_32_mask = Q6_Q_vsetq_R(32 * sizeof(__fp16));
    const HVX_Vector     one       = hvx_vec_splat_f32(1.0f);
@@ -2010,7 +2027,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {

    // ======== Initialize HMX output scales ========
    hmx_init_column_scales(factx.vtcm_hmx_scales_id, Q6_V_vsplat_R(0x3c00)); // 1.0
-    hmx_init_column_scales(factx.vtcm_hmx_scales_qk, hvx_vec_splat_f16(factx.scale));
+    {
+        union { __fp16 f; uint16_t u; } sc = { .f = factx.scale };
+        hmx_init_column_scales(factx.vtcm_hmx_scales_qk, Q6_V_vsplat_R(sc.u));
+    }

    // ======== DMA setup ========
    dma_queue * const dma_q = ctx->dma[0];
@@ -2250,7 +2270,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {

                        // Overlapped: run HVX build diag inv L while HMX is busy executing the update
                        htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
-                        fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
+                        fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br, kv_head, n_rows_g);
                        htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
                        hmx_queue_pop(hmx_q);

@@ -2355,7 +2375,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                            if (kv_blk + 1 == factx.n_kv_blocks) {
                                // Overlapped: run HVX build diag inv L while HMX is busy executing the update
                                htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
-                                fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
+                                fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br, kv_head, n_rows_g);
                                htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
                            }
                            hmx_queue_pop(ctx->hmx_queue);
@@ -17,21 +17,32 @@
 #include "htp-tensor.h"
 #include "hvx-utils.h"
 #include "hvx-quant.h"
-#include "matmul-ops.h"
 #include "get-rows-ops.h"
 #include "work-queue.h"

+struct get_rows_context;
+
+typedef uint32_t (*get_rows_read_index_fn_t)(const void * src1, uint32_t idx, int idx_type);
+typedef void (*get_rows_flat_compute_fn_t)(float * dst_spad, const void * src_spad, uint32_t cur_elems);
+typedef void (*get_rows_tiled_compute_fn_t)(float * dst_spad, const uint8_t * src_spad, uint32_t n_k_tiles, uint32_t tile_stride, uint32_t row);
+
 struct get_rows_context {
    struct htp_ops_context * octx;
    const struct htp_get_rows_kernel_params * kparams;
    struct htp_get_rows_vtcm_layout vtcm_layout;
-    uint8_t * vtcm_base;
+    uint8_t    * vtcm_base;
+    const void * src1_data;
+    get_rows_read_index_fn_t    read_index;
+    get_rows_flat_compute_fn_t  compute_flat;
+    get_rows_tiled_compute_fn_t compute_tiled;
    uint32_t task_start;
    uint32_t tasks;
    uint32_t tasks_per_thread;
    uint32_t tile_size;
    uint32_t tile_stride;
-    bool index_i32;
+    uint32_t src0_row_size;
+    uint32_t dst_row_size;
+    int idx_type;
 };

 #define get_rows_preamble                      \
@@ -64,229 +75,169 @@ struct get_rows_context {
                                               \
    const uint32_t nr = ne10 * ne11 * ne12;

-#define GET_ROWS_THREAD_ST_FN(IDX_TYPE)                                                                                 \
-static void get_rows_thread_st_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) {                             \
-    struct get_rows_context * grctx = (struct get_rows_context *)data;                                                  \
-    struct htp_ops_context * octx = grctx->octx;                                                                        \
-    const struct htp_get_rows_kernel_params * kparams = grctx->kparams;                                                 \
-    get_rows_preamble;                                                                                                  \
-    const uint32_t dr  = grctx->tasks_per_thread;                                                                       \
-    const uint32_t ir0 = grctx->task_start + dr * ith;                                                                  \
-    if (ir0 >= grctx->task_start + grctx->tasks) {                                                                      \
-        return;                                                                                                         \
-    }                                                                                                                   \
-    const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks);                                               \
-    const uint32_t row_size_bytes = htp_tensor_get_row_size(octx->src[0]->type, ne00);                                  \
-    dma_queue * dma_q = octx->ctx->dma[ith];                                                                            \
-    for (uint32_t i = ir0; i < ir1; ++i) {                                                                              \
-        const uint32_t i12 = fastdiv(i, &kparams->div_ne10_ne11);                                                       \
-        const uint32_t rem = i - i12 * ne11 * ne10;                                                                     \
-        const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);                                                          \
-        const uint32_t i10 = rem - i11 * ne10;                                                                          \
-        const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(uintptr_t)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12); \
-        const uint32_t i01 = (uint32_t)*src1_ptr;                                                                       \
-        assert(i01 < ne01);                                                                                             \
-        const uint32_t q02 = fastdiv(i11, &kparams->div_ne02);                                                          \
-        const uint32_t i02 = i11 - q02 * ne02;                                                                          \
-        const uint32_t q03 = fastdiv(i12, &kparams->div_ne03);                                                          \
-        const uint32_t i03 = i12 - q03 * ne03;                                                                          \
-        const dma_addr_t src0_data = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03;                               \
-        const dma_addr_t dst_data  = octx->dst->data    + i10*nb1  + i11*nb2  + i12*nb3;                                \
-        while (!dma_queue_push(dma_q, dma_make_data(dst_data, src0_data), nb1, nb01,                                    \
-                               row_size_bytes, 1)) {                                                                    \
-            dma_queue_pop(dma_q);                                                                                       \
-        }                                                                                                               \
-    }                                                                                                                   \
-    dma_queue_flush(dma_q);                                                                                             \
+static uint32_t get_rows_read_index_ddr(
+    const void * src1,
+    uint32_t idx,
+    int idx_type
+) {
+    const uint32_t off = idx * ((idx_type == HTP_TYPE_I64) ? sizeof(int64_t) : sizeof(int32_t));
+    return *(const uint32_t *) ((const uint8_t *) src1 + off);
 }

-GET_ROWS_THREAD_ST_FN(int32_t)
-GET_ROWS_THREAD_ST_FN(int64_t)
-
-#define GET_ROWS_THREAD_DT_FN(TYPE_NAME, SRC0_SIZE_EXPR, IDX_TYPE, COMPUTE_EXPR)                                                \
-static void get_rows_thread_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) {                          \
-    struct get_rows_context * grctx = (struct get_rows_context *)data;                                                          \
-    struct htp_ops_context * octx = grctx->octx;                                                                                \
-    const struct htp_get_rows_kernel_params * kparams = grctx->kparams;                                                         \
-    get_rows_preamble;                                                                                                          \
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                                      \
-    const uint32_t dr  = grctx->tasks_per_thread;                                                                               \
-    const uint32_t ir0 = grctx->task_start + dr * ith;                                                                          \
-    if (ir0 >= grctx->task_start + grctx->tasks) {                                                                              \
-        return;                                                                                                                 \
-    }                                                                                                                           \
-    const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks);                                                       \
-    const uint32_t chunks_per_row = kparams->chunks_per_row;                                                                    \
-    const uint32_t chunk_size     = kparams->chunk_size;                                                                        \
-    dma_queue * dma_q = octx->ctx->dma[ith];                                                                                    \
-    const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout;                                                  \
-    uint8_t * vtcm_src0 = grctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread;                  \
-    uint8_t * vtcm_dst  = grctx->vtcm_base + vtcm_layout->off_dst  + ith * vtcm_layout->dst_bytes_per_thread;                   \
-    for (uint32_t step = 0, spad_idx = 0; step < ir1 - ir0 && spad_idx < 2; ++step, spad_idx++) {                               \
-        const uint32_t i = ir0 + step;                                                                                          \
-        const uint32_t row_idx   = fastdiv(i, &kparams->div_chunks_per_row);                                                    \
-        const uint32_t chunk_idx = i - row_idx * chunks_per_row;                                                                \
-        const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11);                                                         \
-        const uint32_t rem = row_idx - i12 * ne11 * ne10;                                                                       \
-        const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);                                                                  \
-        const uint32_t i10 = rem - i11 * ne10;                                                                                  \
-        const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(uintptr_t)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12);         \
-        const uint32_t i01 = (uint32_t)*src1_ptr;                                                                               \
-        assert(i01 < ne01);                                                                                                     \
-        const uint32_t q02 = fastdiv(i11, &kparams->div_ne02);                                                                  \
-        const uint32_t i02 = i11 - q02 * ne02;                                                                                  \
-        const uint32_t q03 = fastdiv(i12, &kparams->div_ne03);                                                                  \
-        const uint32_t i03 = i12 - q03 * ne03;                                                                                  \
-        const uint32_t offset = chunk_idx * chunk_size;                                                                         \
-        const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0;                                        \
-        const uint32_t cur_src0_bytes = SRC0_SIZE_EXPR(cur_elems);                                                              \
-        const uint32_t cur_dst_bytes  = cur_elems * sizeof(float);                                                              \
-        const dma_addr_t src0_data = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03 + SRC0_SIZE_EXPR(offset);              \
-        dma_queue_push(dma_q,                                                                                                   \
-                       dma_make_data(octx->dst->data,                                                                           \
-                                     vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size),                                    \
-                       cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 0);                                       \
-        dma_queue_push(dma_q,                                                                                                   \
-                       dma_make_data(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size, src0_data),                       \
-                       vtcm_layout->src0_spad_half_size, cur_src0_bytes, cur_src0_bytes, 1);                                    \
-    }                                                                                                                           \
-    for (uint32_t step = 0; step < ir1 - ir0; ++step) {                                                                         \
-        const uint32_t i = ir0 + step;                                                                                          \
-        void * dst_spad = (void *) dma_queue_pop(dma_q).src;                                                                    \
-        void * src_spad = (void *) dma_queue_pop(dma_q).dst;                                                                    \
-        const uint32_t row_idx   = fastdiv(i, &kparams->div_chunks_per_row);                                                    \
-        const uint32_t chunk_idx = i - row_idx * chunks_per_row;                                                                \
-        const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11);                                                         \
-        const uint32_t rem = row_idx - i12 * ne11 * ne10;                                                                       \
-        const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);                                                                  \
-        const uint32_t i10 = rem - i11 * ne10;                                                                                  \
-        const uint32_t offset = chunk_idx * chunk_size;                                                                         \
-        const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0;                                        \
-        const uint32_t cur_dst_bytes  = cur_elems * sizeof(float);                                                              \
-        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, i);                                                                   \
-        COMPUTE_EXPR;                                                                                                           \
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, i);                                                                    \
-        const dma_addr_t dst_data  = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3 + offset * sizeof(float);                    \
-        dma_queue_push(dma_q,                                                                                                   \
-                       dma_make_data(dst_data, dst_spad),                                                                       \
-                       cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 1);                                       \
-        const uint32_t next_step = step + 2;                                                                                    \
-        if (next_step < ir1 - ir0) {                                                                                            \
-            const uint32_t pi = ir0 + next_step;                                                                                \
-            const uint32_t prow_idx   = fastdiv(pi, &kparams->div_chunks_per_row);                                              \
-            const uint32_t pchunk_idx = pi - prow_idx * chunks_per_row;                                                         \
-            const uint32_t pi12 = fastdiv(prow_idx, &kparams->div_ne10_ne11);                                                   \
-            const uint32_t prem = prow_idx - pi12 * ne11 * ne10;                                                                \
-            const uint32_t pi11 = fastdiv(prem, &kparams->div_ne10);                                                            \
-            const uint32_t pi10 = prem - pi11 * ne10;                                                                           \
-            const IDX_TYPE * psrc1_ptr = (const IDX_TYPE *)(uintptr_t)(octx->src[1]->data + pi10*nb10 + pi11*nb11 + pi12*nb12); \
-            const uint32_t pi01 = (uint32_t)*psrc1_ptr;                                                                         \
-            assert(pi01 < ne01);                                                                                                \
-            const uint32_t pq02 = fastdiv(pi11, &kparams->div_ne02);                                                            \
-            const uint32_t pi02 = pi11 - pq02 * ne02;                                                                           \
-            const uint32_t pq03 = fastdiv(pi12, &kparams->div_ne03);                                                            \
-            const uint32_t pi03 = pi12 - pq03 * ne03;                                                                           \
-            const uint32_t poffset = pchunk_idx * chunk_size;                                                                   \
-            const uint32_t pcur_elems = (poffset < ne00) ? MIN(chunk_size, ne00 - poffset) : 0;                                 \
-            const uint32_t pcur_src0_bytes = SRC0_SIZE_EXPR(pcur_elems);                                                        \
-            const dma_addr_t psrc0_data =                                                                                       \
-                octx->src[0]->data + pi01*nb01 + pi02*nb02 + pi03*nb03 + SRC0_SIZE_EXPR(poffset);                               \
-            dma_queue_push(dma_q,                                                                                               \
-                           dma_make_data(src_spad, psrc0_data),                                                                 \
-                           vtcm_layout->src0_spad_half_size, pcur_src0_bytes, pcur_src0_bytes, 1);                              \
-        }                                                                                                                       \
-    }                                                                                                                           \
-    dma_queue_flush(dma_q);                                                                                                     \
+static __attribute__((noinline)) uint32_t get_rows_read_index_vtcm(
+    const void * src1,
+    uint32_t idx,
+    int idx_type
+) {
+    const uint32_t off = idx * ((idx_type == HTP_TYPE_I64) ? sizeof(int64_t) : sizeof(int32_t));
+    const HVX_Vector v = hvx_vmem((const uint8_t *) src1 + off); // off & ~127u in HW
+    // vextract is normally slow, but here it avoids an HVX->L2 store and L1D miss
+    return (uint32_t) Q6_R_vextract_VR(v, off);
 }

-#define F16_BYTES(n)  ((n) * sizeof(__fp16))
-#define Q8_0_BYTES(n) (((n) / 32) * sizeof(block_q8_0))
+static void get_rows_thread_sametype(unsigned int nth, unsigned int ith, void * data) {
+    struct get_rows_context * grctx = (struct get_rows_context *) data;
+    struct htp_ops_context * octx = grctx->octx;
+    const struct htp_get_rows_kernel_params * kparams = grctx->kparams;
+    get_rows_preamble;

-static __attribute__((noinline)) void compute_get_rows_f16(float * dst_spad, const void * src_spad, uint32_t cur_elems) {
-    hvx_dequantize_row_f16_f32(dst_spad, src_spad, cur_elems);
-}
+    const uint32_t dr  = grctx->tasks_per_thread;
+    const uint32_t ir0 = grctx->task_start + dr * ith;
+    if (ir0 >= grctx->task_start + grctx->tasks) {
+        return;
+    }
+    const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks);
+    const uint32_t row_size_bytes = grctx->src0_row_size;
+    dma_queue * dma_q = octx->ctx->dma[ith];

-static __attribute__((noinline)) void compute_get_rows_q8_0(float * dst_spad, const void * src_spad, uint32_t cur_elems) {
-    hvx_dequantize_row_q8_0_f32(dst_spad, src_spad, cur_elems);
-}
+    for (uint32_t i = ir0; i < ir1; ++i) {
+        const uint32_t i12 = fastdiv(i, &kparams->div_ne10_ne11);
+        const uint32_t rem = i - i12 * ne11 * ne10;
+        const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);
+        const uint32_t i10 = rem - i11 * ne10;

-GET_ROWS_THREAD_DT_FN(f16,  F16_BYTES,  int32_t, { compute_get_rows_f16((float *)dst_spad, src_spad, cur_elems); })
-GET_ROWS_THREAD_DT_FN(f16,  F16_BYTES,  int64_t, { compute_get_rows_f16((float *)dst_spad, src_spad, cur_elems); })
+        const uint32_t i01 = grctx->read_index(grctx->src1_data, i, grctx->idx_type);
+        assert(i01 < ne01);

-GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int32_t, { compute_get_rows_q8_0((float *)dst_spad, src_spad, cur_elems); })
-GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int64_t, { compute_get_rows_q8_0((float *)dst_spad, src_spad, cur_elems); })
+        const uint32_t q02 = fastdiv(i11, &kparams->div_ne02);
+        const uint32_t i02 = i11 - q02 * ne02;
+        const uint32_t q03 = fastdiv(i12, &kparams->div_ne03);
+        const uint32_t i03 = i12 - q03 * ne03;

+        const dma_addr_t src0_data = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03;
+        const dma_addr_t dst_data  = octx->dst->data    + i10*nb1  + i11*nb2  + i12*nb3;

-static __attribute__((noinline)) void compute_get_rows_tiled(float * dst, const uint8_t * tile, uint32_t row, bool q4, bool q4_k) {
-    const HVX_VectorPred first2 = Q6_Q_vsetq_R(2);
-    const HVX_VectorPred first4 = Q6_Q_vsetq_R(4);
-    HVX_Vector vq = Q6_V_vzero();
-    if (q4) {
-        const HVX_VectorPred first1 = Q6_Q_vsetq_R(1);
-        const HVX_VectorPred first3 = Q6_Q_vsetq_R(3);
-        for (int group = 3; group >= 0; --group) {
-            const HVX_Vector v = Q6_V_vror_VR(hvx_vmem(tile + group * VLEN), row);
-            // Four planes contribute bytes at 0, 32, 64 and 96 after rotation.
-            HVX_Vector packed = Q6_V_vmux_QVV(first1, v, Q6_V_vror_VR(v, 31));
-            packed = Q6_V_vmux_QVV(first2, packed, Q6_V_vror_VR(v, 62));
-            packed = Q6_V_vmux_QVV(first3, packed, Q6_V_vror_VR(v, 93));
-            vq = Q6_V_vmux_QVV(first4, packed, Q6_V_vror_VR(vq, VLEN - 4));
-        }
-        const HVX_Vector lo = Q6_V_vand_VV(vq, Q6_Vb_vsplat_R(0x0F));
-        const HVX_Vector hi = Q6_Vub_vlsr_VubR(vq, 4);
-        vq = Q6_V_lo_W(Q6_W_vshuff_VVR(hi, lo, -1));
-        if (!q4_k) {
-            vq = Q6_Vb_vsub_VbVb(vq, Q6_Vb_vsplat_R(8));
-        }
-    } else {
-        for (int group = 7; group >= 0; --group) {
-            const HVX_Vector v = Q6_V_vror_VR(hvx_vmem(tile + group * VLEN), 2 * row);
-            // Two planes contribute halfwords at 0 and 64 after rotation.
-            const HVX_Vector packed = Q6_V_vmux_QVV(first2, v, Q6_V_vror_VR(v, 62));
-            vq = Q6_V_vmux_QVV(first4, packed, Q6_V_vror_VR(vq, VLEN - 4));
+        while (!dma_queue_push(dma_q, dma_make_data(dst_data, src0_data), nb1, nb01, row_size_bytes, 1)) {
+            dma_queue_pop(dma_q);
        }
    }
-    const HVX_Vector scales = hvx_vmem(tile + (q4 ? 512 : 1024));
-    const HVX_Vector scale_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, (q4_k ? 4 : 2) * row));
-    const HVX_Vector scale = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hf));
-    const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq);
-    const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16));
-    HVX_Vector values = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale);
-    if (q4_k) {
-        const HVX_Vector offset_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 4 * row + 2));
-        const HVX_Vector offset = Q6_V_lo_W(hvx_vec_f16_to_f32(offset_hf));
-        values = hvx_vec_add_f32_f32(values, offset);
-    }
-    *(HVX_Vector *) dst = values;
+    dma_queue_flush(dma_q);
 }

-static __attribute__((noinline)) void compute_get_rows_q6_k(float * dst, const uint8_t * tile, uint32_t row) {
-    const HVX_VectorPred first4 = Q6_Q_vsetq_R(4);
-    const HVX_VectorPred first16 = Q6_Q_vsetq_R(16 * sizeof(float));
-    const HVX_Vector mask_0f = Q6_Vb_vsplat_R(0x0F);
-    const HVX_Vector mask_03 = Q6_Vb_vsplat_R(0x03);
-    HVX_Vector vq = Q6_V_vzero();

-    for (int group = 7; group >= 0; --group) {
-        const HVX_Vector lo_plane = Q6_V_vror_VR(hvx_vmem(tile + (group >> 1) * VLEN), 4 * row);
-        const HVX_Vector hi_plane = Q6_V_vror_VR(hvx_vmem(tile + 512 + (group >> 2) * VLEN), 4 * row);
-        const HVX_Vector lo = (group & 1) ? Q6_Vub_vlsr_VubR(lo_plane, 4) : Q6_V_vand_VV(lo_plane, mask_0f);
-        const HVX_Vector hi = Q6_Vub_vlsr_VubR(hi_plane, 2 * (group & 3));
-        const HVX_Vector packed = Q6_V_vor_VV(lo, Q6_Vw_vasl_VwR(Q6_V_vand_VV(hi, mask_03), 4));
-        vq = Q6_V_vmux_QVV(first4, packed, Q6_V_vror_VR(vq, VLEN - 4));
+
+struct get_rows_flat_task {
+    dma_addr_t src0_data;
+    dma_addr_t dst_data;
+};
+
+static inline struct get_rows_flat_task get_rows_flat_calc_task(
+    const struct htp_ops_context * octx,
+    const struct get_rows_context * grctx,
+    uint32_t i,
+    bool need_src0
+) {
+    const struct htp_get_rows_kernel_params * kparams = grctx->kparams;
+    get_rows_preamble;
+
+    const uint32_t i12 = fastdiv(i, &kparams->div_ne10_ne11);
+    const uint32_t rem = i - i12 * ne11 * ne10;
+    const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);
+    const uint32_t i10 = rem - i11 * ne10;
+
+    struct get_rows_flat_task task;
+    task.src0_data = 0;
+    task.dst_data  = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3;
+
+    if (need_src0) {
+        const uint32_t i01 = grctx->read_index(grctx->src1_data, i, grctx->idx_type);
+        assert(i01 < ne01);
+
+        const uint32_t q02 = fastdiv(i11, &kparams->div_ne02);
+        const uint32_t i02 = i11 - q02 * ne02;
+        const uint32_t q03 = fastdiv(i12, &kparams->div_ne03);
+        const uint32_t i03 = i12 - q03 * ne03;
+
+        task.src0_data = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03;
    }

-    const HVX_Vector scales = hvx_vmem(tile + 768);
-    const HVX_Vector scale_lo_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 2 * row));
-    const HVX_Vector scale_hi_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 64 + 2 * row));
-    const HVX_Vector scale_lo = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_lo_hf));
-    const HVX_Vector scale_hi = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hi_hf));
-    const HVX_Vector scale = Q6_V_vmux_QVV(first16, scale_lo, scale_hi);
-    const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(Q6_Vb_vsub_VbVb(vq, Q6_Vb_vsplat_R(32)));
-    const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16));
-    *(HVX_Vector *) dst = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale);
+    return task;
 }

+static void get_rows_thread_flat(unsigned int nth, unsigned int ith, void * data) {
+    struct get_rows_context * grctx = (struct get_rows_context *) data;
+    struct htp_ops_context * octx = grctx->octx;
+    get_rows_preamble;
+
+    const uint32_t dr  = grctx->tasks_per_thread;
+    const uint32_t ir0 = grctx->task_start + dr * ith;
+    if (ir0 >= grctx->task_start + grctx->tasks) {
+        return;
+    }
+    const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks);
+
+    dma_queue * dma_q = octx->ctx->dma[ith];
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout;
+    uint8_t * vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_src0) + ith * vtcm_layout->src0_bytes_per_thread;
+    uint8_t * vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_dst)  + ith * vtcm_layout->dst_bytes_per_thread;
+    const uint32_t src0_bytes = grctx->src0_row_size;
+    const uint32_t dst_bytes  = grctx->dst_row_size;
+    get_rows_flat_compute_fn_t compute = grctx->compute_flat;
+
+    for (uint32_t step = 0, spad_idx = 0; step < ir1 - ir0 && spad_idx < 2; ++step, spad_idx++) {
+        const uint32_t i = ir0 + step;
+        struct get_rows_flat_task task = get_rows_flat_calc_task(octx, grctx, i, true);
+
+        dma_queue_push(dma_q,
+                       dma_make_data(octx->dst->data,
+                                     vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size),
+                       dst_bytes, vtcm_layout->dst_spad_half_size, dst_bytes, 0);
+        dma_queue_push(dma_q,
+                       dma_make_data(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size, task.src0_data),
+                       vtcm_layout->src0_spad_half_size, src0_bytes, src0_bytes, 1);
+    }
+
+    for (uint32_t step = 0; step < ir1 - ir0; ++step) {
+        const uint32_t i = ir0 + step;
+        void * dst_spad = (void *) dma_queue_pop(dma_q).src;
+        void * src_spad = (void *) dma_queue_pop(dma_q).dst;
+
+        struct get_rows_flat_task task = get_rows_flat_calc_task(octx, grctx, i, false);
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
+        compute((float *) dst_spad, src_spad, ne00);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
+
+        dma_queue_push(dma_q,
+                       dma_make_data(task.dst_data, dst_spad),
+                       dst_bytes, vtcm_layout->dst_spad_half_size, dst_bytes, 1);
+
+        const uint32_t next_step = step + 2;
+        if (next_step < ir1 - ir0) {
+            const uint32_t pi = ir0 + next_step;
+            struct get_rows_flat_task next_task = get_rows_flat_calc_task(octx, grctx, pi, true);
+            dma_queue_push(dma_q,
+                           dma_make_data(src_spad, next_task.src0_data),
+                           vtcm_layout->src0_spad_half_size, src0_bytes, src0_bytes, 1);
+        }
+    }
+
+    dma_queue_flush(dma_q);
+}
+
+
+
 struct get_rows_tiled_task {
    dma_addr_t tile_src_base;
    dma_addr_t dst_data;
@@ -307,16 +258,15 @@ static inline struct get_rows_tiled_task get_rows_tiled_calc_task(
    const uint32_t rem = i - i12 * ne11 * ne10;
    const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);
    const uint32_t i10 = rem - i11 * ne10;
-    const dma_addr_t src1_data = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12;
-    const uint32_t i01 = grctx->index_i32 ? *(const int32_t *)(uintptr_t) src1_data : (uint32_t) *(const int64_t *)(uintptr_t) src1_data;
+    const uint32_t i01 = grctx->read_index(grctx->src1_data, i, grctx->idx_type);
    assert(i01 < ne01);

    const uint32_t q02 = fastdiv(i11, &kparams->div_ne02);
    const uint32_t i02 = i11 - q02 * ne02;
    const uint32_t q03 = fastdiv(i12, &kparams->div_ne03);
    const uint32_t i03 = i12 - q03 * ne03;
-    const uint32_t column_tile = i01 / HTP_MM_HMX_TILE_N_ROWS;
-    const uint32_t row = i01 % HTP_MM_HMX_TILE_N_ROWS;
+    const uint32_t column_tile = i01 / HTP_TENSOR_TILE_N_ROWS;
+    const uint32_t row = i01 % HTP_TENSOR_TILE_N_ROWS;
    const dma_addr_t matrix = octx->src[0]->data + i02*nb02 + i03*nb03;

    struct get_rows_tiled_task task;
@@ -329,7 +279,6 @@ static inline struct get_rows_tiled_task get_rows_tiled_calc_task(
 static void get_rows_thread_tiled(unsigned int nth, unsigned int ith, void * data) {
    struct get_rows_context * grctx = (struct get_rows_context *) data;
    struct htp_ops_context * octx = grctx->octx;
-    const struct htp_get_rows_kernel_params * kparams = grctx->kparams;
    get_rows_preamble;

    const uint32_t dr  = grctx->tasks_per_thread;
@@ -339,19 +288,17 @@ static void get_rows_thread_tiled(unsigned int nth, unsigned int ith, void * dat
    }

    const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks);
-    const uint32_t n_k_tiles = ne00 / HTP_MM_HMX_TILE_N_COLS;
+    const uint32_t n_k_tiles = ne00 / HTP_TENSOR_TILE_N_COLS;
    const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout;
-    uint8_t * src_spad_base = grctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread;
-    uint8_t * dst_spad_base = grctx->vtcm_base + vtcm_layout->off_dst + ith * vtcm_layout->dst_bytes_per_thread;
+    uint8_t * src_spad_base = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_src0) + ith * vtcm_layout->src0_bytes_per_thread;
+    uint8_t * dst_spad_base = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_dst)  + ith * vtcm_layout->dst_bytes_per_thread;
    dma_queue * dma_q = octx->ctx->dma[ith];
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];

    const uint32_t tile_size   = grctx->tile_size;
    const uint32_t tile_stride = grctx->tile_stride;
-    const uint32_t dst_bytes   = ne00 * sizeof(float);
-    const bool is_q4 = octx->src[0]->type == HTP_TYPE_Q4_0 || octx->src[0]->type == HTP_TYPE_Q4_K;
-    const bool is_q4_k = octx->src[0]->type == HTP_TYPE_Q4_K;
-    const bool is_q6_k = octx->src[0]->type == HTP_TYPE_Q6_K;
+    const uint32_t dst_bytes   = grctx->dst_row_size;
+    get_rows_tiled_compute_fn_t compute_tiled = grctx->compute_tiled;

    for (uint32_t step = 0, spad_idx = 0; step < ir1 - ir0 && spad_idx < 2; ++step, ++spad_idx) {
        const uint32_t i = ir0 + step;
@@ -376,15 +323,7 @@ static void get_rows_thread_tiled(unsigned int nth, unsigned int ith, void * dat
        struct get_rows_tiled_task task = get_rows_tiled_calc_task(octx, grctx, i, n_k_tiles, tile_size);

        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
-        for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) {
-            const uint8_t * tile = src_spad + k_tile * tile_stride;
-            float * dst_block = dst_spad + k_tile * HTP_MM_HMX_TILE_N_COLS;
-            if (is_q6_k) {
-                compute_get_rows_q6_k(dst_block, tile, task.row);
-            } else {
-                compute_get_rows_tiled(dst_block, tile, task.row, is_q4, is_q4_k);
-            }
-        }
+        compute_tiled(dst_spad, src_spad, n_k_tiles, tile_stride, task.row);
        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);

        // Real writeback of dst_spad
@@ -405,16 +344,40 @@ static void get_rows_thread_tiled(unsigned int nth, unsigned int ith, void * dat
    dma_queue_flush(dma_q);
 }

+static __attribute__((noinline)) void get_rows_dma_src1(
+    struct htp_ops_context * octx,
+    uint8_t * vtcm_src1,
+    uint32_t elem_size
+) {
+    const uint32_t ne10 = octx->src[1]->ne[0];
+    const uint32_t ne11 = octx->src[1]->ne[1];
+    const uint32_t ne12 = octx->src[1]->ne[2];
+    const uint32_t ne13 = octx->src[1]->ne[3];
+    const uint32_t nb11 = octx->src[1]->nb[1];
+    const uint32_t nb12 = octx->src[1]->nb[2];
+    const uint32_t nb13 = octx->src[1]->nb[3];
+
+    dma_queue * dma_q = octx->ctx->dma[0];
+    uint8_t * cur_dst = vtcm_src1;
+    for (uint32_t i13 = 0; i13 < ne13; ++i13) {
+        for (uint32_t i12 = 0; i12 < ne12; ++i12) {
+            const dma_addr_t src_plane = octx->src[1]->data + i12 * nb12 + i13 * nb13;
+            while (!dma_queue_push(dma_q, dma_make_data(cur_dst, src_plane),
+                                   ne10 * elem_size, nb11, ne10 * elem_size, ne11)) {
+                dma_queue_pop(dma_q);
+            }
+            cur_dst += ne10 * ne11 * elem_size;
+        }
+    }
+    dma_queue_flush(dma_q);
+}
+
 int op_get_rows(struct htp_ops_context * octx) {
    const struct htp_get_rows_kernel_params * kparams = (const struct htp_get_rows_kernel_params *) octx->kernel_params;

-    if (octx->src[0]->type != HTP_TYPE_F32 &&
-        octx->src[0]->type != HTP_TYPE_F16 &&
-        octx->src[0]->type != HTP_TYPE_Q4_0 &&
-        octx->src[0]->type != HTP_TYPE_Q4_K &&
-        octx->src[0]->type != HTP_TYPE_Q6_K &&
-        octx->src[0]->type != HTP_TYPE_Q8_0 &&
-        octx->src[0]->type != HTP_TYPE_I32) {
+    const uint32_t ne00          = octx->src[0]->ne[0];
+    const uint32_t src0_row_size = htp_tensor_get_row_size(octx->src[0]->type, ne00);
+    if (src0_row_size == 0) {
        return HTP_STATUS_NO_SUPPORT;
    }

@@ -432,20 +395,16 @@ int op_get_rows(struct htp_ops_context * octx) {
        return HTP_STATUS_NO_SUPPORT;
    }

-    if (htp_tensor_is_extended(octx->src[1])) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
    const struct htp_tensor * dst = octx->dst;
    const uint32_t total_tasks    = kparams->total_tasks;
-    const size_t dst_row_size     = htp_tensor_get_row_size(dst->type, dst->ne[0]);
+    const uint32_t dst_row_size   = (kparams->kernel_type == HTP_GET_ROWS_KERNEL_SAMETYPE) ? src0_row_size : ne00 * sizeof(float);

    uint32_t task_start = 0;
    uint32_t tasks      = total_tasks;

    if (octx->ctx->mdev.count > 1) {
        uint32_t tasks_per_chunk = 1;
-        htp_tensor_mdev_rows_per_chunk(dst, dst_row_size / dst->ne[0], (uint32_t) dst_row_size, &tasks_per_chunk);
+        htp_tensor_mdev_rows_per_chunk(dst, dst_row_size / dst->ne[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;
@@ -465,15 +424,18 @@ int op_get_rows(struct htp_ops_context * octx) {
    grctx.octx = octx;
    grctx.kparams = kparams;
    grctx.vtcm_base = (uint8_t *)octx->ctx->vtcm_base;
+    grctx.compute_flat = NULL;
+    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.tile_size = htp_mm_get_weight_tile_size(octx->src[0]->type);
-    grctx.tile_stride = (grctx.tile_size + 127) & ~127;
-    grctx.index_i32 = octx->src[1]->type == HTP_TYPE_I32;
+    grctx.tile_size   = 0;
+    grctx.tile_stride = 0;
+    grctx.src0_row_size = src0_row_size;
+    grctx.dst_row_size  = dst_row_size;
+    grctx.idx_type      = octx->src[1]->type;

-    const uint32_t ne00 = octx->src[0]->ne[0];
-    htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, kparams->kernel_type, octx->src[0]->type, ne00, n_threads);
+    htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, kparams->kernel_type, octx->src[0]->type, ne00, total_tasks, grctx.idx_type, n_threads);

    if (grctx.vtcm_layout.total_bytes > octx->ctx->vtcm_size) {
        FARF(ERROR, "get-rows: VTCM reservation %zu is too small, needed %zu\n",
@@ -481,27 +443,56 @@ int op_get_rows(struct htp_ops_context * octx) {
        return HTP_STATUS_INVAL_PARAMS;
    }

-    const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32);
-
    work_queue_func_t q_func = NULL;
    switch (kparams->kernel_type) {
        case HTP_GET_ROWS_KERNEL_SAMETYPE:
-            q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_st_int32_t : get_rows_thread_st_int64_t);
+            q_func = get_rows_thread_sametype;
            break;
        case HTP_GET_ROWS_KERNEL_TILED:
+            switch (octx->src[0]->type) {
+                case HTP_TYPE_F16:  grctx.compute_tiled = hvx_dequantize_row_tiled_f16_f32; break;
+                case HTP_TYPE_Q8_0: grctx.compute_tiled = hvx_dequantize_row_tiled_q8_0_f32; break;
+                case HTP_TYPE_Q4_0: grctx.compute_tiled = hvx_dequantize_row_tiled_q4_0_f32; break;
+                case HTP_TYPE_Q4_K: grctx.compute_tiled = hvx_dequantize_row_tiled_q4_k_f32; break;
+                case HTP_TYPE_Q5_K: grctx.compute_tiled = hvx_dequantize_row_tiled_q5_k_f32; break;
+                case HTP_TYPE_Q6_K: grctx.compute_tiled = hvx_dequantize_row_tiled_q6_k_f32; break;
+                default:            return HTP_STATUS_NO_SUPPORT;
+            }
+            grctx.tile_size   = htp_tensor_get_tile_size(octx->src[0]->type);
+            grctx.tile_stride = htp_tensor_get_tile_padded_size(octx->src[0]->type);
            q_func = get_rows_thread_tiled;
            break;
        case HTP_GET_ROWS_KERNEL_FLAT:
            switch (octx->src[0]->type) {
-                case HTP_TYPE_F16:  q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_f16_int32_t  : get_rows_thread_f16_int64_t);  break;
-                case HTP_TYPE_Q8_0: q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_q8_0_int32_t : get_rows_thread_q8_0_int64_t); break;
+                case HTP_TYPE_F16:  grctx.compute_flat = hvx_dequantize_row_flat_f16_f32;  break;
+                case HTP_TYPE_BF16: grctx.compute_flat = hvx_dequantize_row_flat_bf16_f32; break;
+                case HTP_TYPE_Q8_0: grctx.compute_flat = hvx_dequantize_row_flat_q8_0_f32; break;
+                case HTP_TYPE_Q4_0: grctx.compute_flat = hvx_dequantize_row_flat_q4_0_f32; break;
+                case HTP_TYPE_Q4_K: grctx.compute_flat = hvx_dequantize_row_flat_q4_k_f32; break;
+                case HTP_TYPE_Q5_K: grctx.compute_flat = hvx_dequantize_row_flat_q5_k_f32; break;
+                case HTP_TYPE_Q6_K: grctx.compute_flat = hvx_dequantize_row_flat_q6_k_f32; break;
                default:            return HTP_STATUS_NO_SUPPORT;
            }
+            q_func = get_rows_thread_flat;
            break;
        default:
            return HTP_STATUS_NO_SUPPORT;
    }

+    const uint32_t elem_size = (grctx.idx_type == HTP_TYPE_I64) ? sizeof(int64_t) : sizeof(int32_t);
+    const size_t src1_bytes  = (size_t) total_tasks * elem_size;
+
+    if (!htp_tensor_is_extended(octx->src[1]) && src1_bytes <= 128 &&
+        htp_tensor_is_contiguous(octx->src[1], elem_size)) {
+        grctx.src1_data  = (const void *) (uintptr_t) octx->src[1]->data;
+        grctx.read_index = get_rows_read_index_ddr;
+    } else {
+        uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, grctx.vtcm_base, grctx.vtcm_layout.off_src1);
+        get_rows_dma_src1(octx, vtcm_src1, elem_size);
+        grctx.src1_data  = vtcm_src1;
+        grctx.read_index = get_rows_read_index_vtcm;
+    }
+
    FARF(HIGH, "get-rows: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-vtcm-size %zu dst-vtcm-size %zu kernel-type %d n-threads %d\n",
         octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3],
         octx->src[1]->ne[0], octx->src[1]->ne[1], octx->src[1]->ne[2], octx->src[1]->ne[3],
@@ -5,7 +5,8 @@
 #include <string.h>

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

 enum htp_get_rows_kernel_type {
    HTP_GET_ROWS_KERNEL_SAMETYPE = 0,
@@ -16,8 +17,6 @@ enum htp_get_rows_kernel_type {
 struct htp_get_rows_kernel_params {
    int32_t  n_threads;
    int32_t  kernel_type;
-    int32_t  chunks_per_row;
-    int32_t  chunk_size;
    int32_t  total_tasks;
    int32_t  tasks_per_thread;
    int32_t  vtcm_size;
@@ -25,13 +24,14 @@ struct htp_get_rows_kernel_params {
    // Fastdiv helpers
    struct fastdiv_values div_ne10;
    struct fastdiv_values div_ne10_ne11;
-    struct fastdiv_values div_chunks_per_row;
    struct fastdiv_values div_ne02;
    struct fastdiv_values div_ne03;
 };

 struct htp_get_rows_vtcm_layout {
    size_t total_bytes;
+    size_t off_src1;
+    size_t src1_bytes;
    size_t off_src0;
    size_t off_dst;

@@ -43,57 +43,44 @@ struct htp_get_rows_vtcm_layout {
 };

 static inline void htp_get_rows_vtcm_layout_build(
-    struct htp_get_rows_vtcm_layout * vtcm_layout,
+    struct htp_get_rows_vtcm_layout * L,
    int kernel_type,
    int type,
    uint32_t ne00,
+    uint32_t total_tasks,
+    int idx_type,
    uint32_t n_threads) {

-    if (kernel_type == HTP_GET_ROWS_KERNEL_SAMETYPE) {
-        memset(vtcm_layout, 0, sizeof(*vtcm_layout));
-        return;
+    memset(L, 0, sizeof(*L));
+
+    const size_t elem_size = (idx_type == HTP_TYPE_I64) ? sizeof(int64_t) : sizeof(int32_t);
+    const size_t src1_raw_bytes = total_tasks * elem_size;
+    const size_t src1_bytes = ((src1_raw_bytes + 127) & ~127) + 128;
+
+    L->src1_bytes = src1_bytes;
+    const size_t src1_alloc_bytes = (src1_bytes + 255) & ~255;
+
+    if (kernel_type != HTP_GET_ROWS_KERNEL_SAMETYPE) {
+        size_t src0_row_size = 0;
+        if (kernel_type == HTP_GET_ROWS_KERNEL_TILED) {
+            const size_t tile_stride = htp_tensor_get_tile_padded_size(type);
+            const uint32_t n_k_tiles = ne00 / HTP_TENSOR_TILE_N_COLS;
+            src0_row_size = n_k_tiles > 0 ? (n_k_tiles * tile_stride) : tile_stride;
+        } else {
+            src0_row_size = htp_tensor_get_row_size(type, ne00);
+        }
+
+        L->src0_spad_half_size   = (src0_row_size + 255) & ~255;
+        L->dst_spad_half_size    = (ne00 * sizeof(float) + 255) & ~255;
+        L->src0_bytes_per_thread = 2 * L->src0_spad_half_size;
+        L->dst_bytes_per_thread  = 2 * L->dst_spad_half_size;
    }

-    if (kernel_type == HTP_GET_ROWS_KERNEL_TILED) {
-        const size_t tile_size   = htp_mm_get_weight_tile_size(type);
-        const size_t tile_stride = (tile_size + 127) & ~127;
-        const uint32_t n_k_tiles = ne00 / HTP_MM_HMX_TILE_N_COLS;
-        const size_t row_tiles_size = n_k_tiles > 0 ? (n_k_tiles * tile_stride) : tile_stride;
-        vtcm_layout->src0_spad_half_size = (row_tiles_size + 255) & ~255;
-        vtcm_layout->dst_spad_half_size  = (ne00 * sizeof(float) + 255) & ~255;
-        vtcm_layout->src0_bytes_per_thread = 2 * vtcm_layout->src0_spad_half_size;
-        vtcm_layout->dst_bytes_per_thread  = 2 * vtcm_layout->dst_spad_half_size;
-        vtcm_layout->off_src0 = 0;
-        vtcm_layout->off_dst  = vtcm_layout->src0_bytes_per_thread * n_threads;
-        vtcm_layout->total_bytes = vtcm_layout->off_dst + vtcm_layout->dst_bytes_per_thread * n_threads;
-        return;
-    }
-
-    uint32_t src0_row_size = 0;
-    switch (type) {
-        case HTP_TYPE_F16:
-            src0_row_size = ne00 * 2;
-            break;
-        case HTP_TYPE_Q8_0:
-            src0_row_size = (ne00 / 32) * 34;
-            break;
-        default:
-            src0_row_size = 0;
-            break;
-    }
-
-    size_t src0_row_size_aligned = (src0_row_size + 255) & ~255;
-    size_t dst_row_size_aligned  = (ne00 * sizeof(float) + 255) & ~255;
-
-    vtcm_layout->src0_spad_half_size = src0_row_size_aligned;
-    vtcm_layout->dst_spad_half_size  = dst_row_size_aligned;
-
-    vtcm_layout->src0_bytes_per_thread = src0_row_size_aligned * 2;
-    vtcm_layout->dst_bytes_per_thread  = dst_row_size_aligned * 2;
-
-    vtcm_layout->off_src0 = 0;
-    vtcm_layout->off_dst  = vtcm_layout->off_src0 + vtcm_layout->src0_bytes_per_thread * n_threads;
-    vtcm_layout->total_bytes = vtcm_layout->off_dst + vtcm_layout->dst_bytes_per_thread * n_threads;
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src1, src1_alloc_bytes);
+    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;
 }

 #if defined(__cplusplus)
@@ -14,11 +14,15 @@ static const __fp16 iq4_nl_to_fp16_lut[64] __attribute__((aligned(VLEN))) = {

 // --- tiled format dequantizers ---

-typedef struct {
+typedef struct tiled_dequantize_state tiled_dequantize_state_t;
+typedef void (*tiled_dequant_task_fn_t)(const tiled_dequantize_state_t * state, uint32_t start_tile, uint32_t end_tile);
+
+struct tiled_dequantize_state {
    struct htp_context      * ctx;
    struct htp_thread_trace * traces;
    __fp16                  * dst;
    const uint8_t           * src;
+    tiled_dequant_task_fn_t   task_fn;

    struct fastdiv_values     n_k_tiles_div;
    uint32_t                  n_k_tiles;
@@ -31,7 +35,7 @@ typedef struct {
    uint32_t                  k_block;
    size_t                    row_stride;
    uint32_t                  weight_type;
-} tiled_dequantize_state_t;
+};

 // Dequantize a single tile from tiled weight data (already in VTCM) to tile-major FP16.
 static void dequantize_tiled_weight_to_fp16_task_q4_0(
@@ -611,7 +615,7 @@ static void dequantize_tiled_weight_to_fp16_task_q5_k(
    }
 }

-// Q6_K stores 6-bit weights and one fp16 scale per 16 k, see HTP_MM_WEIGHT_TILE_SIZE_Q6_K.
+// Q6_K stores 6-bit weights and one fp16 scale per 16 k, see HTP_TENSOR_TILE_SIZE_Q6_K.
 // A k-group holds 4 k per row, the HMX tile holds 2, so each group is dealt into two tiles.
 static void dequantize_tiled_weight_to_fp16_task_q6_k(
        const tiled_dequantize_state_t *state,
@@ -669,7 +673,7 @@ static void dequantize_tiled_weight_to_fp16_task_q6_k(
    }
 }

-// Q3_K stores 3-bit weights and one fp16 scale per 16 k, see HTP_MM_WEIGHT_TILE_SIZE_Q3_K.
+// Q3_K stores 3-bit weights and one fp16 scale per 16 k, see HTP_TENSOR_TILE_SIZE_Q3_K.
 static void dequantize_tiled_weight_to_fp16_task_q3_k(
        const tiled_dequantize_state_t *state,
        uint32_t start_tile, uint32_t end_tile) {
@@ -701,7 +705,7 @@ static void dequantize_tiled_weight_to_fp16_task_q3_k(
    }
 }

-// Q2_K stores 2-bit weights, the ggml scale bytes and d / -dmin per super-block, see HTP_MM_WEIGHT_TILE_SIZE_Q2_K.
+// Q2_K stores 2-bit weights, the ggml scale bytes and d / -dmin per super-block, see HTP_TENSOR_TILE_SIZE_Q2_K.
 static void dequantize_tiled_weight_to_fp16_task_q2_k(
        const tiled_dequantize_state_t *state,
        uint32_t start_tile, uint32_t end_tile) {
@@ -837,6 +841,56 @@ void quantize_f32_weight_to_fp16_tiles_task(
    }
 }

+static __attribute__((noinline))
+void convert_bf16_weight_to_fp16_tiles_task(
+        const tiled_dequantize_state_t *state,
+        uint32_t start_tile, uint32_t end_tile) {
+
+    const uint32_t n_k_tiles = state->n_k_tiles;
+    const struct fastdiv_values n_k_tiles_div = state->n_k_tiles_div;
+
+    const HVX_Vector v_scat_base  = hvx_vmem(hmx_transpose_scatter_offsets);
+    const HVX_Vector v_scat_step  = Q6_V_vsplat_R(4);
+    const HVX_VectorPred q_mask64 = Q6_Q_vsetq_R(64);
+
+    unsigned ct = fastdiv((unsigned)start_tile, &n_k_tiles_div);
+    unsigned kt = fastmodulo((unsigned)start_tile, n_k_tiles, &n_k_tiles_div);
+
+    for (unsigned t = start_tile; t < (unsigned)end_tile; ) {
+        if (kt >= (unsigned)n_k_tiles) { kt = 0; ct++; }
+
+        __fp16 *tile_base = state->dst + t * HTP_MM_HMX_TILE_N_ELMS;
+        {
+            uint32_t byte_off = kt * 32 * sizeof(uint16_t);
+
+            HVX_Vector v_off = v_scat_base;
+            for (uint32_t r = 0; r < HTP_MM_HMX_TILE_N_ROWS; r += 2) {
+                uint32_t row0 = ct * HTP_MM_HMX_TILE_N_COLS + r;
+                uint32_t row1 = row0 + 1;
+
+                const uint8_t *r0 = state->src + row0 * state->row_stride;
+                const uint8_t *r1 = state->src + row1 * state->row_stride;
+
+                HVX_Vector v0_raw = (row0 < state->n_cols) ? hvx_vmemu((const void *)(r0 + byte_off)) : Q6_V_vzero();
+                HVX_Vector v1_raw = (row1 < state->n_cols) ? hvx_vmemu((const void *)(r1 + byte_off)) : Q6_V_vzero();
+
+                HVX_Vector v0_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(Q6_Ww_vunpack_Vh(v0_raw)), 16);
+                HVX_Vector v1_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(Q6_Ww_vunpack_Vh(v1_raw)), 16);
+
+                HVX_Vector v_out = hvx_vec_f32_to_f16(v0_f32, v1_f32);
+
+                Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v_out);
+                v_off = Q6_Vw_vadd_VwVw(v_off, v_scat_step);
+
+                HVX_Vector v_out_hi = Q6_V_vror_VR(v_out, 64);
+                Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v_out_hi);
+                v_off = Q6_Vw_vadd_VwVw(v_off, v_scat_step);
+            }
+        }
+        ++t; ++kt;
+    }
+}
+
 // --- End tiled dequantizers ---

 // dot-chunk functions require external HMX lock
@@ -3,6 +3,7 @@

 #include "dma-queue.h"
 #include "hmx-queue.h"
+#include "htp-dcache.h"
 #include "htp-ops.h"
 #include "hex-profile.h"
 #include "work-queue.h"
@@ -18,8 +19,7 @@
 #define HTP_MAX_NTHREADS 10
 #endif

-#define HTP_MAX_MMAPS        64
-#define HTP_MAX_DIRTY_RANGES 64
+#define HTP_MAX_MMAPS 64

 // Memory mapping
 struct htp_mmap {
@@ -29,11 +29,6 @@ struct htp_mmap {
    uint32_t flags;
 };

-struct htp_dirty_range {
-    uint32_t start;
-    uint32_t end;
-};
-
 // Scratchpad state
 struct htp_spad {
    const struct htp_tensor * src;             // original src of the data (for reuse)
@@ -1,4 +1,4 @@
-#include "htp-tensor.h"
+#include "htp-dcache.h"

 #include <qurt.h>
 #include <qurt_memory.h>
@@ -94,7 +94,7 @@ static void merge_dirty_ranges(struct htp_context * ctx) {
    }
 }

-void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
+void htp_dcache_mark_dirty_tensors(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
    const struct htp_tensor * pending[HTP_OP_MAX_OUTPUTS];
    uint32_t n_pending = 0;

@@ -301,7 +301,7 @@ static void flush_dirty_ranges(struct htp_context * ctx, const struct htp_dirty_
    }
 }

-void htp_flush_dirty_ranges(struct htp_context * ctx) {
+void htp_dcache_flush_dirty_ranges(struct htp_context * ctx) {
    struct htp_dirty_range ranges[HTP_MAX_DIRTY_RANGES];
    uint32_t n_ranges = 0;
    uint64_t total_dirty = 0;
@@ -328,7 +328,7 @@ void htp_flush_dirty_ranges(struct htp_context * ctx) {
    memset(ctx->dirty_ranges, 0, sizeof(ctx->dirty_ranges));
 }

-void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
+void htp_dcache_flush_dirty_tensors(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
    const struct htp_tensor * dirty_tensors[HTP_OP_MAX_INPUTS];
    struct htp_dirty_range ranges[HTP_OP_MAX_INPUTS];
    uint32_t n_dirty = 0;
@@ -0,0 +1,20 @@
+#ifndef HTP_DCACHE_H
+#define HTP_DCACHE_H
+
+#include <stdint.h>
+#include "htp-ops.h"
+
+#define HTP_MAX_DIRTY_RANGES 64
+
+struct htp_dirty_range {
+    uint32_t start;
+    uint32_t end;
+};
+
+struct htp_context;
+
+void htp_dcache_flush_dirty_ranges(struct htp_context * ctx);
+void htp_dcache_flush_dirty_tensors(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
+void htp_dcache_mark_dirty_tensors(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
+
+#endif // HTP_DCACHE_H
@@ -30,21 +30,14 @@ enum htp_data_type {
    HTP_TYPE_IQ4_NL = 20,
    HTP_TYPE_I32    = 26,
    HTP_TYPE_I64    = 27,
+    HTP_TYPE_BF16   = 30,
    HTP_TYPE_MXFP4  = 39,

-    // types used internally for repack, dyn.quant, etc
-    HTP_TYPE_Q4_0_TILED = 200,
-    HTP_TYPE_Q4_1_TILED,
-    HTP_TYPE_Q8_0_TILED,
-    HTP_TYPE_MXFP4_TILED,
-
    HTP_TYPE_INVALID
 };

-// Constats for internal types
-#define QK_Q4_0_TILED  256  // 32x32 Q4_0 tiled layout
+// Constants for internal types
 #define QK_Q8_0_TILED  128  // 32x32 Q8_0 tiled layout
-#define QK_MXFP4_TILED 256  // 32x32 MXFP4 tiled layout



@@ -141,7 +134,7 @@ enum htp_op_code {

 enum htp_tensor_flags {
    HTP_TENSOR_WEIGHT  = (1U << 0), // Tensor buffer model weight data (not compute)
-    HTP_TENSOR_REPACK  = (1U << 1), // Tensor is in repacked tiled format
+    HTP_TENSOR_TILED   = (1U << 1), // Tensor is in tiled format
    HTP_TENSOR_FENCE   = (1U << 2)  // Tensor is synchronization fence (explicitly managed)
 };

@@ -8,27 +8,59 @@
 #include "hex-common.h"
 #include "hex-fastdiv.h"

-enum {
-    HTP_TENSOR_MDEV_LINE_SIZE = 128,
-};
+#define HTP_TENSOR_MDEV_LINE_SIZE   128
+
+// Currently HMX requires 32x32 tile formats (do not change)
+#define HTP_TENSOR_TILE_N_COLS      32
+#define HTP_TENSOR_TILE_N_ROWS      32
+
+#define HTP_TENSOR_TILE_SIZE_F16    2048
+#define HTP_TENSOR_TILE_SIZE_Q4_0   576
+#define HTP_TENSOR_TILE_SIZE_Q4_1   640
+#define HTP_TENSOR_TILE_SIZE_Q8_0   1088
+#define HTP_TENSOR_TILE_SIZE_IQ4_NL 576
+#define HTP_TENSOR_TILE_SIZE_MXFP4  544
+#define HTP_TENSOR_TILE_SIZE_Q4_K   640
+#define HTP_TENSOR_TILE_SIZE_Q5_K   768
+#define HTP_TENSOR_TILE_SIZE_Q6_K   896
+#define HTP_TENSOR_TILE_SIZE_Q3_K   512
+#define HTP_TENSOR_TILE_SIZE_Q2_K   336
+
+// Q2_K, Q3_K, Q6_K use vrmpy byte order within each 128B HVX vector:
+//   k = 0..31 is split into 8 groups g = k/4 of 4 k each; byte 4*r + (k%4) holds row r
+//
+// Q2_K (336B, lossless tiling with same size as block_q2_K):
+//   bytes   0..255 (vec 0..1): unsigned 2-bit quants, vec m holds groups 4m..4m+3 at bit offsets 0,2,4,6
+//   bytes 256..319:            raw scales[] bytes (min << 4 | sc): [0..31] for k 0..15, [32..63] for k 16..31
+//   bytes 320..335:            8 fp16 in tile j (0..7) of super-block: d (j < 4) or -dmin (j >= 4) for rows 8*(j%4)..8*(j%4)+7
+//
+// Q3_K (512B):
+//   bytes   0..255 (vec 0..1): low 2 bits, vec m holds groups 4m..4m+3 at bit offsets 0,2,4,6
+//   bytes 256..383 (vec 2):    sign mask, bit g set when hmask is 0 (quant = low2 - 4)
+//   bytes 384..511 (vec 3):    fp16 scales d * (scales[] - 32): fp16[0..31] for k 0..15, fp16[32..63] for k 16..31
+//
+// Q4_K (640B, same layout as Q4_1 tile; each 256-k super-block splits into 8 tiles of 32 k):
+//   bytes   0..511 (vec 0..3): 4-bit quants, byte 32*cp + r holds k = 2*cp (low nibble) and 2*cp + 1 (high nibble)
+//   bytes 512..639 (vec 4):    32x2 fp16 per row: (d * sc, -dmin * m)
+//
+// Q5_K (768B):
+//   bytes   0..639 (vec 0..4): Q4_K tile (low 4 bits + fp16 scale/min pairs)
+//   bytes 640..767 (vec 5):    5th-bit plane; byte l holds the 8 high bits for lane l across vec i = 0..3 (bit 2i = lo, 2i+1 = hi)
+//
+// Q6_K (896B):
+//   bytes   0..511 (vec 0..3): low 4 bits, vec i holds group 2i (low nibble) and 2i+1 (high nibble)
+//   bytes 512..767 (vec 4..5): high 2 bits, vec 4+m holds groups 4m..4m+3 at bit offsets 0,2,4,6
+//   bytes 768..895 (vec 6):    fp16 scales d * scales[]: fp16[0..31] for k 0..15, fp16[32..63] for k 16..31

 struct htp_tensor_mdev_range {
    uint32_t start;
    uint32_t count;
 };

-static inline void * htp_tensor_data(const struct htp_tensor * t) {
-    return (void *) (uintptr_t) t->data;
-}
-
 static inline bool htp_tensor_is_extended(const struct htp_tensor * t) {
    return t && (t->data >> 32) != 0;
 }

-static inline uint32_t * htp_tensor_flags(const struct htp_tensor * t) {
-    return (uint32_t *) &t->flags;
-}
-
 static inline bool htp_tensor_is_contiguous(const struct htp_tensor * t, uint32_t type_size) {
    uint32_t next_nb = type_size;
    if (t->ne[0] != 1 && t->nb[0] != next_nb) {
@@ -125,19 +157,50 @@ static inline struct htp_tensor_mdev_range htp_tensor_mdev_partition(uint32_t to
    return range;
 }

-static inline uint32_t htp_tensor_get_row_size(int type, uint32_t ne00) {
+static inline uint32_t htp_tensor_get_tile_size(int type) {
    switch (type) {
-        case HTP_TYPE_F32:  return ne00 * 4;
-        case HTP_TYPE_F16:  return ne00 * 2;
-        case HTP_TYPE_Q8_0: return (ne00 / 32) * 34;
-        case HTP_TYPE_I32:  return ne00 * 4;
-        default:            return 0;
+        case HTP_TYPE_Q4_0:   return HTP_TENSOR_TILE_SIZE_Q4_0;
+        case HTP_TYPE_IQ4_NL: return HTP_TENSOR_TILE_SIZE_IQ4_NL;
+        case HTP_TYPE_Q4_1:   return HTP_TENSOR_TILE_SIZE_Q4_1;
+        case HTP_TYPE_Q4_K:   return HTP_TENSOR_TILE_SIZE_Q4_K;
+        case HTP_TYPE_Q8_0:   return HTP_TENSOR_TILE_SIZE_Q8_0;
+        case HTP_TYPE_Q5_K:   return HTP_TENSOR_TILE_SIZE_Q5_K;
+        case HTP_TYPE_Q6_K:   return HTP_TENSOR_TILE_SIZE_Q6_K;
+        case HTP_TYPE_Q3_K:   return HTP_TENSOR_TILE_SIZE_Q3_K;
+        case HTP_TYPE_Q2_K:   return HTP_TENSOR_TILE_SIZE_Q2_K;
+        case HTP_TYPE_MXFP4:  return HTP_TENSOR_TILE_SIZE_MXFP4;
+        case HTP_TYPE_F16:    return HTP_TENSOR_TILE_SIZE_F16;
+        default:              return 0;
    }
 }

-struct htp_context;
-void htp_flush_dirty_ranges(struct htp_context * ctx);
-void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
-void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
+static inline uint32_t htp_tensor_get_tile_padded_size(int type) {
+    return (htp_tensor_get_tile_size(type) + 127u) & ~127u;
+}
+
+static inline bool htp_tensor_is_tiled(const struct htp_tensor * t) {
+    return t && (t->flags & HTP_TENSOR_TILED) != 0;
+}
+
+static inline uint32_t htp_tensor_get_row_size(int type, uint32_t ne00) {
+    switch (type) {
+        case HTP_TYPE_F32:    return ne00 * 4;
+        case HTP_TYPE_F16:
+        case HTP_TYPE_BF16:   return ne00 * 2;
+        case HTP_TYPE_Q8_0:   return (ne00 / 32) * 34;
+        case HTP_TYPE_Q4_0:
+        case HTP_TYPE_IQ4_NL: return (ne00 / 32) * 18;
+        case HTP_TYPE_Q4_1:   return (ne00 / 32) * 20;
+        case HTP_TYPE_MXFP4:  return (ne00 / 32) * 17;
+        case HTP_TYPE_Q2_K:   return (ne00 / 256) * 84;
+        case HTP_TYPE_Q3_K:   return (ne00 / 256) * 110;
+        case HTP_TYPE_Q4_K:   return (ne00 / 256) * 144;
+        case HTP_TYPE_Q5_K:   return (ne00 / 256) * 176;
+        case HTP_TYPE_Q6_K:   return (ne00 / 256) * 210;
+        case HTP_TYPE_I32:    return ne00 * 4;
+        case HTP_TYPE_I64:    return ne00 * 8;
+        default:              return 0;
+    }
+}

 #endif // HTP_TENSOR_H
@@ -76,7 +76,7 @@ static inline void quantize_f16_f16_kernel(
 #define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b)
 #endif

-static inline void vec_dot_f32_f32_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
+static inline void flat_vec_dot_f32_f32_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
    const HVX_Vector * restrict x = (const HVX_Vector *) vx;
    const HVX_Vector * restrict y = (const HVX_Vector *) vy;

@@ -104,7 +104,7 @@ static inline void vec_dot_f32_f32_aa_1x1(const uint32_t n, float * restrict s,
    *s = hvx_vec_get_f32(hvx_vec_reduce_sum_f32(rsum));
 }

-static inline void vec_dot_f32_f32_aa_2x1(const uint32_t n, float * restrict s0,
+static inline void flat_vec_dot_f32_f32_aa_2x1(const uint32_t n, float * restrict s0,
                                const void * restrict vx0, const void * restrict vx1,
                                const void * restrict vy0) {
    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
@@ -143,7 +143,7 @@ static inline void vec_dot_f32_f32_aa_2x1(const uint32_t n, float * restrict s0,
    hvx_vec_store_u(s0, 8, rsum);
 }

-static inline void vec_dot_f32_f32_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
+static inline void flat_vec_dot_f32_f32_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
                                const void * restrict vx0, const void * restrict vx1,
                                const void * restrict vy0, const void * restrict vy1) {
    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
@@ -196,10 +196,203 @@ static inline void vec_dot_f32_f32_aa_2x2(const uint32_t n, float * restrict s0,
    hvx_vec_store_u(s1, 8, r0_r1_c1_sum);
 }

+static inline void flat_vec_dot_bf16_f32_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
+    const HVX_Vector * restrict x = (const HVX_Vector *) vx;
+    const HVX_Vector * restrict y = (const HVX_Vector *) vy;
+
+    uint32_t nvec = n / 64;
+    uint32_t rem  = n % 64;
+
+    HVX_Vector rsum0 = Q6_V_vzero();
+    HVX_Vector rsum1 = Q6_V_vzero();
+
+    uint32_t i = 0;
+    #pragma unroll(2)
+    for (i = 0; i < nvec; i++) {
+        HVX_Vector v_bf16 = x[i];
+        HVX_VectorPair p = Q6_Ww_vunpack_Vh(v_bf16);
+        HVX_Vector w0_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p), 16);
+        HVX_Vector w1_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p), 16);
+
+        rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w0_f32, y[2 * i + 0]));
+        rsum1 = HVX_OP_ADD_F32(rsum1, HVX_OP_MUL_F32(w1_f32, y[2 * i + 1]));
+    }
+
+    if (rem) {
+        HVX_VectorPred bmask = Q6_Q_vsetq2_R(rem * 2);
+        HVX_Vector v_bf16 = Q6_V_vand_QV(bmask, x[i]);
+        HVX_VectorPair p = Q6_Ww_vunpack_Vh(v_bf16);
+        HVX_Vector w0_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p), 16);
+        HVX_Vector w1_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p), 16);
+
+        if (rem <= 32) {
+            HVX_VectorPred a_mask = Q6_Q_vsetq2_R(rem * 4);
+            HVX_Vector y0 = Q6_V_vand_QV(a_mask, y[2 * i + 0]);
+            rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w0_f32, y0));
+        } else {
+            HVX_Vector y0 = y[2 * i + 0];
+            rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w0_f32, y0));
+
+            HVX_VectorPred a_mask = Q6_Q_vsetq2_R((rem - 32) * 4);
+            HVX_Vector y1 = Q6_V_vand_QV(a_mask, y[2 * i + 1]);
+            rsum1 = HVX_OP_ADD_F32(rsum1, HVX_OP_MUL_F32(w1_f32, y1));
+        }
+    }
+
+    HVX_Vector rsum = HVX_OP_ADD_F32(rsum0, rsum1);
+    *s = hvx_vec_get_f32(hvx_vec_reduce_sum_f32(rsum));
+}
+
+static inline void flat_vec_dot_bf16_f32_aa_2x1(const uint32_t n, float * restrict s0,
+                                const void * restrict vx0, const void * restrict vx1,
+                                const void * restrict vy0) {
+    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
+    const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
+    const HVX_Vector * restrict y  = (const HVX_Vector *) vy0;
+
+    uint32_t nvec = n / 64;
+    uint32_t rem  = n % 64;
+
+    HVX_Vector rsum0 = Q6_V_vzero();
+    HVX_Vector rsum1 = Q6_V_vzero();
+
+    uint32_t i = 0;
+    #pragma unroll(2)
+    for (i = 0; i < nvec; i++) {
+        HVX_VectorPair p0 = Q6_Ww_vunpack_Vh(x0[i]);
+        HVX_VectorPair p1 = Q6_Ww_vunpack_Vh(x1[i]);
+        HVX_Vector w00_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p0), 16);
+        HVX_Vector w01_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p0), 16);
+        HVX_Vector w10_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p1), 16);
+        HVX_Vector w11_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p1), 16);
+
+        HVX_Vector y0 = y[2 * i + 0];
+        HVX_Vector y1 = y[2 * i + 1];
+
+        rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w00_f32, y0));
+        rsum1 = HVX_OP_ADD_F32(rsum1, HVX_OP_MUL_F32(w10_f32, y0));
+        rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w01_f32, y1));
+        rsum1 = HVX_OP_ADD_F32(rsum1, HVX_OP_MUL_F32(w11_f32, y1));
+    }
+
+    if (rem) {
+        HVX_VectorPred bmask = Q6_Q_vsetq2_R(rem * 2);
+        HVX_VectorPair p0 = Q6_Ww_vunpack_Vh(Q6_V_vand_QV(bmask, x0[i]));
+        HVX_VectorPair p1 = Q6_Ww_vunpack_Vh(Q6_V_vand_QV(bmask, x1[i]));
+        HVX_Vector w00_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p0), 16);
+        HVX_Vector w01_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p0), 16);
+        HVX_Vector w10_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p1), 16);
+        HVX_Vector w11_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p1), 16);
+
+        if (rem <= 32) {
+            HVX_VectorPred a_mask = Q6_Q_vsetq2_R(rem * 4);
+            HVX_Vector y0 = Q6_V_vand_QV(a_mask, y[2 * i + 0]);
+            rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w00_f32, y0));
+            rsum1 = HVX_OP_ADD_F32(rsum1, HVX_OP_MUL_F32(w10_f32, y0));
+        } else {
+            HVX_Vector y0 = y[2 * i + 0];
+            rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w00_f32, y0));
+            rsum1 = HVX_OP_ADD_F32(rsum1, HVX_OP_MUL_F32(w10_f32, y0));
+
+            HVX_VectorPred a_mask = Q6_Q_vsetq2_R((rem - 32) * 4);
+            HVX_Vector y1 = Q6_V_vand_QV(a_mask, y[2 * i + 1]);
+            rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w01_f32, y1));
+            rsum1 = HVX_OP_ADD_F32(rsum1, HVX_OP_MUL_F32(w11_f32, y1));
+        }
+    }
+
+    HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1);
+    hvx_vec_store_u(s0, 8, rsum);
+}
+
+static inline void flat_vec_dot_bf16_f32_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
+                                const void * restrict vx0, const void * restrict vx1,
+                                const void * restrict vy0, const void * restrict vy1) {
+    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
+    const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
+    const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0;
+    const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1;
+
+    uint32_t nvec = n / 64;
+    uint32_t rem  = n % 64;
+
+    HVX_Vector r0_c0_sum = Q6_V_vzero();
+    HVX_Vector r0_c1_sum = Q6_V_vzero();
+    HVX_Vector r1_c0_sum = Q6_V_vzero();
+    HVX_Vector r1_c1_sum = Q6_V_vzero();
+
+    uint32_t i = 0;
+    #pragma unroll(2)
+    for (i = 0; i < nvec; i++) {
+        HVX_VectorPair p0 = Q6_Ww_vunpack_Vh(x0[i]);
+        HVX_VectorPair p1 = Q6_Ww_vunpack_Vh(x1[i]);
+        HVX_Vector w00_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p0), 16);
+        HVX_Vector w01_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p0), 16);
+        HVX_Vector w10_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p1), 16);
+        HVX_Vector w11_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p1), 16);
+
+        HVX_Vector c0_0 = y0[2 * i + 0];
+        HVX_Vector c0_1 = y0[2 * i + 1];
+        HVX_Vector c1_0 = y1[2 * i + 0];
+        HVX_Vector c1_1 = y1[2 * i + 1];
+
+        r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(w00_f32, c0_0));
+        r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(w00_f32, c1_0));
+        r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(w10_f32, c0_0));
+        r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(w10_f32, c1_0));
+
+        r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(w01_f32, c0_1));
+        r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(w01_f32, c1_1));
+        r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(w11_f32, c0_1));
+        r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(w11_f32, c1_1));
+    }
+
+    if (rem) {
+        HVX_VectorPred bmask = Q6_Q_vsetq2_R(rem * 2);
+        HVX_VectorPair p0 = Q6_Ww_vunpack_Vh(Q6_V_vand_QV(bmask, x0[i]));
+        HVX_VectorPair p1 = Q6_Ww_vunpack_Vh(Q6_V_vand_QV(bmask, x1[i]));
+        HVX_Vector w00_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p0), 16);
+        HVX_Vector w01_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p0), 16);
+        HVX_Vector w10_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p1), 16);
+        HVX_Vector w11_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p1), 16);
+
+        if (rem <= 32) {
+            HVX_VectorPred a_mask = Q6_Q_vsetq2_R(rem * 4);
+            HVX_Vector c0_0 = Q6_V_vand_QV(a_mask, y0[2 * i + 0]);
+            HVX_Vector c1_0 = Q6_V_vand_QV(a_mask, y1[2 * i + 0]);
+            r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(w00_f32, c0_0));
+            r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(w00_f32, c1_0));
+            r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(w10_f32, c0_0));
+            r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(w10_f32, c1_0));
+        } else {
+            HVX_Vector c0_0 = y0[2 * i + 0];
+            HVX_Vector c1_0 = y1[2 * i + 0];
+            r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(w00_f32, c0_0));
+            r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(w00_f32, c1_0));
+            r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(w10_f32, c0_0));
+            r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(w10_f32, c1_0));
+
+            HVX_VectorPred a_mask = Q6_Q_vsetq2_R((rem - 32) * 4);
+            HVX_Vector c0_1 = Q6_V_vand_QV(a_mask, y0[2 * i + 1]);
+            HVX_Vector c1_1 = Q6_V_vand_QV(a_mask, y1[2 * i + 1]);
+            r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(w01_f32, c0_1));
+            r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(w01_f32, c1_1));
+            r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(w11_f32, c0_1));
+            r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(w11_f32, c1_1));
+        }
+    }
+
+    HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum);
+    HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum);
+
+    hvx_vec_store_u(s0, 8, r0_r1_c0_sum);
+    hvx_vec_store_u(s1, 8, r0_r1_c1_sum);
+}
+
 #undef HVX_OP_ADD_F32
 #undef HVX_OP_MUL_F32

-static inline void vec_dot_f16_f16_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
+static inline void flat_vec_dot_f16_f16_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
    const HVX_Vector * restrict x = (const HVX_Vector *) vx;
    const HVX_Vector * restrict y = (const HVX_Vector *) vy;

@@ -226,7 +419,7 @@ static inline void vec_dot_f16_f16_aa_1x1(const uint32_t n, float * restrict s,
    hvx_vec_store_u(s, 4, hvx_vec_reduce_sum_f32(rsum));
 }

-static inline void vec_dot_f16_f16_aa_2x1(const uint32_t n, float * restrict s0,
+static inline void flat_vec_dot_f16_f16_aa_2x1(const uint32_t n, float * restrict s0,
                                const void * restrict vx0, const void * restrict vx1,
                                const void * restrict vy0) {
    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
@@ -263,7 +456,7 @@ static inline void vec_dot_f16_f16_aa_2x1(const uint32_t n, float * restrict s0,
    hvx_vec_store_u(s0, 8, rsum);
 }

-static inline void vec_dot_f16_f16_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
+static inline void flat_vec_dot_f16_f16_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
                                const void * restrict vx0, const void * restrict vx1,
                                const void * restrict vy0, const void * restrict vy1) {
    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
@@ -274,6 +274,72 @@ static void quantize_row_f32_q8_1_s16_tiled(float * restrict x, uint8_t * restri
    }
 }

+static inline void quantize_half_f16_tiled_splat(HVX_Vector v_src, HVX_Vector * restrict dst) {
+    static const uint8_t __attribute__((aligned(128))) repl[128] = {
+        0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
+        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
+        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
+        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
+        0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
+        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
+        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
+        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
+    };
+    HVX_Vector v_repl_ctrl = *(const HVX_Vector *) repl;
+
+    #pragma unroll(4)
+    for (int cp = 0; cp < 32; cp++) {
+        dst[cp] = Q6_V_vdelta_VV(cp == 0 ? v_src : Q6_V_vror_VR(v_src, cp * 4), v_repl_ctrl);
+    }
+}
+
+static inline void quantize_block_f32_f16_tiled(const float * restrict x, uint8_t * restrict y_block) {
+    assert((unsigned long) x % 128 == 0);
+    assert((unsigned long) y_block % 128 == 0);
+
+    const HVX_Vector * vx = (const HVX_Vector *) x;
+    HVX_Vector * restrict dst = (HVX_Vector *) y_block;
+
+    HVX_Vector vx01_hf = hvx_vec_f32_to_f16(vx[0], vx[1]);
+    quantize_half_f16_tiled_splat(vx01_hf, dst);
+
+    HVX_Vector vx23_hf = hvx_vec_f32_to_f16(vx[2], vx[3]);
+    quantize_half_f16_tiled_splat(vx23_hf, dst + 32);
+}
+
+static inline void quantize_block_f16_f16_tiled(const __fp16 * restrict x, uint8_t * restrict y_block) {
+    assert((unsigned long) x % 128 == 0);
+    assert((unsigned long) y_block % 128 == 0);
+
+    const HVX_Vector * vx = (const HVX_Vector *) x;
+    HVX_Vector * restrict dst = (HVX_Vector *) y_block;
+
+    quantize_half_f16_tiled_splat(vx[0], dst);
+    quantize_half_f16_tiled_splat(vx[1], dst + 32);
+}
+
+static void quantize_row_f32_f16_tiled(const float * restrict x, uint8_t * restrict y, uint32_t k) {
+    assert(k % 32 == 0);
+    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t nb = (k + qk - 1) / qk;
+
+    for (uint32_t i = 0; i < nb; i++) {
+        uint8_t * restrict y_block = y + i * 4 * 2048;
+        quantize_block_f32_f16_tiled(x + i * qk, y_block);
+    }
+}
+
+static void quantize_row_f16_f16_tiled(const __fp16 * restrict x, uint8_t * restrict y, uint32_t k) {
+    assert(k % 32 == 0);
+    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t nb = (k + qk - 1) / qk;
+
+    for (uint32_t i = 0; i < nb; i++) {
+        uint8_t * restrict y_block = y + i * 4 * 2048;
+        quantize_block_f16_f16_tiled(x + i * qk, y_block);
+    }
+}
+
 // Dot kernels & helpers that consume tiled activations

 static inline HVX_Vector hvx_vec_mul_f16_f16_to_f32_lower32(HVX_Vector v1, HVX_Vector v2) {
@@ -454,7 +520,7 @@ static inline HVX_VectorPair accum_q8_0_32x2(
    return Q6_W_vcombine_VV(v_sum1, v_sum0);
 }

-// Q5_K: OR 0x10 into every lane of v whose flag j is set in the plane (see HTP_MM_WEIGHT_TILE_SIZE_Q5_K)
+// Q5_K: OR 0x10 into every lane of v whose flag j is set in the plane (see HTP_TENSOR_TILE_SIZE_Q5_K)
 static inline HVX_Vector hvx_q5k_or_hibit(HVX_Vector v, HVX_Vector v_plane, int j) {
    HVX_VectorPred q = Q6_Q_vand_VR(v_plane, 0x01010101u << j);
    return Q6_V_vandor_VQR(v, q, 0x10101010);
@@ -526,7 +592,7 @@ static inline HVX_VectorPair accum_5bit_32x2(
    return Q6_W_vcombine_VV(v_sum1, v_sum0);
 }

-// Q6_K weights are stored unsigned (0..63), see HTP_MM_WEIGHT_TILE_SIZE_Q6_K. Unpack k-group g of a tile to signed bytes (q - 32)
+// Q6_K weights are stored unsigned (0..63), see HTP_TENSOR_TILE_SIZE_Q6_K. Unpack k-group g of a tile to signed bytes (q - 32)
 static inline HVX_Vector unpack_q6_k_group(const HVX_Vector * restrict vptr, int g, HVX_Vector mask_0f, HVX_Vector mask_03, HVX_Vector i32) {
    HVX_Vector v_lo = (g & 1) ? Q6_Vub_vlsr_VubR(vptr[g >> 1], 4) : Q6_V_vand_VV(vptr[g >> 1], mask_0f);
    HVX_Vector v_hi = (g & 3) ? Q6_Vub_vlsr_VubR(vptr[4 + (g >> 2)], 2 * (g & 3)) : vptr[4 + (g >> 2)];
@@ -593,7 +659,7 @@ static inline HVX_Vector scale_q6_k_32x1(HVX_VectorPair v_sums, HVX_Vector v_sca
    return hvx_vec_add_f32_f32(v_lo, v_hi);
 }

-// Q3_K / Q2_K: low 2 bits of k-group g, see HTP_MM_WEIGHT_TILE_SIZE_Q3_K
+// Q3_K / Q2_K: low 2 bits of k-group g, see HTP_TENSOR_TILE_SIZE_Q3_K
 static inline HVX_Vector unpack_q3_k_low2(const HVX_Vector * restrict vptr, int g, HVX_Vector mask_03) {
    HVX_Vector v = vptr[g >> 2];
    if ((g & 3) == 3) {
@@ -1482,6 +1548,59 @@ static void tiled_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, floa
    }
 }

+static void tiled_vec_dot_f16_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
+    const HVX_Vector * restrict vptr  = (const HVX_Vector *) vx;
+    const HVX_Vector * restrict v_act = (const HVX_Vector *) vy;
+    const uint32_t n_vecs = ((n + 31) / 32) * 16;
+
+    HVX_VectorPair acc = Q6_W_vzero();
+
+    #pragma unroll(4)
+    for (uint32_t i = 0; i < n_vecs; i++) {
+        acc = hvx_vec_mpyacc_f32_f16(acc, vptr[i], v_act[i]);
+    }
+
+    HVX_Vector v_sum = hvx_vec_add_f32_f32(Q6_V_lo_W(acc), Q6_V_hi_W(acc));
+
+    if (sz) {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum, hvx_vmemu(sz)));
+    } else {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum);
+    }
+}
+
+static void tiled_vec_dot_f16_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
+    const HVX_Vector * restrict vptr   = (const HVX_Vector *) vx;
+    const HVX_Vector * restrict v_act0 = (const HVX_Vector *) vy0;
+    const HVX_Vector * restrict v_act1 = (const HVX_Vector *) vy1;
+    const uint32_t n_vecs = ((n + 31) / 32) * 16;
+
+    HVX_VectorPair acc0 = Q6_W_vzero();
+    HVX_VectorPair acc1 = Q6_W_vzero();
+
+    #pragma unroll(4)
+    for (uint32_t i = 0; i < n_vecs; i++) {
+        HVX_Vector v_w = vptr[i];
+        acc0 = hvx_vec_mpyacc_f32_f16(acc0, v_w, v_act0[i]);
+        acc1 = hvx_vec_mpyacc_f32_f16(acc1, v_w, v_act1[i]);
+    }
+
+    HVX_Vector v_sum0 = hvx_vec_add_f32_f32(Q6_V_lo_W(acc0), Q6_V_hi_W(acc0));
+    HVX_Vector v_sum1 = hvx_vec_add_f32_f32(Q6_V_lo_W(acc1), Q6_V_hi_W(acc1));
+
+    if (sz0) {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum0, hvx_vmemu(sz0)));
+    } else {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum0);
+    }
+    if (sz1) {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum1, hvx_vmemu(sz1)));
+    } else {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum1);
+    }
+}
+
+
 static inline void quantize_f32_q8_0_tiled_kernel(
    const uint8_t * restrict src_data,
    uint8_t * restrict dst_data,
@@ -1533,6 +1652,40 @@ static inline void quantize_f32_q8_1_s16_tiled_kernel(
    }
 }

+static inline void quantize_f32_f16_tiled_kernel(
+    const uint8_t * restrict src_data,
+    uint8_t * restrict dst_data,
+    uint8_t * restrict tmp_data,
+    uint32_t ne0,
+    uint32_t nrows,
+    size_t src_row_size,
+    size_t dst_row_size
+) {
+    (void) tmp_data;
+    for (uint32_t i = 0; i < nrows; ++i) {
+        quantize_row_f32_f16_tiled((const float *) src_data, dst_data, ne0);
+        dst_data += dst_row_size;
+        src_data += src_row_size;
+    }
+}
+
+static inline void quantize_f16_f16_tiled_kernel(
+    const uint8_t * restrict src_data,
+    uint8_t * restrict dst_data,
+    uint8_t * restrict tmp_data,
+    uint32_t ne0,
+    uint32_t nrows,
+    size_t src_row_size,
+    size_t dst_row_size
+) {
+    (void) tmp_data;
+    for (uint32_t i = 0; i < nrows; ++i) {
+        quantize_row_f16_f16_tiled((const __fp16 *) src_data, dst_data, ne0);
+        dst_data += dst_row_size;
+        src_data += src_row_size;
+    }
+}
+
 static inline void quantize_f32_q8_0_tiled_block_kernel(
    const float * restrict src,
    uint8_t * restrict dst,
@@ -1622,3 +1775,63 @@ static inline void quantize_f32_q8_1_s16_tiled_block_kernel(
        }
    }
 }
+
+static inline void quantize_f32_f16_tiled_block_kernel(
+    const float * restrict src,
+    uint8_t * restrict dst,
+    uint8_t * restrict tmp_data,
+    uint32_t ne0,
+    uint32_t ib_first,
+    uint32_t ib_last,
+    size_t src_row_size,
+    size_t dst_row_size,
+    uint32_t r,
+    uint32_t c
+) {
+    (void) tmp_data;
+    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t nb = (ne0 + qk - 1) / qk;
+
+    for (uint32_t ib = ib_first; ib < ib_last; ++ib) {
+        const float * restrict src_ptr = (const float *) ((const uint8_t *) src + r * src_row_size + c * qk * sizeof(float));
+        uint8_t * restrict dst_ptr = dst + r * dst_row_size + c * 4 * 2048;
+
+        quantize_block_f32_f16_tiled(src_ptr, dst_ptr);
+
+        c++;
+        if (c == nb) {
+            c = 0;
+            r++;
+        }
+    }
+}
+
+static inline void quantize_f16_f16_tiled_block_kernel(
+    const __fp16 * restrict src,
+    uint8_t * restrict dst,
+    uint8_t * restrict tmp_data,
+    uint32_t ne0,
+    uint32_t ib_first,
+    uint32_t ib_last,
+    size_t src_row_size,
+    size_t dst_row_size,
+    uint32_t r,
+    uint32_t c
+) {
+    (void) tmp_data;
+    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t nb = (ne0 + qk - 1) / qk;
+
+    for (uint32_t ib = ib_first; ib < ib_last; ++ib) {
+        const __fp16 * restrict src_ptr = (const __fp16 *) ((const uint8_t *) src + r * src_row_size + c * qk * sizeof(__fp16));
+        uint8_t * restrict dst_ptr = dst + r * dst_row_size + c * 4 * 2048;
+
+        quantize_block_f16_f16_tiled(src_ptr, dst_ptr);
+
+        c++;
+        if (c == nb) {
+            c = 0;
+            r++;
+        }
+    }
+}
@@ -5,6 +5,7 @@
 #include <stdint.h>
 #include <string.h>

+#include "htp-tensor.h"
 #include "hvx-arith.h"
 #include "hvx-base.h"
 #include "hvx-reduce.h"
@@ -96,7 +97,7 @@ static inline void hvx_quantize_row_q8_0_f32(void * restrict dst_ptr, const floa
    }
 }

-static inline void hvx_dequantize_row_q8_0_f32(float * restrict dst_ptr, const void * restrict src_ptr, int n) {
+static inline void hvx_dequantize_row_flat_q8_0_f32(float * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) {
    const int nb = n / QK8_0;
    const block_q8_0 * src = (const block_q8_0 *) src_ptr;

@@ -120,7 +121,7 @@ static inline void hvx_dequantize_row_q8_0_f32(float * restrict dst_ptr, const v
    }
 }

-static inline void hvx_dequantize_row_q8_0_f16(__fp16 * restrict dst_ptr, const void * restrict src_ptr, int n) {
+static inline void hvx_dequantize_row_flat_q8_0_f16(__fp16 * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) {
    const int nb = n / QK8_0;
    const block_q8_0 * src = (const block_q8_0 *) src_ptr;

@@ -146,20 +147,460 @@ static inline void hvx_dequantize_row_q8_0_f16(__fp16 * restrict dst_ptr, const
    }
 }

-static inline void hvx_dequantize_row_f16_f32(float * restrict dst_ptr, const void * restrict src_ptr, int n) {
-    const int nb = n / 32;
+static inline void hvx_dequantize_row_flat_f16_f32(float * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) {
+    const uint32_t nvec = n / VLEN_FP32;
+    const uint32_t nloe = n % VLEN_FP32;
    const _Float16 * src = (const _Float16 *) src_ptr;

-    for (int i = 0; i < nb; i++) {
-        HVX_Vector v_f16 = *(const HVX_UVector *) (src + i * 32);
+    for (uint32_t i = 0; i < nvec; i++) {
+        HVX_Vector v_f16 = *(const HVX_UVector *) (src + i * VLEN_FP32);
        HVX_VectorPair vp_f32 = hvx_vec_f16_to_f32(v_f16);
        HVX_Vector res = Q6_V_lo_W(vp_f32);

-        float * block_dst = dst_ptr + i * 32;
-        hvx_vmem(block_dst) = res;
+        hvx_vmem(dst_ptr + i * VLEN_FP32) = res;
+    }
+
+    if (nloe) {
+        HVX_Vector v_f16 = *(const HVX_UVector *) (src + nvec * VLEN_FP32);
+        HVX_VectorPair vp_f32 = hvx_vec_f16_to_f32(v_f16);
+        HVX_Vector res = Q6_V_lo_W(vp_f32);
+
+        hvx_vec_store_a(dst_ptr + nvec * VLEN_FP32, nloe * sizeof(float), res);
    }
 }

+static inline void hvx_dequantize_row_flat_bf16_f32(float * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) {
+    const uint32_t nvec = n / VLEN_FP32;
+    const uint32_t nloe = n % VLEN_FP32;
+    const uint16_t * src = (const uint16_t *) src_ptr;

+    for (uint32_t i = 0; i < nvec; i++) {
+        HVX_Vector v_bf16 = *(const HVX_UVector *) (src + i * VLEN_FP32);
+        HVX_VectorPair p = Q6_Ww_vunpack_Vh(v_bf16);
+        HVX_Vector res = Q6_Vw_vasl_VwR(Q6_V_lo_W(p), 16);
+
+        hvx_vmem(dst_ptr + i * VLEN_FP32) = res;
+    }
+
+    if (nloe) {
+        HVX_Vector v_bf16 = *(const HVX_UVector *) (src + nvec * VLEN_FP32);
+        HVX_VectorPair p = Q6_Ww_vunpack_Vh(v_bf16);
+        HVX_Vector res = Q6_Vw_vasl_VwR(Q6_V_lo_W(p), 16);
+
+        hvx_vec_store_a(dst_ptr + nvec * VLEN_FP32, nloe * sizeof(float), res);
+    }
+}
+
+static inline void hvx_dequantize_row_flat_q4_0_f32(float * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) {
+    const int nb = n / QK4_0;
+    const block_q4_0 * src = (const block_q4_0 *) src_ptr;
+
+    const HVX_VectorPred q16 = Q6_Q_vsetq_R(16);
+    const HVX_Vector v_eight = Q6_Vb_vsplat_R(8);
+    const HVX_Vector v_zero  = Q6_V_vzero();
+
+    for (int i = 0; i < nb; i++) {
+        HVX_Vector vd_f16     = Q6_Vh_vsplat_R(*(const int16_t *) &src[i].d);
+        HVX_VectorPair vp_f32 = hvx_vec_f16_to_f32(vd_f16);
+        HVX_Vector vd         = Q6_V_lo_W(vp_f32);
+
+        HVX_Vector vq = *(const HVX_UVector *) src[i].qs;
+        HVX_Vector v_lo = Q6_V_vand_VV(vq, Q6_Vb_vsplat_R(0x0F));
+        HVX_Vector v_hi = Q6_Vub_vlsr_VubR(vq, 4);
+
+        HVX_Vector v_lo_s8 = Q6_Vb_vsub_VbVb(v_lo, v_eight);
+        HVX_Vector v_hi_s8 = Q6_Vb_vsub_VbVb(v_hi, v_eight);
+
+        HVX_Vector v32 = Q6_V_vmux_QVV(q16, v_lo_s8, Q6_V_vlalign_VVR(v_hi_s8, v_zero, 16));
+
+        HVX_VectorPair p16   = Q6_Wh_vunpack_Vb(v32);
+        HVX_Vector     v_i16 = Q6_V_lo_W(p16);
+        HVX_VectorPair p32   = Q6_Ww_vunpack_Vh(v_i16);
+        HVX_Vector     v_i32 = Q6_V_lo_W(p32);
+
+        HVX_Vector v_f32 = Q6_Vsf_equals_Vw(v_i32);
+        HVX_Vector res   = hvx_vec_mul_f32_f32(v_f32, vd);
+
+        hvx_vmem(dst_ptr + i * QK4_0) = res;
+    }
+}
+
+static inline void hvx_dequantize_row_flat_q4_k_f32(float * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) {
+    const int nb = n / QK_K;
+    const block_q4_K * src = (const block_q4_K *) src_ptr;
+
+    for (int i = 0; i < nb; i++) {
+        const float d   = GGML_FP16_TO_FP32(src[i].d);
+        const float min = GGML_FP16_TO_FP32(src[i].dmin);
+        const uint8_t * sc = src[i].scales;
+
+        float * y = dst_ptr + i * QK_K;
+
+        for (int step = 0; step < 4; ++step) {
+            uint8_t sc1, m1_u, sc2, m2_u;
+            if (step < 2) {
+                sc1  = sc[2 * step + 0] & 63;
+                m1_u = sc[2 * step + 4] & 63;
+                sc2  = sc[2 * step + 1] & 63;
+                m2_u = sc[2 * step + 5] & 63;
+            } else {
+                const int s = step - 2;
+                sc1  = (sc[8 + 2 * s + 0] & 0xF) | ((sc[2 * s + 0] >> 6) << 4);
+                m1_u = (sc[8 + 2 * s + 0] >>  4) | ((sc[4 + 2 * s + 0] >> 6) << 4);
+                sc2  = (sc[8 + 2 * s + 1] & 0xF) | ((sc[2 * s + 1] >> 6) << 4);
+                m2_u = (sc[8 + 2 * s + 1] >>  4) | ((sc[4 + 2 * s + 1] >> 6) << 4);
+            }
+
+            HVX_Vector vq = *(const HVX_UVector *) (src[i].qs + step * 32);
+            HVX_Vector v_lo = Q6_V_vand_VV(vq, Q6_Vb_vsplat_R(0x0F));
+            HVX_Vector v_hi = Q6_Vub_vlsr_VubR(vq, 4);
+
+            HVX_VectorPair p16_lo = Q6_Wuh_vunpack_Vub(v_lo);
+            HVX_VectorPair p32_lo = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_lo));
+            HVX_Vector     v_i32_lo = Q6_V_lo_W(p32_lo);
+            HVX_Vector     v_f32_lo = Q6_Vsf_equals_Vw(v_i32_lo);
+
+            HVX_Vector vd1 = hvx_vec_splat_f32(d * sc1);
+            HVX_Vector vm1 = hvx_vec_splat_f32(min * m1_u);
+            HVX_Vector v_out1 = hvx_vec_sub_f32_f32(hvx_vec_mul_f32_f32(v_f32_lo, vd1), vm1);
+            hvx_vmem(y + 0) = v_out1;
+
+            HVX_VectorPair p16_hi = Q6_Wuh_vunpack_Vub(v_hi);
+            HVX_VectorPair p32_hi = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_hi));
+            HVX_Vector     v_i32_hi = Q6_V_lo_W(p32_hi);
+            HVX_Vector     v_f32_hi = Q6_Vsf_equals_Vw(v_i32_hi);
+
+            HVX_Vector vd2 = hvx_vec_splat_f32(d * sc2);
+            HVX_Vector vm2 = hvx_vec_splat_f32(min * m2_u);
+            HVX_Vector v_out2 = hvx_vec_sub_f32_f32(hvx_vec_mul_f32_f32(v_f32_hi, vd2), vm2);
+            hvx_vmem(y + 32) = v_out2;
+
+            y += 64;
+        }
+    }
+}
+
+static inline void hvx_dequantize_row_flat_q5_k_f32(float * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) {
+    const int nb = n / QK_K;
+    const block_q5_K * src = (const block_q5_K *) src_ptr;
+
+    for (int i = 0; i < nb; i++) {
+        const float d   = GGML_FP16_TO_FP32(src[i].d);
+        const float min = GGML_FP16_TO_FP32(src[i].dmin);
+        const uint8_t * sc = src[i].scales;
+
+        HVX_Vector v_qh = *(const HVX_UVector *) src[i].qh;
+        float * y = dst_ptr + i * QK_K;
+
+        for (int step = 0; step < 4; ++step) {
+            uint8_t sc1, m1_u, sc2, m2_u;
+            if (step < 2) {
+                sc1  = sc[2 * step + 0] & 63;
+                m1_u = sc[2 * step + 4] & 63;
+                sc2  = sc[2 * step + 1] & 63;
+                m2_u = sc[2 * step + 5] & 63;
+            } else {
+                const int s = step - 2;
+                sc1  = (sc[8 + 2 * s + 0] & 0xF) | ((sc[2 * s + 0] >> 6) << 4);
+                m1_u = (sc[8 + 2 * s + 0] >>  4) | ((sc[4 + 2 * s + 0] >> 6) << 4);
+                sc2  = (sc[8 + 2 * s + 1] & 0xF) | ((sc[2 * s + 1] >> 6) << 4);
+                m2_u = (sc[8 + 2 * s + 1] >>  4) | ((sc[4 + 2 * s + 1] >> 6) << 4);
+            }
+
+            HVX_Vector v_ql = *(const HVX_UVector *) (src[i].qs + step * 32);
+            HVX_Vector v_lo = Q6_V_vand_VV(v_ql, Q6_Vb_vsplat_R(0x0F));
+            HVX_Vector v_hi = Q6_Vub_vlsr_VubR(v_ql, 4);
+
+            HVX_Vector v_h1 = Q6_Vw_vasl_VwR(Q6_V_vand_VV(Q6_Vub_vlsr_VubR(v_qh, 2 * step), Q6_Vb_vsplat_R(1)), 4);
+            HVX_Vector v_val1 = Q6_V_vor_VV(v_lo, v_h1);
+
+            HVX_VectorPair p16_1 = Q6_Wuh_vunpack_Vub(v_val1);
+            HVX_VectorPair p32_1 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_1));
+            HVX_Vector     v_i32_1 = Q6_V_lo_W(p32_1);
+            HVX_Vector     v_f32_1 = Q6_Vsf_equals_Vw(v_i32_1);
+
+            HVX_Vector vd1 = hvx_vec_splat_f32(d * sc1);
+            HVX_Vector vm1 = hvx_vec_splat_f32(min * m1_u);
+            HVX_Vector v_out1 = hvx_vec_sub_f32_f32(hvx_vec_mul_f32_f32(v_f32_1, vd1), vm1);
+            hvx_vmem(y + 0) = v_out1;
+
+            HVX_Vector v_h2 = Q6_Vw_vasl_VwR(Q6_V_vand_VV(Q6_Vub_vlsr_VubR(v_qh, 2 * step + 1), Q6_Vb_vsplat_R(1)), 4);
+            HVX_Vector v_val2 = Q6_V_vor_VV(v_hi, v_h2);
+
+            HVX_VectorPair p16_2 = Q6_Wuh_vunpack_Vub(v_val2);
+            HVX_VectorPair p32_2 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_2));
+            HVX_Vector     v_i32_2 = Q6_V_lo_W(p32_2);
+            HVX_Vector     v_f32_2 = Q6_Vsf_equals_Vw(v_i32_2);
+
+            HVX_Vector vd2 = hvx_vec_splat_f32(d * sc2);
+            HVX_Vector vm2 = hvx_vec_splat_f32(min * m2_u);
+            HVX_Vector v_out2 = hvx_vec_sub_f32_f32(hvx_vec_mul_f32_f32(v_f32_2, vd2), vm2);
+            hvx_vmem(y + 32) = v_out2;
+
+            y += 64;
+        }
+    }
+}
+
+static inline void hvx_dequantize_row_flat_q6_k_f32(float * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) {
+    const int nb = n / QK_K;
+    const block_q6_K * src = (const block_q6_K *) src_ptr;
+
+    const HVX_VectorPred q64 = Q6_Q_vsetq_R(64);
+    const HVX_Vector v32_b   = Q6_Vb_vsplat_R(32);
+    const HVX_Vector mask_0f = Q6_Vb_vsplat_R(0x0F);
+    const HVX_Vector mask_03 = Q6_Vb_vsplat_R(0x03);
+
+    for (int i = 0; i < nb; i++) {
+        const float d = GGML_FP16_TO_FP32(src[i].d);
+
+        const uint8_t * restrict ql = src[i].ql;
+        const uint8_t * restrict qh = src[i].qh;
+        const int8_t  * restrict sc = src[i].scales;
+
+        float * y = dst_ptr + i * QK_K;
+
+        for (int half = 0; half < 2; ++half) {
+            float s[8];
+            for (int j = 0; j < 8; ++j) {
+                s[j] = d * sc[j];
+            }
+
+            HVX_Vector vscale1 = Q6_V_vmux_QVV(q64, hvx_vec_splat_f32(s[0]), hvx_vec_splat_f32(s[1]));
+            HVX_Vector vscale2 = Q6_V_vmux_QVV(q64, hvx_vec_splat_f32(s[2]), hvx_vec_splat_f32(s[3]));
+            HVX_Vector vscale3 = Q6_V_vmux_QVV(q64, hvx_vec_splat_f32(s[4]), hvx_vec_splat_f32(s[5]));
+            HVX_Vector vscale4 = Q6_V_vmux_QVV(q64, hvx_vec_splat_f32(s[6]), hvx_vec_splat_f32(s[7]));
+
+            HVX_Vector v_ql0 = *(const HVX_UVector *) (ql);
+            HVX_Vector v_ql1 = *(const HVX_UVector *) (ql + 32);
+            HVX_Vector v_qh  = *(const HVX_UVector *) (qh);
+
+            HVX_Vector ql0_lo = Q6_V_vand_VV(v_ql0, mask_0f);
+            HVX_Vector qh_0   = Q6_Vw_vasl_VwR(Q6_V_vand_VV(v_qh, mask_03), 4);
+            HVX_Vector v_q1   = Q6_Vb_vsub_VbVb(Q6_V_vor_VV(ql0_lo, qh_0), v32_b);
+
+            HVX_Vector ql1_lo = Q6_V_vand_VV(v_ql1, mask_0f);
+            HVX_Vector qh_1   = Q6_Vw_vasl_VwR(Q6_V_vand_VV(Q6_Vub_vlsr_VubR(v_qh, 2), mask_03), 4);
+            HVX_Vector v_q2   = Q6_Vb_vsub_VbVb(Q6_V_vor_VV(ql1_lo, qh_1), v32_b);
+
+            HVX_Vector ql0_hi = Q6_Vub_vlsr_VubR(v_ql0, 4);
+            HVX_Vector qh_2   = Q6_Vw_vasl_VwR(Q6_V_vand_VV(Q6_Vub_vlsr_VubR(v_qh, 4), mask_03), 4);
+            HVX_Vector v_q3   = Q6_Vb_vsub_VbVb(Q6_V_vor_VV(ql0_hi, qh_2), v32_b);
+
+            HVX_Vector ql1_hi = Q6_Vub_vlsr_VubR(v_ql1, 4);
+            HVX_Vector qh_3   = Q6_Vw_vasl_VwR(Q6_V_vand_VV(Q6_Vub_vlsr_VubR(v_qh, 6), mask_03), 4);
+            HVX_Vector v_q4   = Q6_Vb_vsub_VbVb(Q6_V_vor_VV(ql1_hi, qh_3), v32_b);
+
+            HVX_VectorPair p16_1 = Q6_Wh_vunpack_Vb(v_q1);
+            HVX_VectorPair p32_1 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_1));
+            HVX_Vector     v_f32_1 = Q6_Vsf_equals_Vw(Q6_V_lo_W(p32_1));
+            hvx_vmem(y + 0) = hvx_vec_mul_f32_f32(v_f32_1, vscale1);
+
+            HVX_VectorPair p16_2 = Q6_Wh_vunpack_Vb(v_q2);
+            HVX_VectorPair p32_2 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_2));
+            HVX_Vector     v_f32_2 = Q6_Vsf_equals_Vw(Q6_V_lo_W(p32_2));
+            hvx_vmem(y + 32) = hvx_vec_mul_f32_f32(v_f32_2, vscale2);
+
+            HVX_VectorPair p16_3 = Q6_Wh_vunpack_Vb(v_q3);
+            HVX_VectorPair p32_3 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_3));
+            HVX_Vector     v_f32_3 = Q6_Vsf_equals_Vw(Q6_V_lo_W(p32_3));
+            hvx_vmem(y + 64) = hvx_vec_mul_f32_f32(v_f32_3, vscale3);
+
+            HVX_VectorPair p16_4 = Q6_Wh_vunpack_Vb(v_q4);
+            HVX_VectorPair p32_4 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_4));
+            HVX_Vector     v_f32_4 = Q6_Vsf_equals_Vw(Q6_V_lo_W(p32_4));
+            hvx_vmem(y + 96) = hvx_vec_mul_f32_f32(v_f32_4, vscale4);
+
+            y  += 128;
+            ql += 64;
+            qh += 32;
+            sc += 8;
+        }
+    }
+}
+
+// bytes row, row + 32, row + 64 and row + 96 of v, to bytes 0..3
+static inline HVX_Vector hvx_dequantize_gather4(HVX_Vector v, uint32_t row) {
+    v = Q6_V_vror_VR(v, row);
+    HVX_Vector packed = Q6_V_vmux_QVV(Q6_Q_vsetq_R(1), v, Q6_V_vror_VR(v, 31));
+    packed = Q6_V_vmux_QVV(Q6_Q_vsetq_R(2), packed, Q6_V_vror_VR(v, 62));
+    return Q6_V_vmux_QVV(Q6_Q_vsetq_R(3), packed, Q6_V_vror_VR(v, 93));
+}
+
+static inline void hvx_dequantize_row_tiled_q8_0_f32(float * dst_spad, const uint8_t * src_spad, uint32_t n_k_tiles, uint32_t tile_stride, uint32_t row) {
+    const HVX_VectorPred first2 = Q6_Q_vsetq_R(2);
+    const HVX_VectorPred first4 = Q6_Q_vsetq_R(4);
+
+    for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) {
+        const uint8_t * tile = src_spad + k_tile * tile_stride;
+        float * dst = dst_spad + k_tile * HTP_TENSOR_TILE_N_COLS;
+
+        HVX_Vector vq = Q6_V_vzero();
+        for (int group = 7; group >= 0; --group) {
+            const HVX_Vector v = Q6_V_vror_VR(hvx_vmem(tile + group * VLEN), 2 * row);
+            const HVX_Vector packed = Q6_V_vmux_QVV(first2, v, Q6_V_vror_VR(v, 62));
+            vq = Q6_V_vmux_QVV(first4, packed, Q6_V_vror_VR(vq, VLEN - 4));
+        }
+
+        const HVX_Vector scales = hvx_vmem(tile + 1024);
+        const HVX_Vector scale_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 2 * row));
+        const HVX_Vector scale = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hf));
+        const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq);
+        const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16));
+        HVX_Vector values = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale);
+        *(HVX_Vector *) dst = values;
+    }
+}
+
+static inline void hvx_dequantize_row_tiled_q4_0_f32(float * dst_spad, const uint8_t * src_spad, uint32_t n_k_tiles, uint32_t tile_stride, uint32_t row) {
+    const HVX_VectorPred first4 = Q6_Q_vsetq_R(4);
+
+    for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) {
+        const uint8_t * tile = src_spad + k_tile * tile_stride;
+        float * dst = dst_spad + k_tile * HTP_TENSOR_TILE_N_COLS;
+
+        HVX_Vector vq = Q6_V_vzero();
+        for (int group = 3; group >= 0; --group) {
+            vq = Q6_V_vmux_QVV(first4, hvx_dequantize_gather4(hvx_vmem(tile + group * VLEN), row), Q6_V_vror_VR(vq, VLEN - 4));
+        }
+
+        const HVX_Vector lo = Q6_V_vand_VV(vq, Q6_Vb_vsplat_R(0x0F));
+        const HVX_Vector hi = Q6_Vub_vlsr_VubR(vq, 4);
+        vq = Q6_V_lo_W(Q6_W_vshuff_VVR(hi, lo, -1));
+        vq = Q6_Vb_vsub_VbVb(vq, Q6_Vb_vsplat_R(8));
+
+        const HVX_Vector scales = hvx_vmem(tile + 512);
+        const HVX_Vector scale_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 2 * row));
+        const HVX_Vector scale = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hf));
+        const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq);
+        const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16));
+        HVX_Vector values = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale);
+        *(HVX_Vector *) dst = values;
+    }
+}
+
+static inline void hvx_dequantize_row_tiled_q4_k_f32(float * dst_spad, const uint8_t * src_spad, uint32_t n_k_tiles, uint32_t tile_stride, uint32_t row) {
+    const HVX_VectorPred first4 = Q6_Q_vsetq_R(4);
+
+    for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) {
+        const uint8_t * tile = src_spad + k_tile * tile_stride;
+        float * dst = dst_spad + k_tile * HTP_TENSOR_TILE_N_COLS;
+
+        HVX_Vector vq = Q6_V_vzero();
+        for (int group = 3; group >= 0; --group) {
+            vq = Q6_V_vmux_QVV(first4, hvx_dequantize_gather4(hvx_vmem(tile + group * VLEN), row), Q6_V_vror_VR(vq, VLEN - 4));
+        }
+
+        const HVX_Vector lo = Q6_V_vand_VV(vq, Q6_Vb_vsplat_R(0x0F));
+        const HVX_Vector hi = Q6_Vub_vlsr_VubR(vq, 4);
+        vq = Q6_V_lo_W(Q6_W_vshuff_VVR(hi, lo, -1));
+
+        const HVX_Vector scales = hvx_vmem(tile + 512);
+        const HVX_Vector scale_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 4 * row));
+        const HVX_Vector scale = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hf));
+        const HVX_Vector offset_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 4 * row + 2));
+        const HVX_Vector offset = Q6_V_lo_W(hvx_vec_f16_to_f32(offset_hf));
+
+        const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq);
+        const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16));
+        HVX_Vector values = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale);
+        *(HVX_Vector *) dst = hvx_vec_add_f32_f32(values, offset);
+    }
+}
+
+static inline void hvx_dequantize_row_tiled_q5_k_f32(float * dst_spad, const uint8_t * src_spad, uint32_t n_k_tiles, uint32_t tile_stride, uint32_t row) {
+    const HVX_VectorPred first4 = Q6_Q_vsetq_R(4);
+    const HVX_Vector one = Q6_Vb_vsplat_R(1);
+
+    for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) {
+        const uint8_t * tile = src_spad + k_tile * tile_stride;
+        float * dst = dst_spad + k_tile * HTP_TENSOR_TILE_N_COLS;
+
+        const HVX_Vector p5 = hvx_dequantize_gather4(hvx_vmem(tile + 640), row);
+        HVX_Vector vq = Q6_V_vzero();
+        HVX_Vector vh = Q6_V_vzero();
+        for (int group = 3; group >= 0; --group) {
+            vq = Q6_V_vmux_QVV(first4, hvx_dequantize_gather4(hvx_vmem(tile + group * VLEN), row), Q6_V_vror_VR(vq, VLEN - 4));
+            vh = Q6_V_vmux_QVV(first4, Q6_Vub_vlsr_VubR(p5, 2 * group), Q6_V_vror_VR(vh, VLEN - 4));
+        }
+
+        const HVX_Vector lo = Q6_V_vor_VV(Q6_V_vand_VV(vq, Q6_Vb_vsplat_R(0x0F)), Q6_Vw_vasl_VwR(Q6_V_vand_VV(vh, one), 4));
+        const HVX_Vector hi = Q6_V_vor_VV(Q6_Vub_vlsr_VubR(vq, 4), Q6_Vw_vasl_VwR(Q6_V_vand_VV(Q6_Vub_vlsr_VubR(vh, 1), one), 4));
+        vq = Q6_V_lo_W(Q6_W_vshuff_VVR(hi, lo, -1));
+
+        const HVX_Vector scales = hvx_vmem(tile + 512);
+        const HVX_Vector scale_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 4 * row));
+        const HVX_Vector scale = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hf));
+        const HVX_Vector offset_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 4 * row + 2));
+        const HVX_Vector offset = Q6_V_lo_W(hvx_vec_f16_to_f32(offset_hf));
+
+        const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq);
+        const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16));
+        HVX_Vector values = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale);
+        *(HVX_Vector *) dst = hvx_vec_add_f32_f32(values, offset);
+    }
+}
+
+static inline void hvx_dequantize_row_tiled_q6_k_f32(float * dst_spad, const uint8_t * src_spad, uint32_t n_k_tiles, uint32_t tile_stride, uint32_t row) {
+    const HVX_VectorPred first4 = Q6_Q_vsetq_R(4);
+    const HVX_VectorPred first16 = Q6_Q_vsetq_R(16 * sizeof(float));
+    const HVX_Vector mask_0f = Q6_Vb_vsplat_R(0x0F);
+    const HVX_Vector mask_03 = Q6_Vb_vsplat_R(0x03);
+
+    for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) {
+        const uint8_t * tile = src_spad + k_tile * tile_stride;
+        float * dst = dst_spad + k_tile * HTP_TENSOR_TILE_N_COLS;
+
+        HVX_Vector vq = Q6_V_vzero();
+        for (int group = 7; group >= 0; --group) {
+            const HVX_Vector lo_plane = Q6_V_vror_VR(hvx_vmem(tile + (group >> 1) * VLEN), 4 * row);
+            const HVX_Vector hi_plane = Q6_V_vror_VR(hvx_vmem(tile + 512 + (group >> 2) * VLEN), 4 * row);
+            const HVX_Vector lo = (group & 1) ? Q6_Vub_vlsr_VubR(lo_plane, 4) : Q6_V_vand_VV(lo_plane, mask_0f);
+            const HVX_Vector hi = Q6_Vub_vlsr_VubR(hi_plane, 2 * (group & 3));
+            const HVX_Vector packed = Q6_V_vor_VV(lo, Q6_Vw_vasl_VwR(Q6_V_vand_VV(hi, mask_03), 4));
+            vq = Q6_V_vmux_QVV(first4, packed, Q6_V_vror_VR(vq, VLEN - 4));
+        }
+
+        const HVX_Vector scales = hvx_vmem(tile + 768);
+        const HVX_Vector scale_lo_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 2 * row));
+        const HVX_Vector scale_hi_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 64 + 2 * row));
+        const HVX_Vector scale_lo = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_lo_hf));
+        const HVX_Vector scale_hi = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hi_hf));
+        const HVX_Vector scale = Q6_V_vmux_QVV(first16, scale_lo, scale_hi);
+        const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(Q6_Vb_vsub_VbVb(vq, Q6_Vb_vsplat_R(32)));
+        const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16));
+        *(HVX_Vector *) dst = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale);
+    }
+}
+
+static inline HVX_Vector hvx_dequantize_tiled_f16_gather4(const uint8_t * p, uint32_t r_byte) {
+    HVX_Vector w0  = Q6_V_vror_VR(hvx_vmem(p + 0 * VLEN), r_byte);
+    HVX_Vector w1  = Q6_V_vror_VR(hvx_vmem(p + 1 * VLEN), r_byte);
+    HVX_Vector w01 = Q6_V_lo_W(Q6_W_vshuff_VVR(w1, w0, 4));
+    HVX_Vector w2  = Q6_V_vror_VR(hvx_vmem(p + 2 * VLEN), r_byte);
+    HVX_Vector w3  = Q6_V_vror_VR(hvx_vmem(p + 3 * VLEN), r_byte);
+    HVX_Vector w23 = Q6_V_lo_W(Q6_W_vshuff_VVR(w3, w2, 4));
+    return Q6_V_lo_W(Q6_W_vshuff_VVR(w23, w01, 8));
+}
+
+static inline void hvx_dequantize_row_tiled_f16_f32(float * dst_spad, const uint8_t * src_spad, uint32_t n_k_tiles, uint32_t tile_stride, uint32_t row) {
+    const uint32_t r_byte = 4 * row;
+
+    for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) {
+        const uint8_t * tile = src_spad + k_tile * tile_stride;
+        float * dst = dst_spad + k_tile * HTP_TENSOR_TILE_N_COLS;
+
+        HVX_Vector g0 = hvx_dequantize_tiled_f16_gather4(tile + 0 * VLEN, r_byte);
+        HVX_Vector g1 = hvx_dequantize_tiled_f16_gather4(tile + 4 * VLEN, r_byte);
+        HVX_Vector g01 = Q6_V_lo_W(Q6_W_vshuff_VVR(g1, g0, 16));
+
+        HVX_Vector g2 = hvx_dequantize_tiled_f16_gather4(tile + 8 * VLEN, r_byte);
+        HVX_Vector g3 = hvx_dequantize_tiled_f16_gather4(tile + 12 * VLEN, r_byte);
+        HVX_Vector g23 = Q6_V_lo_W(Q6_W_vshuff_VVR(g3, g2, 16));
+
+        HVX_Vector g = Q6_V_lo_W(Q6_W_vshuff_VVR(g23, g01, 32));
+        *(HVX_Vector *) dst = Q6_V_lo_W(hvx_vec_f16_to_f32(g));
+    }
+}

 #endif // HVX_QUANT_H
@@ -28,8 +28,8 @@
 #include "ggml-common.h"
 #include "hex-bitmap.h"
 #include "htp-ctx.h"
+#include "htp-dcache.h"
 #include "htp-ops.h"
-#include "htp-tensor.h"
 #include "htp_iface.h"
 #include "work-queue.h"
 #include "hex-profile.h"
@@ -835,7 +835,7 @@ static int op_fence(struct htp_ops_context * octx) {
    atomic_uint * sync_fence = (atomic_uint *) (uintptr_t) sync->data;

    if (mode == 1) {
-        htp_flush_dirty_ranges(ctx);
+        htp_dcache_flush_dirty_ranges(ctx);

        htp_mdev_group_barrier(octx);

@@ -1211,7 +1211,7 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_buf_desc * bufs
            src->ne[0], src->ne[1], src->ne[2], src->ne[3]);
    }

-    htp_tensor_flush_all(octx->ctx, octx->src, HTP_OP_MAX_INPUTS);
+    htp_dcache_flush_dirty_tensors(octx->ctx, octx->src, HTP_OP_MAX_INPUTS);

    // Prep output tensors
    for (uint32_t i = 0; i < HTP_OP_MAX_OUTPUTS; i++) {
@@ -1227,7 +1227,7 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_buf_desc * bufs
            dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
    }

-    htp_tensor_dirty_all(octx->ctx, octx->dsts, HTP_OP_MAX_OUTPUTS);
+    htp_dcache_mark_dirty_tensors(octx->ctx, octx->dsts, HTP_OP_MAX_OUTPUTS);

    htp_mdev_group_barrier(octx);

@@ -4,6 +4,7 @@
 #include <stdint.h>
 #include <stddef.h>
 #include "htp-ops.h"
+#include "htp-tensor.h"
 #include "hex-fastdiv.h"
 #include "hex-common.h"
 #include "htp-vtcm.h"
@@ -19,45 +20,10 @@ extern "C" {
 #define HTP_MM_HMX_TILE_N_ELMS 1024
 #define HTP_MM_HMX_MIN_NROWS   4

-// --- Weight Repacked Tile Sizes ---
-#define HTP_MM_WEIGHT_TILE_SIZE_Q4_0   576
-#define HTP_MM_WEIGHT_TILE_SIZE_Q4_1   640
-#define HTP_MM_WEIGHT_TILE_SIZE_Q8_0   1088
-#define HTP_MM_WEIGHT_TILE_SIZE_IQ4_NL 576
-#define HTP_MM_WEIGHT_TILE_SIZE_MXFP4  544
-// Q5_K: the Q4_1 tile (640) followed by a 128-byte plane with the 5th bit of every quant, transposed so that
-//   plane byte l holds the eight flags of lane l: bit 2i = low nibble of nibble vector i, bit 2i+1 = high nibble
-#define HTP_MM_WEIGHT_TILE_SIZE_Q5_K   768
-// Q6_K native 6-bit tile (32 rows x 32 k), vrmpy-ready: byte 4*row+b of a vector holds k = 4*group+b
-//   vectors 0..3: low nibbles, vector i holds group 2i (low nibble) and group 2i+1 (high nibble)
-//   vectors 4..5: high 2 bits, vector m holds groups 4m..4m+3 at bit offsets 0,2,4,6
-//   vector 6: fp16 scales per row, d * scales[]: k 0..15 in lanes 0..31, k 16..31 in lanes 32..63
-#define HTP_MM_WEIGHT_TILE_SIZE_Q6_K   896
-// Q3_K native 3-bit tile, vrmpy-ready like Q6_K
-//   vectors 0..1: low 2 bits, vector m holds groups 4m..4m+3 at bit offsets 0,2,4,6
-//   vector 2: bit g set where the hmask bit of group g is clear (quant = low 2 bits - 4)
-//   vector 3: fp16 scales per row, d * (scales[] - 32): k 0..15 in lanes 0..31, k 16..31 in lanes 32..63
-#define HTP_MM_WEIGHT_TILE_SIZE_Q3_K   512
-// Q2_K native 2-bit tile with the ggml scale bytes, vrmpy-ready like Q6_K (same bits per weight as block_q2_K)
-//   bytes   0..255: unsigned 2-bit quants, vector m holds groups 4m..4m+3 at bit offsets 0,2,4,6
-//   bytes 256..319: scales[] bytes (min << 4 | scale), sub-block s of row r at 32 * s + r
-//   bytes 320..335: 8 fp16, tile j of a super-block: d (j < 4) or -dmin (j >= 4) of rows 8 * (j % 4) .. 8 * (j % 4) + 7
-#define HTP_MM_WEIGHT_TILE_SIZE_Q2_K   336
-
-// --- Weight Repacked Aligned Tile Sizes ---
-#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q4_0   640
-#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q4_1   640
-#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q8_0   1152
-#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_IQ4_NL 640
-#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_MXFP4  640
-#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q5_K   768
-#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q6_K   896
-#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q3_K   512
-#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q2_K   384
-
 // --- Activation Tiled Block Sizes (including padding) ---
 #define HTP_MM_ACT_TILE_SIZE_Q8_0      1152
 #define HTP_MM_ACT_TILE_SIZE_Q8_1      1280
+#define HTP_MM_ACT_TILE_SIZE_F16       2048

 #define HTP_MM_MAX_PREFETCH 16

@@ -74,11 +40,12 @@ enum htp_mm_kernel_type {

    // HMX paths
    HTP_MM_KERNEL_HMX_2D,
-    HTP_MM_KERNEL_HMX_F16_BATCHED,
+    HTP_MM_KERNEL_HMX_BATCHED,

    // HVX floating-point paths
    HTP_MM_KERNEL_HVX_F16_F16_VTCM,
    HTP_MM_KERNEL_HVX_F32_F32_VTCM,
+    HTP_MM_KERNEL_HVX_BF16_F32_VTCM,

    // HVX quantized paths
    HTP_MM_KERNEL_HVX_QUANT_ROW,      // standard row-wise parallel quantization
@@ -212,57 +179,6 @@ next_nc:
    return 0;
 }

-// --- Tile Size Helpers ---
-static inline uint32_t htp_mm_get_weight_tile_size(int weight_type) {
-    switch (weight_type) {
-        case HTP_TYPE_Q4_0:
-        case HTP_TYPE_IQ4_NL:
-            return HTP_MM_WEIGHT_TILE_SIZE_Q4_0;
-        case HTP_TYPE_Q4_1:
-        case HTP_TYPE_Q4_K:
-            return HTP_MM_WEIGHT_TILE_SIZE_Q4_1;
-        case HTP_TYPE_Q8_0:
-            return HTP_MM_WEIGHT_TILE_SIZE_Q8_0;
-        case HTP_TYPE_Q5_K:
-            return HTP_MM_WEIGHT_TILE_SIZE_Q5_K;
-        case HTP_TYPE_Q6_K:
-            return HTP_MM_WEIGHT_TILE_SIZE_Q6_K;
-        case HTP_TYPE_Q3_K:
-            return HTP_MM_WEIGHT_TILE_SIZE_Q3_K;
-        case HTP_TYPE_Q2_K:
-            return HTP_MM_WEIGHT_TILE_SIZE_Q2_K;
-        case HTP_TYPE_MXFP4:
-            return HTP_MM_WEIGHT_TILE_SIZE_MXFP4;
-        default:
-            return 0;
-    }
-}
-
-static inline uint32_t htp_mm_get_weight_aligned_tile_size(int weight_type) {
-    switch (weight_type) {
-        case HTP_TYPE_Q4_0:
-        case HTP_TYPE_IQ4_NL:
-            return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q4_0;
-        case HTP_TYPE_Q4_1:
-        case HTP_TYPE_Q4_K:
-            return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q4_1;
-        case HTP_TYPE_Q8_0:
-            return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q8_0;
-        case HTP_TYPE_Q5_K:
-            return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q5_K;
-        case HTP_TYPE_Q6_K:
-            return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q6_K;
-        case HTP_TYPE_Q3_K:
-            return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q3_K;
-        case HTP_TYPE_Q2_K:
-            return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q2_K;
-        case HTP_TYPE_MXFP4:
-            return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_MXFP4;
-        default:
-            return 0;
-    }
-}
-
 // weight types whose tiles carry a per-block offset (x = d * q + m): the activations need block sums (q8_1)
 // (Q2_K: per-16 k sums, q8_1_s16)
 static inline bool htp_mm_weight_has_offset(int weight_type) {
@@ -283,27 +199,17 @@ static inline size_t htp_mm_q8_1_tiled_row_size(uint32_t ne) {
    return nb_32 * HTP_MM_ACT_TILE_SIZE_Q8_1;
 }

-static inline size_t htp_mm_get_tiled_row_stride(int weight_type, uint32_t k) {
-    uint32_t nb = (k + QK_Q4_0_TILED - 1) / QK_Q4_0_TILED;
-    switch (weight_type) {
-        case HTP_TYPE_Q4_0:
-        case HTP_TYPE_IQ4_NL:
-        case HTP_TYPE_Q4_1:
-        case HTP_TYPE_Q4_K:
-        case HTP_TYPE_Q8_0:
-        case HTP_TYPE_Q5_K:
-        case HTP_TYPE_Q6_K:
-        case HTP_TYPE_Q3_K:
-        case HTP_TYPE_Q2_K:
-        case HTP_TYPE_MXFP4:
-            return (size_t) nb * htp_mm_get_weight_tile_size(weight_type);
-        case HTP_TYPE_F16:
-            return (size_t) k * sizeof(__fp16);
-        case HTP_TYPE_F32:
-            return (size_t) k * sizeof(float);
-        default:
-            return 0;
+static inline size_t htp_mm_f16_tiled_row_size(uint32_t ne) {
+    const uint32_t ne_padded = ((ne + 127) / 128) * 128;
+    const uint32_t nb_32 = ne_padded / 32;
+    return nb_32 * HTP_MM_ACT_TILE_SIZE_F16;
+}
+
+static inline size_t htp_mm_tiled_act_row_size(int weight_type, uint32_t ne) {
+    if (weight_type == HTP_TYPE_F16) {
+        return htp_mm_f16_tiled_row_size(ne);
    }
+    return htp_mm_weight_has_offset(weight_type) ? htp_mm_q8_1_tiled_row_size(ne) : htp_mm_q8_0_tiled_row_size(ne);
 }

 static inline size_t htp_mm_round_up(size_t n, size_t m) {
@@ -318,8 +224,8 @@ static inline void htp_mm_hmx_get_2d_chunk_costs(
    int wtype, uint32_t k, bool pipeline, uint32_t aligned_tile_size,
    size_t * size_per_n_out, size_t * size_per_m_out, size_t * size_per_mn_out
 ) {
-    const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32);
-    const size_t row_stride = htp_mm_get_tiled_row_stride(wtype, k);
+    const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32 && wtype != HTP_TYPE_BF16);
+    const size_t row_stride = htp_tensor_get_row_size(wtype, k);
    const size_t vec_dot_size = k * sizeof(uint16_t);
    const uint32_t n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS;
    const size_t qweight_row_stride = is_quant ? (size_t)(n_k_tiles * aligned_tile_size) / 32 : 0;
@@ -404,7 +310,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
 ) {
    size_t off = 0;

-    if (kernel_type == HTP_MM_KERNEL_HMX_F16_BATCHED) {
+    if (kernel_type == HTP_MM_KERNEL_HMX_BATCHED) {
        const size_t vec_dot_size     = k * sizeof(uint16_t);
        const size_t act_head_stride   = mc * k;
        const size_t weight_area_size  = hex_align_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
@@ -450,8 +356,8 @@ static inline void htp_mm_hmx_vtcm_layout_build(
        off = off_group_a + hex_smax(group_b_size, group_c_size);
    } else {
        // HTP_MM_KERNEL_HMX_2D
-        const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32);
-        const size_t row_stride = htp_mm_get_tiled_row_stride(wtype, k);
+        const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32 && wtype != HTP_TYPE_BF16);
+        const size_t row_stride = htp_tensor_get_row_size(wtype, k);
        const size_t vec_dot_size = k * sizeof(uint16_t);
        const uint32_t n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS;

@@ -527,28 +433,13 @@ static inline void htp_mm_hvx_vtcm_layout_build(
    size_t dst_sz     = 0;
    size_t act_raw_sz = 0;

-    const bool is_repack = (wtype == HTP_TYPE_Q4_0 || wtype == HTP_TYPE_Q4_1 ||
-                            wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL ||
-                            wtype == HTP_TYPE_MXFP4 || wtype == HTP_TYPE_Q6_K ||
-                            wtype == HTP_TYPE_Q4_K || wtype == HTP_TYPE_Q5_K ||
-                            wtype == HTP_TYPE_Q3_K || wtype == HTP_TYPE_Q2_K);
-
    if (is_fused_nx) {
-        const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128);
+        uint32_t aligned_tile_size = htp_tensor_get_tile_padded_size(wtype);
+        uint32_t n_k_tiles = hex_round_up(ne10, 32) / 32;
+        uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
+        size_t weight_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128);

-        size_t weight_sz_per_thread = 0;
-
-        if (is_repack) {
-            uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
-            uint32_t n_k_tiles = hex_round_up(ne10, 32) / 32;
-            uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
-
-            weight_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128);
-        } else {
-            weight_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128);
-        }
-
-        size_t tiled_act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
+        size_t tiled_act_row_size = htp_mm_tiled_act_row_size(wtype, ne10);
        size_t q_act_sz = hex_round_up(tiled_act_row_size * act_nrows, 128);
        size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));

@@ -558,20 +449,13 @@ static inline void htp_mm_hvx_vtcm_layout_build(
        dst_sz     = 0;
        act_raw_sz = hex_round_up(raw_row_size * act_nrows, 128);
    } else if (is_matmul_id) {
-        const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128);
-        const size_t act_row_size_tiled = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10)
-                                                                          : htp_mm_q8_0_tiled_row_size(ne10);
+        const size_t act_row_size_tiled = htp_mm_tiled_act_row_size(wtype, ne10);
+        const uint32_t aligned_tile_size = htp_tensor_get_tile_padded_size(wtype);
+        const uint32_t n_k_tiles         = ne10 / 32;
+        const uint32_t tile_row_size     = n_k_tiles * aligned_tile_size;
+        size_t src0_sz_per_thread        = htp_mm_round_up(n_prefetch * tile_row_size, 256);

-        size_t src0_sz_per_thread = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
-        act_sz                    = htp_mm_round_up(act_row_size_tiled * act_nrows, 256);
-
-        if (is_repack) {
-            const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
-            const uint32_t n_k_tiles         = ne10 / 32;
-            const uint32_t tile_row_size     = n_k_tiles * aligned_tile_size;
-            size_t repacked_vtcm_size        = htp_mm_round_up(n_prefetch * tile_row_size, 256);
-            src0_sz_per_thread               = repacked_vtcm_size;
-        }
+        act_sz = htp_mm_round_up(act_row_size_tiled * act_nrows, 256);

        size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));

@@ -592,7 +476,8 @@ static inline void htp_mm_hvx_vtcm_layout_build(
                act_raw_sz = hex_round_up(hex_round_up(ne10 * sizeof(float), 128) * act_nrows, 128);
                break;
            }
-            case HTP_MM_KERNEL_HVX_F32_F32_VTCM: {
+            case HTP_MM_KERNEL_HVX_F32_F32_VTCM:
+            case HTP_MM_KERNEL_HVX_BF16_F32_VTCM: {
                size_t f32_act_row_size = htp_mm_round_up(ne10 * 4, 128);
                act_sz     = htp_mm_round_up(f32_act_row_size * act_nrows, 256);
                src0_sz    = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads;
@@ -602,21 +487,15 @@ static inline void htp_mm_hvx_vtcm_layout_build(
            }
            case HTP_MM_KERNEL_HVX_QUANT_BLOCK:
            case HTP_MM_KERNEL_HVX_QUANT_ROW: {
-                size_t q_act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
+                size_t q_act_row_size = htp_mm_tiled_act_row_size(wtype, ne10);
+                uint32_t aligned_tile_size = htp_tensor_get_tile_padded_size(wtype);
+                uint32_t n_k_tiles = (ne10 + 31) / 32;
+                uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
+                size_t tiled_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);

-                src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
+                src0_sz = tiled_vtcm_size * n_threads;
                act_sz  = htp_mm_round_up(q_act_row_size * act_nrows, 256);

-                src0_sz = src0_sz * n_threads;
-
-                if (is_repack) {
-                    uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
-                    uint32_t n_k_tiles = ne10 / 32;
-                    uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
-                    size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
-                    src0_sz = repacked_vtcm_size * n_threads;
-                }
-
                size_t dst_slice_per_thread = (dst_nrows > 0 && act_nrows == 1) ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0;
                dst_sz = dst_slice_per_thread * n_threads;
                size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
@@ -687,9 +566,7 @@ static inline bool htp_mm_hvx_solve_vtcm_params(
    const size_t avail_act = vtcm_budget - fixed_bytes;
    size_t row_size = 0;
    if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
-        row_size = htp_mm_weight_has_offset(wtype)
-                 ? htp_mm_q8_1_tiled_row_size(ne10)
-                 : htp_mm_q8_0_tiled_row_size(ne10);
+        row_size = htp_mm_tiled_act_row_size(wtype, ne10);
    } else if (kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM) {
        row_size = hex_round_up(ne10 * 2, 128);
    } else {
@@ -749,7 +626,7 @@ static inline size_t htp_mm_hmx_get_batched_vtcm_size(
    int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool pipeline, uint32_t act_threads, size_t bias_size) {
    (void)pipeline;
    struct htp_mm_hmx_vtcm_layout L;
-    htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, false, act_threads, 0, bias_size);
+    htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_BATCHED, wtype, k, mc, nc, group_size, false, act_threads, 0, bias_size);
    return L.total_bytes;
 }

@@ -45,6 +45,7 @@ set(METALLIB_KERNEL_SOURCES
    kernels/fa_q5_0.metal
    kernels/fa_q5_1.metal
    kernels/fa_q8_0.metal
+    kernels/fa_tensor.metal
    kernels/fa_vec_f16.metal
    kernels/fa_vec_f32.metal
    kernels/fa_vec_q4_0.metal
@@ -233,12 +234,12 @@ else()
            VERBATIM
        )

-        set(AIR_FA_TENSOR "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/fa_f16_tensor.air")
+        set(AIR_FA_TENSOR "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/fa_tensor_tensor.air")
        add_custom_command(
            OUTPUT ${AIR_FA_TENSOR}
-            COMMAND xcrun -sdk ${METAL_SDK} metal ${XC_FLAGS_TENSOR} -DGGML_METAL_HAS_TENSOR -I ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/fa_f16.metal -o ${AIR_FA_TENSOR}
-            DEPENDS kernels/fa_f16.metal ${METALLIB_KERNELS_FA_SHARED} kernels/common.h kernels/dequantize.h ${METALLIB_COMMON} ggml-metal-impl.h
-            COMMENT "Compiling kernels/fa_f16.metal (tensor API)"
+            COMMAND xcrun -sdk ${METAL_SDK} metal ${XC_FLAGS_TENSOR} -DGGML_METAL_HAS_TENSOR -I ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/fa_tensor.metal -o ${AIR_FA_TENSOR}
+            DEPENDS kernels/fa_tensor.metal kernels/common.h ${METALLIB_COMMON} ggml-metal-impl.h
+            COMMENT "Compiling kernels/fa_tensor.metal (tensor API)"
            VERBATIM
        )

@@ -1866,7 +1866,9 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_t
    const int32_t dk = (int32_t) op->src[1]->ne[0];
    const int32_t dv = (int32_t) op->src[2]->ne[0];

-    snprintf(base, 256, "kernel_flash_attn_ext_tensor_f16_dk%d_dv%d", dk, dv);
+    const char * type = op->src[1]->type == GGML_TYPE_BF16 ? "bf16" : "f16";
+
+    snprintf(base, 256, "kernel_flash_attn_ext_tensor_%s_dk%d_dv%d", type, dk, dv);
    snprintf(name, 256, "%s_mask=%d_sinks=%d_bias=%d_scap=%d", base, has_mask, has_sinks, has_bias, has_scap);

    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
@@ -118,6 +118,7 @@ int ggml_metal_pipeline_max_theads_per_threadgroup(struct ggml_metal_pipeline_wi
    X(FA_Q5_0,         fa_q5_0)        \
    X(FA_Q5_1,         fa_q5_1)        \
    X(FA_Q8_0,         fa_q8_0)        \
+    X(FA_TENSOR,       fa_tensor)      \
    X(FA_VEC_F16,      fa_vec_f16)     \
    X(FA_VEC_F32,      fa_vec_f32)     \
    X(FA_VEC_Q4_0,     fa_vec_q4_0)    \
@@ -3161,7 +3161,7 @@ static bool ggml_metal_op_flash_attn_ext_use_tensor(const ggml_tensor * op, bool
        return false;
    }

-    if (op->src[1]->type != GGML_TYPE_F16 && !ggml_metal_op_flash_attn_ext_use_kv_f16(op)) {
+    if (op->src[1]->type != GGML_TYPE_F16 && op->src[1]->type != GGML_TYPE_BF16 && !ggml_metal_op_flash_attn_ext_use_kv_f16(op)) {
        return false;
    }

@@ -3678,7 +3678,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
            /*.logit_softcap =*/ logit_softcap,
        };

-        // shared memory layout: queries (half), scores (float), probabilities (half), row scale (float), rescale flag (int)
+        // shared memory layout: queries (16-bit), scores (float), probabilities (16-bit), row scale (float), rescale flag (int)
        const size_t smem = GGML_PAD(nqptg*ne00*sizeof(ggml_fp16_t) + nqptg*ncpsg*(sizeof(float) + sizeof(ggml_fp16_t)) + nqptg*sizeof(float) + sizeof(int32_t), 16);

        auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_tensor(lib, op, has_mask, has_sinks, has_bias, has_scap);
@@ -75,291 +75,3 @@ template [[host_name("kernel_flash_attn_ext_bf16_dk576_dv512")]] kernel flash_at
 #undef FA_TYPES
 #undef FA_TYPES_BF
 #undef FA_TYPES_F32
-
-#ifdef GGML_METAL_HAS_TENSOR
-
-constant bool FC_flash_attn_ext_tensor_has_mask  [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 0)]];
-constant bool FC_flash_attn_ext_tensor_has_sinks [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 1)]];
-constant bool FC_flash_attn_ext_tensor_has_bias  [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 2)]];
-constant bool FC_flash_attn_ext_tensor_has_scap  [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 3)]];
-
-// ref: https://arxiv.org/pdf/2307.08691.pdf
-template<
-    short DK,                                   // K head size
-    short DV,                                   // V head size
-    short Q   = OP_FLASH_ATTN_EXT_TENSOR_NQPSG, // queries per threadgroup
-    short C   = OP_FLASH_ATTN_EXT_TENSOR_NCPSG, // cache items per threadgroup
-    short NSG = OP_FLASH_ATTN_EXT_TENSOR_NSG>   // number of simd groups
-kernel void kernel_flash_attn_ext_tensor(
-        constant ggml_metal_kargs_flash_attn_ext & args,
-        device const char * q,
-        device const char * k,
-        device const char * v,
-        device const char * mask,
-        device const char * sinks,
-        device const char * blk,
-        device       char * dst,
-        threadgroup  char * shmem [[threadgroup(0)]],
-        uint3  tgpig [[threadgroup_position_in_grid]],
-        ushort tiisg [[thread_index_in_simdgroup]],
-        ushort sgitg [[simdgroup_index_in_threadgroup]]) {
-    constexpr short NW = N_SIMDWIDTH;
-    constexpr short NT = NW*NSG;
-    constexpr short NQ = Q/NSG;
-    constexpr short NC = C/NW; // columns per thread
-
-    static_assert(DK % 4 == 0,  "DK must be divisible by 4");
-    static_assert(Q % NSG == 0, "Q must be divisible by NSG");
-    static_assert(C % NW == 0,  "C must be divisible by NW");
-
-    const int iq3 = tgpig[2];
-    const int iq2 = tgpig[1];
-    const int iq1 = tgpig[0]*Q;
-
-    const short tiitg = sgitg*NW + tiisg;
-
-    threadgroup half  * sq = (threadgroup half  *) shmem;       // [Q, DK] queries
-    threadgroup float * ss = (threadgroup float *) (sq + Q*DK); // [Q, C]  scores
-    threadgroup half  * sp = (threadgroup half  *) (ss + Q*C);  // [Q, C]  probabilities
-    threadgroup float * sr = (threadgroup float *) (sp + Q*C);  // [Q]     per-row scale of O
-    threadgroup int   * sf = (threadgroup int   *) (sr + Q);    // [1]     last iteration (ic0 + 1) that rescaled O
-
-    q += iq1*args.nb01 + iq2*args.nb02 + iq3*args.nb03;
-
-    {
-        const int ikv2 = iq2/(args.ne02/args.ne_12_2);
-        const int ikv3 = iq3/(args.ne03/args.ne_12_3);
-
-        k += ikv2*args.nb12 + ikv3*args.nb13;
-        v += ikv2*args.nb22 + ikv3*args.nb23;
-    }
-
-    // with softcap the scale is small (scale/softcap), so it is applied to the scores to keep the precision of Q
-    const float qscale = FC_flash_attn_ext_tensor_has_scap ? 1.0f : args.scale;
-
-    // load the queries, with the scale folded in
-    for (int i = tiitg; i < Q*DK/4; i += NT) {
-        const int j = i/(DK/4);
-
-        float4 q4 = 0.0f;
-        if (iq1 + j < args.ne01) {
-            q4 = ((device const float4 *) (q + j*args.nb01))[i%(DK/4)];
-        }
-
-        ((threadgroup half4 *) sq)[i] = (half4) (q4*qscale);
-    }
-
-    device const half * pm[NQ];
-
-    FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
-        const short j = jj*NSG + sgitg;
-
-        pm[jj] = (device const half *) (mask + (iq1 + j)*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33);
-    }
-
-    {
-        const int nblk1 = (args.ne01 + Q - 1)/Q;
-        const int nblk0 = (args.ne11 + C - 1)/C;
-
-        blk += (((iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*nblk1 + iq1/Q)*nblk0;
-    }
-
-    float M[NQ];
-    float S[NQ];
-
-    FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
-        M[jj] = -FLT_MAX/2;
-        S[jj] = 0.0f;
-    }
-
-    float slope = 1.0f;
-
-    // ALiBi
-    if (FC_flash_attn_ext_tensor_has_bias) {
-        const short h = iq2;
-
-        const float base = h < args.n_head_log2 ? args.m0 : args.m1;
-        const short exph = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
-
-        slope = pow(base, exph);
-    }
-
-    const int sk = args.ns10;
-    const int sv = args.ns20;
-
-    auto tq = tensor(sq, dextents<int32_t, 2>(DK, Q));
-    auto ts = tensor(ss, dextents<int32_t, 2>(C,  Q));
-    auto tp = tensor(sp, dextents<int32_t, 2>(C,  Q));
-
-    mpp::tensor_ops::matmul2d<
-        mpp::tensor_ops::matmul2d_descriptor(Q, C, DK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply),
-        execution_simdgroups<NSG>> mm_qk;
-
-    mpp::tensor_ops::matmul2d<
-        mpp::tensor_ops::matmul2d_descriptor(Q, DV, C, false, false, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
-        execution_simdgroups<NSG>> mm_pv;
-
-    auto tv0 = tensor((device half *) v, dextents<int32_t, 2>(DV, C), array<int, 2>({1, sv}));
-
-    // the O matrix from the paper
-    auto co = mm_pv.template get_destination_cooperative_tensor<decltype(tp), decltype(tv0), float>();
-
-    FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) {
-        if (co.is_valid_element(i)) {
-            co[i] = 0.0f;
-        }
-    }
-
-    if (tiitg == 0) {
-        sf[0] = 0;
-    }
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    // the host guarantees ne11 % C == 0
-    for (int ic0 = 0, ic = 0; ic < args.ne11; ++ic0, ic += C) {
-        char blk_cur = 1;
-
-        if (FC_flash_attn_ext_tensor_has_mask) {
-            blk_cur = blk[ic0];
-
-            if (blk_cur == 0) {
-                continue;
-            }
-        }
-
-        // Q*K^T
-        {
-            auto tk = tensor((device half *) (k + (uint64_t) ic*args.nb11), dextents<int32_t, 2>(DK, C), array<int, 2>({1, sk}));
-
-            mm_qk.run(tq, tk, ts);
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        // online softmax
-        FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
-            const short j = jj*NSG + sgitg;
-
-            float s[NC];
-
-            FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
-                s[ii] = ss[j*C + ii*NW + tiisg];
-            }
-
-            if (FC_flash_attn_ext_tensor_has_scap) {
-                FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
-                    s[ii] = args.logit_softcap*precise::tanh(s[ii]*args.scale);
-                }
-            }
-
-            if (FC_flash_attn_ext_tensor_has_mask && blk_cur != 2 && iq1 + j < args.ne31) {
-                FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
-                    s[ii] += slope*(float) pm[jj][ic + ii*NW + tiisg];
-                }
-            }
-
-            float m = M[jj];
-
-            FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
-                m = max(m, s[ii]);
-            }
-
-            m = simd_max(m);
-
-            // lazy rescaling: move the running max only when it grows by more than 8 (e^8 fits in half)
-            float ms = 1.0f;
-
-            if (m > M[jj] + 8.0f) {
-                ms    = exp(M[jj] - m);
-                M[jj] = m;
-
-                if (tiisg == 0) {
-                    sf[0] = ic0 + 1;
-                }
-            }
-
-            float sum = 0.0f;
-
-            FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
-                // the sum uses the same rounded values as P*V
-                const half p = (half) exp(s[ii] - M[jj]);
-
-                sp[j*C + ii*NW + tiisg] = p;
-
-                sum += (float) p;
-            }
-
-            S[jj] = S[jj]*ms + simd_sum(sum);
-
-            if (tiisg == 0) {
-                sr[j] = ms;
-            }
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        // O = diag(ms)*O + P*V
-        if (sf[0] == ic0 + 1) {
-            FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) {
-                if (co.is_valid_element(i)) {
-                    co[i] *= sr[co.get_multidimensional_index(i)[1]];
-                }
-            }
-        }
-
-        {
-            auto tv = tensor((device half *) (v + (uint64_t) ic*args.nb21), dextents<int32_t, 2>(DV, C), array<int, 2>({1, sv}));
-
-            mm_pv.run(tp, tv, co);
-        }
-
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-    }
-
-    FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
-        const short j = jj*NSG + sgitg;
-
-        // the sink only adds to the denominator - its rescale of O is folded into the final scale
-        float ms = 1.0f;
-
-        if (FC_flash_attn_ext_tensor_has_sinks) {
-            const float s = ((device const float *) sinks)[iq2];
-            const float m = max(M[jj], s);
-
-            ms = exp(M[jj] - m);
-
-            S[jj] = S[jj]*ms + exp(s - m);
-        }
-
-        if (tiisg == 0) {
-            sr[j] = S[jj] == 0.0f ? 0.0f : ms/S[jj];
-        }
-    }
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) {
-        if (co.is_valid_element(i)) {
-            co[i] *= sr[co.get_multidimensional_index(i)[1]];
-        }
-    }
-
-    // store to global memory - rows past ne01 are clipped by the tensor extents
-    device float * pdst = (device float *) dst + ((uint64_t) iq3*args.ne2*args.ne1 + iq2 + (uint64_t) iq1*args.ne1)*DV;
-
-    auto td = tensor(pdst, dextents<int32_t, 2>(DV, args.ne01 - iq1), array<int, 2>({1, args.ne1*DV}));
-
-    co.store(td);
-}
-
-typedef decltype(kernel_flash_attn_ext_tensor<64, 64>) flash_attn_ext_tensor_t;
-
-template [[host_name("kernel_flash_attn_ext_tensor_f16_dk64_dv64"  )]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<64,  64>;
-template [[host_name("kernel_flash_attn_ext_tensor_f16_dk128_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<128, 128>;
-template [[host_name("kernel_flash_attn_ext_tensor_f16_dk192_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<192, 128>;
-template [[host_name("kernel_flash_attn_ext_tensor_f16_dk256_dv256")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<256, 256>;
-template [[host_name("kernel_flash_attn_ext_tensor_f16_dk512_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<512, 512, OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE>;
-template [[host_name("kernel_flash_attn_ext_tensor_f16_dk576_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<576, 512, OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE>;
-
-#endif // GGML_METAL_HAS_TENSOR
@@ -0,0 +1,298 @@
+#include "common.h"
+
+#ifdef GGML_METAL_HAS_TENSOR
+
+constant bool FC_flash_attn_ext_tensor_has_mask  [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 0)]];
+constant bool FC_flash_attn_ext_tensor_has_sinks [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 1)]];
+constant bool FC_flash_attn_ext_tensor_has_bias  [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 2)]];
+constant bool FC_flash_attn_ext_tensor_has_scap  [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 3)]];
+
+// ref: https://arxiv.org/pdf/2307.08691.pdf
+template<
+    typename T,                                 // Q, K, V and P types
+    typename T4,
+    short DK,                                   // K head size
+    short DV,                                   // V head size
+    short Q   = OP_FLASH_ATTN_EXT_TENSOR_NQPSG, // queries per threadgroup
+    short C   = OP_FLASH_ATTN_EXT_TENSOR_NCPSG, // cache items per threadgroup
+    short NSG = OP_FLASH_ATTN_EXT_TENSOR_NSG>   // number of simd groups
+kernel void kernel_flash_attn_ext_tensor(
+        constant ggml_metal_kargs_flash_attn_ext & args,
+        device const char * q,
+        device const char * k,
+        device const char * v,
+        device const char * mask,
+        device const char * sinks,
+        device const char * blk,
+        device       char * dst,
+        threadgroup  char * shmem [[threadgroup(0)]],
+        uint3  tgpig [[threadgroup_position_in_grid]],
+        ushort tiisg [[thread_index_in_simdgroup]],
+        ushort sgitg [[simdgroup_index_in_threadgroup]]) {
+    constexpr short NW = N_SIMDWIDTH;
+    constexpr short NT = NW*NSG;
+    constexpr short NQ = Q/NSG;
+    constexpr short NC = C/NW; // columns per thread
+
+    static_assert(DK % 4 == 0,  "DK must be divisible by 4");
+    static_assert(Q % NSG == 0, "Q must be divisible by NSG");
+    static_assert(C % NW == 0,  "C must be divisible by NW");
+
+    const int iq3 = tgpig[2];
+    const int iq2 = tgpig[1];
+    const int iq1 = tgpig[0]*Q;
+
+    const short tiitg = sgitg*NW + tiisg;
+
+    threadgroup T     * sq = (threadgroup T     *) shmem;       // [Q, DK] queries
+    threadgroup float * ss = (threadgroup float *) (sq + Q*DK); // [Q, C]  scores
+    threadgroup T     * sp = (threadgroup T     *) (ss + Q*C);  // [Q, C]  probabilities
+    threadgroup float * sr = (threadgroup float *) (sp + Q*C);  // [Q]     per-row scale of O
+    threadgroup int   * sf = (threadgroup int   *) (sr + Q);    // [1]     last iteration (ic0 + 1) that rescaled O
+
+    q += iq1*args.nb01 + iq2*args.nb02 + iq3*args.nb03;
+
+    {
+        const int ikv2 = iq2/(args.ne02/args.ne_12_2);
+        const int ikv3 = iq3/(args.ne03/args.ne_12_3);
+
+        k += ikv2*args.nb12 + ikv3*args.nb13;
+        v += ikv2*args.nb22 + ikv3*args.nb23;
+    }
+
+    // with softcap the scale is small (scale/softcap), so it is applied to the scores to keep the precision of Q
+    const float qscale = FC_flash_attn_ext_tensor_has_scap ? 1.0f : args.scale;
+
+    // load the queries, with the scale folded in
+    for (int i = tiitg; i < Q*DK/4; i += NT) {
+        const int j = i/(DK/4);
+
+        float4 q4 = 0.0f;
+        if (iq1 + j < args.ne01) {
+            q4 = ((device const float4 *) (q + j*args.nb01))[i%(DK/4)];
+        }
+
+        ((threadgroup T4 *) sq)[i] = (T4) (q4*qscale);
+    }
+
+    device const half * pm[NQ];
+
+    FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
+        const short j = jj*NSG + sgitg;
+
+        pm[jj] = (device const half *) (mask + (iq1 + j)*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33);
+    }
+
+    {
+        const int nblk1 = (args.ne01 + Q - 1)/Q;
+        const int nblk0 = (args.ne11 + C - 1)/C;
+
+        blk += (((iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*nblk1 + iq1/Q)*nblk0;
+    }
+
+    float M[NQ];
+    float S[NQ];
+
+    FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
+        M[jj] = -FLT_MAX/2;
+        S[jj] = 0.0f;
+    }
+
+    float slope = 1.0f;
+
+    // ALiBi
+    if (FC_flash_attn_ext_tensor_has_bias) {
+        const short h = iq2;
+
+        const float base = h < args.n_head_log2 ? args.m0 : args.m1;
+        const short exph = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
+
+        slope = pow(base, exph);
+    }
+
+    const int sk = args.ns10;
+    const int sv = args.ns20;
+
+    auto tq = tensor(sq, dextents<int32_t, 2>(DK, Q));
+    auto ts = tensor(ss, dextents<int32_t, 2>(C,  Q));
+    auto tp = tensor(sp, dextents<int32_t, 2>(C,  Q));
+
+    mpp::tensor_ops::matmul2d<
+        mpp::tensor_ops::matmul2d_descriptor(Q, C, DK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply),
+        execution_simdgroups<NSG>> mm_qk;
+
+    mpp::tensor_ops::matmul2d<
+        mpp::tensor_ops::matmul2d_descriptor(Q, DV, C, false, false, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
+        execution_simdgroups<NSG>> mm_pv;
+
+    auto tv0 = tensor((device T *) v, dextents<int32_t, 2>(DV, C), array<int, 2>({1, sv}));
+
+    // the O matrix from the paper
+    auto co = mm_pv.template get_destination_cooperative_tensor<decltype(tp), decltype(tv0), float>();
+
+    FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) {
+        if (co.is_valid_element(i)) {
+            co[i] = 0.0f;
+        }
+    }
+
+    if (tiitg == 0) {
+        sf[0] = 0;
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // the host guarantees ne11 % C == 0
+    for (int ic0 = 0, ic = 0; ic < args.ne11; ++ic0, ic += C) {
+        char blk_cur = 1;
+
+        if (FC_flash_attn_ext_tensor_has_mask) {
+            blk_cur = blk[ic0];
+
+            if (blk_cur == 0) {
+                continue;
+            }
+        }
+
+        // Q*K^T
+        {
+            auto tk = tensor((device T *) (k + (uint64_t) ic*args.nb11), dextents<int32_t, 2>(DK, C), array<int, 2>({1, sk}));
+
+            mm_qk.run(tq, tk, ts);
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // online softmax
+        FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
+            const short j = jj*NSG + sgitg;
+
+            float s[NC];
+
+            FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
+                s[ii] = ss[j*C + ii*NW + tiisg];
+            }
+
+            if (FC_flash_attn_ext_tensor_has_scap) {
+                FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
+                    s[ii] = args.logit_softcap*precise::tanh(s[ii]*args.scale);
+                }
+            }
+
+            if (FC_flash_attn_ext_tensor_has_mask && blk_cur != 2 && iq1 + j < args.ne31) {
+                FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
+                    s[ii] += slope*(float) pm[jj][ic + ii*NW + tiisg];
+                }
+            }
+
+            float m = M[jj];
+
+            FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
+                m = max(m, s[ii]);
+            }
+
+            m = simd_max(m);
+
+            // lazy rescaling: move the running max only when it grows by more than 8 (e^8 fits in half)
+            float ms = 1.0f;
+
+            if (m > M[jj] + 8.0f) {
+                ms    = exp(M[jj] - m);
+                M[jj] = m;
+
+                if (tiisg == 0) {
+                    sf[0] = ic0 + 1;
+                }
+            }
+
+            float sum = 0.0f;
+
+            FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
+                // the sum uses the same rounded values as P*V
+                const T p = (T) exp(s[ii] - M[jj]);
+
+                sp[j*C + ii*NW + tiisg] = p;
+
+                sum += (float) p;
+            }
+
+            S[jj] = S[jj]*ms + simd_sum(sum);
+
+            if (tiisg == 0) {
+                sr[j] = ms;
+            }
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // O = diag(ms)*O + P*V
+        if (sf[0] == ic0 + 1) {
+            FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) {
+                if (co.is_valid_element(i)) {
+                    co[i] *= sr[co.get_multidimensional_index(i)[1]];
+                }
+            }
+        }
+
+        {
+            auto tv = tensor((device T *) (v + (uint64_t) ic*args.nb21), dextents<int32_t, 2>(DV, C), array<int, 2>({1, sv}));
+
+            mm_pv.run(tp, tv, co);
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
+        const short j = jj*NSG + sgitg;
+
+        // the sink only adds to the denominator - its rescale of O is folded into the final scale
+        float ms = 1.0f;
+
+        if (FC_flash_attn_ext_tensor_has_sinks) {
+            const float s = ((device const float *) sinks)[iq2];
+            const float m = max(M[jj], s);
+
+            ms = exp(M[jj] - m);
+
+            S[jj] = S[jj]*ms + exp(s - m);
+        }
+
+        if (tiisg == 0) {
+            sr[j] = S[jj] == 0.0f ? 0.0f : ms/S[jj];
+        }
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) {
+        if (co.is_valid_element(i)) {
+            co[i] *= sr[co.get_multidimensional_index(i)[1]];
+        }
+    }
+
+    // store to global memory - rows past ne01 are clipped by the tensor extents
+    device float * pdst = (device float *) dst + ((uint64_t) iq3*args.ne2*args.ne1 + iq2 + (uint64_t) iq1*args.ne1)*DV;
+
+    auto td = tensor(pdst, dextents<int32_t, 2>(DV, args.ne01 - iq1), array<int, 2>({1, args.ne1*DV}));
+
+    co.store(td);
+}
+
+typedef decltype(kernel_flash_attn_ext_tensor<half, half4, 64, 64>) flash_attn_ext_tensor_t;
+
+template [[host_name("kernel_flash_attn_ext_tensor_f16_dk64_dv64"  )]]  kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<half, half4, 64,  64>;
+template [[host_name("kernel_flash_attn_ext_tensor_f16_dk128_dv128")]]  kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<half, half4, 128, 128>;
+template [[host_name("kernel_flash_attn_ext_tensor_f16_dk192_dv128")]]  kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<half, half4, 192, 128>;
+template [[host_name("kernel_flash_attn_ext_tensor_f16_dk256_dv256")]]  kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<half, half4, 256, 256>;
+template [[host_name("kernel_flash_attn_ext_tensor_f16_dk512_dv512")]]  kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<half, half4, 512, 512, OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE>;
+template [[host_name("kernel_flash_attn_ext_tensor_f16_dk576_dv512")]]  kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<half, half4, 576, 512, OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE>;
+
+template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk64_dv64"  )]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<bfloat, bfloat4, 64,  64>;
+template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk128_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<bfloat, bfloat4, 128, 128>;
+template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk192_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<bfloat, bfloat4, 192, 128>;
+template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk256_dv256")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<bfloat, bfloat4, 256, 256>;
+template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk512_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<bfloat, bfloat4, 512, 512, OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE>;
+template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk576_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<bfloat, bfloat4, 576, 512, OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE>;
+
+#endif // GGML_METAL_HAS_TENSOR
@@ -175,6 +175,12 @@ __kernel void flash_attn_f16(
            ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale;
            ACC_TYPE s2 = (dot_acc2.s0 + dot_acc2.s1 + dot_acc2.s2 + dot_acc2.s3) * scale;
            ACC_TYPE s3 = (dot_acc3.s0 + dot_acc3.s1 + dot_acc3.s2 + dot_acc3.s3) * scale;
+            if (logit_softcap > 0.0f) {
+                s0 = logit_softcap * tanh(s0 / logit_softcap);
+                s1 = logit_softcap * tanh(s1 / logit_softcap);
+                s2 = logit_softcap * tanh(s2 / logit_softcap);
+                s3 = logit_softcap * tanh(s3 / logit_softcap);
+            }

            if (is_causal) {
                const int causal_limit = n_kv - n_q + my_query_row;
@@ -196,13 +202,6 @@ __kernel void flash_attn_f16(
                if (k_row3 < n_kv) s3 += slope * (ACC_TYPE)mask_ptr[k_row3];
            }

-            if (logit_softcap > 0.0f) {
-                s0 = logit_softcap * tanh(s0 / logit_softcap);
-                s1 = logit_softcap * tanh(s1 / logit_softcap);
-                s2 = logit_softcap * tanh(s2 / logit_softcap);
-                s3 = logit_softcap * tanh(s3 / logit_softcap);
-            }
-
            const ACC_TYPE m_new      = max(m_i, max(max(s0, s1), max(s2, s3)));
            const ACC_TYPE scale_prev = native_exp(m_i - m_new);
            const ACC_TYPE p0         = native_exp(s0 - m_new);
@@ -331,13 +330,13 @@ __kernel void flash_attn_f16_q1(
            dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc);
        }
        ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }
        if (mask_base != NULL) {
            const global DATA_TYPE* mask_ptr = (const global DATA_TYPE*)(mask_base);
            score += slope * (ACC_TYPE)mask_ptr[k_idx];
        }
-        if (logit_softcap > 0.0f) {
-            score = logit_softcap * tanh(score / logit_softcap);
-        }
        m_i = max(m_i, score);
    }

@@ -367,13 +366,13 @@ __kernel void flash_attn_f16_q1(
            dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc);
        }
        ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }
        if (mask_base != NULL) {
            const global DATA_TYPE* mask_ptr = (const global DATA_TYPE*)(mask_base);
            score += slope * (ACC_TYPE)mask_ptr[k_idx];
        }
-        if (logit_softcap > 0.0f) {
-            score = logit_softcap * tanh(score / logit_softcap);
-        }
        const ACC_TYPE p = exp(score - m_final);
        l_i += p;
        FA_UNROLL
@@ -176,6 +176,12 @@ __kernel void flash_attn_f32(
            ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale;
            ACC_TYPE s2 = (dot_acc2.s0 + dot_acc2.s1 + dot_acc2.s2 + dot_acc2.s3) * scale;
            ACC_TYPE s3 = (dot_acc3.s0 + dot_acc3.s1 + dot_acc3.s2 + dot_acc3.s3) * scale;
+            if (logit_softcap > 0.0f) {
+                s0 = logit_softcap * tanh(s0 / logit_softcap);
+                s1 = logit_softcap * tanh(s1 / logit_softcap);
+                s2 = logit_softcap * tanh(s2 / logit_softcap);
+                s3 = logit_softcap * tanh(s3 / logit_softcap);
+            }

            if (is_causal) {
                const int causal_limit = n_kv - n_q + my_query_row;
@@ -197,13 +203,6 @@ __kernel void flash_attn_f32(
                if (k_row3 < n_kv) s3 += slope * (ACC_TYPE)mask_ptr[k_row3];
            }

-            if (logit_softcap > 0.0f) {
-                s0 = logit_softcap * tanh(s0 / logit_softcap);
-                s1 = logit_softcap * tanh(s1 / logit_softcap);
-                s2 = logit_softcap * tanh(s2 / logit_softcap);
-                s3 = logit_softcap * tanh(s3 / logit_softcap);
-            }
-
            const ACC_TYPE m_new      = max(m_i, max(max(s0, s1), max(s2, s3)));
            const ACC_TYPE scale_prev = native_exp(m_i - m_new);
            const ACC_TYPE p0         = native_exp(s0 - m_new);
@@ -332,13 +331,13 @@ __kernel void flash_attn_f32_q1(
            dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc);
        }
        ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }
        if (mask_base != NULL) {
            const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base);
            score += slope * (ACC_TYPE)mask_ptr[k_idx];
        }
-        if (logit_softcap > 0.0f) {
-            score = logit_softcap * tanh(score / logit_softcap);
-        }
        m_i = max(m_i, score);
    }

@@ -368,13 +367,13 @@ __kernel void flash_attn_f32_q1(
            dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc);
        }
        ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }
        if (mask_base != NULL) {
            const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base);
            score += slope * (ACC_TYPE)mask_ptr[k_idx];
        }
-        if (logit_softcap > 0.0f) {
-            score = logit_softcap * tanh(score / logit_softcap);
-        }
        const ACC_TYPE p = exp(score - m_final);
        l_i += p;
        FA_UNROLL
@@ -336,6 +336,10 @@ __kernel void FA_TILE_NAME(

                ACC_TYPE score0 = partial0 * scale;
                ACC_TYPE score1 = partial1 * scale;
+                if (logit_softcap > 0.0f) {
+                    score0 = logit_softcap * tanh(score0 / logit_softcap);
+                    score1 = logit_softcap * tanh(score1 / logit_softcap);
+                }

                if (!query_valid) { score0 = FA_M_INIT; score1 = FA_M_INIT; }
                if (is_causal) {
@@ -359,11 +363,6 @@ __kernel void FA_TILE_NAME(
                    }
                }

-                if (logit_softcap > 0.0f) {
-                    score0 = logit_softcap * tanh(score0 / logit_softcap);
-                    score1 = logit_softcap * tanh(score1 / logit_softcap);
-                }
-
                const ACC_TYPE m_new = max(m_i, max(score0, score1));
                // Whole tile masked (m_new == FA_M_INIT): force the exp() args
                // far negative so the tile contributes 0, not exp(0)=1.
@@ -408,6 +407,9 @@ __kernel void FA_TILE_NAME(
                        score += local_partial[j][q_lane * N_SPLIT + s];
                    }
                    score *= scale;
+                    if (logit_softcap > 0.0f) {
+                        score = logit_softcap * tanh(score / logit_softcap);
+                    }

                    if (is_causal && k_row > (n_kv - n_q + my_query_row)) score = FA_M_INIT;
                    if (k_row >= n_kv) score = FA_M_INIT;
@@ -424,10 +426,6 @@ __kernel void FA_TILE_NAME(
                        }
                    }

-                    if (logit_softcap > 0.0f) {
-                        score = logit_softcap * tanh(score / logit_softcap);
-                    }
-
                    m_new = max(m_new, score);
                    local_p[q_lane][j] = score;
                }
@@ -502,6 +500,12 @@ __kernel void FA_TILE_NAME(
                ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale;
                ACC_TYPE s2 = (dot_acc2.s0 + dot_acc2.s1 + dot_acc2.s2 + dot_acc2.s3) * scale;
                ACC_TYPE s3 = (dot_acc3.s0 + dot_acc3.s1 + dot_acc3.s2 + dot_acc3.s3) * scale;
+                if (logit_softcap > 0.0f) {
+                    s0 = logit_softcap * tanh(s0 / logit_softcap);
+                    s1 = logit_softcap * tanh(s1 / logit_softcap);
+                    s2 = logit_softcap * tanh(s2 / logit_softcap);
+                    s3 = logit_softcap * tanh(s3 / logit_softcap);
+                }

                if (is_causal) {
                    const int causal_limit = n_kv - n_q + my_query_row;
@@ -531,13 +535,6 @@ __kernel void FA_TILE_NAME(
                    }
                }

-                if (logit_softcap > 0.0f) {
-                    s0 = logit_softcap * tanh(s0 / logit_softcap);
-                    s1 = logit_softcap * tanh(s1 / logit_softcap);
-                    s2 = logit_softcap * tanh(s2 / logit_softcap);
-                    s3 = logit_softcap * tanh(s3 / logit_softcap);
-                }
-
                const ACC_TYPE m_new      = max(m_i, max(max(s0, s1), max(s2, s3)));
                // Whole tile masked (m_new == FA_M_INIT): force the exp() args
                // far negative so the tile contributes 0, not exp(0)=1.
@@ -748,13 +745,13 @@ __kernel void flash_attn_f32_f16_q1(
            dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc);
        }
        ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }
        if (mask_base != NULL) {
            const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base);
            score += slope * (ACC_TYPE)mask_ptr[k_idx];
        }
-        if (logit_softcap > 0.0f) {
-            score = logit_softcap * tanh(score / logit_softcap);
-        }
        m_i = max(m_i, score);
    }

@@ -776,13 +773,13 @@ __kernel void flash_attn_f32_f16_q1(
            dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc);
        }
        ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }
        if (mask_base != NULL) {
            const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base);
            score += slope * (ACC_TYPE)mask_ptr[k_idx];
        }
-        if (logit_softcap > 0.0f) {
-            score = logit_softcap * tanh(score / logit_softcap);
-        }
        const ACC_TYPE p = exp(score - m_final);
        l_i += p;
        FA_UNROLL
@@ -1122,13 +1119,13 @@ __kernel void flash_attn_f32_f16_q1_local_tile(

            if (tid == 0) {
                float s = red[0] * scale;
+                if (logit_softcap > 0.0f) {
+                    s = logit_softcap * tanh(s / logit_softcap);
+                }
                if (mask_base != NULL) {
                    const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base;
                    s += slope * (float) mask_ptr[kv_idx];
                }
-                if (logit_softcap > 0.0f) {
-                    s = logit_softcap * tanh(s / logit_softcap);
-                }
                score_shared = s;
            }
            barrier(CLK_LOCAL_MEM_FENCE);
@@ -1323,13 +1320,13 @@ __kernel void flash_attn_f32_f16_q1_local_mq_split(
                    contrib += q_shared[h * DK + d] * (float) k_tile[j * DK + d];
                }
                float s = sub_group_reduce_add(contrib) * scale;
+                if (logit_softcap > 0.0f) {
+                    s = logit_softcap * tanh(s / logit_softcap);
+                }
                if (mask_base[h] != NULL) {
                    const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
                    s += slope[h] * (float) mask_ptr[kv_idx];
                }
-                if (logit_softcap > 0.0f) {
-                    s = logit_softcap * tanh(s / logit_softcap);
-                }
                score[h] = s;
            }

@@ -1510,13 +1507,13 @@ __kernel void flash_attn_f32_f16_q1_vec_mq(
        for (int h = 0; h < MQ_GQA; ++h) {
            const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
            ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale;
+            if (logit_softcap > 0.0f) {
+                s = logit_softcap * tanh(s / logit_softcap);
+            }
            if (mask_base[h] != NULL) {
                const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
                s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
            }
-            if (logit_softcap > 0.0f) {
-                s = logit_softcap * tanh(s / logit_softcap);
-            }
            score[h] = s;
        }

@@ -2448,13 +2445,13 @@ __kernel void flash_attn_f32_f16_q1_vec_mq_split_k_img(
        for (int h = 0; h < MQ_GQA; ++h) {
            const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
            ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale;
+            if (logit_softcap > 0.0f) {
+                s = logit_softcap * tanh(s / logit_softcap);
+            }
            if (mask_base[h] != NULL) {
                const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
                s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
            }
-            if (logit_softcap > 0.0f) {
-                s = logit_softcap * tanh(s / logit_softcap);
-            }
            score[h] = s;
        }

@@ -2633,13 +2630,13 @@ __kernel void flash_attn_f32_f16_q1_split(
            dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc);
        }
        ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }
        if (mask_base != NULL) {
            const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base);
            score += slope * (ACC_TYPE) mask_ptr[k_idx];
        }
-        if (logit_softcap > 0.0f) {
-            score = logit_softcap * tanh(score / logit_softcap);
-        }
        m_i = max(m_i, score);
    }

@@ -2662,13 +2659,13 @@ __kernel void flash_attn_f32_f16_q1_split(
            dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc);
        }
        ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }
        if (mask_base != NULL) {
            const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base);
            score += slope * (ACC_TYPE) mask_ptr[k_idx];
        }
-        if (logit_softcap > 0.0f) {
-            score = logit_softcap * tanh(score / logit_softcap);
-        }
        const ACC_TYPE p = exp(score - m_c);
        l_i += p;
        #pragma unroll
@@ -287,14 +287,14 @@ __kernel void flash_attn_f32_q4_0_q1(
 #endif
        }
        score *= scale;
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }

        if (mask_base != NULL) {
            const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base);
            score += slope * (ACC_TYPE)mask_ptr[k_idx];
        }
-        if (logit_softcap > 0.0f) {
-            score = logit_softcap * tanh(score / logit_softcap);
-        }

        // Online softmax step.
        const ACC_TYPE m_new = max(m_i, score);
@@ -554,14 +554,14 @@ __kernel void flash_attn_f32_q4_0_q1_vec(
        ACC_TYPE dot_partial = dot4.s0 + dot4.s1 + dot4.s2 + dot4.s3;
        ACC_TYPE score = sub_group_reduce_add(dot_partial) * scale;
 #endif
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }

        if (mask_base != NULL) {
            const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base;
            score += slope * (ACC_TYPE) mask_ptr[k_idx];
        }
-        if (logit_softcap > 0.0f) {
-            score = logit_softcap * tanh(score / logit_softcap);
-        }

        const ACC_TYPE m_new      = max(m_i, score);
        const ACC_TYPE scale_prev = native_exp(m_i - m_new);
@@ -745,13 +745,13 @@ __kernel void flash_attn_f32_q4_0_q1_split(
 #endif
        }
        score *= scale;
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }
        if (mask_base != NULL) {
            const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base);
            score += slope * (ACC_TYPE) mask_ptr[k_idx];
        }
-        if (logit_softcap > 0.0f) {
-            score = logit_softcap * tanh(score / logit_softcap);
-        }

        // Online softmax step.
        const ACC_TYPE m_new = max(m_i, score);
@@ -1037,13 +1037,13 @@ __kernel void flash_attn_f32_q4_0_q1_vec_mq_split(
        #pragma unroll
        for (int h = 0; h < MQ_GQA; ++h) {
            ACC_TYPE s = sub_group_reduce_add(lane_contrib[h]) * scale;
+            if (logit_softcap > 0.0f) {
+                s = logit_softcap * tanh(s / logit_softcap);
+            }
            if (mask_base[h] != NULL) {
                const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
                s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
            }
-            if (logit_softcap > 0.0f) {
-                s = logit_softcap * tanh(s / logit_softcap);
-            }
            score[h] = s;
        }
 #else
@@ -1067,13 +1067,13 @@ __kernel void flash_attn_f32_q4_0_q1_vec_mq_split(
        for (int h = 0; h < MQ_GQA; ++h) {
            const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
            ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale;
+            if (logit_softcap > 0.0f) {
+                s = logit_softcap * tanh(s / logit_softcap);
+            }
            if (mask_base[h] != NULL) {
                const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
                s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
            }
-            if (logit_softcap > 0.0f) {
-                s = logit_softcap * tanh(s / logit_softcap);
-            }
            score[h] = s;
        }
 #endif
@@ -1385,13 +1385,13 @@ __kernel void flash_attn_f32_q4_0_q1_vec_mq_split_c8(
                s += sub_group_shuffle_xor(s, step);
            }
            s *= scale;
+            if (logit_softcap > 0.0f) {
+                s = logit_softcap * tanh(s / logit_softcap);
+            }
            if (mask_base[h] != NULL) {
                const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
                s += slope[h] * (ACC_TYPE) mask_ptr[k_safe];
            }
-            if (logit_softcap > 0.0f) {
-                s = logit_softcap * tanh(s / logit_softcap);
-            }
            score[h] = valid ? s : FA_M_INIT;
        }

@@ -1852,9 +1852,15 @@ __kernel void flash_attn_f32_q4_0(
                    s2 = sub_group_shuffle(s2, tri_base + 0) + sub_group_shuffle(s2, tri_base + 1) + sub_group_shuffle(s2, tri_base + 2);
                    s3 = sub_group_shuffle(s3, tri_base + 0) + sub_group_shuffle(s3, tri_base + 1) + sub_group_shuffle(s3, tri_base + 2);
                #endif
-                if (!query_valid) { s0 = FA_M_INIT; s1 = FA_M_INIT; s2 = FA_M_INIT; s3 = FA_M_INIT; }
 #endif
                s0 *= scale; s1 *= scale; s2 *= scale; s3 *= scale;
+                if (logit_softcap > 0.0f) {
+                    s0 = logit_softcap * tanh(s0 / logit_softcap);
+                    s1 = logit_softcap * tanh(s1 / logit_softcap);
+                    s2 = logit_softcap * tanh(s2 / logit_softcap);
+                    s3 = logit_softcap * tanh(s3 / logit_softcap);
+                }
+                if (!query_valid) { s0 = FA_M_INIT; s1 = FA_M_INIT; s2 = FA_M_INIT; s3 = FA_M_INIT; }

                if (is_causal) {
                    const int causal_limit = n_kv - n_q + my_query_row;
@@ -1876,12 +1882,6 @@ __kernel void flash_attn_f32_q4_0(
                    if (k_row2 < n_kv) s2 += slope * (ACC_TYPE) mask_ptr[k_row2];
                    if (k_row3 < n_kv) s3 += slope * (ACC_TYPE) mask_ptr[k_row3];
                }
-                if (logit_softcap > 0.0f) {
-                    s0 = logit_softcap * tanh(s0 / logit_softcap);
-                    s1 = logit_softcap * tanh(s1 / logit_softcap);
-                    s2 = logit_softcap * tanh(s2 / logit_softcap);
-                    s3 = logit_softcap * tanh(s3 / logit_softcap);
-                }

                const ACC_TYPE m_new      = max(m_i, max(max(s0, s1), max(s2, s3)));
                // Whole tile masked (m_new == FA_M_INIT): force the exp() args
@@ -232,14 +232,14 @@ __kernel void flash_attn_f32_q8_0_q1(
 #endif
        }
        score *= scale;
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }

        if (mask_base != NULL) {
            const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base);
            score += slope * (ACC_TYPE)mask_ptr[k_idx];
        }
-        if (logit_softcap > 0.0f) {
-            score = logit_softcap * tanh(score / logit_softcap);
-        }

        // Online softmax step.
        const ACC_TYPE m_new = max(m_i, score);
@@ -437,14 +437,14 @@ __kernel void flash_attn_f32_q8_0_q1_vec(
        }
        ACC_TYPE dot_partial = dot4.s0 + dot4.s1 + dot4.s2 + dot4.s3;
        ACC_TYPE score = sub_group_reduce_add(dot_partial) * scale;
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }

        if (mask_base != NULL) {
            const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base;
            score += slope * (ACC_TYPE) mask_ptr[k_idx];
        }
-        if (logit_softcap > 0.0f) {
-            score = logit_softcap * tanh(score / logit_softcap);
-        }

        const ACC_TYPE m_new      = max(m_i, score);
        const ACC_TYPE scale_prev = native_exp(m_i - m_new);
@@ -627,13 +627,13 @@ __kernel void flash_attn_f32_q8_0_q1_split(
 #endif
        }
        score *= scale;
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }
        if (mask_base != NULL) {
            const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base);
            score += slope * (ACC_TYPE) mask_ptr[k_idx];
        }
-        if (logit_softcap > 0.0f) {
-            score = logit_softcap * tanh(score / logit_softcap);
-        }

        // Online softmax step.
        const ACC_TYPE m_new = max(m_i, score);
@@ -882,13 +882,13 @@ __kernel void flash_attn_f32_q8_0_q1_vec_mq_split(
        for (int h = 0; h < MQ_GQA; ++h) {
            const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
            ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale;
+            if (logit_softcap > 0.0f) {
+                s = logit_softcap * tanh(s / logit_softcap);
+            }
            if (mask_base[h] != NULL) {
                const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
                s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
            }
-            if (logit_softcap > 0.0f) {
-                s = logit_softcap * tanh(s / logit_softcap);
-            }
            score[h] = s;
        }

@@ -1151,13 +1151,13 @@ __kernel void flash_attn_f32_q8_0_q1_vec_mq_split_c8(
                s += sub_group_shuffle_xor(s, step);
            }
            s *= scale;
+            if (logit_softcap > 0.0f) {
+                s = logit_softcap * tanh(s / logit_softcap);
+            }
            if (mask_base[h] != NULL) {
                const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
                s += slope[h] * (ACC_TYPE) mask_ptr[k_safe];
            }
-            if (logit_softcap > 0.0f) {
-                s = logit_softcap * tanh(s / logit_softcap);
-            }
            score[h] = valid ? s : FA_M_INIT;
        }

@@ -1652,9 +1652,15 @@ __kernel void flash_attn_f32_q8_0(
                    s2 = sub_group_shuffle(s2, tri_base + 0) + sub_group_shuffle(s2, tri_base + 1) + sub_group_shuffle(s2, tri_base + 2);
                    s3 = sub_group_shuffle(s3, tri_base + 0) + sub_group_shuffle(s3, tri_base + 1) + sub_group_shuffle(s3, tri_base + 2);
                #endif
-                if (!query_valid) { s0 = FA_M_INIT; s1 = FA_M_INIT; s2 = FA_M_INIT; s3 = FA_M_INIT; }
 #endif
                s0 *= scale; s1 *= scale; s2 *= scale; s3 *= scale;
+                if (logit_softcap > 0.0f) {
+                    s0 = logit_softcap * tanh(s0 / logit_softcap);
+                    s1 = logit_softcap * tanh(s1 / logit_softcap);
+                    s2 = logit_softcap * tanh(s2 / logit_softcap);
+                    s3 = logit_softcap * tanh(s3 / logit_softcap);
+                }
+                if (!query_valid) { s0 = FA_M_INIT; s1 = FA_M_INIT; s2 = FA_M_INIT; s3 = FA_M_INIT; }

                if (is_causal) {
                    const int causal_limit = n_kv - n_q + my_query_row;
@@ -1676,12 +1682,6 @@ __kernel void flash_attn_f32_q8_0(
                    if (k_row2 < n_kv) s2 += slope * (ACC_TYPE) mask_ptr[k_row2];
                    if (k_row3 < n_kv) s3 += slope * (ACC_TYPE) mask_ptr[k_row3];
                }
-                if (logit_softcap > 0.0f) {
-                    s0 = logit_softcap * tanh(s0 / logit_softcap);
-                    s1 = logit_softcap * tanh(s1 / logit_softcap);
-                    s2 = logit_softcap * tanh(s2 / logit_softcap);
-                    s3 = logit_softcap * tanh(s3 / logit_softcap);
-                }

                const ACC_TYPE m_new      = max(m_i, max(max(s0, s1), max(s2, s3)));
                // Whole tile masked (m_new == FA_M_INIT): force the exp() args
@@ -230,6 +230,8 @@ kernel void kernel_gemv_noshuffle_iq4_nl_f32(
    uint M = ne01;

    uint LINE_STRIDE_A = M / 2;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, LINE_STRIDE_A - 1);
    uint BLOCK_STRIDE_A = NSUBGROUPS * M;

    private uint4     regA;
@@ -240,7 +242,7 @@ kernel void kernel_gemv_noshuffle_iq4_nl_f32(

    // loop along K in block granularity, skip 4 blocks every iter
    for (uint k = groupId; k < (K / QK4_NL); k += NSUBGROUPS) {
-        regS = src0_d[gid + k * LINE_STRIDE_A]; // each fiber loads scale of two rows
+        regS = src0_d[gid_s + k * LINE_STRIDE_A]; // each fiber loads scale of two rows
        // first 4 fibers in each wave load 8 B values to its private scope
        if (slid < 4) {
            regB.s0123 = read_imagef(src1, (slid * 2 + k * 8));
@@ -248,20 +250,20 @@ kernel void kernel_gemv_noshuffle_iq4_nl_f32(
        }

        // load half weights for two blocks in consecutive rows
-        regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
-        regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
-        regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
-        regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
+        regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
+        regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
+        regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
+        regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
 #ifdef VECTOR_SUB_GROUP_BROADCAST
        dequantizeBlockAccum_ns_sgbroadcast_8_hi(totalSum, as_ushort8(regA), regS, regB);
 #else
        dequantizeBlockAccum_ns_sgbroadcast_1_hi(totalSum, as_ushort8(regA), regS, regB);
 #endif // VECTOR_SUB_GROUP_BROADCAST

-        regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
-        regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
-        regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
-        regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
+        regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
+        regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
+        regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
+        regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
 #ifdef VECTOR_SUB_GROUP_BROADCAST
        dequantizeBlockAccum_ns_sgbroadcast_8_lo(totalSum, as_ushort8(regA), regS, regB);
 #else
@@ -73,6 +73,8 @@ __kernel void kernel_gemv_noshuffle_q1_0_f32(
    uint M = ne01;

    uint LINE_STRIDE_A  = M;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, LINE_STRIDE_A - 1);
    uint BLOCK_STRIDE_A = 4 * M;

    uint4  regA;
@@ -83,7 +85,7 @@ __kernel void kernel_gemv_noshuffle_q1_0_f32(

    #pragma unroll 1
    for (uint kb = groupId; kb < (K / QK1_0); kb += N_SIMDGROUP) {
-        regS = src0_d[gid + kb * LINE_STRIDE_A]; // each fiber loads its row's scale
+        regS = src0_d[gid_s + kb * LINE_STRIDE_A]; // each fiber loads its row's scale

        // first 16 fibers load 8 B values each -> 128 activations for this block
        if (slid < 16) {
@@ -92,10 +94,10 @@ __kernel void kernel_gemv_noshuffle_q1_0_f32(
        }

        // load this row's 4 uint32 (128 sign bits)
-        regA.s0 = read_imageui(src0_q, (gid + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
-        regA.s1 = read_imageui(src0_q, (gid + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
-        regA.s2 = read_imageui(src0_q, (gid + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
-        regA.s3 = read_imageui(src0_q, (gid + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
+        regA.s0 = read_imageui(src0_q, (gid_s + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
+        regA.s1 = read_imageui(src0_q, (gid_s + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
+        regA.s2 = read_imageui(src0_q, (gid_s + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
+        regA.s3 = read_imageui(src0_q, (gid_s + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;

        float scale = (float)regS;
        dequantizeBlockAccum_q1(totalSum, regA.s0, scale, regB, 0);
@@ -224,6 +224,8 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32(
    uint M = ne01;

    uint LINE_STRIDE_A = M / 2;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, LINE_STRIDE_A - 1);
    uint BLOCK_STRIDE_A = N_SIMDGROUP * M;

    __private uint4     regA;
@@ -234,7 +236,7 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32(

    // loop along K in block granularity, skip 4 blocks every iter
    for (uint k = groupId; k < (K / QK4_0); k += N_SIMDGROUP) {
-        regS = src0_d[gid + k * LINE_STRIDE_A]; // each fiber loads scale of two rows
+        regS = src0_d[gid_s + k * LINE_STRIDE_A]; // each fiber loads scale of two rows
        // first 4 fibers in each wave load 8 B values to its private scope
        if (slid < 4) {
            regB.s0123 = read_imagef(src1, (slid * 2 + k * 8));
@@ -242,20 +244,20 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32(
        }

        // load half weights for two blocks in consecutive rows
-        regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
-        regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
-        regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
-        regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
+        regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
+        regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
+        regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
+        regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
 #ifdef VECTOR_SUB_GROUP_BROADCAST
        dequantizeBlockAccum_ns_sgbroadcast_8_hi(totalSum, as_ushort8(regA), regS, regB);
 #else
        dequantizeBlockAccum_ns_sgbroadcast_1_hi(totalSum, as_ushort8(regA), regS, regB);
 #endif // VECTOR_SUB_GROUP_BROADCAST

-        regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
-        regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
-        regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
-        regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
+        regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
+        regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
+        regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
+        regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
 #ifdef VECTOR_SUB_GROUP_BROADCAST
        dequantizeBlockAccum_ns_sgbroadcast_8_lo(totalSum, as_ushort8(regA), regS, regB);
 #else
@@ -329,6 +331,8 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32_mc3(
    uint M = ne01;

    uint LINE_STRIDE_A  = M / 2;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, LINE_STRIDE_A - 1);
    // BLOCK_STRIDE_A is the LAYOUT stride between consecutive K-blocks = 4 uints
    // per q4_0 block * M (set by the trans4_ns convert). The "4" is uints/block, NOT
    // the subgroup count — keep it fixed so the K-split count (nsg) can vary.
@@ -347,17 +351,17 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32_mc3(
    __private float2 ts3 = (float2)(0.0f);

    for (uint k = groupId; k < (K / QK4_0); k += nsg) {
-        regS = src0_d[gid + k * LINE_STRIDE_A];
+        regS = src0_d[gid_s + k * LINE_STRIDE_A];

        // weights loaded ONCE, reused across the columns
-        regA_hi.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
-        regA_hi.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
-        regA_hi.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
-        regA_hi.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
-        regA_lo.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
-        regA_lo.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
-        regA_lo.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
-        regA_lo.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
+        regA_hi.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
+        regA_hi.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
+        regA_hi.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
+        regA_hi.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
+        regA_lo.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
+        regA_lo.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
+        regA_lo.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
+        regA_lo.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;

        MC_COL_Q40(ts0, 0);
        MC_COL_Q40(ts1, 1);
@@ -381,10 +385,12 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32_mc3(
        }
        dst = (global float*)((global char*)dst + offsetd);
        // dst is column-major [M rows x n_cols cols]: (row, col) at col*M + row
-        vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2]));
-        vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2]));
-        if (n_cols > 2) vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2]));
-        if (n_cols > 3) vstore2((float2)(acc.s6, acc.s7), 0, &(dst[3 * M + gid * 2]));
+        if (gid < LINE_STRIDE_A) {
+            vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2]));
+            vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2]));
+            if (n_cols > 2) vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2]));
+            if (n_cols > 3) vstore2((float2)(acc.s6, acc.s7), 0, &(dst[3 * M + gid * 2]));
+        }
    }
 }
 #undef MC_COL_Q40
@@ -210,6 +210,8 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32(
 {
    uint groupId = get_local_id(1);
    uint gid     = get_global_id(0);
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, (uint)LINE_STRIDE_A - 1u);
    ushort slid    = get_sub_group_local_id();

    __private uint4     regA;
@@ -220,7 +222,7 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32(

    // loop along K in block granularity, skip 4 blocks every iter
    for (uint k = groupId; k < (K / QK4_0); k += N_SIMDGROUP) {
-        regS = src0_d[gid + k * LINE_STRIDE_A]; // each fiber loads scale of two rows
+        regS = src0_d[gid_s + k * LINE_STRIDE_A]; // each fiber loads scale of two rows
        // first 4 fibers in each wave load 8 B values to its private scope
        if (slid < 4) {
            regB.s0123 = read_imagef(src1, (slid * 2 + k * 8));
@@ -228,20 +230,20 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32(
        }

        // load half weights for two blocks in consecutive rows
-        regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
-        regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
-        regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
-        regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
+        regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
+        regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
+        regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
+        regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
 #ifdef VECTOR_SUB_GROUP_BROADCAST
        dequantizeBlockAccum_ns_sgbroadcast_8_hi(totalSum, as_ushort8(regA), regS, regB);
 #else
        dequantizeBlockAccum_ns_sgbroadcast_1_hi(totalSum, as_ushort8(regA), regS, regB);
 #endif // VECTOR_SUB_GROUP_BROADCAST

-        regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
-        regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
-        regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
-        regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
+        regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
+        regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
+        regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
+        regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
 #ifdef VECTOR_SUB_GROUP_BROADCAST
        dequantizeBlockAccum_ns_sgbroadcast_8_lo(totalSum, as_ushort8(regA), regS, regB);
 #else
@@ -209,6 +209,8 @@ kernel void kernel_gemv_noshuffle_q4_1_f32(
    uint M = ne01;

    uint LINE_STRIDE_A = M / 2;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, LINE_STRIDE_A - 1);
    uint BLOCK_STRIDE_A = NSUBGROUPS * M;

    private uint4     regA;
@@ -220,8 +222,8 @@ kernel void kernel_gemv_noshuffle_q4_1_f32(

    // loop along K in block granularity, skip 4 blocks every iter
    for (uint k = groupId; k < (K / QK4_0); k += NSUBGROUPS) {
-        regS = src0_d[gid + k * LINE_STRIDE_A]; // each fiber loads scale of two rows
-        regM = src0_m[gid + k * LINE_STRIDE_A]; // each fiber loads min of two rows
+        regS = src0_d[gid_s + k * LINE_STRIDE_A]; // each fiber loads scale of two rows
+        regM = src0_m[gid_s + k * LINE_STRIDE_A]; // each fiber loads min of two rows
        // first 4 fibers in each wave load 8 B values to its private scope
        if (slid < 4) {
            regB.s0123 = read_imagef(src1, (slid * 2 + k * 8));
@@ -229,20 +231,20 @@ kernel void kernel_gemv_noshuffle_q4_1_f32(
        }

        // load half weights for two blocks in consecutive rows
-        regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
-        regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
-        regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
-        regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
+        regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
+        regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
+        regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
+        regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
 #ifdef VECTOR_SUB_GROUP_BROADCAT
        dequantizeBlockAccum_ns_sgbroadcast_8_hi(totalSum, as_ushort8(regA), regS, regM, regB);
 #else
        dequantizeBlockAccum_ns_sgbroadcast_1_hi(totalSum, as_ushort8(regA), regS, regM, regB);
 #endif // VECTOR_SUB_GROUP_BROADCAT

-        regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
-        regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
-        regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
-        regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
+        regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
+        regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
+        regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
+        regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
 #ifdef VECTOR_SUB_GROUP_BROADCAT
        dequantizeBlockAccum_ns_sgbroadcast_8_lo(totalSum, as_ushort8(regA), regS, regM, regB);
 #else
@@ -326,6 +328,8 @@ kernel void kernel_gemv_noshuffle_q4_1_f32_mc3(
    uint M = ne01;

    uint LINE_STRIDE_A  = M / 2;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, LINE_STRIDE_A - 1);
    uint BLOCK_STRIDE_A = NSUBGROUPS * M;
    uint COL_STRIDE     = K / 4;   // float4 pixels per activation column

@@ -339,18 +343,18 @@ kernel void kernel_gemv_noshuffle_q4_1_f32_mc3(
    private float2 ts3 = (float2)(0.0f);

    for (uint k = groupId; k < (K / QK4_0); k += NSUBGROUPS) {
-        regS = src0_d[gid + k * LINE_STRIDE_A];
-        regM = src0_m[gid + k * LINE_STRIDE_A];
+        regS = src0_d[gid_s + k * LINE_STRIDE_A];
+        regM = src0_m[gid_s + k * LINE_STRIDE_A];

        // weights loaded ONCE, reused across the columns
-        regA_hi.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
-        regA_hi.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
-        regA_hi.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
-        regA_hi.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
-        regA_lo.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
-        regA_lo.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
-        regA_lo.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
-        regA_lo.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
+        regA_hi.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
+        regA_hi.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
+        regA_hi.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
+        regA_hi.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
+        regA_lo.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
+        regA_lo.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
+        regA_lo.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
+        regA_lo.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;

        MC_COL_Q41(ts0, 0);
        MC_COL_Q41(ts1, 1);
@@ -373,10 +377,12 @@ kernel void kernel_gemv_noshuffle_q4_1_f32_mc3(
        acc += reduceLM[SUBGROUP_SIZE * 2 + slid];
        dst = (global float*)((global char*)dst + offsetd);
        // dst is column-major [M rows x n_cols cols]: (row, col) at col*M + row
-        vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2]));
-        vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2]));
-        if (n_cols > 2) vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2]));
-        if (n_cols > 3) vstore2((float2)(acc.s6, acc.s7), 0, &(dst[3 * M + gid * 2]));
+        if (gid < LINE_STRIDE_A) {
+            vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2]));
+            vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2]));
+            if (n_cols > 2) vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2]));
+            if (n_cols > 3) vstore2((float2)(acc.s6, acc.s7), 0, &(dst[3 * M + gid * 2]));
+        }
    }
 }
 #undef MC_COL_Q41
@@ -415,6 +415,8 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_glu(
    uint M = ne01;

    uint LINE_STRIDE_A  = M / 2;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, LINE_STRIDE_A - 1);
    uint BLOCK_STRIDE_A = 4 * M;

    private uint4  regA;
@@ -434,9 +436,9 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_glu(
    for (uint k = groupId; k < (K / 32); k += nsg) {                                           \
        uint sb = k / 8;                                                                       \
        uint j  = k % 8;                                                                       \
-        half2 d   = DD[gid + sb * LINE_STRIDE_A];                                              \
-        half2 dm  = MM[gid + sb * LINE_STRIDE_A];                                              \
-        global const uchar * sc0 = SS + sb * 12 * M + 2 * gid;                                 \
+        half2 d   = DD[gid_s + sb * LINE_STRIDE_A];                                            \
+        half2 dm  = MM[gid_s + sb * LINE_STRIDE_A];                                            \
+        global const uchar * sc0 = SS + sb * 12 * M + 2 * gid_s;                               \
        global const uchar * sc1 = sc0 + 1;                                                    \
        uchar sv0, mn0, sv1, mn1;                                                              \
        get_scale_min_k4(j, sc0, M, &sv0, &mn0, mask_d6, mask_d4, mask_hi2);                      \
@@ -447,15 +449,15 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_glu(
            regB.s0123 = read_imagef(src1, (slid * 2 + k * 8));                                \
            regB.s4567 = read_imagef(src1, (1 + slid * 2 + k * 8));                            \
        }                                                                                      \
-        regA.s0 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;           \
-        regA.s1 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;           \
-        regA.s2 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;           \
-        regA.s3 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;           \
+        regA.s0 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;           \
+        regA.s1 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;           \
+        regA.s2 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;           \
+        regA.s3 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;           \
        DEQ_HI(SUM, as_ushort8(regA), regS, regM, regB);                                       \
-        regA.s0 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;           \
-        regA.s1 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;           \
-        regA.s2 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;           \
-        regA.s3 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;           \
+        regA.s0 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;           \
+        regA.s1 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;           \
+        regA.s2 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;           \
+        regA.s3 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;           \
        DEQ_LO(SUM, as_ushort8(regA), regS, regM, regB);                                       \
    }

@@ -489,8 +491,10 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_glu(
            upSum   += p.zw;
        }
        dst = (global float*)((global char*)dst + offsetd);
-        dst[gid * 2 + 0] = glu_apply(glu_op, gateSum.s0, upSum.s0);
-        dst[gid * 2 + 1] = glu_apply(glu_op, gateSum.s1, upSum.s1);
+        if (gid < LINE_STRIDE_A) {
+            dst[gid * 2 + 0] = glu_apply(glu_op, gateSum.s0, upSum.s0);
+            dst[gid * 2 + 1] = glu_apply(glu_op, gateSum.s1, upSum.s1);
+        }
    }
 }

@@ -533,6 +537,9 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_splitk(
    uint LINE_STRIDE_A  = M / 2;
    uint BLOCK_STRIDE_A = 4 * M;      // physical, independent of the K-split

+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, LINE_STRIDE_A - 1);
+
    private uint4  regA;
    private half2  regS, regM;
    private float8 regB;
@@ -542,9 +549,9 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_splitk(
    for (uint k = kslice * nsg + groupId; k < (K / 32); k += ksplit * nsg) {
        uint sb = k / 8;
        uint j  = k % 8;
-        half2 d   = src0_d[gid + sb * LINE_STRIDE_A];
-        half2 dm  = src0_m[gid + sb * LINE_STRIDE_A];
-        global const uchar * sc0 = src0_s + sb * 12 * M + 2 * gid;
+        half2 d   = src0_d[gid_s + sb * LINE_STRIDE_A];
+        half2 dm  = src0_m[gid_s + sb * LINE_STRIDE_A];
+        global const uchar * sc0 = src0_s + sb * 12 * M + 2 * gid_s;
        global const uchar * sc1 = sc0 + 1;
        uchar sv0, mn0, sv1, mn1;
        get_scale_min_k4(j, sc0, M, &sv0, &mn0, mask_d6, mask_d4, mask_hi2);
@@ -555,19 +562,19 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_splitk(
            regB.s0123 = read_imagef(src1, (slid * 2 + k * 8));
            regB.s4567 = read_imagef(src1, (1 + slid * 2 + k * 8));
        }
-        regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
-        regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
-        regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
-        regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
+        regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
+        regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
+        regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
+        regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
 #ifdef VECTOR_SUB_GROUP_BROADCAST
        dequantizeBlockAccum_ns_sgbroadcast_8_hi(totalSum, as_ushort8(regA), regS, regM, regB);
 #else
        dequantizeBlockAccum_ns_sgbroadcast_1_hi(totalSum, as_ushort8(regA), regS, regM, regB);
 #endif
-        regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
-        regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
-        regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
-        regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
+        regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
+        regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
+        regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
+        regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
 #ifdef VECTOR_SUB_GROUP_BROADCAST
        dequantizeBlockAccum_ns_sgbroadcast_8_lo(totalSum, as_ushort8(regA), regS, regM, regB);
 #else
@@ -584,7 +591,9 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_splitk(
        for (uint i = 0; i < nsg - 1; ++i) {
            totalSum += reduceLM[SUBGROUP_SIZE * i + slid];
        }
-        vstore2(totalSum, 0, &(partial[kslice * M + gid * 2]));
+        // Keep tail stores inside this K-slice.
+        if (gid * 2 + 0 < M) partial[kslice * M + gid * 2 + 0] = totalSum.s0;
+        if (gid * 2 + 1 < M) partial[kslice * M + gid * 2 + 1] = totalSum.s1;
    }
 }

@@ -689,6 +698,8 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_mc3(
    uint M = ne01;

    uint LINE_STRIDE_A  = M / 2;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, LINE_STRIDE_A - 1);
    uint BLOCK_STRIDE_A = NSUBGROUPS * M;
    uint COL_STRIDE     = K / 4;   // float4 pixels per activation column

@@ -713,10 +724,10 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_mc3(
        uint sb = k / 8;
        uint j  = k % 8;

-        half2 d   = src0_d[gid + sb * LINE_STRIDE_A];
-        half2 dm  = src0_m[gid + sb * LINE_STRIDE_A];
+        half2 d   = src0_d[gid_s + sb * LINE_STRIDE_A];
+        half2 dm  = src0_m[gid_s + sb * LINE_STRIDE_A];

-        global const uchar * sc0 = src0_s + sb * 12 * M + 2 * gid;
+        global const uchar * sc0 = src0_s + sb * 12 * M + 2 * gid_s;
        global const uchar * sc1 = sc0 + 1;

        uchar sv0, mn0, sv1, mn1;
@@ -727,14 +738,14 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_mc3(
        regM = convert_half2(convert_float2(dm) * convert_float2((uchar2)(mn0, mn1)));

        // weights loaded ONCE, reused across the 3 columns
-        regA_hi.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
-        regA_hi.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
-        regA_hi.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
-        regA_hi.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
-        regA_lo.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
-        regA_lo.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
-        regA_lo.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
-        regA_lo.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
+        regA_hi.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
+        regA_hi.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
+        regA_hi.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
+        regA_hi.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
+        regA_lo.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
+        regA_lo.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
+        regA_lo.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
+        regA_lo.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;

 #ifdef Q4K_MC3_DEQUANT_ONCE
        // Dequant the 32 weights/row (16 hi + 16 lo) ONCE into half2[] (byte-
@@ -827,8 +838,10 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_mc3(
        acc += reduceLM[SUBGROUP_SIZE * 2 + slid];
        dst = (global float*)((global char*)dst + offsetd);
        // dst is column-major [M rows x 3 cols]: (row, col) at col*M + row
-        vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2]));
-        vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2]));
-        vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2]));
+        if (gid < LINE_STRIDE_A) {
+            vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2]));
+            vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2]));
+            vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2]));
+        }
    }
 }
@@ -231,8 +231,9 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_o4(

    // Two consecutive pair-indices (each the same access pattern the 2-output
    // kernel uses); together they cover 4 consecutive output rows.
-    uint gid_a = gid * 2;
-    uint gid_b = gid * 2 + 1;
+    // Keep tail lanes active for subgroup ops, but read valid weight pairs.
+    uint gid_a = min(gid * 2,     (uint)(ne01 / 2) - 1u);
+    uint gid_b = min(gid * 2 + 1, (uint)(ne01 / 2) - 1u);

    uint K = ne00;
    uint M = ne01;
@@ -209,6 +209,8 @@ __kernel void kernel_gemv_noshuffle_q5_0_f32(
    uint M = ne01;

    uint LINE_STRIDE_A    = M / 2;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, LINE_STRIDE_A - 1);
    uint BLOCK_STRIDE_A  = NSUBGROUPS * M;

    private uint4     regA;
@@ -218,13 +220,13 @@ __kernel void kernel_gemv_noshuffle_q5_0_f32(
    private float2 totalSum = (float2)(0.0f);

    for (uint k = groupId; k < (K / QK5_0); k += NSUBGROUPS) {
-        regS = src0_d[gid + k * LINE_STRIDE_A];
+        regS = src0_d[gid_s + k * LINE_STRIDE_A];

        ushort4 qh_raw;
-        qh_raw.s0 = src0_qh[gid + (4*k + 0) * LINE_STRIDE_A];
-        qh_raw.s1 = src0_qh[gid + (4*k + 1) * LINE_STRIDE_A];
-        qh_raw.s2 = src0_qh[gid + (4*k + 2) * LINE_STRIDE_A];
-        qh_raw.s3 = src0_qh[gid + (4*k + 3) * LINE_STRIDE_A];
+        qh_raw.s0 = src0_qh[gid_s + (4*k + 0) * LINE_STRIDE_A];
+        qh_raw.s1 = src0_qh[gid_s + (4*k + 1) * LINE_STRIDE_A];
+        qh_raw.s2 = src0_qh[gid_s + (4*k + 2) * LINE_STRIDE_A];
+        qh_raw.s3 = src0_qh[gid_s + (4*k + 3) * LINE_STRIDE_A];

        uchar8 raw = as_uchar8(qh_raw);
        uchar8 qh_bytes = (uchar8)(raw.s0, raw.s2, raw.s4, raw.s6,
@@ -236,10 +238,10 @@ __kernel void kernel_gemv_noshuffle_q5_0_f32(
            regB.s4567 = read_imagef(src1, (1 + slid * 2 + k * 8));
        }

-        regA.s0 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
-        regA.s1 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
-        regA.s2 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
-        regA.s3 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
+        regA.s0 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
+        regA.s1 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
+        regA.s2 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
+        regA.s3 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;

 #ifdef VECTOR_SUB_GROUP_BROADCAST
        dequantizeBlockAccum_ns_q5_0_sgbroadcast_8_hi(totalSum, as_ushort8(regA), qh_bytes, regS, regB);
@@ -247,10 +249,10 @@ __kernel void kernel_gemv_noshuffle_q5_0_f32(
        dequantizeBlockAccum_ns_q5_0_sgbroadcast_1_hi(totalSum, as_ushort8(regA), qh_bytes, regS, regB);
 #endif // VECTOR_SUB_GROUP_BROADCAST

-        regA.s0 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
-        regA.s1 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
-        regA.s2 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
-        regA.s3 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
+        regA.s0 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
+        regA.s1 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
+        regA.s2 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
+        regA.s3 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
 #ifdef VECTOR_SUB_GROUP_BROADCAST
        dequantizeBlockAccum_ns_q5_0_sgbroadcast_8_lo(totalSum, as_ushort8(regA), qh_bytes, regS, regB);
 #else
@@ -210,6 +210,8 @@ __kernel void kernel_gemv_noshuffle_q5_1_f32(
    uint M = ne01;

    uint LINE_STRIDE_A    = M / 2;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, LINE_STRIDE_A - 1);
   uint BLOCK_STRIDE_A  = NSUBGROUPS * M;

    __private uint4     regA;
@@ -220,14 +222,14 @@ __kernel void kernel_gemv_noshuffle_q5_1_f32(
    __private float2 totalSum = (float2)(0.0f);

    for (uint k = groupId; k < (K / QK5_1); k += NSUBGROUPS) {
-        regS = src0_d[gid + k * LINE_STRIDE_A];
-        regM = src0_m[gid + k * LINE_STRIDE_A];
+        regS = src0_d[gid_s + k * LINE_STRIDE_A];
+        regM = src0_m[gid_s + k * LINE_STRIDE_A];

        ushort4 qh_raw;
-        qh_raw.s0 = src0_qh[gid + (4*k + 0) * LINE_STRIDE_A];
-        qh_raw.s1 = src0_qh[gid + (4*k + 1) * LINE_STRIDE_A];
-        qh_raw.s2 = src0_qh[gid + (4*k + 2) * LINE_STRIDE_A];
-        qh_raw.s3 = src0_qh[gid + (4*k + 3) * LINE_STRIDE_A];
+        qh_raw.s0 = src0_qh[gid_s + (4*k + 0) * LINE_STRIDE_A];
+        qh_raw.s1 = src0_qh[gid_s + (4*k + 1) * LINE_STRIDE_A];
+        qh_raw.s2 = src0_qh[gid_s + (4*k + 2) * LINE_STRIDE_A];
+        qh_raw.s3 = src0_qh[gid_s + (4*k + 3) * LINE_STRIDE_A];

        uchar8 raw = as_uchar8(qh_raw);
        uchar8 qh_bytes = (uchar8)(raw.s0, raw.s2, raw.s4, raw.s6,
@@ -239,10 +241,10 @@ __kernel void kernel_gemv_noshuffle_q5_1_f32(
            regB.s4567 = read_imagef(src1, (1 + slid * 2 + k * 8));
        }

-        regA.s0 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
-        regA.s1 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
-        regA.s2 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
-        regA.s3 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
+        regA.s0 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
+        regA.s1 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
+        regA.s2 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
+        regA.s3 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;

 #ifdef VECTOR_SUB_GROUP_BROADCAST
        dequantizeBlockAccum_ns_q5_1_sgbroadcast_8_hi(totalSum, as_ushort8(regA), qh_bytes, regS, regM, regB);
@@ -250,10 +252,10 @@ __kernel void kernel_gemv_noshuffle_q5_1_f32(
        dequantizeBlockAccum_ns_q5_1_sgbroadcast_1_hi(totalSum, as_ushort8(regA), qh_bytes, regS, regM, regB);
 #endif // VECTOR_SUB_GROUP_BROADCAST

-        regA.s0 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
-        regA.s1 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
-        regA.s2 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
-        regA.s3 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
+        regA.s0 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
+        regA.s1 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
+        regA.s2 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
+        regA.s3 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
 #ifdef VECTOR_SUB_GROUP_BROADCAST
        dequantizeBlockAccum_ns_q5_1_sgbroadcast_8_lo(totalSum, as_ushort8(regA), qh_bytes, regS, regM, regB);
 #else
@@ -232,6 +232,8 @@ kernel void kernel_gemv_noshuffle_q5_k_f32(
    uint M = ne01;

    uint LINE_STRIDE_A     = M / 2;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, LINE_STRIDE_A - 1);
    uint BLOCK_STRIDE_A    = NSUBGROUPS * M;

    uint LINE_STRIDE_A_QH  = M / 2;
@@ -250,11 +252,11 @@ kernel void kernel_gemv_noshuffle_q5_k_f32(
        uint sb = k / 8;
        uint j  = k % 8;

-        half2 d   = src0_d[gid + sb * LINE_STRIDE_A];
-        half2 dm  = src0_m[gid + sb * LINE_STRIDE_A];
+        half2 d   = src0_d[gid_s + sb * LINE_STRIDE_A];
+        half2 dm  = src0_m[gid_s + sb * LINE_STRIDE_A];

-        global const uchar * sc0 = src0_s + 2 * gid * scales_per_row + sb * 12;
-        global const uchar * sc1 = src0_s + (2 * gid + 1) * scales_per_row + sb * 12;
+        global const uchar * sc0 = src0_s + 2 * gid_s * scales_per_row + sb * 12;
+        global const uchar * sc1 = src0_s + (2 * gid_s + 1) * scales_per_row + sb * 12;

        uchar sv0, mn0, sv1, mn1;
        get_scale_min_k4(j, sc0, &sv0, &mn0, mask_d6, mask_d4, mask_hi2);
@@ -268,25 +270,25 @@ kernel void kernel_gemv_noshuffle_q5_k_f32(
            regB.s4567 = read_imagef(src1, (1 + slid * 2 + k * 8));
        }

-        regH.s0 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 0)).x);
-        regH.s1 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 1)).x);
-        regH.s2 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 2)).x);
-        regH.s3 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 3)).x);
+        regH.s0 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 0)).x);
+        regH.s1 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 1)).x);
+        regH.s2 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 2)).x);
+        regH.s3 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 3)).x);

-        regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
-        regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
-        regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
-        regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
+        regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
+        regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
+        regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
+        regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
 #ifdef VECTOR_SUB_GROUP_BROADCAST
        dequantizeBlockAccum_ns_sgbroadcast_8_hi(totalSum, as_ushort8(regA), as_uchar8(regH), regS, regM, regB);
 #else
        dequantizeBlockAccum_ns_sgbroadcast_1_hi(totalSum, as_ushort8(regA), as_uchar8(regH), regS, regM, regB);
 #endif // VECTOR_SUB_GROUP_BROADCAST

-        regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
-        regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
-        regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
-        regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
+        regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
+        regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
+        regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
+        regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
 #ifdef VECTOR_SUB_GROUP_BROADCAST
        dequantizeBlockAccum_ns_sgbroadcast_8_lo(totalSum, as_ushort8(regA), as_uchar8(regH), regS, regM, regB);
 #else
@@ -373,6 +375,8 @@ kernel void kernel_gemv_noshuffle_q5_k_f32_mc3(
    uint M = ne01;

    uint LINE_STRIDE_A     = M / 2;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, LINE_STRIDE_A - 1);
    uint BLOCK_STRIDE_A    = NSUBGROUPS * M;
    uint LINE_STRIDE_A_QH  = M / 2;
    uint BLOCK_STRIDE_A_QH = NSUBGROUPS * M / 2;
@@ -393,11 +397,11 @@ kernel void kernel_gemv_noshuffle_q5_k_f32_mc3(
        uint sb = k / 8;
        uint j  = k % 8;

-        half2 d   = src0_d[gid + sb * LINE_STRIDE_A];
-        half2 dm  = src0_m[gid + sb * LINE_STRIDE_A];
+        half2 d   = src0_d[gid_s + sb * LINE_STRIDE_A];
+        half2 dm  = src0_m[gid_s + sb * LINE_STRIDE_A];

-        global const uchar * sc0 = src0_s + 2 * gid * scales_per_row + sb * 12;
-        global const uchar * sc1 = src0_s + (2 * gid + 1) * scales_per_row + sb * 12;
+        global const uchar * sc0 = src0_s + 2 * gid_s * scales_per_row + sb * 12;
+        global const uchar * sc1 = src0_s + (2 * gid_s + 1) * scales_per_row + sb * 12;

        uchar sv0, mn0, sv1, mn1;
        get_scale_min_k4(j, sc0, &sv0, &mn0, mask_d6, mask_d4, mask_hi2);
@@ -407,19 +411,19 @@ kernel void kernel_gemv_noshuffle_q5_k_f32_mc3(
        regM = convert_half2(convert_float2(dm) * convert_float2((uchar2)(mn0, mn1)));

        // high-bit plane + weights loaded ONCE, reused across the columns
-        regH.s0 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 0)).x);
-        regH.s1 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 1)).x);
-        regH.s2 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 2)).x);
-        regH.s3 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 3)).x);
+        regH.s0 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 0)).x);
+        regH.s1 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 1)).x);
+        regH.s2 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 2)).x);
+        regH.s3 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 3)).x);

-        regA_hi.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
-        regA_hi.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
-        regA_hi.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
-        regA_hi.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
-        regA_lo.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
-        regA_lo.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
-        regA_lo.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
-        regA_lo.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
+        regA_hi.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
+        regA_hi.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
+        regA_hi.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
+        regA_hi.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
+        regA_lo.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
+        regA_lo.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
+        regA_lo.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
+        regA_lo.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;

        MC_COL_Q5K(ts0, 0);
        MC_COL_Q5K(ts1, 1);
@@ -442,10 +446,12 @@ kernel void kernel_gemv_noshuffle_q5_k_f32_mc3(
        acc += reduceLM[SUBGROUP_SIZE * 2 + slid];
        dst = (global float*)((global char*)dst + offsetd);
        // dst is column-major [M rows x n_cols cols]: (row, col) at col*M + row
-        vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2]));
-        vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2]));
-        if (n_cols > 2) vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2]));
-        if (n_cols > 3) vstore2((float2)(acc.s6, acc.s7), 0, &(dst[3 * M + gid * 2]));
+        if (gid < LINE_STRIDE_A) {
+            vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2]));
+            vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2]));
+            if (n_cols > 2) vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2]));
+            if (n_cols > 3) vstore2((float2)(acc.s6, acc.s7), 0, &(dst[3 * M + gid * 2]));
+        }
    }
 }
 #undef MC_COL_Q5K
@@ -219,26 +219,28 @@ kernel void kernel_gemv_noshuffle_q6_K_f32(
    float2  total_sum = 0.0f;

    int line_stride_a = ne01 / 2;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    int gid_s = min(gid, line_stride_a - 1);
    int block_stride_a = NSUBGROUPS * ne01;

    for (int k = grp; k < nb; k += NSUBGROUPS) {
-        reg_d = src0_d[gid + k/8 * line_stride_a];
-        reg_s = as_char4(src0_s[gid + k * line_stride_a]);
+        reg_d = src0_d[gid_s + k/8 * line_stride_a];
+        reg_s = as_char4(src0_s[gid_s + k * line_stride_a]);

        if (slid < 4) {
            reg_b.s0123 = read_imagef(src1, 0 + slid*2 + k*8);
            reg_b.s4567 = read_imagef(src1, 1 + slid*2 + k*8);
        }

-        reg_a_l.s0 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*0).x;
-        reg_a_l.s1 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*1).x;
-        reg_a_l.s2 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*2).x;
-        reg_a_l.s3 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*3).x;
+        reg_a_l.s0 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*0).x;
+        reg_a_l.s1 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*1).x;
+        reg_a_l.s2 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*2).x;
+        reg_a_l.s3 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*3).x;

-        reg_a_h.s0 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*0).x);
-        reg_a_h.s1 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*1).x);
-        reg_a_h.s2 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*2).x);
-        reg_a_h.s3 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*3).x);
+        reg_a_h.s0 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*0).x);
+        reg_a_h.s1 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*1).x);
+        reg_a_h.s2 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*2).x);
+        reg_a_h.s3 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*3).x);

 #ifdef VECTOR_SUB_GROUP_BROADCAT
        dequantize_block_acc_bcast_8_hi(total_sum, as_ushort8(reg_a_l), as_uchar8(reg_a_h), reg_d, reg_s, reg_b);
@@ -246,15 +248,15 @@ kernel void kernel_gemv_noshuffle_q6_K_f32(
        dequantize_block_acc_bcast_1_hi(total_sum, as_ushort8(reg_a_l), as_uchar8(reg_a_h), reg_d, reg_s, reg_b);
 #endif // VECTOR_SUB_GROUP_BROADCAT

-        reg_a_l.s0 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*4).x;
-        reg_a_l.s1 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*5).x;
-        reg_a_l.s2 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*6).x;
-        reg_a_l.s3 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*7).x;
+        reg_a_l.s0 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*4).x;
+        reg_a_l.s1 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*5).x;
+        reg_a_l.s2 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*6).x;
+        reg_a_l.s3 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*7).x;

-        reg_a_h.s0 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*4).x);
-        reg_a_h.s1 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*5).x);
-        reg_a_h.s2 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*6).x);
-        reg_a_h.s3 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*7).x);
+        reg_a_h.s0 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*4).x);
+        reg_a_h.s1 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*5).x);
+        reg_a_h.s2 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*6).x);
+        reg_a_h.s3 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*7).x);

 #ifdef VECTOR_SUB_GROUP_BROADCAT
        dequantize_block_acc_bcast_8_lo(total_sum, as_ushort8(reg_a_l), as_uchar8(reg_a_h), reg_d, reg_s, reg_b);
@@ -323,6 +325,8 @@ kernel void kernel_gemv_noshuffle_q6_K_f32_mc3(

    int nb = ne00 / 32;
    int line_stride_a  = ne01 / 2;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    int gid_s = min(gid, line_stride_a - 1);
    int block_stride_a = NSUBGROUPS * ne01;
    int COL_STRIDE     = ne00 / 4;   // float4 pixels per activation column

@@ -335,27 +339,27 @@ kernel void kernel_gemv_noshuffle_q6_K_f32_mc3(
    float2  ts0 = 0.0f, ts1 = 0.0f, ts2 = 0.0f;

    for (int k = grp; k < nb; k += NSUBGROUPS) {
-        reg_d = src0_d[gid + k/8 * line_stride_a];
-        reg_s = as_char4(src0_s[gid + k * line_stride_a]);
+        reg_d = src0_d[gid_s + k/8 * line_stride_a];
+        reg_s = as_char4(src0_s[gid_s + k * line_stride_a]);

        // weights loaded ONCE (hi: blocks 0-3, lo: blocks 4-7), reused x3 cols
-        ql_hi.s0 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*0).x;
-        ql_hi.s1 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*1).x;
-        ql_hi.s2 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*2).x;
-        ql_hi.s3 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*3).x;
-        qh_hi.s0 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*0).x);
-        qh_hi.s1 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*1).x);
-        qh_hi.s2 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*2).x);
-        qh_hi.s3 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*3).x);
+        ql_hi.s0 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*0).x;
+        ql_hi.s1 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*1).x;
+        ql_hi.s2 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*2).x;
+        ql_hi.s3 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*3).x;
+        qh_hi.s0 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*0).x);
+        qh_hi.s1 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*1).x);
+        qh_hi.s2 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*2).x);
+        qh_hi.s3 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*3).x);

-        ql_lo.s0 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*4).x;
-        ql_lo.s1 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*5).x;
-        ql_lo.s2 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*6).x;
-        ql_lo.s3 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*7).x;
-        qh_lo.s0 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*4).x);
-        qh_lo.s1 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*5).x);
-        qh_lo.s2 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*6).x);
-        qh_lo.s3 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*7).x);
+        ql_lo.s0 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*4).x;
+        ql_lo.s1 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*5).x;
+        ql_lo.s2 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*6).x;
+        ql_lo.s3 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*7).x;
+        qh_lo.s0 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*4).x);
+        qh_lo.s1 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*5).x);
+        qh_lo.s2 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*6).x);
+        qh_lo.s3 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*7).x);

        // Per-column: load only this column's activation (single reg_b live) ->
        // 1/3 the activation register pressure, cutting the private-mem spill.
@@ -402,8 +406,10 @@ kernel void kernel_gemv_noshuffle_q6_K_f32_mc3(
        acc += reduce_lm[SUBGROUP_SIZE*2 + slid];
        dst = (global float*)((global char*)dst + offsetd);
        // dst column-major [ne01 rows x 3 cols]: (row, col) at col*ne01 + row
-        vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0*ne01 + gid*2]));
-        vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1*ne01 + gid*2]));
-        vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2*ne01 + gid*2]));
+        if (gid < line_stride_a) {
+            vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0*ne01 + gid*2]));
+            vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1*ne01 + gid*2]));
+            vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2*ne01 + gid*2]));
+        }
    }
 }
@@ -249,8 +249,9 @@ kernel void Q6K_O4_NAME(
    // ne01/4 apart) is slower because two distant cache-line streams have worse
    // locality than the adjacent pair whose reads interleave into the same lines
    // each iteration.
-    int gid_a = gid * 2;
-    int gid_b = gid * 2 + 1;
+    // Keep tail lanes active for subgroup ops, but read valid weight pairs.
+    int gid_a = min(gid * 2,     ne01 / 2 - 1);
+    int gid_b = min(gid * 2 + 1, ne01 / 2 - 1);

    int nb = ne00 / 32;

@@ -155,6 +155,8 @@ __kernel void kernel_gemv_noshuffle_q8_0_f32_splitk(
    uint M = ne01;

    uint LINE_STRIDE_A  = M;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, LINE_STRIDE_A - 1);
    uint BLOCK_STRIDE_A = 8 * M;   // physical, independent of the K-split

    __private uint8  regA;
@@ -164,19 +166,19 @@ __kernel void kernel_gemv_noshuffle_q8_0_f32_splitk(

    #pragma unroll 1
    for (uint k = kslice * nsg + groupId; k < (K / QK8_0); k += ksplit * nsg) {
-        regS = src0_d[gid + k * LINE_STRIDE_A];
+        regS = src0_d[gid_s + k * LINE_STRIDE_A];
        if (slid < 4) {
            regB.s0123 = read_imagef(src1, (slid * 2 + k * 8));
            regB.s4567 = read_imagef(src1, (1 + slid * 2 + k * 8));
        }
-        regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
-        regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
-        regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
-        regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
-        regA.s4 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
-        regA.s5 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
-        regA.s6 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
-        regA.s7 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
+        regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
+        regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
+        regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
+        regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
+        regA.s4 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
+        regA.s5 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
+        regA.s6 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
+        regA.s7 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;

        dequantizeBlockAccum_ns_sgbroadcast_1(totalSum, regA, convert_float(regS), regB);
    }
@@ -227,6 +229,8 @@ __kernel void kernel_gemv_noshuffle_q8_0_f32(
    uint M = ne01;

    uint LINE_STRIDE_A = M;
+    // Keep tail lanes active for subgroup ops, but read a valid weight row.
+    uint gid_s = min(gid, LINE_STRIDE_A - 1);
    uint BLOCK_STRIDE_A = 8 * M;   // 32 / 4 = 8

    __private uint8     regA;
@@ -238,7 +242,7 @@ __kernel void kernel_gemv_noshuffle_q8_0_f32(
    // loop along K in block granularity, skip 4 blocks every iter
    #pragma unroll 1 /* tell compiler not to unroll */
    for (uint k = groupId; k < (K / QK8_0); k += N_SIMDGROUP) {
-        regS = src0_d[gid + k * LINE_STRIDE_A]; // each fiber loads scale of one rows
+        regS = src0_d[gid_s + k * LINE_STRIDE_A]; // each fiber loads the scale of one row
        // first 4 fibers in each wave load 8 B values to its private scope
        if (slid < 4) {
            regB.s0123 = read_imagef(src1, (slid * 2 + k * 8));
@@ -246,14 +250,14 @@ __kernel void kernel_gemv_noshuffle_q8_0_f32(
        }

        // load weights for one block in consecutive rows
-        regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
-        regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
-        regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
-        regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
-        regA.s4 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
-        regA.s5 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
-        regA.s6 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
-        regA.s7 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;
+        regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x;
+        regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x;
+        regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x;
+        regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x;
+        regA.s4 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x;
+        regA.s5 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x;
+        regA.s6 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x;
+        regA.s7 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x;

        dequantizeBlockAccum_ns_sgbroadcast_1(totalSum, regA, convert_float(regS), regB);
    }
@@ -75,6 +75,15 @@ OutputVector translate_glu_swiglu_oai(const NodeContext & context) {
    const float alpha = reinterpret_cast<const float *>(params)[2];
    const float limit = reinterpret_cast<const float *>(params)[3];

+    // Compute in f32: alpha_const/one are f32, and Swish/Add need matching input types.
+    auto output_type = context.get_output_type();
+    if (src0.get_element_type() != ov::element::f32) {
+        src0 = std::make_shared<ov::op::v0::Convert>(src0, ov::element::f32);
+    }
+    if (src1.get_element_type() != ov::element::f32) {
+        src1 = std::make_shared<ov::op::v0::Convert>(src1, ov::element::f32);
+    }
+
    auto gate = std::make_shared<ov::op::v0::Clamp>(src0, -std::numeric_limits<float>::infinity(), limit);
    auto alpha_const = ov::op::v0::Constant::create(ov::element::f32, {}, {alpha});
    auto out_glu = std::make_shared<ov::op::v4::Swish>(gate, alpha_const);
@@ -82,7 +91,10 @@ OutputVector translate_glu_swiglu_oai(const NodeContext & context) {
    auto up = std::make_shared<ov::op::v0::Clamp>(src1, -limit, limit);
    auto one = ov::op::v0::Constant::create(ov::element::f32, {}, {1.0f});
    auto up_plus_one = std::make_shared<ov::op::v1::Add>(up, one);
-    auto res = std::make_shared<ov::op::v1::Multiply>(out_glu, up_plus_one);
+    ov::Output<ov::Node> res = std::make_shared<ov::op::v1::Multiply>(out_glu, up_plus_one);
+    if (res.get_element_type() != output_type) {
+        res = std::make_shared<ov::op::v0::Convert>(res, output_type);
+    }

    return rename_outputs_with_suffix({res}, context.get_name());
 }
@@ -115,7 +115,8 @@ void * ggml_sycl_malloc_device(size_t size, sycl::queue &q, ggml_sycl_mem_type t
 #else
        ze_device_mem_alloc_desc_t alloc_desc = {ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC, nullptr, 0, 0};
 #endif
-        ze_result_t r = zeMemAllocDevice(ze_ctx, &alloc_desc, size, 64, ze_dev, &ptr);
+        // request SYCL_BUFFER_ALIGNMENT so tensor offsets aligned to it stay aligned in absolute terms
+        ze_result_t r = zeMemAllocDevice(ze_ctx, &alloc_desc, size, SYCL_BUFFER_ALIGNMENT, ze_dev, &ptr);
        if (r == ZE_RESULT_SUCCESS && ptr) {
            ggml_sycl_memtrace_add(type, ptr, size);
            return ptr;
@@ -212,20 +212,10 @@ template <typename dst_t>
 static void dequantize_row_q4_K_sycl(const void *vx, dst_t *y, const int64_t k,
                                     dpct::queue_ptr stream) {
    const int64_t nb = k / QK_K;
-    {
-        dpct::has_capability_or_fail(stream->get_device(),
-                                     {sycl::aspect::fp16});
-
-        stream->submit([&](sycl::handler &cgh) {
-            sycl::local_accessor<uint8_t, 1> scale_local_acc(sycl::range<1>(12), cgh);
-            cgh.parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, nb) *
-                                                   sycl::range<3>(1, 1, 32),
-                                               sycl::range<3>(1, 1, 32)),
-                             [=](sycl::nd_item<3> item_ct1) {
-                                 dequantize_block_q4_K(vx, y, get_pointer(scale_local_acc), item_ct1);
-                             });
-        });
-    }
+    dpct::has_capability_or_fail(stream->get_device(), { sycl::aspect::fp16 });
+    stream->parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t) nb * SYCL_DEQK_WI_PER_BLOCK),
+                                           sycl::range<1>((size_t) SYCL_DEQK_WI_PER_BLOCK)),
+                         [=](sycl::nd_item<1> it) { dequantize_block_q4_K_wide(vx, y, nb, it); });
 }

 template <typename dst_t>
@@ -250,32 +240,10 @@ template <typename dst_t>
 static void dequantize_row_q5_K_sycl(const void *vx, dst_t *y, const int64_t k,
                                     dpct::queue_ptr stream) {
    const int64_t nb = k / QK_K;
-#if QK_K == 256
-    {
-        dpct::has_capability_or_fail(stream->get_device(),
-                                     {sycl::aspect::fp16});
-
-        stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, nb) *
-                                                   sycl::range<3>(1, 1, 64),
-                                               sycl::range<3>(1, 1, 64)),
-                             [=](sycl::nd_item<3> item_ct1) {
-                                 dequantize_block_q5_K(vx, y, item_ct1);
-                             });
-    }
-#else
-    {
-        dpct::has_capability_or_fail(stream->get_device(),
-                                     {sycl::aspect::fp16});
-
-        stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, nb) *
-                                                   sycl::range<3>(1, 1, 32),
-                                               sycl::range<3>(1, 1, 32)),
-                             [=](sycl::nd_item<3> item_ct1) {
-                                 dequantize_block_q5_K(vx, y, item_ct1);
-                             });
-    }
-
-#endif
+    dpct::has_capability_or_fail(stream->get_device(), { sycl::aspect::fp16 });
+    stream->parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t) nb * SYCL_DEQK_WI_PER_BLOCK),
+                                           sycl::range<1>((size_t) SYCL_DEQK_WI_PER_BLOCK)),
+                         [=](sycl::nd_item<1> it) { dequantize_block_q5_K_wide(vx, y, nb, it); });
 }

 template <typename dst_t>
@@ -1255,6 +1255,8 @@ void ggml_sycl_cpy(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, co
    scope_op_debug_print scope_dbg_print(__func__, src1, /*num_src=*/0, debug_get_tensor_str("\tsrc0", src0));
    const int64_t ne = ggml_nelements(src0);
    GGML_ASSERT(ne == ggml_nelements(src1));
+    // the cpy kernels index in 32-bit
+    GGML_ASSERT(ne <= INT32_MAX);

    GGML_TENSOR_BINARY_OP_LOCALS01;

@@ -1108,6 +1108,53 @@ static inline void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8
 }
 #endif

+// Work-items that cover one 256-weight K-quant block: four 64-element chunks, 8 work-items
+// per chunk, 4 qs bytes per work-item.
+static constexpr int SYCL_DEQK_WI_PER_BLOCK = 32;
+
+template <typename dst_t>
+static inline void dequantize_q4_K_wide_one(dst_t * __restrict__ y, const uint8_t * __restrict__ qs,
+                                            const uint8_t * __restrict__ scales, const float dall,
+                                            const float dmin, const int tid) {
+    const int il = tid >> 3;        // 64-element chunk, 0..3
+    const int ir = (tid & 7) * 4;   // 4 qs bytes within the chunk
+
+    uint8_t sc, m;
+    get_scale_min_k4(2 * il + 0, scales, sc, m);
+    const float d1 = dall * sc;
+    const float m1 = dmin * m;
+    get_scale_min_k4(2 * il + 1, scales, sc, m);
+    const float d2 = dall * sc;
+    const float m2 = dmin * m;
+
+    const sycl::vec<uint8_t, 4> q = vec_aligned_load<uint8_t, 4>(qs + 32 * il + ir);
+
+    sycl::vec<dst_t, 4> lo, hi;
+#pragma unroll
+    for (int l = 0; l < 4; ++l) {
+        lo[l] = d1 * (q[l] & 0xF) - m1;
+        hi[l] = d2 * (q[l] >> 4)  - m2;
+    }
+    dst_t * yb = y + 64 * il + ir;
+    *reinterpret_cast<sycl::vec<dst_t, 4> *>(yb +  0) = lo;
+    *reinterpret_cast<sycl::vec<dst_t, 4> *>(yb + 32) = hi;
+}
+
+template <typename dst_t>
+static void dequantize_block_q4_K_wide(const void * __restrict__ vx, dst_t * __restrict__ yy, int64_t nb,
+                                       const sycl::nd_item<1> & it) {
+    static_assert(QK_K == 256, "wide K-quant dequant assumes QK_K == 256");
+    const int64_t g   = it.get_global_id(0);
+    const int64_t b   = g / SYCL_DEQK_WI_PER_BLOCK;
+    const int     tid = (int) (g % SYCL_DEQK_WI_PER_BLOCK);
+    if (b >= nb) {
+        return;
+    }
+    const block_q4_K * x = (const block_q4_K *) vx + b;
+    dequantize_q4_K_wide_one(yy + b * QK_K, x->qs, x->scales, (float) x->dm[0], (float) x->dm[1], tid);
+}
+
+
 template <typename dst_t>
 inline void dequantize_q4_K_common(dst_t * __restrict__ y, const uint8_t * __restrict__ qs_ptr, const float dall,
                                   const float dmin, uint8_t * __restrict__ scales_local, int il, int ir) {
@@ -1130,41 +1177,6 @@ inline void dequantize_q4_K_common(dst_t * __restrict__ y, const uint8_t * __res
    }
 }

-template<typename dst_t>
-static void dequantize_block_q4_K(const void * __restrict__ vx, dst_t * __restrict__ yy,
-                                  uint8_t* scales_local, const sycl::nd_item<3> &item_ct1) {
-    const block_q4_K * x = (const block_q4_K *) vx;
-
-    const int64_t i = item_ct1.get_group(2);
-
-#if QK_K == 256
-    const int64_t tid = item_ct1.get_local_id(2);
-    const int64_t il  = tid / 8;
-    const int64_t ir  = tid % 8;
-
-    dst_t * y = yy + i * QK_K + 64 * il + 4 * ir;
-
-    const sycl::half2 dm = x[i].dm;
-    const float dall = dm[0];
-    const float dmin = dm[1];
-
-    if (tid < 12) {
-        scales_local[tid] = x[i].scales[tid];
-    }
-
-    item_ct1.barrier(sycl::access::fence_space::local_space);
-    dequantize_q4_K_common(y, x[i].qs, dall, dmin, scales_local, il, ir);
-#else
-    const int64_t tid = item_ct1.get_local_id(2);
-    const uint8_t * q = x[i].qs;
-    dst_t * y = yy + i*QK_K;
-    const float d = (float)x[i].dm[0];
-    const float m = (float)x[i].dm[1];
-    y[tid+ 0] = d * (x[i].scales[0] & 0xF) * (q[tid] & 0xF) - m * (x[i].scales[0] >> 4);
-    y[tid+32] = d * (x[i].scales[1] & 0xF) * (q[tid] >>  4) - m * (x[i].scales[1] >> 4);
-#endif
-}
-
 template <typename dst_t>
 static void dequantize_block_q4_K_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy, uint8_t * scales_local,
                                          const sycl::nd_item<1> & item_ct1, int64_t nb) {
@@ -1195,52 +1207,51 @@ static void dequantize_block_q4_K_reorder(const void * __restrict__ vx, dst_t *
    dequantize_q4_K_common(y, qs_ptr, dall, dmin, scales_local, il, ir);
 }

-template<typename dst_t>
-static void dequantize_block_q5_K(const void * __restrict__ vx, dst_t * __restrict__ yy,
-                                  const sycl::nd_item<3> &item_ct1) {
-    const block_q5_K * x = (const block_q5_K *) vx;
-
-    const int64_t i = item_ct1.get_group(2);
-
-#if QK_K == 256
-    // assume 64 threads - this is very slightly better than the one below
-    const int64_t tid = item_ct1.get_local_id(2);
-    const int64_t il  = tid/16;   // il is in 0...3
-    const int64_t ir  = tid%16;   // ir is in 0...15
-    const int64_t is  = 2*il;     // is is in 0...6
-
-    dst_t * y = yy + i*QK_K + 64*il + 2*ir;
-
-    const float dall = x[i].dm[0];
-    const float dmin = x[i].dm[1];
-
-    const uint8_t * ql = x[i].qs + 32*il + 2*ir;
-    const uint8_t * qh = x[i].qh + 2*ir;
+template <typename dst_t>
+static inline void dequantize_q5_K_wide_one(dst_t * __restrict__ y, const uint8_t * __restrict__ qs,
+                                            const uint8_t * __restrict__ qh_ptr,
+                                            const uint8_t * __restrict__ scales, const float dall,
+                                            const float dmin, const int tid) {
+    const int il = tid >> 3;        // 64-element chunk, 0..3
+    const int ir = (tid & 7) * 4;   // 4 qs bytes within the chunk; qh holds the 5th bit

    uint8_t sc, m;
-    get_scale_min_k4(is + 0, x[i].scales, sc, m);
-    const float d1 = dall * sc; const float m1 = dmin * m;
-    get_scale_min_k4(is + 1, x[i].scales, sc, m);
-    const float d2 = dall * sc; const float m2 = dmin * m;
+    get_scale_min_k4(2 * il + 0, scales, sc, m);
+    const float d1 = dall * sc;
+    const float m1 = dmin * m;
+    get_scale_min_k4(2 * il + 1, scales, sc, m);
+    const float d2 = dall * sc;
+    const float m2 = dmin * m;

-    uint8_t   hm  = 1 << (2*il);
-    y[ 0] = d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1;
-    y[ 1] = d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1;
-    hm <<= 1;
-    y[32] = d2 * ((ql[ 0] >>  4) + (qh[ 0] & hm ? 16 : 0)) - m2;
-    y[33] = d2 * ((ql[ 1] >>  4) + (qh[ 1] & hm ? 16 : 0)) - m2;
-#else
-    const int64_t tid = item_ct1.get_local_id(2);
-    const uint8_t q = x[i].qs[tid];
-    const int64_t im = tid/8;  // 0...3
-    const int64_t in = tid%8;  // 0...7
-    const int64_t is = tid/16; // 0 or 1
-    const uint8_t h = x[i].qh[in] >> im;
-    const float d = x[i].d;
-    dst_t * y = yy + i*QK_K + tid;
-    y[ 0] = d * x[i].scales[is+0] * ((q & 0xF) - ((h >> 0) & 1 ? 0 : 16));
-    y[32] = d * x[i].scales[is+2] * ((q >>  4) - ((h >> 4) & 1 ? 0 : 16));
-#endif
+    const sycl::vec<uint8_t, 4> q = vec_aligned_load<uint8_t, 4>(qs + 32 * il + ir);
+    const sycl::vec<uint8_t, 4> h = vec_aligned_load<uint8_t, 4>(qh_ptr + ir);
+
+    const uint8_t hm1 = 1 << (2 * il);
+    const uint8_t hm2 = hm1 << 1;
+
+    sycl::vec<dst_t, 4> lo, hi;
+#pragma unroll
+    for (int l = 0; l < 4; ++l) {
+        lo[l] = d1 * ((q[l] & 0xF) + ((h[l] & hm1) ? 16 : 0)) - m1;
+        hi[l] = d2 * ((q[l] >>  4) + ((h[l] & hm2) ? 16 : 0)) - m2;
+    }
+    dst_t * yb = y + 64 * il + ir;
+    *reinterpret_cast<sycl::vec<dst_t, 4> *>(yb +  0) = lo;
+    *reinterpret_cast<sycl::vec<dst_t, 4> *>(yb + 32) = hi;
+}
+
+template <typename dst_t>
+static void dequantize_block_q5_K_wide(const void * __restrict__ vx, dst_t * __restrict__ yy, int64_t nb,
+                                       const sycl::nd_item<1> & it) {
+    static_assert(QK_K == 256, "wide K-quant dequant assumes QK_K == 256");
+    const int64_t g   = it.get_global_id(0);
+    const int64_t b   = g / SYCL_DEQK_WI_PER_BLOCK;
+    const int     tid = (int) (g % SYCL_DEQK_WI_PER_BLOCK);
+    if (b >= nb) {
+        return;
+    }
+    const block_q5_K * x = (const block_q5_K *) vx + b;
+    dequantize_q5_K_wide_one(yy + b * QK_K, x->qs, x->qh, x->scales, (float) x->dm[0], (float) x->dm[1], tid);
 }

 template <typename dst_t>
--- a/Show More
+++ b/Show More