vendor: attempt to ignore warnings from vendored files

Signed-off-by: Aaron Teo <[email protected]>

.devops/musa.Dockerfile

@@ -1,10 +1,9 @@
 ARG UBUNTU_VERSION=22.04
 # This needs to generally match the container host's environment.
-ARG MUSA_VERSION=rc4.3.0
 # Target the MUSA build image
-ARG BASE_MUSA_DEV_CONTAINER=docker.io/mthreads/musa:${MUSA_VERSION}-devel-ubuntu${UBUNTU_VERSION}-amd64
+ARG BASE_MUSA_DEV_CONTAINER=registry.mthreads.com/mcconline/inference/pytorch:2.9.1.post1-py3.10-musa5.2.0-mp31-devel-ubuntu${UBUNTU_VERSION}-amd64
 
-ARG BASE_MUSA_RUN_CONTAINER=docker.io/mthreads/musa:${MUSA_VERSION}-runtime-ubuntu${UBUNTU_VERSION}-amd64
+ARG BASE_MUSA_RUN_CONTAINER=${BASE_MUSA_DEV_CONTAINER}
 
 ARG BUILD_DATE=N/A
 ARG APP_VERSION=N/A

.github/workflows/build-apple.yml

@@ -33,6 +33,7 @@ concurrency:
 env:
   GGML_NLOOP: 3
   GGML_N_THREADS: 1
+  GGML_SCHED_DEBUG_REALLOC: 1
   LLAMA_ARG_LOG_COLORS: 1
   LLAMA_ARG_LOG_PREFIX: 1
   LLAMA_ARG_LOG_TIMESTAMPS: 1
@@ -98,7 +99,8 @@ jobs:
         id: cmake_test
         run: |
           cd build
-          ctest -L main -E "test-llama-archs" --verbose --timeout 900
+          # ref: https://github.com/ggml-org/llama.cpp/pull/19802#issuecomment-4013704023
+          ctest -L main -E "test-llama-archs|test-save-load-state" --verbose --timeout 900
 
   macos-latest-x64:
     runs-on: macos-15-intel

.github/workflows/build-cpu.yml

@@ -37,6 +37,7 @@ concurrency:
 env:
   GGML_NLOOP: 3
   GGML_N_THREADS: 1
+  GGML_SCHED_DEBUG_REALLOC: 1
   LLAMA_ARG_LOG_COLORS: 1
   LLAMA_ARG_LOG_PREFIX: 1
   LLAMA_ARG_LOG_TIMESTAMPS: 1

.github/workflows/build-ibm.yml

@@ -100,6 +100,46 @@ jobs:
           wget https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K-be.gguf
           ./bin/llama-completion -m stories260K-be.gguf -p "One day, Lily met a Shoggoth" -n 500 -c 256
 
+  ubuntu-26-zdnn-s390x:
+    name: ubuntu-26-zdnn-s390x
+    runs-on: ubuntu-24.04-s390x
+    container: ubuntu:26.04  # required to get GCC 15.1 and binutils 2.44
+    defaults:
+      run:
+        shell: bash
+
+    steps:
+      - name: Build Dependencies
+        id: build_depends
+        run: |
+          apt-get update
+          apt-get install -y --no-install-recommends \
+            build-essential cmake git ca-certificates \
+            libssl-dev libzdnn-dev
+
+      - name: Clone
+        id: checkout
+        uses: actions/checkout@v6
+
+      - name: Toolchain workaround (GCC 15)
+        run: |
+          apt-get install -y gcc-15 g++-15
+          echo "CC=gcc-15" >> "$GITHUB_ENV"
+          echo "CXX=g++-15" >> "$GITHUB_ENV"
+
+      - name: Build with zDNN Backend
+        id: cmake_build
+        run: |
+          cmake -B build \
+            -DLLAMA_FATAL_WARNINGS=ON \
+            -DGGML_NATIVE=OFF \
+            -DGGML_VXE=ON \
+            -DGGML_ZDNN=ON \
+            -DGGML_RPC=ON \
+            -DCMAKE_C_FLAGS="-march=arch15" \
+            -DCMAKE_CXX_FLAGS="-march=arch15"
+          time cmake --build build --config Release -j $(nproc)
+
   ubuntu-24-ppc64le:
     runs-on: ubuntu-24.04-ppc64le
 

.github/workflows/build-vulkan.yml

@@ -31,6 +31,7 @@ concurrency:
 env:
   GGML_NLOOP: 3
   GGML_N_THREADS: 1
+  GGML_SCHED_DEBUG_REALLOC: 1
   LLAMA_ARG_LOG_COLORS: 1
   LLAMA_ARG_LOG_PREFIX: 1
   LLAMA_ARG_LOG_TIMESTAMPS: 1

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

@@ -31,7 +31,7 @@ jobs:
         uses: actions/setup-python@v6
         with:
           python-version: "3.11"
-          pip-install: -r requirements/requirements-all.txt ty==0.0.78
+          pip-install: -r requirements/requirements-all.txt ty==0.0.84
       # - name: Type-check with Pyright
       #   uses: jakebailey/pyright-action@v2
       #   with:

CODEOWNERS

@@ -57,7 +57,7 @@
 /ggml/src/ggml-cann/                    @ggml-org/ggml-cann
 /ggml/src/ggml-common.h                 @ggerganov
 /ggml/src/ggml-cpu/                     @ggerganov
-/ggml/src/ggml-cpu/iqp.*                @bartowski1182
+/ggml/src/ggml-cpu/tiled/               @jbooth @bartowski1182
 /ggml/src/ggml-cpu/spacemit/            @alex-spacemit
 /ggml/src/ggml-cuda/                    @ggml-org/ggml-cuda
 /ggml/src/ggml-cuda/vendors/hip.h       @IMbackK

common/arg.cpp

@@ -2673,16 +2673,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.video_ffmpeg_bin_dir = value;
         }
     ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_FFMPEG_DIR"));
-    if (params.is_gen_docs || llama_supports_rpc()) {
-        add_opt(common_arg(
-            {"--rpc"}, "SERVERS",
-            "comma-separated list of RPC servers (host:port)",
-            [](common_params & params, const std::string & value) {
-                add_rpc_devices(value);
-                GGML_UNUSED(params);
+    add_opt(common_arg(
+        {"--rpc"}, "SERVERS",
+        "comma-separated list of RPC servers (host:port)",
+        [](common_params & params, const std::string & value) {
+            if (!llama_supports_rpc()) {
+                throw std::invalid_argument("RPC not supported in this build");
             }
-        ).set_env("LLAMA_ARG_RPC"));
-    }
+            add_rpc_devices(value);
+            GGML_UNUSED(params);
+        }
+    ).set_env("LLAMA_ARG_RPC"));
     add_opt(common_arg(
         {"-lm", "--load-mode"}, "MODE",
         "model loading mode (default: auto)\n"

common/common.h

@@ -809,7 +809,9 @@ static std::vector<T> string_split(const std::string & str, char delim) {
     while (std::getline(str_stream, token, delim)) {
         T value;
         std::istringstream token_stream(token);
-        token_stream >> value;
+        if (!(token_stream >> value)) {
+            throw std::invalid_argument("invalid value: \"" + token + "\"");
+        }
         values.push_back(value);
     }
     return values;

common/fit.cpp

@@ -192,9 +192,9 @@ static void common_params_fit_impl(
     uint32_t hp_nct = 0; // hparams.n_ctx_train
     uint32_t hp_nex = 0; // hparams.n_expert
 
-    // size the context for all sequences, but keep minimums and alignment per KV stream
-    const uint32_t n_seq_max  = std::max<uint32_t>(1, cparams->n_seq_max);
-    const uint32_t n_streams  = cparams->kv_unified ? 1 : n_seq_max;
+    // with non-unified kv, we need to take into account n_streams
+    // for example, if memory can hold more than model's trained context size, we must extend the n_ctx to hold enough n_streams
+    const uint32_t n_streams  = cparams->kv_unified ? 1 : std::max<uint32_t>(1, cparams->n_seq_max);
     const bool     n_ctx_auto = cparams->n_ctx == 0;
 
     dmds_t   dmds_extra;       // memory of the extra model, laid out on the devices of the main model
@@ -264,15 +264,15 @@ static void common_params_fit_impl(
     dmds_t dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
 
     // saturate instead of overflowing, this also preserves the UINT32_MAX sentinel of n_ctx_min:
-    const uint32_t n_ctx_max       = (uint32_t) std::min<uint64_t>(uint64_t(hp_nct)    * n_seq_max, UINT32_MAX);
+    const uint32_t n_ctx_max       = (uint32_t) std::min<uint64_t>(uint64_t(hp_nct)    * n_streams, UINT32_MAX);
     const uint32_t n_ctx_min_total = (uint32_t) std::min<uint64_t>(uint64_t(n_ctx_min) * n_streams, UINT32_MAX);
 
     // llama_context would use only hp_nct in total for n_ctx == 0, resolve the context before measuring anything else:
     if (n_ctx_auto) {
         cparams->n_ctx = n_ctx_max;
-        if (n_seq_max > 1) {
-            LOG_TRC("%s: context size unset -> using %" PRIu32 " for %" PRIu32 " sequences:\n",
-                __func__, n_ctx_max, n_seq_max);
+        if (n_streams > 1) {
+            LOG_TRC("%s: context size unset and KV cache not unified -> using %" PRIu32 " for %" PRIu32 " sequences:\n",
+                __func__, n_ctx_max, n_streams);
             dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
         }
     }

common/jinja/parser.cpp

@@ -429,8 +429,15 @@ private:
             bool negate = false;
             if (is_identifier("not")) { ++current; negate = true; }
             auto test_id = parse_primary_expression();
-            // FIXME: tests can also be expressed like this: if x is eq 3
-            if (is(token::open_paren)) test_id = parse_call_expression(std::move(test_id));
+            if (is(token::open_paren)) {
+                test_id = parse_call_expression(std::move(test_id));
+            } else if (is(token::numeric_literal) || is(token::string_literal) || is(token::open_curly_bracket) || is(token::open_square_bracket) ||
+                    (is(token::identifier) && !is_identifier("and") && !is_identifier("or") && !is_identifier("else"))) {
+                size_t call_pos = current;
+                statements args;
+                args.push_back(parse_unary_expression());
+                test_id = mk_stmt<call_expression>(call_pos, std::move(test_id), std::move(args));
+            }
             operand = mk_stmt<test_expression>(start_pos, std::move(operand), negate, std::move(test_id));
         }
         return operand;

common/jinja/value.cpp

@@ -348,6 +348,43 @@ static value default_value(const func_args & args) {
     return no_value ? args.get_pos(1) : args.get_pos(0);
 }
 
+static value toobject(const func_args & args) {
+    auto out = mk_val<value_object>();
+    value iter = args.get_pos(0, mk_val<value_undefined>());
+    bool iter_first = false;
+    if (is_val<value_array>(iter)) {
+        iter_first = true;
+        for (const auto & it : iter->as_array()) {
+            if (is_val<value_array>(it) && it->as_array().size() == 2) {
+                auto tuple = it->as_array();
+                auto key = tuple[0];
+                auto val = tuple[1];
+                JJ_DEBUG("namespace/dict: adding key '%s'", key->as_string().str().c_str());
+                out->insert(key, val);
+            } else {
+                throw raised_exception("namespace/dict() iterable argument must consist of tuples, not " + it->type());
+            }
+        }
+    } else if (is_val<value_object>(iter)) {
+        iter_first = true;
+        for (const auto & pair : iter->as_ordered_object()) {
+            JJ_DEBUG("namespace/dict: adding key '%s'", pair.first->as_string().str().c_str());
+            out->insert(pair.first, pair.second);
+        }
+    }
+    for (const auto & arg : args.get_args()) {
+        if (is_val<value_kwarg>(arg)) {
+            auto kwarg = cast_val<value_kwarg>(arg);
+            JJ_DEBUG("namespace/dict: adding key '%s'", kwarg->key.c_str());
+            out->insert(kwarg->key, kwarg->val);
+        } else if (!iter_first) {
+            throw raised_exception("namespace/dict() arguments must be kwargs, dict and/or iterable of tuples, not " + arg->type());
+        }
+        iter_first = false;
+    }
+    return out;
+}
+
 const func_builtins & global_builtins() {
     static const func_builtins builtins = {
         {"raise_exception", [](const func_args & args) -> value {
@@ -355,18 +392,8 @@ const func_builtins & global_builtins() {
             std::string msg = args.get_pos(0)->as_string().str();
             throw raised_exception("Jinja Exception: " + msg);
         }},
-        {"namespace", [](const func_args & args) -> value {
-            auto out = mk_val<value_object>();
-            for (const auto & arg : args.get_args()) {
-                if (!is_val<value_kwarg>(arg)) {
-                    throw raised_exception("namespace() arguments must be kwargs");
-                }
-                auto kwarg = cast_val<value_kwarg>(arg);
-                JJ_DEBUG("namespace: adding key '%s'", kwarg->key.c_str());
-                out->insert(kwarg->key, kwarg->val);
-            }
-            return out;
-        }},
+        {"dict", toobject},
+        {"namespace", toobject},
         {"strftime_now", [](const func_args & args) -> value {
             args.ensure_vals<value_string>();
             std::string format = args.get_pos(0)->as_string().str();
@@ -515,8 +542,28 @@ const func_builtins & global_builtins() {
         }},
         {"test_is_sameas", [](const func_args & args) -> value {
             // Check if an object points to the same memory address as another object
-            (void)args;
-            throw not_implemented_exception("sameas test not implemented");
+            args.ensure_count(2);
+            auto a = args.get_pos(0);
+            auto b = args.get_pos(1);
+            bool res = false;
+            if (!is_val<value_undefined>(a) && !is_val<value_undefined>(b)) {
+                if (is_val<value_none>(a) && is_val<value_none>(b)) {
+                    res = true;
+                } else if (is_val<value_bool>(a) && is_val<value_bool>(b)) {
+                    if (a->as_bool() == b->as_bool()) {
+                        res = true;
+                    }
+                } else if (is_val<value_int>(a) && is_val<value_int>(b)) {
+                    const int64_t x = a->as_int();
+                    // Allow comparison within small-int cache range
+                    if (x >= -5 && x <= 256 && x == b->as_int()) {
+                        res = true;
+                    }
+                } else if (a == b) {
+                    res = true;
+                }
+            }
+            return mk_val<value_bool>(res);
         }},
         {"test_is_escaped", [](const func_args & args) -> value {
             (void)args;

common/sampling.cpp

@@ -214,7 +214,7 @@ struct common_sampler * common_sampler_init(
 #ifdef LLAMA_USE_LLGUIDANCE
         grmr = llama_sampler_init_llg(vocab, "lark", grammar_str.c_str());
 #else
-        GGML_ABORT("llguidance (cmake -DLLAMA_LLGUIDANCE=ON) is not enabled");
+        throw std::runtime_error("failed to parse grammar: llguidance is not enabled");
 #endif // LLAMA_USE_LLGUIDANCE
     } else {
         std::vector<std::string> trigger_patterns;

conversion/bert.py

@@ -341,7 +341,7 @@ class NomicBertModel(BertModel):
         else:
             raise ValueError(f"unrecognized parameters: n_positions={npos}, max_trained_positions={mtp}")
 
-        assert self.hparams["activation_function"] == "gelu" if self.is_moe else "swiglu"
+        assert self.hparams["activation_function"] == ("gelu" if self.is_moe else "swiglu")
 
         # this doesn't do anything in the HF version
         assert self.hparams["causal"] is False

conversion/plamo.py

@@ -154,6 +154,21 @@ class Plamo2Model(TextModel):
 class Plamo3Model(TextModel):
     model_arch = gguf.MODEL_ARCH.PLAMO3
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # PLaMo-3 builds rope_parameters from flat config keys at runtime; mirror the YaRN settings for GGUF.
+        rope_scaling_factor = self.hparams.get("rope_scaling_factor", 1)
+        if rope_scaling_factor != 1 and "rope_type" not in self.rope_parameters:
+            self.rope_parameters.update({
+                "rope_type": "yarn",
+                "factor": float(rope_scaling_factor),
+                "original_max_position_embeddings": int(self.hparams["initial_context_length"]),
+                "beta_fast": 32.0,
+                "beta_slow": 1.0,
+                "truncate": False,
+            })
+
     def set_vocab(self):
         self._set_vocab_plamo()
 

docs/docker.md

@@ -123,7 +123,7 @@ You may want to pass in some different `ARGS`, depending on the MUSA environment
 
 The defaults are:
 
-- `MUSA_VERSION` set to `rc4.3.0`
+- the base image is the MUSA 5.2.0 image from the Moore Threads registry
 
 The resulting images, are essentially the same as the non-MUSA images:
 

docs/ops.md

@@ -14,16 +14,16 @@ Legend:
 
 | Operation | BLAS | CANN | CPU | CUDA | ET | HTP | MTL | OpenCL | SYCL | Vulkan | WebGPU | ZenDNN | zDNN |
 |-----------|------|------|------|------|------|------|------|------|------|------|------|------|------|
-|                              ABS | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              ABS | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              ACC | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              ADD | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             ADD1 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                           ADD_ID | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           ARANGE | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                           ARGMAX | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                          ARGSORT | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           ARGMAX | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                          ARGSORT | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             CEIL | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                            CLAMP | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            CLAMP | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        COL2IM_1D | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           CONCAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                             CONT | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
@@ -55,7 +55,7 @@ Legend:
 |                GATED_LINEAR_ATTN | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                            GEGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        GEGLU_ERF | ❌ | ✅ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                      GEGLU_QUICK | ❌ | ✅ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                      GEGLU_QUICK | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             GELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         GELU_ERF | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                       GELU_QUICK | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -64,17 +64,21 @@ Legend:
 |                       GROUP_NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                      HARDSIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        HARDSWISH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                           IM2COL | ❌ | ✅ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           IM2COL | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        IM2COL_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                          L2_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | 🟡 | ❌ | ❌ |
-|                       LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | 🟡 | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                       LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                LIGHTNING_INDEXER | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                              LOG | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              LOG | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ | ❌ |
 |                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
-|                 MUL_MAT_HADAMARD | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                 MUL_MAT_HADAMARD | ❌ | ❌ | ❌ | ❌ | ✅ | 🟡 | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | 🟡 | ❌ |
+|                  MUL_MAT_ID_W4A4 | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A8 | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A4 | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A8 | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -85,12 +89,12 @@ Legend:
 |                          POOL_1D | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                          POOL_2D | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                            REGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             RELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             RELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           REPEAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                      REPEAT_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                         RMS_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    RMS_NORM_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                             ROLL | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                             ROLL | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                             ROPE | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        ROPE_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                            ROUND | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -108,20 +112,20 @@ Legend:
 |                         SOFT_MAX | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    SOFT_MAX_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ |
 |                        SOLVE_TRI | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ |
-|                              SQR | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             SQRT | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              SQR | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             SQRT | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         SSM_CONV | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         SSM_SCAN | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | 🟡 | 🟡 | ✅ | ❌ | ❌ |
-|                             STEP | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             STEP | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              SUB | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                              SUM | ❌ | 🟡 | ✅ | 🟡 | ❌ | ❌ | 🟡 | ❌ | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
+|                              SUM | ❌ | 🟡 | ✅ | 🟡 | ❌ | 🟡 | 🟡 | ❌ | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
 |                         SUM_ROWS | ❌ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ |
 |                           SWIGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                     SWIGLU_CLAMP | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                       SWIGLU_OAI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             TANH | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |               TIMESTEP_EMBEDDING | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            TOP_K | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ✅ | ❌ | ✅ | 🟡 | ✅ | ❌ | ❌ |
+|                            TOP_K | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | 🟡 | ✅ | ❌ | ❌ |
 |                              TRI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                            TRUNC | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          UPSCALE | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |

ggml/include/ggml-zdnn.h

@@ -10,6 +10,8 @@ extern "C" {
 // device buffer
 GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_zdnn_buffer_type(void);
 
+GGML_BACKEND_API bool ggml_backend_is_zdnn(ggml_backend_t backend);
+
 GGML_BACKEND_API ggml_backend_reg_t ggml_backend_zdnn_reg(void);
 
 #ifdef __cplusplus

ggml/src/ggml-cpu/CMakeLists.txt

@@ -31,8 +31,6 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
         ggml-cpu/ggml-cpu.cpp
         ggml-cpu/repack.cpp
         ggml-cpu/repack.h
-        ggml-cpu/iqp.cpp
-        ggml-cpu/iqp.h
         ggml-cpu/hbm.cpp
         ggml-cpu/hbm.h
         ggml-cpu/quants.c
@@ -43,6 +41,10 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
         ggml-cpu/amx/amx.h
         ggml-cpu/amx/mmq.cpp
         ggml-cpu/amx/mmq.h
+        ggml-cpu/tiled/tiled.cpp
+        ggml-cpu/tiled/tiled.h
+        ggml-cpu/tiled/tiled-kernel.cpp
+        ggml-cpu/tiled/tiled-kernel.h
         ggml-cpu/ggml-cpu-impl.h
         ggml-cpu/common.h
         ggml-cpu/binary-ops.h

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

@@ -4,7 +4,6 @@
 #include "ggml-backend-impl.h"
 #include "ggml-backend.h"
 #include "traits.h"
-#include "iqp.h"
 #include "ggml-cpu-impl.h"
 #include "ggml-impl.h"
 #include "quants.h"
@@ -15,6 +14,7 @@
 #include "ops.h"
 #include "ggml.h"
 #include "common.h"
+#include "tiled/tiled.h"
 
 #if defined(_MSC_VER) || defined(__MINGW32__)
 #include <malloc.h> // using malloc.h with MSC/MINGW
@@ -1265,6 +1265,11 @@ void ggml_compute_forward_mul_mat(
         return;
     }
 
+    // If tiled is supported, it will execute the full op here and we return
+    if (ggml_compute_forward_mul_mat_tiled(params, dst)) {
+        return;
+    }
+
     GGML_TENSOR_BINARY_OP_LOCALS
 
     const int ith = params->ith;
@@ -1372,13 +1377,6 @@ UseGgmlGemm1:;
 
     ggml_barrier(params->threadpool);
 
-    // IQ panel gemm (see iqp.h) - must come after the barrier above, it consumes the q8_K rows
-    // of src1 from the work buffer
-    if (ggml_cpu_iqp_supports_mul_mat(dst) && !params->use_ref) {
-        ggml_compute_forward_mul_mat_iqp(params, dst);
-        return;
-    }
-
 #if GGML_USE_LLAMAFILE
     if (src1->type != vec_dot_type) {
         const void* wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata;
@@ -1596,15 +1594,9 @@ static void ggml_compute_forward_mul_mat_id(
     char (*atomic_current_chunk)[CACHE_LINE_SIZE] = // [n_as]
         incr_ptr_aligned(&wdata_cur, CACHE_LINE_SIZE * n_as, CACHE_LINE_SIZE);
 
-    // IQ panel gemm (see iqp.h); per expert eligibility is decided below, but the work buffer is
-    // reserved for the whole node (ggml_graph_plan sizes it without params, use_ref only skips the dispatch)
-    const bool iqp = ggml_cpu_iqp_supports_mul_mat_id(dst) && !params->use_ref;
-
-    char * iqp_panels = NULL;
-
-    if (iqp) {
-        iqp_panels = incr_ptr_aligned(&wdata_cur, nth * ggml_cpu_iqp_scratch_size(dst), 64);
-    }
+    // Tiled matmul (see tiled.h); per-thread work buffers, 0 bytes when disabled. The
+    // reservation is unconditional, the per expert eligibility is decided at dispatch time
+    char * tiled_scratch = incr_ptr_aligned(&wdata_cur, ggml_tiled_wdata_size(nth, dst), 64);
 
     GGML_ASSERT(params->wsize >= (size_t)((char *) wdata_cur - (char *) params->wdata));
 
@@ -1677,10 +1669,8 @@ static void ggml_compute_forward_mul_mat_id(
             continue;
         }
 
-        if (iqp && ggml_cpu_iqp_mul_mat_id_min_batch(cne1)) {
-            ggml_compute_forward_mul_mat_id_iqp(params, dst, cur_a, cne1, (const int32_t *) &MMID_MATRIX_ROW(cur_a, 0),
-                                                iqp_panels);
-
+        // tiled takes over if profitable for this expert (see tiled.h)
+        if (ggml_compute_forward_mul_mat_id_tiled(params, dst, cur_a, cne1, (const int32_t *) &MMID_MATRIX_ROW(cur_a, 0), tiled_scratch)) {
             continue;
         }
 
@@ -2887,15 +2877,12 @@ struct ggml_cplan ggml_graph_plan(
                 case GGML_OP_MUL_MAT:
                     {
                         const enum ggml_type vec_dot_type = type_traits_cpu[node->src[0]->type].vec_dot_type;
-
                         if (node->src[1]->type != vec_dot_type) {
                             cur = ggml_row_size(vec_dot_type, ggml_nelements(node->src[1]));
                         }
-
-                        // the IQ panel path needs one scratch panel per thread past the q8_K rows
-                        if (ggml_cpu_iqp_supports_mul_mat(node)) {
-                            cur = GGML_PAD(cur, 64) + n_tasks * ggml_cpu_iqp_scratch_size(node);
-                        }
+                        // Workspace for tiled (see tiled.h)
+                        cur = GGML_PAD(cur, 64);
+                        cur += ggml_tiled_wdata_size(n_tasks, node);
                     } break;
                 case GGML_OP_MUL_MAT_ID:
                     {
@@ -2915,10 +2902,9 @@ struct ggml_cplan ggml_graph_plan(
                         cur += n_as*ids->ne[0]*ids->ne[1]*sizeof(struct mmid_row_mapping) + sizeof(int64_t);
                         // atomic_current_chunk
                         cur += CACHE_LINE_SIZE*n_as + CACHE_LINE_SIZE;
-                        // the IQ panel path needs one scratch panel per thread on top of that
-                        if (ggml_cpu_iqp_supports_mul_mat_id(node)) {
-                            cur += n_tasks * ggml_cpu_iqp_scratch_size(node) + 64;
-                        }
+                        // Workspace for tiled (see tiled.h)
+                        cur = GGML_PAD(cur, 64);
+                        cur += ggml_tiled_wdata_size(n_tasks, node);
                     } break;
                 case GGML_OP_OUT_PROD:
                     {

ggml/src/ggml-cpu/iqp.h

@@ -1,39 +0,0 @@
-#pragma once
-
-#include "ggml-cpu-impl.h"
-#include "ggml.h"
-
-// GGML internal header
-
-// batched mul_mat path for the grid based IQ types: decode 8 src0 rows at a time into per thread scratch
-// (block_iqp_x8, see iqp.cpp) and run an integer gemm over them against all src1 columns
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-
-// whether cne1 rows of src1 are enough for the decode to pay for itself, per expert, for MUL_MAT_ID
-bool ggml_cpu_iqp_mul_mat_id_min_batch(int64_t cne1);
-
-bool ggml_cpu_iqp_supports_mul_mat(const struct ggml_tensor * dst);
-
-// node level test only - per expert eligibility is decided with ggml_cpu_iqp_mul_mat_id_min_batch
-bool ggml_cpu_iqp_supports_mul_mat_id(const struct ggml_tensor * dst);
-
-// per thread panel scratch bytes, padded
-size_t ggml_cpu_iqp_scratch_size(const struct ggml_tensor * dst);
-
-// must be called after src1 has been converted to q8_K into params->wdata and the threads have synchronized on it
-void ggml_compute_forward_mul_mat_iqp(const struct ggml_compute_params * params, struct ggml_tensor * dst);
-
-// one expert: expert_rows points at its row of the matrix_rows table of (i1, i2) int32 pairs, panels at the base of the per thread panel scratches
-void ggml_compute_forward_mul_mat_id_iqp(const struct ggml_compute_params * params,
-                                         struct ggml_tensor *               dst,
-                                         int64_t                            cur_a,
-                                         int64_t                            cne1,
-                                         const int32_t *                    expert_rows,
-                                         void *                             panels);
-
-#ifdef __cplusplus
-}
-#endif

ggml/src/ggml-cpu/tiled/tiled-kernel.cpp

@@ -0,0 +1,719 @@
+// mulmat microtile kernels
+
+#include "tiled-kernel.h"
+
+#include "ggml.h"
+
+#include <cstring>
+
+#if defined(__AVX512VNNI__) || defined(__AVX2__) || defined(__AVX__)
+#include <immintrin.h>
+#endif
+
+
+// Reference implementation, slower than existing vec_dot approach
+template <int SUBBLK, bool HAS_MIN, int BIAS, int NK>
+static void tiled_run_microtile_scalar(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1,
+                                       int i0, int j0, int num_k, int slab, float * buf, int buf_stride) {
+    constexpr int NB = TILED_TILE_K / SUBBLK;    // subblocks per 256-K block
+    constexpr int NS = SUBBLK / 16; // per-16 bsums per subblock
+
+    // num_k/slab: the narrow path holds num_k slabs at row stride num_k*256 (num_k=1, slab=0 = standard)
+    // NK: 1 = standard single-slab (offsets 0, strides compile-time); 0 = runtime
+    // (num_k/slab from the args). The narrow path is memory-bound, so one runtime
+    // version serves all of it.
+    const int nkr  = (NK > 0) ? NK : num_k;
+    const int seff = (NK > 0) ? 0 : slab;
+    const int qk_stride = nkr * TILED_TILE_K;
+    const int qk_off = seff * TILED_TILE_K;
+    const int nb_stride = NB * nkr;
+    const int nb_off = seff * NB;
+    const int bs_stride = (nkr == 1) ? TILED_TILE_ROWS : TILED_MICRO;
+    const int bs_off = seff * TILED_MICRO;
+
+    float acc[TILED_MICRO][TILED_MICRO];
+    memset(acc, 0, sizeof(acc));
+
+    // subdots at subblock granularity over the 256-K block, exact integer math
+    for (int s = 0; s < NB; s++) {
+        for (int j = 0; j < TILED_MICRO; j++) {
+            const int br = j0 + j;
+            const int8_t * q1 = &src1.q[br * qk_stride + qk_off + s * SUBBLK];
+            int32_t bsum = 0;
+            for (int u = 0; u < NS; u++) {
+                bsum += src1.bsums[(bs_off + s * NS + u) * bs_stride + br];
+            }
+
+            for (int i = 0; i < TILED_MICRO; i++) {
+                const int ar = i0 + i;
+                const int d_off = seff * TILED_MICRO + ar;
+                const uint8_t * q0 = &src0.q[ar * qk_stride + qk_off + s * SUBBLK];
+
+                int32_t raw = 0;
+                for (int e = 0; e < SUBBLK; e++) {
+                    raw += (int32_t) q0[e] * (int32_t) q1[e];
+                }
+
+                // BIAS: subtract BIAS*bsum (src1's per-subblock code sum) from the exact int raw
+                int32_t corr = raw;
+                if constexpr (BIAS != 0) {
+                    corr -= BIAS * bsum;
+                }
+                const int32_t scales_raw = (int32_t) src0.scales[ar * nb_stride + nb_off + s] * corr;
+                // d is NOT applied here: it is constant over the s-loop, so we apply it last before write-out
+                if constexpr (HAS_MIN) {
+                    const int32_t mins_bsum = (int32_t) src0.mins[ar * nb_stride + nb_off + s] * bsum;
+                    acc[i][j] += (float) src0.d[d_off] * (float) scales_raw
+                              - (float) src0.dmin[d_off] * (float) mins_bsum;
+                } else {
+                    acc[i][j] += (float) src0.d[d_off] * (float) scales_raw;
+                }
+            }
+        }
+    }
+
+    // Apply d and write out to buf
+    for (int i = 0; i < TILED_MICRO; i++) {
+        for (int j = 0; j < TILED_MICRO; j++) {
+            buf[(i0 + i) * buf_stride + (j0 + j)] += src1.d[seff * TILED_MICRO + j0 + j] * acc[i][j];
+        }
+    }
+}
+
+#if defined(__AVX512VNNI__) && defined(__AVX512VL__) && defined(__AVX512DQ__)
+
+
+// VNNI microkernel
+// 8x16 band pass: 8 src0 rows (i0, i0+8) x 16 src1 cols (j0, j0+16)
+// Math: s1_acc += scales_s*raw_s, s2_acc += mins_s*bsums_s per subblock; the row result
+// is d0*s1_acc - dmin*s2_acc (f32, exact in range), applied once per row.
+//
+// Register pressure: the band pass holds acc16 + s1_acc = 16 zmm for the
+// 8-row band (s2_acc has no dpbusd dependency, so it is computed after the
+// band pass and adds no register pressure);
+template <int SUBBLK, bool HAS_MIN, int BIAS, int NK>
+static void tiled_run_micro_vnni_8x16(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1,
+                                          int i0, int j0, int num_k, int slab, float * buf, int buf_stride) {
+    constexpr int NB = TILED_TILE_K / SUBBLK;
+    constexpr int NS = SUBBLK / 16;
+    constexpr int NG = SUBBLK / 4;
+
+    // band width, see the register-pressure note above
+    constexpr int NUM_ROWS = 8;
+
+    // num_k = K-blocks per row (the narrow path holds num_k slabs at row stride num_k*256 so
+    // each weight row is one long stream); slab = this call's slab. num_k=1, slab=0 = standard.
+    // NK: 1 = standard single-slab (offsets 0, strides compile-time); 0 = runtime
+    // (num_k/slab from the args). The narrow path is memory-bound, so one runtime
+    // version serves all of it.
+    const int nkr  = (NK > 0) ? NK : num_k;
+    const int seff = (NK > 0) ? 0 : slab;
+    const int qk_stride = nkr * TILED_TILE_K;
+    const int qk_off = seff * TILED_TILE_K;
+    const int nb_stride = NB * nkr;
+    const int nb_off = seff * NB;
+    const int bs_stride = (nkr == 1) ? TILED_TILE_ROWS : TILED_MICRO;
+    const int bs_off = seff * TILED_MICRO;
+
+    const __m512 d1_vec = _mm512_load_ps(&src1.d[seff * TILED_MICRO + j0]);
+
+    __m512i s1_acc[NUM_ROWS];
+    for (int t = 0; t < NUM_ROWS; t++) { s1_acc[t] = _mm512_setzero_si512(); }
+    __m512i s2_acc[NUM_ROWS];
+    if constexpr (HAS_MIN) {
+        for (int t = 0; t < NUM_ROWS; t++) { s2_acc[t] = _mm512_setzero_si512(); }
+    }
+
+    const int32_t * bsums = src1.bsums; // int32 per-16 sums; NS > 1 combines NS lanes per subblock
+    for (int s = 0; s < NB; s++) {
+        __m512i bsums32 = _mm512_load_si512((const __m512i *) &bsums[(bs_off + s * NS) * bs_stride + j0]);
+        for (int u = 1; u < NS; u++) {
+            bsums32 = _mm512_add_epi32(bsums32, _mm512_load_si512((const __m512i *) &bsums[(bs_off + s * NS + u) * bs_stride + j0]));
+        }
+
+        __m512i bias32 = _mm512_setzero_si512();
+        if constexpr (BIAS != 0) {
+            bias32 = _mm512_mullo_epi32(bsums32, _mm512_set1_epi32(BIAS));
+        }
+
+        __m512i acc16[NUM_ROWS];
+        for (int t = 0; t < NUM_ROWS; t++) { acc16[t] = _mm512_setzero_si512(); }
+
+#ifdef __GNUC__
+        #pragma GCC unroll 8 // pragma unrolled justified by measuing with/without
+#endif
+        for (int g = 0; g < NG; g++) {
+            const int kg = s * NG + g;
+            // in-place interleave layout: [kg%16 @ TILED_TILE_K][kg/16 @ 64][row @ 4]
+            const __m512i codes = _mm512_load_si512((const __m512i *) &src1.q[(j0 / TILED_MICRO) * (TILED_MICRO * qk_stride)
+                + qk_off + (kg % TILED_MICRO) * qk_stride + (kg / TILED_MICRO) * (TILED_MICRO * 4)]);
+            for (int t = 0; t < NUM_ROWS; t++) {
+                const uint32_t u4 = *(const uint32_t *) &src0.q[(i0 + t) * qk_stride + qk_off + kg * 4];
+                const __m512i u4b = _mm512_set1_epi32((int) u4);
+                acc16[t] = _mm512_dpbusd_epi32(acc16[t], u4b, codes);
+            }
+        }
+
+        // int correction: s1_acc += scales*(raw-BIAS*bsums)
+        for (int t = 0; t < NUM_ROWS; t++) {
+            const int ar = i0 + t;
+            __m512i rawi = acc16[t];
+            if constexpr (BIAS != 0) {
+                rawi = _mm512_sub_epi32(rawi, bias32);
+            }
+            s1_acc[t] = _mm512_add_epi32(s1_acc[t],
+                _mm512_mullo_epi32(rawi, _mm512_set1_epi32(src0.scales[ar * nb_stride + nb_off + s])));
+        }
+    }
+
+    // s2_acc += mins*bsums; independent of the dpbusd results, so it runs here instead of in the
+    // band pass: the band pass keeps the register budget for the 8-row band
+    if constexpr (HAS_MIN) {
+        for (int s = 0; s < NB; s++) {
+            __m512i bsums32 = _mm512_load_si512((const __m512i *) &bsums[(bs_off + s * NS) * bs_stride + j0]);
+            for (int u = 1; u < NS; u++) {
+                bsums32 = _mm512_add_epi32(bsums32, _mm512_load_si512((const __m512i *) &bsums[(bs_off + s * NS + u) * bs_stride + j0]));
+            }
+            for (int t = 0; t < NUM_ROWS; t++) {
+                s2_acc[t] = _mm512_add_epi32(s2_acc[t],
+                    _mm512_mullo_epi32(bsums32, _mm512_set1_epi32(src0.mins[(i0 + t) * nb_stride + nb_off + s])));
+            }
+        }
+    }
+
+    // epilogue: int->float, apply per-row scales, store to buf
+    for (int t = 0; t < NUM_ROWS; t++) {
+        const int ar = i0 + t;
+        const int d_off = seff * TILED_MICRO + ar;
+        __m512 f1 = _mm512_cvtepi32_ps(s1_acc[t]);
+        __m512 result = _mm512_mul_ps(f1, _mm512_set1_ps(src0.d[d_off]));
+        if constexpr (HAS_MIN) {
+            __m512 f2 = _mm512_cvtepi32_ps(s2_acc[t]);
+            result = _mm512_fnmadd_ps(_mm512_set1_ps(src0.dmin[d_off]), f2, result);
+        }
+        float * p = &buf[(i0 + t) * buf_stride + j0];
+        _mm512_store_ps(p, _mm512_add_ps(_mm512_load_ps(p), _mm512_mul_ps(result, d1_vec)));
+    }
+}
+
+// 16x16 microtile as two explicit 8x16 band passes
+template <int SUBBLK, bool HAS_MIN, int BIAS, int NK>
+static void tiled_run_microtile_vnni(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1,
+                                         int i0, int j0, int num_k, int slab, float * buf, int buf_stride) {
+    tiled_run_micro_vnni_8x16<SUBBLK, HAS_MIN, BIAS, NK>(src0, src1, i0,      j0, num_k, slab, buf, buf_stride);
+    tiled_run_micro_vnni_8x16<SUBBLK, HAS_MIN, BIAS, NK>(src0, src1, i0 + 8,  j0, num_k, slab, buf, buf_stride);
+}
+
+#endif // __AVX512VNNI__ && __AVX512VL__
+
+#if defined(__AVX2__)
+
+// AVX2 kernel.
+// We're effectively applying the existing vec_dot algorithms to an 8x16 block here.
+// Different paths based on subblock size as it affects when/where we multiply in scales and apply mins
+// SUBBLK=32: one 256-bit maddubs per (s, column), a 256-bit set1 scale,
+// one 8-lane i32 accumulator per column.
+// SUBBLK=16: two subblocks per 32B load. The 256-bit maddubs product
+// still gives 16 i16 lanes (0..7 = subblock sp, 8..15 = sp+1), but the
+// scale is applied per 128-bit half.
+// With BIAS != 0 the q0 codes are biased (up to 241) and a maddubs i16 pair
+// overflows, so the weight is debiased (w = q0 - BIAS, small). Two MACs:
+//   ACTBIAS: activation pre-biased (+128) by repack, one maddubs(q1b, w); the +128 is
+//     corrected out in the epilogue (128*sum_s scales[s]*w_bsum[s], precomputed in repack_src0).
+//     Fits only when the debiased weight magnitude Wmax <= 64 (2*255*Wmax < 32767).
+//   else (iq4_xs, Wmax = 127): the signed-int8 sign trick, ax = |w|, sy = q1*sign(w),
+//     one maddubs (2*127*127 < 32767) so it always fits.
+template <int SUBBLK, bool HAS_MIN, int BIAS, int NK, bool ACTBIAS>
+static void tiled_run_microtile_avx2(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1,
+                                     int i0, int j0, int num_k, int slab, float * buf, int buf_stride) {
+    constexpr int NB = TILED_TILE_K / SUBBLK;
+    constexpr int GROUP = 8; // Process 8 src1 columns concurrently (1 YMM vector)
+
+    static_assert(SUBBLK == 16 || SUBBLK == 32, "unsupported SUBBLK");
+
+    // num_k/slab: the narrow path holds num_k slabs at row stride num_k*256 (num_k=1, slab=0 = standard)
+    // NK: 1 = standard single-slab (offsets 0, strides compile-time); 0 = runtime
+    // (num_k/slab from the args). The narrow path is memory-bound, so one runtime
+    // version serves all of it.
+    const int nkr  = (NK > 0) ? NK : num_k;
+    const int seff = (NK > 0) ? 0 : slab;
+    const int qk_stride = nkr * TILED_TILE_K;
+    const int qk_off = seff * TILED_TILE_K;
+    const int nb_stride = NB * nkr;
+    const int nb_off = seff * NB;
+    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;
+        const float d0 = src0.d[d_off];
+        const float dmin0 = src0.dmin[d_off];
+        const uint8_t * q0 = &src0.q[ar * qk_stride + qk_off];
+        const int32_t * scales_row = &src0.scales[ar * nb_stride + nb_off];
+        const int32_t * mins_row = &src0.mins[ar * nb_stride + nb_off];
+
+        for (int g = 0; g < TILED_MICRO; g += GROUP) {
+            // Pre-calculate weight pointers for this column group
+            const int8_t * q1_ptr[GROUP];
+            for (int t = 0; t < GROUP; t++) {
+                q1_ptr[t] = &src1.q[(j0 + g + t) * qk_stride + qk_off];
+            }
+
+            __m256i s2   = _mm256_setzero_si256();
+
+            __m256i acc[GROUP];
+            for (int t = 0; t < GROUP; t++) {
+                acc[t] = _mm256_setzero_si256();
+            }
+
+            if constexpr (SUBBLK == 32) {
+                for (int s = 0; s < NB; s++) {
+                    // OPTIMIZATION 1: Use aligned loads (_mm256_load_si256)
+                    const __m256i q0_32 = _mm256_load_si256((const __m256i *) &q0[s * SUBBLK]);
+                    const __m256i scales16 = _mm256_set1_epi16(scales_row[s]);
+
+                    if constexpr (BIAS != 0) {
+                        const __m256i w = _mm256_sub_epi8(q0_32, _mm256_set1_epi8((int8_t) BIAS)); // debiased weight
+                        if constexpr (ACTBIAS) {
+                            // activation pre-biased (+128) by repack; q1b unsigned, w signed
+                            #pragma GCC unroll 8
+                            for (int t = 0; t < GROUP; t++) {
+                                const __m256i q1b_32 = _mm256_load_si256((const __m256i *) &q1_ptr[t][s * SUBBLK]);
+                                acc[t] = _mm256_add_epi32(acc[t],
+                                    _mm256_madd_epi16(scales16, _mm256_maddubs_epi16(q1b_32, w)));
+                            }
+                        } else {
+                            // iq4_xs (Wmax > 64 overflows act-bias): sign trick, both operands <= 127
+                            const __m256i ax = _mm256_sign_epi8(w, w);
+                            #pragma GCC unroll 8
+                            for (int t = 0; t < GROUP; t++) {
+                                const __m256i q1_32 = _mm256_load_si256((const __m256i *) &q1_ptr[t][s * SUBBLK]);
+                                acc[t] = _mm256_add_epi32(acc[t],
+                                    _mm256_madd_epi16(scales16, _mm256_maddubs_epi16(ax, _mm256_sign_epi8(q1_32, w))));
+                            }
+                        }
+                    } else {
+                        #pragma GCC unroll 8
+                        for (int t = 0; t < GROUP; t++) {
+                            const __m256i q1_32 = _mm256_load_si256((const __m256i *) &q1_ptr[t][s * SUBBLK]);
+                            acc[t] = _mm256_add_epi32(acc[t],
+                                _mm256_madd_epi16(scales16, _mm256_maddubs_epi16(q0_32, q1_32)));
+                        }
+                    }
+
+                    if constexpr (HAS_MIN) {
+                        const __m256i bsums_v = _mm256_add_epi32(
+                            _mm256_load_si256((const __m256i *) &src1.bsums[(bs_off + s * 2) * bs_stride + j0 + g]),
+                            _mm256_load_si256((const __m256i *) &src1.bsums[(bs_off + s * 2 + 1) * bs_stride + j0 + g]));
+                        s2 = _mm256_add_epi32(s2, _mm256_mullo_epi32(bsums_v, _mm256_set1_epi32(mins_row[s])));
+                    }
+                }
+            } else { // SUBBLK == 16
+                for (int sp = 0; sp < NB; sp += 2) {
+                    const __m256i q0_32 = _mm256_load_si256((const __m256i *) &q0[sp * SUBBLK]);
+                    const __m256i scalesv = _mm256_set_m128i(_mm_set1_epi16(scales_row[sp + 1]), _mm_set1_epi16(scales_row[sp]));
+                    if constexpr (BIAS != 0) {
+                        const __m256i w = _mm256_sub_epi8(q0_32, _mm256_set1_epi8((int8_t) BIAS)); // debiased weight
+                        if constexpr (ACTBIAS) {
+                            // activation pre-biased (+128) by repack; q1b unsigned, w signed
+                            #pragma GCC unroll 8
+                            for (int t = 0; t < GROUP; t++) {
+                                const __m256i q1b_32 = _mm256_load_si256((const __m256i *) &q1_ptr[t][sp * SUBBLK]);
+                                acc[t] = _mm256_add_epi32(acc[t], _mm256_madd_epi16(
+                                    scalesv, _mm256_maddubs_epi16(q1b_32, w)));
+                            }
+                        } else {
+                            // iq4_xs (Wmax > 64 overflows act-bias): sign trick, both operands <= 127
+                            const __m256i ax = _mm256_sign_epi8(w, w);
+                            #pragma GCC unroll 8
+                            for (int t = 0; t < GROUP; t++) {
+                                const __m256i q1_32 = _mm256_load_si256((const __m256i *) &q1_ptr[t][sp * SUBBLK]);
+                                acc[t] = _mm256_add_epi32(acc[t], _mm256_madd_epi16(
+                                    scalesv, _mm256_maddubs_epi16(ax, _mm256_sign_epi8(q1_32, w))));
+                            }
+                        }
+                    } else {
+                        #pragma GCC unroll 8
+                        for (int t = 0; t < GROUP; t++) {
+                            const __m256i q1_32 = _mm256_load_si256((const __m256i *) &q1_ptr[t][sp * SUBBLK]);
+                            acc[t] = _mm256_add_epi32(acc[t], _mm256_madd_epi16(
+                                scalesv, _mm256_maddubs_epi16(q0_32, q1_32)));
+                        }
+                    }
+
+                    if constexpr (HAS_MIN) {
+                        const __m256i bsums0_v = _mm256_load_si256((const __m256i *) &src1.bsums[(bs_off + sp) * bs_stride + j0 + g]);
+                        const __m256i bsums1_v = _mm256_load_si256((const __m256i *) &src1.bsums[(bs_off + sp + 1) * bs_stride + j0 + g]);
+                        s2 = _mm256_add_epi32(s2, _mm256_add_epi32(
+                            _mm256_mullo_epi32(bsums0_v, _mm256_set1_epi32(mins_row[sp])),
+                            _mm256_mullo_epi32(bsums1_v, _mm256_set1_epi32(mins_row[sp + 1]))));
+                    }
+                }
+            }
+
+            // 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);
+
+            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]));
+            }
+
+            const __m256 src1_d_vec = _mm256_load_ps(&src1.d[seff * TILED_MICRO + j0 + g]);
+            __m256 res_vec = _mm256_mul_ps(_mm256_set1_ps(d0), _mm256_cvtepi32_ps(s1_vec));
+
+            if constexpr (HAS_MIN) {
+#if defined(__FMA__)
+                // FMA3 implementation: res_vec = res_vec - (dmin0 * s2)
+                res_vec = _mm256_fnmadd_ps(_mm256_set1_ps(dmin0), _mm256_cvtepi32_ps(s2), res_vec);
+#else
+                // Non-FMA fallback
+                res_vec = _mm256_sub_ps(res_vec, _mm256_mul_ps(_mm256_set1_ps(dmin0), _mm256_cvtepi32_ps(s2)));
+#endif
+            }
+
+            float * buf_ptr = &buf[ar * buf_stride + j0 + g];
+            __m256 current_buf = _mm256_loadu_ps(buf_ptr);
+
+#if defined(__FMA__)
+            // FMA3 accumulation: current_buf + (res_vec * src1_d_vec)
+            __m256 updated_buf = _mm256_fmadd_ps(res_vec, src1_d_vec, current_buf);
+#else
+            __m256 updated_buf = _mm256_add_ps(current_buf, _mm256_mul_ps(res_vec, src1_d_vec));
+#endif
+
+            _mm256_storeu_ps(buf_ptr, updated_buf);
+        }
+    }
+}
+
+#endif // __AVX2__
+
+#if defined(__AVX__) && !defined(__AVX2__)
+template <int SUBBLK, bool HAS_MIN, int BIAS, int NK>
+static void tiled_run_microtile_avx(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1,
+                                    int i0, int j0, int num_k, int slab, float * buf, int buf_stride) {
+    constexpr int NB = TILED_TILE_K / SUBBLK;
+    constexpr int NS = SUBBLK / 16;
+    constexpr int GROUP = 4; // src1 columns per group: one acc32 per column
+
+    // num_k/slab: the narrow path holds num_k slabs at row stride num_k*256 (num_k=1, slab=0 = standard)
+    // NK: 1 = standard single-slab (offsets 0, strides compile-time); 0 = runtime
+    // (num_k/slab from the args). The narrow path is memory-bound, so one runtime
+    // version serves all of it.
+    const int nkr  = (NK > 0) ? NK : num_k;
+    const int seff = (NK > 0) ? 0 : slab;
+    const int qk_stride = nkr * TILED_TILE_K;
+    const int qk_off = seff * TILED_TILE_K;
+    const int nb_stride = NB * nkr;
+    const int nb_off = seff * NB;
+    const int bs_stride = (nkr == 1) ? TILED_TILE_ROWS : TILED_MICRO;
+    const int bs_off = seff * TILED_MICRO;
+
+    for (int i = 0; i < TILED_MICRO; i++) {
+        const int ar = i0 + i;
+        const int d_off = seff * TILED_MICRO + ar;
+        const float d0 = src0.d[d_off];
+        const float dmin0 = src0.dmin[d_off];
+        const uint8_t * q0 = &src0.q[ar * qk_stride + qk_off];
+        const int32_t * scales_row = &src0.scales[ar * nb_stride + nb_off];
+        const int32_t * mins_row = &src0.mins[ar * nb_stride + nb_off];
+
+        for (int g = 0; g < TILED_MICRO; g += GROUP) {
+            const int8_t * q1g[GROUP];
+            __m128i acc[GROUP];
+            __m128i s2 = _mm_setzero_si128();   // sum_s mins_s * bsums_s
+            for (int t = 0; t < GROUP; t++) {
+                q1g[t] = &src1.q[(j0 + g + t) * qk_stride + qk_off];
+                acc[t] = _mm_setzero_si128();
+            }
+
+            for (int s = 0; s < NB; s++) {
+                const __m128i scales16 = _mm_set1_epi16(scales_row[s]);
+                if constexpr (HAS_MIN) {
+                    __m128i bsums_v = _mm_setzero_si128(); // per-col bsums as a 4-lane i32 vector
+                    for (int u = 0; u < NS; u++) {
+                        bsums_v = _mm_add_epi32(bsums_v, _mm_loadu_si128(
+                            (const __m128i *) &src1.bsums[(bs_off + s * NS + u) * bs_stride + j0 + g]));
+                    }
+                    s2 = _mm_add_epi32(s2, _mm_mullo_epi32(bsums_v, _mm_set1_epi32(mins_row[s])));
+                }
+                for (int u = 0; u < NS; u++) {
+                    const __m128i a16 = _mm_loadu_si128((const __m128i *) &q0[s * SUBBLK + u * 16]);
+                    if constexpr (BIAS != 0) {
+                        // debiased weight w = a16 - BIAS; signed dot via the sign trick (see header)
+                        const __m128i a16s = _mm_sub_epi8(a16, _mm_set1_epi8((int8_t) BIAS));
+                        const __m128i ax    = _mm_sign_epi8(a16s, a16s);
+                        for (int t = 0; t < GROUP; t++) {
+                            const __m128i q1_16 = _mm_loadu_si128((const __m128i *) &q1g[t][s * SUBBLK + u * 16]);
+                            acc[t] = _mm_add_epi32(acc[t],
+                                _mm_madd_epi16(scales16, _mm_maddubs_epi16(ax, _mm_sign_epi8(q1_16, a16s))));
+                        }
+                    } else {
+                        for (int t = 0; t < GROUP; t++) {
+                            acc[t] = _mm_add_epi32(acc[t], _mm_madd_epi16(scales16,
+                                _mm_maddubs_epi16(a16,
+                                    _mm_loadu_si128((const __m128i *) &q1g[t][s * SUBBLK + u * 16]))));
+                        }
+                    }
+                }
+            }
+
+            int32_t s2_s[GROUP] = { 0, 0, 0, 0 };
+            if constexpr (HAS_MIN) {
+                _mm_storeu_si128((__m128i *) s2_s, s2);
+            }
+
+            for (int t = 0; t < GROUP; t++) {
+                // 4 i32 lanes -> scalar
+                __m128i v = _mm_shuffle_epi32(acc[t], _MM_SHUFFLE(2, 3, 0, 1));
+                acc[t] = _mm_add_epi32(acc[t], v);
+                v = _mm_shuffle_epi32(acc[t], _MM_SHUFFLE(1, 0, 3, 2));
+                acc[t] = _mm_add_epi32(acc[t], v);
+                const int32_t s1 = _mm_cvtsi128_si32(acc[t]);
+
+                float res = d0 * (float) s1;
+                if constexpr (HAS_MIN) {
+                    res -= dmin0 * (float) s2_s[t];
+                }
+                buf[ar * buf_stride + j0 + g + t] += src1.d[seff * TILED_MICRO + j0 + g + t] * res;
+            }
+        }
+    }
+}
+
+#endif // __AVX__ && !__AVX2__
+
+// main microtile entry point. num_k = K-blocks per row (the narrow path reads num_k slabs at row
+// stride num_k*256 so each weight row is one long stream); slab = this call's slab. num_k=1,
+// slab=0 is the standard single-slab path.
+// num_k==1 dispatches to the NK=1 instantiation, whose strides/offsets are compile-time
+// constants (shifts/scaled-leas). num_k>1 (the narrow path) uses the NK=0 runtime-strided
+// version; it is memory-bound, so one version serves all of it.
+template <int SUBBLK, bool HAS_MIN, int BIAS, bool ACTBIAS>
+void tiled_run_microtile(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1,
+                         int i0, int j0, int num_k, int slab, float * buf, int buf_stride) {
+    const bool standard = (num_k == 1);
+#if defined(__AVX512VNNI__) && defined(__AVX512VL__) && defined(__AVX512DQ__)
+    if (standard) {
+        tiled_run_microtile_vnni<SUBBLK, HAS_MIN, BIAS, 1>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
+    } else {
+        tiled_run_microtile_vnni<SUBBLK, HAS_MIN, BIAS, 0>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
+    }
+#elif defined(__AVX2__)
+    if (standard) {
+        tiled_run_microtile_avx2<SUBBLK, HAS_MIN, BIAS, 1, ACTBIAS>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
+    } else {
+        tiled_run_microtile_avx2<SUBBLK, HAS_MIN, BIAS, 0, ACTBIAS>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
+    }
+#elif defined(__AVX__)
+    if (standard) {
+        tiled_run_microtile_avx<SUBBLK, HAS_MIN, BIAS, 1>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
+    } else {
+        tiled_run_microtile_avx<SUBBLK, HAS_MIN, BIAS, 0>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
+    }
+#else
+    if (standard) {
+        tiled_run_microtile_scalar<SUBBLK, HAS_MIN, BIAS, 1>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
+    } else {
+        tiled_run_microtile_scalar<SUBBLK, HAS_MIN, BIAS, 0>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
+    }
+#endif
+}
+
+// explicit instantiations for the in-use formats (q4_K and q5_K share the constants)
+// ACTBIAS selects the act-bias MAC (Wmax <= 64); iq4_xs (Wmax 127) keeps the sign trick
+// q4_K / q5_K: BIAS = 0
+ template void tiled_run_microtile<32, true, 0, false>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1,
+                                               int i0, int j0, int num_k, int slab, float * buf, int buf_stride);
+// iq4_xs: sign trick (Wmax 127 overflows act-bias)
+template void tiled_run_microtile<32, false, 128, false>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1,
+                                                  int i0, int j0, int num_k, int slab, float * buf, int buf_stride);
+// iq2_xxs, iq3_xxs, iq3_s, iq1_s: act-bias (Wmax <= 62)
+template void tiled_run_microtile<32, false, 128, true>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1,
+                                                  int i0, int j0, int num_k, int slab, float * buf, int buf_stride);
+// iq2_xs, iq2_s, iq1_m: act-bias
+template void tiled_run_microtile<16, false, 128, true>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1,
+                                                  int i0, int j0, int num_k, int slab, float * buf, int buf_stride);
+// q6_K: act-bias (Wmax 32)
+template void tiled_run_microtile<16, false, 32, true>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1,
+                                                 int i0, int j0, int num_k, int slab, float * buf, int buf_stride);
+// q3_K: act-bias (Wmax 4)
+template void tiled_run_microtile<16, false, 4, true>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1,
+                                                int i0, int j0, int num_k, int slab, float * buf, int buf_stride);
+// q2_K: BIAS = 0
+template void tiled_run_microtile<16, true, 0, false>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1,
+                                               int i0, int j0, int num_k, int slab, float * buf, int buf_stride);
+
+
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
+
+// block_q8_K in 4-byte words, for the int32 gather indices
+static_assert(sizeof(block_q8_K) == 292 && offsetof(block_q8_K, qs) == 4,
+              "block_q8_K layout changed, fix the src1 repack");
+
+// VNNI: interleave the natural [row][k] codes in-place into the group-local
+// [kg%16][kg/16][row][4] layout. base points at row 0 (row r at base + r*row_stride);
+// n_tiles 16x16 int32 tiles are transposed to the right (tile t occupies bytes
+// [t*64, t*64+64) per row). Each 16-row group is self-contained in row_stride*16 bytes.
+#if defined(__AVX512VNNI__) && defined(__AVX512VL__) && defined(__AVX512DQ__)
+void tiled_repack_src1(tiled_tile_src1 * src1, int row0, int num_k, bool bias) {
+    GGML_UNUSED(bias);
+    const int n_tiles = num_k * (TILED_TILE_K / (TILED_MICRO * 4));
+    const int row_stride = num_k * TILED_TILE_K;
+    uint8_t * base = (uint8_t *) src1->q + row0 * row_stride;
+    for (int t = 0; t < n_tiles; t++) {
+        uint8_t * cbase = base + t * (TILED_MICRO * 4);
+        __m512i v[16];
+        for (int r = 0; r < 16; r++) {
+            v[r] = _mm512_load_si512((const __m512i *) (cbase + r * row_stride));
+        }
+
+        // 16x16 int32 transpose, 4 butterfly phases
+        // Phase 1: 1-element interleave, pairs (0,1), (2,3), ..., (14,15)
+        {
+            const __m512i idx_a = _mm512_setr_epi32(0, 16, 1, 17, 2, 18, 3, 19, 4, 20, 5, 21, 6, 22, 7, 23);
+            const __m512i idx_b = _mm512_setr_epi32(8, 24, 9, 25, 10, 26, 11, 27, 12, 28, 13, 29, 14, 30, 15, 31);
+            for (int i = 0; i < 16; i += 2) {
+                __m512i a = v[i], b = v[i+1];
+                v[i]   = _mm512_permutex2var_epi32(a, idx_a, b);
+                v[i+1] = _mm512_permutex2var_epi32(a, idx_b, b);
+            }
+        }
+        // Phase 2: 2-element interleave, pairs (0,2), (1,3), (4,6), (5,7), ...
+        {
+            const __m512i idx_a = _mm512_setr_epi32(0, 1, 16, 17, 4, 5, 20, 21, 8, 9, 24, 25, 12, 13, 28, 29);
+            const __m512i idx_b = _mm512_setr_epi32(2, 3, 18, 19, 6, 7, 22, 23, 10, 11, 26, 27, 14, 15, 30, 31);
+            for (int i = 0; i < 16; i += 4) {
+                __m512i a = v[i],   b = v[i+2];
+                v[i]   = _mm512_permutex2var_epi32(a, idx_a, b);
+                v[i+2] = _mm512_permutex2var_epi32(a, idx_b, b);
+                a = v[i+1], b = v[i+3];
+                v[i+1] = _mm512_permutex2var_epi32(a, idx_a, b);
+                v[i+3] = _mm512_permutex2var_epi32(a, idx_b, b);
+            }
+        }
+        // Phase 3: 4-element interleave, pairs (0,4), (1,5), ..., (7,11), (8,12), ...
+        {
+            const __m512i idx_a = _mm512_setr_epi32(0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23);
+            const __m512i idx_b = _mm512_setr_epi32(8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31);
+            for (int i = 0; i < 16; i += 8) {
+                for (int j = 0; j < 4; j++) {
+                    __m512i a = v[i+j], b = v[i+4+j];
+                    v[i+j]   = _mm512_permutex2var_epi32(a, idx_a, b);
+                    v[i+4+j] = _mm512_permutex2var_epi32(a, idx_b, b);
+                }
+            }
+        }
+        // Phase 4: 8-element interleave, pairs (0,8), (1,9), ..., (7,15)
+        {
+            const __m512i idx_a = _mm512_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7, 16, 17, 18, 19, 20, 21, 22, 23);
+            const __m512i idx_b = _mm512_setr_epi32(8, 9, 10, 11, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31);
+            for (int i = 0; i < 8; i++) {
+                __m512i a = v[i], b = v[i+8];
+                v[i]   = _mm512_permutex2var_epi32(a, idx_a, b);
+                v[i+8] = _mm512_permutex2var_epi32(a, idx_b, b);
+            }
+        }
+
+        // store: v[g] holds 16 int32s for k-group col_order[g], rows 0..15
+        static const int col_order[16] = {0, 8, 1, 9, 4, 12, 5, 13, 2, 10, 3, 11, 6, 14, 7, 15};
+        for (int g = 0; g < 16; g++) {
+            _mm512_store_si512((void *) (cbase + col_order[g] * row_stride), v[g]);
+        }
+    }
+}
+#elif defined(__AVX2__)
+// act-bias: pre-bias the activation in place (+128) so the MAC's maddubs(q1b, w) sees unsigned bytes;
+// no transpose on AVX2 (natural layout). 16 rows x n_tiles*64 bytes, row r at base + r*row_stride
+void tiled_repack_src1(tiled_tile_src1 * src1, int row0, int num_k, bool bias) {
+    if (!bias) {
+        return;
+    }
+    const int row_stride = num_k * TILED_TILE_K;
+    uint8_t * base = (uint8_t *) src1->q + row0 * row_stride;
+    const int nbytes = row_stride; // full row width
+    const __m256i b128 = _mm256_set1_epi8((int8_t) 0x80);
+    for (int r = 0; r < TILED_MICRO; r++) {
+        uint8_t * row = base + r * row_stride;
+        for (int off = 0; off < nbytes; off += 32) {
+            const __m256i v = _mm256_loadu_si256((const __m256i *) (row + off));
+            _mm256_storeu_si256((__m256i *) (row + off), _mm256_xor_si256(v, b128));
+        }
+    }
+}
+#else
+void tiled_repack_src1(tiled_tile_src1 * src1, int row0, int num_k, bool bias) {
+    GGML_UNUSED(src1); GGML_UNUSED(row0); GGML_UNUSED(num_k); GGML_UNUSED(bias);
+}
+#endif
+
+#if defined(__AVX2__)
+// sum 32 unsigned bytes to int32 (via maddubs -> int16 pairs -> hsum)
+static int32_t tiled_byte_sum_32(const uint8_t * p) {
+    const __m256i dot = _mm256_maddubs_epi16(_mm256_loadu_si256((const __m256i *) p), _mm256_set1_epi8(1));
+    const __m256i sum = _mm256_madd_epi16(_mm256_set1_epi16(1), dot); // 8 int32
+    // reduce 8 lanes: lo+hi -> 4 lanes, hadd -> 2, hadd -> 1 (a single hadd pair only folds 4 lanes)
+    const __m128i t = _mm_add_epi32(_mm256_castsi256_si128(sum), _mm256_extracti128_si256(sum, 1));
+    const __m128i h = _mm_hadd_epi32(t, t);
+    return _mm_cvtsi128_si32(_mm_hadd_epi32(h, h));
+}
+// sum 16 unsigned bytes to int32
+static int32_t tiled_byte_sum_16(const uint8_t * p) {
+    const __m128i dot = _mm_maddubs_epi16(_mm_loadu_si128((const __m128i *) p), _mm_set1_epi8(1));
+    const __m128i sum = _mm_madd_epi16(_mm_set1_epi16(1), dot); // 4 int32
+    const __m128i h = _mm_hadd_epi32(sum, sum);
+    return _mm_cvtsi128_si32(_mm_hadd_epi32(h, h));
+}
+#endif
+
+// act-bias: precompute the per-weight-row correction corr = 128*sum_s scales[s]*w_bsum[s] where
+// w_bsum[s] = sum of the debiased weight bytes in subblock s. Stored in mins[r][0] (unused for
+// HAS_MIN = false, which is exactly the act-bias case). Reused across all act bands and weight groups.
+// No-op off AVX2 (the act-bias MAC is AVX2-only) and when corr is false (iq4_xs / BIAS = 0).
+template <int SUBBLK>
+void tiled_repack_src0(tiled_tile_src0 * tile, int n_rows, int num_k, int BIAS, bool corr) {
+#if defined(__AVX2__)
+    if (!corr) {
+        return;
+    }
+    constexpr int NB = TILED_TILE_K / SUBBLK;
+    const int nb_stride = NB * num_k;
+    for (int slab = 0; slab < num_k; slab++) {
+        for (int r = 0; r < n_rows; r++) {
+            const uint8_t * q = &tile->q[r * (num_k * TILED_TILE_K) + slab * TILED_TILE_K];
+            const int32_t * scales = &tile->scales[r * nb_stride + slab * NB];
+            int32_t corr_val = 0;
+            for (int s = 0; s < NB; s++) {
+                const int32_t qsum = (SUBBLK == 32) ? tiled_byte_sum_32(&q[s * SUBBLK]) : tiled_byte_sum_16(&q[s * SUBBLK]);
+                corr_val += 128 * scales[s] * (qsum - SUBBLK * BIAS);
+            }
+            tile->mins[r * nb_stride + slab * NB] = corr_val;
+        }
+    }
+#else
+    GGML_UNUSED(tile); GGML_UNUSED(n_rows); GGML_UNUSED(num_k); GGML_UNUSED(BIAS); GGML_UNUSED(corr);
+#endif
+}
+
+template void tiled_repack_src0<16>(tiled_tile_src0 * tile, int n_rows, int num_k, int BIAS, bool corr);
+template void tiled_repack_src0<32>(tiled_tile_src0 * tile, int n_rows, int num_k, int BIAS, bool corr);
+
+

ggml/src/ggml-cpu/tiled/tiled-kernel.h

@@ -0,0 +1,180 @@
+#pragma once
+
+// Tiled matmul kernel API: tile structs, kernel definitions
+
+// Currently only optimized for x86, new architectures should implement:
+// tiled_run_microtile:  16x16 microkernel
+// tiled_repack_src0: Optional repack/recalculation of src0, per macrotile
+// tiled_repack_src1: Optional repack/recalculation of src1, per microtile-band
+// bit unpacking routines: tiled_unpk_nib4, tiled_unpk_2bit, tiled_unpk_or
+// LUT value expansion routines: tiled_lut8, tiled_unpk_sign32, tiled_unpk_tern8
+
+#define GGML_COMMON_DECL_C
+#include "ggml-common.h"
+
+#include <stddef.h>
+#include <stdint.h>
+
+#if defined(__AVX2__)
+#include <immintrin.h>
+#endif
+
+#define TILED_TILE_K    256 // one QK_K block
+#define TILED_TILE_ROWS 256 // max window rows, ragged at edges
+#define TILED_MICRO     16  // microtile edge (also the bsums code-sum granularity)
+#define TILED_WS_SLOT   (512 * 1024) // per-thread workspace slot, a clean 512KB (multiple of 64B)
+
+// src0 tile: weight side, shared by all formats.
+// scales/mins are sized for the max subblock count (SUBBLK=16);
+// SUBBLK=32 formats index at stride 8 and leave the slack unused.
+struct tiled_tile_src0 {
+    static constexpr int NB_MAX = TILED_TILE_K / 16; // max subblocks per 256-elem block
+
+    alignas(64) uint8_t q[TILED_TILE_ROWS * TILED_TILE_K]; // unsigned quants, widened to uint8
+    float    d[TILED_TILE_ROWS];  // One d from each input block, widened to f32
+    float    dmin[TILED_TILE_ROWS]; // dmin from each input block (if applicable), widened to F32
+    int32_t   scales[TILED_TILE_ROWS * NB_MAX];  // per-subblock scale, stored as int32_t
+    int32_t   mins[TILED_TILE_ROWS * NB_MAX];    // per-subblock min, used when HAS_MIN
+};
+
+// src1 tile: built from q8_K (wdata)
+struct tiled_tile_src1 {
+    // q8 codes, one byte per element. Note for VNNI these are reshaped + transposed to be suitable for dpbusd.
+    alignas(64) int8_t  q[TILED_TILE_ROWS * TILED_TILE_K];
+    // per-16 code sums from q8_k (int16), widened to int32 so the kernels load them directly, no per-use cvt
+    alignas(64) int32_t bsums[(TILED_TILE_K / 16) * TILED_TILE_ROWS];
+    // f32 (not f16): q8_k stores fp16, the unpack converts once
+    float       d[TILED_TILE_ROWS];
+};
+
+// per-thread workspace: all tiled state lives here, allocated in wdata (one slot per thread)
+struct tiled_ws {
+    tiled_tile_src0 src0;
+    tiled_tile_src1 src1;
+    alignas(64) float acc[TILED_TILE_ROWS * TILED_TILE_ROWS];
+};
+
+static_assert(sizeof(tiled_ws) <= TILED_WS_SLOT, "tiled workspace exceeds the 512KB per-thread slot");
+// the slot base is 64B-aligned and TILED_WS_SLOT is a multiple of 64B, so each per-thread slot
+// is 64B-aligned; these pin the tile fields and the acc buffer at aligned offsets within a slot
+static_assert(offsetof(tiled_ws, src0) % 64 == 0, "src0 not 64B-aligned in the workspace");
+static_assert(offsetof(tiled_ws, src1) % 64 == 0, "src1 not 64B-aligned in the workspace");
+static_assert(offsetof(tiled_ws, acc)  % 64 == 0, "acc not 64B-aligned in the workspace");
+
+// unpack primitives for reading quants, defined as inline here to keep arch-specific code in kernel.h/.cpp
+// If this section gets too hairy later, we can break up into separate includes.
+#if defined(__AVX2__)
+// packed 4-bit codes -> low nibbles (lo) + high nibbles (hi)
+inline void tiled_unpk_nib4(const uint8_t * src, uint8_t * lo, uint8_t * hi) {
+    const __m256i v = _mm256_loadu_si256((const __m256i *) src);
+    // mask before the lane shift so bits do not cross byte boundaries
+    _mm256_storeu_si256((__m256i *) lo, _mm256_and_si256(v, _mm256_set1_epi8(0x0F)));
+    _mm256_storeu_si256((__m256i *) hi, _mm256_srli_epi32(_mm256_and_si256(v, _mm256_set1_epi8((int8_t) 0xF0)), 4));
+}
+// 2-bit values at bit offset S
+template <int S> inline void tiled_unpk_2bit(const uint8_t * src, uint8_t * dst) {
+    _mm256_storeu_si256((__m256i *) dst, _mm256_and_si256(
+        _mm256_srli_epi32(_mm256_loadu_si256((const __m256i *) src), S), _mm256_set1_epi8(0x03)));
+}
+// OR the M-bit value at bit offset S of src into bit offset D of dst
+template <int S, int D, int M>
+inline void tiled_unpk_or(uint8_t * dst, const uint8_t * src) {
+    const __m256i v = _mm256_slli_epi32(_mm256_and_si256(
+        _mm256_srli_epi32(_mm256_loadu_si256((const __m256i *) src), S), _mm256_set1_epi8((uint8_t) M)), D);
+    _mm256_storeu_si256((__m256i *) dst, _mm256_or_si256(_mm256_loadu_si256((const __m256i *) dst), v));
+}
+
+
+// Unpacking kernels for IQ quants
+
+// LUT value expansion for the LUT-based formats (iq4_xs, iq grids): the bit unpackers
+// above give the indices, these expand 8/16 of them to widened codes in one pass
+// 16-entry byte LUT: dst[j] = lut[src[j]] (16 bytes)
+inline void tiled_lut8(const uint8_t * lut, const uint8_t * src, uint8_t * dst) {
+    _mm_storeu_si128((__m128i *) dst, _mm_shuffle_epi8(_mm_loadu_si128((const __m128i *) lut),
+                                                       _mm_loadu_si128((const __m128i *) src)));
+}
+
+
+// 32 grid magnitudes (4 x 64-bit groups g0..g3, 8 values each) + 4 sign bytes
+// (byte l signs values 8*l .. 8*l+7) -> codes stored as (value + 128)
+inline void tiled_unpk_sign32(uint64_t g0, uint64_t g1, uint64_t g2, uint64_t g3,
+                              const uint8_t signs[4], uint8_t * dst32) {
+    const __m256i v  = _mm256_set_epi64x((int64_t) g3, (int64_t) g2, (int64_t) g1, (int64_t) g0);
+    const __m256i sv = _mm256_shuffle_epi8(_mm256_set1_epi32((int32_t) (signs[0] | signs[1] << 8 | signs[2] << 16 | signs[3] << 24)),
+                                           _mm256_setr_epi8(0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
+                                                            2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3));
+    const __m256i sel  = _mm256_set1_epi64x((int64_t) 0x8040201008040201ULL);
+    // 0xFF in every lane whose sign bit is set; GFNI does the and+compare in one instruction
+#ifdef __GFNI__
+    const __m256i mask = _mm256_gf2p8affine_epi64_epi8(sel, sv, 0);
+#else
+    const __m256i mask = _mm256_cmpeq_epi8(_mm256_and_si256(sv, sel), sel);
+#endif
+    // v ^ mask - mask negates the signed lanes (v elsewhere); + 128 gives the biased code
+    const __m256i sgn  = _mm256_sub_epi8(_mm256_xor_si256(v, mask), mask);
+    _mm256_storeu_si256((__m256i *) dst32, _mm256_add_epi8(sgn, _mm256_set1_epi8((int8_t) 128)));
+}
+// 8 ternary grid bytes (0 = 0, 1 = +1, 0xFF = -1): dst[j] = 128 + delta + 8 * (int8_t) src[j]
+inline void tiled_unpk_tern8(const uint8_t * src, int8_t delta, uint8_t * dst) {
+    const __m128i v = _mm_cvtepi8_epi16(_mm_loadl_epi64((const __m128i *) src));
+    const __m128i p = _mm_add_epi16(_mm_slli_epi16(v, 3), _mm_set1_epi16(128 + (int) delta));
+    _mm_storel_epi64((__m128i *) dst, _mm_packus_epi16(p, _mm_setzero_si128()));
+}
+#else
+
+// Scalar definitions for unpackers.
+
+
+inline void tiled_unpk_nib4(const uint8_t * src, uint8_t * lo, uint8_t * hi) {
+    for (int l = 0; l < 32; l++) { lo[l] = (uint8_t) (src[l] & 0xF); hi[l] = (uint8_t) (src[l] >> 4); }
+}
+template <int S>
+inline void tiled_unpk_2bit(const uint8_t * src, uint8_t * dst) {
+    for (int l = 0; l < 32; l++) { dst[l] = (uint8_t) ((src[l] >> S) & 3); }
+}
+template <int S, int D, int M>
+inline void tiled_unpk_or(uint8_t * dst, const uint8_t * src) {
+    for (int l = 0; l < 32; l++) { dst[l] = (uint8_t) (dst[l] | (((src[l] >> S) & M) << D)); }
+}
+inline void tiled_lut8(const uint8_t * lut, const uint8_t * src, uint8_t * dst) {
+    for (int j = 0; j < 16; j++) { dst[j] = lut[src[j]]; }
+}
+inline void tiled_unpk_sign32(uint64_t g0, uint64_t g1, uint64_t g2, uint64_t g3,
+                              const uint8_t signs[4], uint8_t * dst32) {
+    const uint64_t g[4] = { g0, g1, g2, g3 };
+    for (int l = 0; l < 4; l++) {
+        const uint8_t * v = (const uint8_t *) &g[l];
+        const uint8_t s = signs[l];
+        for (int j = 0; j < 8; j++) {
+            dst32[8 * l + j] = (s & (1 << j)) ? (uint8_t) (128 - v[j]) : (uint8_t) (128 + v[j]);
+        }
+    }
+}
+// 8 ternary grid bytes (0 = 0, 1 = +1, 0xFF = -1): dst[j] = 128 + delta + 8 * (int8_t) src[j]
+inline void tiled_unpk_tern8(const uint8_t * src, int8_t delta, uint8_t * dst) {
+    for (int j = 0; j < 8; j++) { dst[j] = (uint8_t) (128 + (int) delta + 8 * (int8_t) src[j]); }
+}
+#endif
+
+// Accumulate one 16x16 microtile (src0 rows [i0, i0+16), src1 cols [j0, j0+16))
+// over one 256-K slab held in the tiles into a j-major float buffer
+// (row width buf_stride): buf[i*buf_stride + j] += partial.
+// SUBBLK/HAS_MIN/BIAS are the src0 format constants (see tiled_tile_src0).
+// ACTBIAS (AVX2 only): the activation is pre-biased +128 by tiled_repack_src1,
+// num_k = K-blocks per row: the tile holds num_k slabs at row stride num_k*256
+// (Default case is num_k=1, 256x256 tiles, we go to longer num_k to improve memory bandwidth when num_rows is small)
+template <int SUBBLK, bool HAS_MIN, int BIAS, bool ACTBIAS>
+void tiled_run_microtile(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1,
+                         int i0, int j0, int num_k, int slab, float * buf, int buf_stride);
+
+// Optional repack, if profitable for the kernel.
+// Repacks one 16-row band of src1 codes, called by driver as we reach each 16-row band in outer loop
+void tiled_repack_src1(tiled_tile_src1 * src1, int row0, int num_k, bool bias);
+
+// Optional repack, if profitable for the kernel
+// Repack the entire src0 panel in place, called by the driver immediately after dequant
+template <int SUBBLK>
+void tiled_repack_src0(tiled_tile_src0 * tile, int n_rows, int num_k, int BIAS, bool corr);
+
+

ggml/src/ggml-cpu/tiled/tiled.h

@@ -0,0 +1,31 @@
+#pragma once
+
+#include <stddef.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Amount of wdata to reserve for tiled workspaces
+size_t ggml_tiled_wdata_size(int n_tasks, struct ggml_tensor * dst);
+
+// tiled K-quant matmul; returns true if the op was computed here,
+// false to fall through to the stock path
+bool ggml_compute_forward_mul_mat_tiled(const struct ggml_compute_params * params,
+                                        struct ggml_tensor * dst);
+
+// MUL_MAT_ID (MoE) path, one expert; returns true if the expert was computed here, per expert
+// eligibility (type gate, batch floor) is decided inside. expert_rows points at the expert's
+// row of the matrix_rows table of (expert slot, batch row) int32 pairs; scratch is the
+// per-thread tiled_ws region reserved in wdata, n_tasks of ggml_tiled_ws_size() bytes
+bool ggml_compute_forward_mul_mat_id_tiled(const struct ggml_compute_params * params,
+                                           struct ggml_tensor *               dst,
+                                           int64_t                            cur_a,
+                                           int64_t                            cne1,
+                                           const int32_t *                    expert_rows,
+                                           char *                             scratch);
+
+#ifdef __cplusplus
+}
+#endif

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

@@ -1875,7 +1875,7 @@ static __global__ void flash_attn_ext_f16(
 #endif // defined(AMD_WMMA_AVAILABLE)
 
 #if defined(AMD_MFMA_AVAILABLE)
-    if (ncols1*ncols2 < 16 || DKQ > 256) {
+    if (ncols1*ncols2 < 16 || (DKQ > 256 && ncols1*ncols2 < 64)) {
         NO_DEVICE_CODE;
         return;
     }

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

@@ -19,41 +19,41 @@
     }                                                                                                 \
 
 static constexpr __host__ __device__ uint32_t ggml_cuda_fattn_tile_get_config_nvidia_fp16(const int DKQ, const int DV, const int ncols) {
-    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 40,  40,  2,  64, 2,  64,  40)
-    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 40,  40,  4, 128, 2,  64,  40)
-    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 40,  40,  8, 256, 2,  64,  40)
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 40,  40,  2, 128, 3, 128,  40)
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 40,  40,  4, 128, 2, 128,  40)
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 40,  40,  8, 128, 3, 128,  40)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 40,  40, 16, 256, 2,  64,  40)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 40,  40, 32, 256, 2,  64,  40)
 
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 64,  64,  2,  64, 2,  64,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 64,  64,  4, 128, 2,  64,  64)
-    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 64,  64,  8, 256, 2,  64,  64)
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 64,  64,  8, 256, 3, 128,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 64,  64, 16, 256, 2,  64,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 64,  64, 32, 256, 2,  64,  64)
 
-    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 72,  72,  2,  64, 2,  64,  72)
-    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 72,  72,  4, 128, 2,  64,  72)
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 72,  72,  2, 128, 2,  64,  72)
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 72,  72,  4, 128, 3, 128,  72)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 72,  72,  8, 256, 2,  64,  72)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 72,  72, 16, 256, 2,  64,  72)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 72,  72, 32, 256, 2,  64,  72)
 
-    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 80,  80,  2,  64, 2,  64,  40)
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 80,  80,  2, 128, 2,  64,  40)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 80,  80,  4, 128, 2,  64,  40)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 80,  80,  8, 256, 2,  64,  40)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 80,  80, 16, 256, 2,  64,  40)
-    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 80,  80, 32, 256, 2,  64,  40)
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 80,  80, 32, 256, 2, 128,  40)
 
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 96,  96,  2,  64, 2,  64,  48)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 96,  96,  4, 128, 2,  64,  48)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 96,  96,  8, 256, 2,  64,  48)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE( 96,  96, 16, 256, 2,  64,  48)
-    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 96,  96, 32, 256, 2,  64,  48)
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE( 96,  96, 32, 256, 2, 128,  24)
 
-    GGML_CUDA_FATTN_TILE_CONFIG_CASE(112, 112,  2,  64, 2,  64,  56)
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE(112, 112,  2, 128, 3,  64,  56)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(112, 112,  4, 128, 2,  64,  56)
-    GGML_CUDA_FATTN_TILE_CONFIG_CASE(112, 112,  8, 256, 2,  64,  56)
-    GGML_CUDA_FATTN_TILE_CONFIG_CASE(112, 112, 16, 256, 2,  64,  56)
-    GGML_CUDA_FATTN_TILE_CONFIG_CASE(112, 112, 32, 256, 2,  64,  56)
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE(112, 112,  8, 256, 3,  64, 112)
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE(112, 112, 16, 256, 2,  32, 112)
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE(112, 112, 32, 256, 2, 128,  56)
 
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(128, 128,  2,  64, 2,  64,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(128, 128,  4, 128, 2,  64,  64)

ggml/src/ggml-cuda/fattn.cu

@@ -681,7 +681,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
     }
 
     // AMD MFMA needs a certain minimum batch size to outscale the tile kernel for large head sizes.
-    if ((amd_mfma_available(cc) && Q->ne[0] <= 256) && Q->ne[0] != 40 && Q->ne[0] != 72) {
+    if (amd_mfma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) {
         if ((Q->ne[0] <= 64 && Q->ne[1] * gqa_ratio_eff > 8)) {
             return BEST_FATTN_KERNEL_MMA_F16;
         }
@@ -691,6 +691,9 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
         if ((Q->ne[0] <= 256 && Q->ne[1] * gqa_ratio_eff > 64)) {
             return BEST_FATTN_KERNEL_MMA_F16;
         }
+        if (Q->ne[0] > 256 && gqa_opt_applies && Q->ne[1] * gqa_ratio_eff > 128) {
+            return BEST_FATTN_KERNEL_MMA_F16;
+        }
     }
 
     // AMD WMMA is faster than the tile kernel if the wide tiles with high arithmetic intensity can be utilized.

ggml/src/ggml-cuda/fwht.cu

@@ -1,9 +1,10 @@
 #include "common.cuh"
+#include "convert.cuh"
 #include "fwht.cuh"
 
-template <int N>
+template <int N, typename T>
 __launch_bounds__(4*ggml_cuda_get_physical_warp_size(), 1)
-__global__ void fwht_cuda(const float * src, float * dst, const int64_t n_rows, const float scale) {
+__global__ void fwht_cuda(const T * src, float * dst, const int64_t n_rows, const float scale) {
     constexpr int warp_size = ggml_cuda_get_physical_warp_size();
 
     const int64_t r = (int64_t) blockIdx.x * blockDim.y + threadIdx.y;
@@ -22,7 +23,7 @@ __global__ void fwht_cuda(const float * src, float * dst, const int64_t n_rows,
     ggml_cuda_pdl_sync();
 #pragma unroll
     for (int i = 0; i < el_w; ++i) {
-        reg[i] = src[i * warp_size + lane] * scale;
+        reg[i] = ggml_cuda_cast<float>(src[i * warp_size + lane]) * scale;
     }
 
 #pragma unroll
@@ -58,15 +59,12 @@ __global__ void fwht_cuda(const float * src, float * dst, const int64_t n_rows,
     }
 }
 
-bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) {
-    GGML_ASSERT(ggml_are_same_shape(src, dst));
-    if (!ggml_is_contiguous(src) || !ggml_is_contiguous(dst)) {
-        return false;
-    }
+template <typename T>
+static bool ggml_cuda_op_fwht_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) {
     const int     n    = src->ne[0];
     const int64_t rows = ggml_nrows(src);
 
-    const float * src_d = (const float *) src->data;
+    const T *     src_d = (const T *) src->data;
     float *       dst_d = (float *) dst->data;
 
     const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
@@ -84,18 +82,47 @@ bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src,
 
     switch (n) {
         case 64:
-            ggml_cuda_kernel_launch(fwht_cuda<64>, launch_params, src_d, dst_d, rows, scale);
+            ggml_cuda_kernel_launch(fwht_cuda<64, T>, launch_params, src_d, dst_d, rows, scale);
             return true;
         case 128:
-            ggml_cuda_kernel_launch(fwht_cuda<128>, launch_params, src_d, dst_d, rows, scale);
+            ggml_cuda_kernel_launch(fwht_cuda<128, T>, launch_params, src_d, dst_d, rows, scale);
             return true;
         case 256:
-            ggml_cuda_kernel_launch(fwht_cuda<256>, launch_params, src_d, dst_d, rows, scale);
+            ggml_cuda_kernel_launch(fwht_cuda<256, T>, launch_params, src_d, dst_d, rows, scale);
             return true;
         case 512:
-            ggml_cuda_kernel_launch(fwht_cuda<512>, launch_params, src_d, dst_d, rows, scale);
+            ggml_cuda_kernel_launch(fwht_cuda<512, T>, launch_params, src_d, dst_d, rows, scale);
             return true;
         default:
             return false;
     }
 }
+
+bool ggml_cuda_op_mul_mat_use_fwht(const struct ggml_tensor * op) {
+    const struct ggml_tensor * a = op->src[0];
+    const struct ggml_tensor * b = op->src[1];
+
+    return op->op == GGML_OP_MUL_MAT && ggml_get_op_params_i32(op, 1) == GGML_HINT_SRC0_IS_HADAMARD &&
+           a->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 &&
+           (b->type == GGML_TYPE_F32 || b->type == GGML_TYPE_F16) && ggml_is_contiguous(b) && ggml_is_contiguous(op) &&
+           ggml_are_same_shape(b, op);
+}
+
+bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) {
+    GGML_ASSERT(ggml_are_same_shape(src, dst));
+    if (!ggml_is_contiguous(src) || !ggml_is_contiguous(dst)) {
+        return false;
+    }
+    if (dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+
+    switch (src->type) {
+        case GGML_TYPE_F32:
+            return ggml_cuda_op_fwht_impl<float>(ctx, src, dst);
+        case GGML_TYPE_F16:
+            return ggml_cuda_op_fwht_impl<half>(ctx, src, dst);
+        default:
+            return false;
+    }
+}

ggml/src/ggml-cuda/fwht.cuh

@@ -1,4 +1,6 @@
 #include "common.cuh"
 
+bool ggml_cuda_op_mul_mat_use_fwht(const struct ggml_tensor * op);
+
 // Returns whether the Fast Walsh-Hadamard transform could be used.
 bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst);

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

@@ -1818,8 +1818,7 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) {
 static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
     GGML_TENSOR_BINARY_OP_LOCALS
 
-    const int32_t hint = ggml_get_op_params_i32(dst, 1);
-    if (hint == GGML_HINT_SRC0_IS_HADAMARD && ggml_cuda_op_fwht(ctx, src1, dst)) {
+    if (ggml_cuda_op_mul_mat_use_fwht(dst) && ggml_cuda_op_fwht(ctx, src1, dst)) {
         return;
     }
 
@@ -5212,7 +5211,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                 if (a->nb[0] != ggml_element_size(a) || b->nb[0] != ggml_element_size(b)) {
                     return false; // TODO this could in principle be implemented though currently there is no use case.
                 }
-                if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16) {
+                if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16 && !ggml_cuda_op_mul_mat_use_fwht(op)) {
                     return false;
                 }
                 if (op->op == GGML_OP_MUL_MAT_ID && ggml_get_op_params_i32(op, 3) == GGML_PREC_F32) {
@@ -5488,8 +5487,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
 
             if (op->src[3]->ne[0] == 1) {
                 // Mamba2
-                // (kernel only supports (d_state == 128 || d_state == 256) && d_head % 16 == 0)
-                return (op->src[0]->ne[0] == 128 || op->src[0]->ne[0] == 256) && op->src[0]->ne[1] % 16 == 0;
+                // (kernel only supports (d_state == 96 || d_state == 128 || d_state == 256) && d_head % 16 == 0)
+                const int64_t d_state = op->src[0]->ne[0];
+                return (d_state == 96 || d_state == 128 || d_state == 256) && op->src[0]->ne[1] % 16 == 0;
             } else {
                 if (K > 1) {
                     return false;

ggml/src/ggml-cuda/ssm-scan.cu

@@ -163,13 +163,19 @@ __global__ void __launch_bounds__(d_state, 1)
     const int lane     = threadIdx.x % WARP_SIZE;
     const int warp_idx = blockIdx.x  * c_factor + warp;
 
+    ggml_cuda_pdl_sync();
+
+    // the last block can have unused warps when n_head*d_head is not a multiple of c_factor
+    if (warp_idx >= n_head * d_head) {
+        return;
+    }
+
     const int head_idx =  warp_idx / d_head;
     const int head_off = (warp_idx % d_head) * sizeof(float);
     const int seq_idx  = blockIdx.y;
 
     const int group_off = (head_idx / (n_head / n_group)) * d_state * sizeof(float);
 
-    ggml_cuda_pdl_sync();
     // TODO: refactor strides to be in elements/floats instead of bytes to be cleaner and consistent with the rest of the codebase
     const float * s0_warp = (const float *) ((const char *) src0 + src6[seq_idx] * src0_nb3 + head_idx * src0_nb2 + head_off * d_state);
     const float * x_warp  = (const float *) ((const char *) src1 + (seq_idx * src1_nb3) + (warp_idx * sizeof(float)));
@@ -246,7 +252,17 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
     // NOTE: if you change conditions here, be sure to update the corresponding supports_op condition!
     if (src3_nb1 == sizeof(float)) {
         // Mamba-2
-        if (d_state == 128) {
+        if (d_state == 96) {
+            constexpr int threads   = 96;
+            constexpr int num_warps = threads/WARP_SIZE;
+
+            const dim3 blocks((n_head * head_dim + (num_warps - 1)) / num_warps, n_seq, 1);
+            const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks, threads, 0, stream);
+            ggml_cuda_kernel_launch(ssm_scan_f32_group<96/WARP_SIZE, 96>, launch_params,
+                    src0, src1, src2, src3, src4, src5, src6, dst,
+                    src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
+                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K);
+        } else if (d_state == 128) {
             constexpr int threads   = 128;
             constexpr int num_warps = threads/WARP_SIZE;
 
@@ -267,7 +283,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
                     src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
                     src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K);
         } else {
-            GGML_ABORT("doesn't support d_state!=(128 or 256).");
+            GGML_ABORT("doesn't support d_state!=(96, 128 or 256).");
         }
     } else {
         // Mamba-1

ggml/src/ggml-hexagon/ggml-hexagon.cpp

@@ -63,6 +63,7 @@
 #include "htp/rope-ops.h"
 #include "htp/ssm-conv.h"
 #include "htp/gated-delta-net-ops.h"
+#include "htp/argsort-ops.h"
 #include "htp_iface.h"
 #include "htp-drv.h"
 
@@ -368,6 +369,13 @@ static void ggml_hexagon_precompute_gated_delta_net_params(
     struct htp_gdn_kernel_params * kparams
 );
 
+static void ggml_hexagon_precompute_sort_params(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    bool is_top_k,
+    struct htp_sort_kernel_params * kparams
+);
+
 static void ggml_hexagon_precompute_fused_mmnx_params(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * src0,
@@ -4998,13 +5006,48 @@ static bool ggml_hexagon_precompute_binary_params(
     const bool is_add_id = op == HTP_OP_ADD_ID;
     const bool is_scalar = !is_add_id && src1->ne[0] == 1;
     const bool is_transposed = src0->nb[1] < src0_row_size || src1->nb[1] < src1_row_size || dst->nb[1] < dst_row_size;
+    const bool is_row_bcast = !is_add_id && !is_scalar && !is_transposed &&
+        src1->ne[0] == src0->ne[0] &&
+        (src0->ne[1] > 1 || src0->ne[2] > 1 || src0->ne[3] > 1) &&
+        src1->ne[1] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1;
     const bool is_same_shape = !is_add_id && !is_scalar && !is_transposed &&
         src1->ne[0] == src0->ne[0] &&
-        (src1->ne[1] == src0->ne[1] || src1->ne[1] == 1) &&
+        src1->ne[1] == src0->ne[1] &&
         (src1->ne[2] == src0->ne[2] || src1->ne[2] == 1) &&
         (src1->ne[3] == src0->ne[3] || src1->ne[3] == 1);
-    const bool is_row_bcast = is_same_shape && src1->ne[1] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1;
-    const bool is_complex   = !is_add_id && !is_scalar && !is_same_shape && (src1->ne[0] == src0->ne[0]);
+    const bool is_complex   = !is_add_id && !is_scalar && !is_same_shape && !is_row_bcast && (src1->ne[0] == src0->ne[0]);
+    const bool is_contig    = ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst);
+    const bool is_scalar_broadcast = !is_add_id && (ggml_nelements(src1) == 1);
+
+    if (!is_add_id && is_contig && (ggml_are_same_shape(src0, src1) || is_scalar_broadcast)) {
+        const uint32_t total_elems = (uint32_t) ggml_nelements(src0);
+        const uint32_t n_threads = sess->n_threads;
+        const uint32_t max_chunk_elems = 32768 / elem_size;
+        const uint32_t min_chunk_elems = 256;
+        const uint32_t target_chunk_elems = hex_round_up((total_elems + (2 * n_threads) - 1) / (2 * n_threads), 32);
+        const uint32_t chunk_size = (std::min)(max_chunk_elems, (std::max)(target_chunk_elems, min_chunk_elems));
+        const uint32_t chunk_bytes = hex_round_up(chunk_size * elem_size, 128);
+
+        kparams->kernel_type           = HTP_BINARY_KERNEL_CHUNKED;
+        kparams->n_threads             = n_threads;
+        kparams->rows_per_buffer       = 1;
+        kparams->src0_row_size_aligned = src0_row_size_aligned;
+        kparams->src1_row_size_aligned = is_scalar_broadcast ? 0 : src1_row_size_aligned;
+        kparams->dst_row_size_aligned  = dst_row_size_aligned;
+        kparams->src1_size             = 0;
+        kparams->chunk_size            = chunk_size;
+        kparams->chunk_bytes           = chunk_bytes;
+        kparams->is_scalar             = is_scalar_broadcast ? 1 : 0;
+
+        struct htp_binary_vtcm_layout L;
+        htp_binary_vtcm_layout_build(&L, kparams, sess->vtcm_size);
+        if (L.total_bytes == 0 || L.total_bytes > sess->vtcm_size) {
+            return false;
+        }
+
+        kparams->vtcm_size = L.total_bytes;
+        return true;
+    }
 
     enum htp_binary_kernel_type kernel_type;
     size_t src1_size = 0;
@@ -5042,6 +5085,32 @@ static bool ggml_hexagon_precompute_binary_params(
     struct htp_binary_vtcm_layout L;
     htp_binary_vtcm_layout_build(&L, kparams, sess->vtcm_size);
     if (L.rows_per_buffer == 0 || L.total_bytes > sess->vtcm_size) {
+        if (!is_add_id && is_contig && (ggml_are_same_shape(src0, src1) || is_scalar_broadcast)) {
+            const uint32_t total_elems = (uint32_t) ggml_nelements(src0);
+            const uint32_t n_threads = sess->n_threads;
+            const uint32_t max_chunk_elems = 32768 / elem_size;
+            const uint32_t min_chunk_elems = 256;
+            const uint32_t target_chunk_elems = hex_round_up((total_elems + (2 * n_threads) - 1) / (2 * n_threads), 32);
+            const uint32_t chunk_size = (std::min)(max_chunk_elems, (std::max)(target_chunk_elems, min_chunk_elems));
+            const uint32_t chunk_bytes = hex_round_up(chunk_size * elem_size, 128);
+
+            kparams->kernel_type           = HTP_BINARY_KERNEL_CHUNKED;
+            kparams->n_threads             = n_threads;
+            kparams->rows_per_buffer       = 1;
+            kparams->src1_row_size_aligned = is_scalar_broadcast ? 0 : src1_row_size_aligned;
+            kparams->src1_size             = 0;
+            kparams->chunk_size            = chunk_size;
+            kparams->chunk_bytes           = chunk_bytes;
+            kparams->is_scalar             = is_scalar_broadcast ? 1 : 0;
+
+            htp_binary_vtcm_layout_build(&L, kparams, sess->vtcm_size);
+            if (L.total_bytes == 0 || L.total_bytes > sess->vtcm_size) {
+                return false;
+            }
+
+            kparams->vtcm_size = L.total_bytes;
+            return true;
+        }
         return false;
     }
 
@@ -5139,49 +5208,51 @@ static void ggml_hexagon_precompute_get_rows_params(
     const uint32_t ne12 = src1->ne[2];
     const uint32_t nr = ne10 * ne11 * ne12;
 
-    const size_t nb01 = src0->nb[1];
-    const size_t nb1 = dst->nb[1];
+    const ggml_tensor * src0_base = src0->view_src ? src0->view_src : src0;
+    const auto * extra = src0_base->buffer && ggml_backend_buffer_is_hexagon(src0_base->buffer) ?
+        (const ggml_hexagon_tensor_extra *) src0_base->extra : nullptr;
+    const bool tiled = src0->type == GGML_TYPE_Q4_0 || (extra && (extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0) ||
+                       sess->needs_repack.count(src0_base) || sess->needs_repack.count(src0);
 
-    const bool can_use_dma = (src0->type == dst->type) && (nb01 == nb1);
-    const bool use_dma = can_use_dma && (ne00 >= 2048);
-
-    kparams->use_dma = use_dma ? 1 : 0;
-
-    uint32_t chunks_per_row = 1;
-    uint32_t chunk_size = ne00;
-    uint32_t total_tasks = nr;
-
-    if (use_dma) {
-        kparams->n_threads = (std::min)((uint32_t)sess->n_threads, nr);
-        kparams->tasks_per_thread = (nr + kparams->n_threads - 1) / kparams->n_threads;
+    if (src0->type == dst->type) {
+        kparams->kernel_type = HTP_GET_ROWS_KERNEL_SAMETYPE;
+    } else if (tiled) {
+        kparams->kernel_type = HTP_GET_ROWS_KERNEL_TILED;
     } else {
-        if (src0->type == GGML_TYPE_F32 && nr < sess->n_threads) {
-            const uint32_t min_chunk_size = 1024;
-            uint32_t max_chunks = ne00 / min_chunk_size;
-            if (max_chunks == 0) {
-                max_chunks = 1;
-            }
-            chunks_per_row = (std::min)((sess->n_threads + nr - 1) / nr, max_chunks);
-            chunk_size = (ne00 + chunks_per_row - 1) / chunks_per_row;
-            total_tasks = nr * chunks_per_row;
-        }
-        kparams->n_threads = (std::min)(total_tasks, (uint32_t)sess->n_threads);
-        kparams->tasks_per_thread = (total_tasks + kparams->n_threads - 1) / kparams->n_threads;
+        kparams->kernel_type = HTP_GET_ROWS_KERNEL_FLAT;
     }
 
+    const uint32_t chunks_per_row = 1;
+    const uint32_t chunk_size = ne00;
+    const uint32_t total_tasks = nr;
+
+    kparams->n_threads = (std::min)((uint32_t)sess->n_threads, total_tasks);
+
+    struct htp_get_rows_vtcm_layout vtcm_layout = {};
+    while (kparams->n_threads > 0) {
+        htp_get_rows_vtcm_layout_build(&vtcm_layout, kparams->kernel_type, src0->type, ne00, kparams->n_threads);
+        if (vtcm_layout.total_bytes <= sess->vtcm_size) {
+            break;
+        }
+        --kparams->n_threads;
+    }
+
+    if (kparams->n_threads == 0 && total_tasks > 0) {
+        htp_get_rows_vtcm_layout_build(&vtcm_layout, kparams->kernel_type, src0->type, ne00, 1);
+    }
+
+    kparams->vtcm_size = (total_tasks == 0) ? 0 : vtcm_layout.total_bytes;
+    kparams->tasks_per_thread = kparams->n_threads > 0 ? (total_tasks + kparams->n_threads - 1) / kparams->n_threads : 0;
+
     kparams->chunks_per_row = chunks_per_row;
     kparams->chunk_size = chunk_size;
     kparams->total_tasks = total_tasks;
 
-    kparams->div_ne10 = init_fastdiv_values(ne10);
-    kparams->div_ne10_ne11 = init_fastdiv_values(ne10 * ne11);
-    kparams->div_chunks_per_row = init_fastdiv_values(chunks_per_row);
-    kparams->div_ne02 = init_fastdiv_values(ne02);
-    kparams->div_ne03 = init_fastdiv_values(ne03);
-
-    struct htp_get_rows_vtcm_layout vtcm_layout;
-    htp_get_rows_vtcm_layout_build(&vtcm_layout, src0->type, ne00, kparams->n_threads);
-    kparams->vtcm_size = vtcm_layout.total_bytes;
+    kparams->div_ne10 = ne10 > 0 ? init_fastdiv_values(ne10) : fastdiv_values{0, 0};
+    kparams->div_ne10_ne11 = (ne10 * ne11) > 0 ? init_fastdiv_values(ne10 * ne11) : fastdiv_values{0, 0};
+    kparams->div_chunks_per_row = chunks_per_row > 0 ? init_fastdiv_values(chunks_per_row) : fastdiv_values{0, 0};
+    kparams->div_ne02 = ne02 > 0 ? init_fastdiv_values(ne02) : fastdiv_values{0, 0};
+    kparams->div_ne03 = ne03 > 0 ? init_fastdiv_values(ne03) : fastdiv_values{0, 0};
 }
 
 static void ggml_hexagon_precompute_set_rows_params(
@@ -5479,6 +5550,52 @@ static void ggml_hexagon_precompute_gated_delta_net_params(
     if (kparams->n_threads > 0) kparams->div_n_threads = init_fastdiv_values(kparams->n_threads);
 }
 
+static void ggml_hexagon_precompute_sort_params(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    bool is_top_k,
+    struct htp_sort_kernel_params * kparams
+) {
+    memset(kparams, 0, sizeof(*kparams));
+
+    const struct ggml_tensor * src0 = op->src[0];
+    const struct ggml_tensor * dst  = op;
+
+    const uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3];
+    const uint32_t ne00       = src0->ne[0];
+    const uint32_t k          = dst->ne[0];
+
+    int32_t order = GGML_SORT_ORDER_DESC;
+    if (!is_top_k) {
+        order = ((const int32_t *) op->op_params)[0];
+    }
+
+    const uint32_t n_threads_max = sess->n_threads > 0 ? sess->n_threads : 4;
+    const size_t vtcm_budget = sess->vtcm_size > 0 ? sess->vtcm_size : (8 * 1024 * 1024);
+
+    struct htp_sort_vtcm_layout layout;
+    bool ok = htp_sort_solve_layout(&layout, ne00, total_rows, k, n_threads_max, vtcm_budget, is_top_k);
+    GGML_ASSERT(ok);
+
+    kparams->n_threads         = (int32_t) layout.n_threads;
+    kparams->total_rows        = (int32_t) total_rows;
+    kparams->row_start         = 0;
+    kparams->row_end           = (int32_t) total_rows;
+    kparams->ne00              = (int32_t) ne00;
+    kparams->k                 = (int32_t) k;
+    kparams->order             = order;
+    kparams->is_top_k          = is_top_k ? 1 : 0;
+    kparams->use_dma           = 1;
+    kparams->chunk_elems       = (int32_t) layout.chunk_elems;
+    kparams->n_chunks          = (int32_t) layout.n_chunks;
+    kparams->vtcm_size         = (int32_t) layout.total_bytes;
+    kparams->phase1_slot_size  = (int32_t) layout.phase1_slot_size;
+    kparams->merge_values_off  = (int32_t) layout.merge_values_off;
+    kparams->merge_indices_off = (int32_t) layout.merge_indices_off;
+    kparams->merge_elems       = (int32_t) layout.merge_elems;
+    kparams->n_slots           = (int32_t) layout.n_slots;
+}
+
 static void ggml_hexagon_precompute_fused_mmnx_params(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * src0, // W0
@@ -5830,7 +5947,8 @@ static bool ggml_hexagon_supported_unary(const struct ggml_hexagon_session * ses
             case GGML_OP_LOG:
                 break;
             case GGML_OP_UNARY:
-                if (ggml_get_unary_op(op) != GGML_UNARY_OP_ABS) {
+                if (ggml_get_unary_op(op) != GGML_UNARY_OP_ABS &&
+                    ggml_get_unary_op(op) != GGML_UNARY_OP_STEP) {
                     return false;
                 }
                 break;
@@ -5853,6 +5971,26 @@ static bool ggml_hexagon_supported_unary(const struct ggml_hexagon_session * ses
     GGML_UNUSED(sess);
 }
 
+static bool ggml_hexagon_supported_sum(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
+    const struct ggml_tensor * src0 = op->src[0];
+    const struct ggml_tensor * dst  = op;
+
+    if (src0->type != GGML_TYPE_F32) {
+        return false;
+    }
+    if (dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+
+    if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(dst)) {
+        return false;
+    }
+
+    return true;
+
+    GGML_UNUSED(sess);
+}
+
 static bool ggml_hexagon_supported_sum_rows(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
     const struct ggml_tensor * src0 = op->src[0];
     const struct ggml_tensor * dst  = op;
@@ -5874,6 +6012,26 @@ static bool ggml_hexagon_supported_sum_rows(const struct ggml_hexagon_session *
     GGML_UNUSED(sess);
 }
 
+static bool ggml_hexagon_supported_argmax(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
+    const struct ggml_tensor * src0 = op->src[0];
+    const struct ggml_tensor * dst  = op;
+
+    if (src0->type != GGML_TYPE_F32) {
+        return false;
+    }
+    if (dst->type != GGML_TYPE_I32) {
+        return false;
+    }
+
+    if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(dst)) {
+        return false;
+    }
+
+    return true;
+
+    GGML_UNUSED(sess);
+}
+
 static bool ggml_hexagon_supported_activations(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
     const struct ggml_tensor * src0 = op->src[0];
     const struct ggml_tensor * src1 = op->src[1];
@@ -6020,19 +6178,36 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session *
     const struct ggml_tensor * src1 = op->src[1]; // indices
     const struct ggml_tensor * dst  = op;
 
-    if (src0->extra) {
-        const auto * extra = (const ggml_hexagon_tensor_extra *) src0->extra;
-        if (extra->flags & GGML_HEXAGON_TENSOR_REPACK) {
+    if (src0->type == GGML_TYPE_Q4_0 && src0->view_src) {
+        return false;
+    }
+
+    const ggml_tensor * src0_base = src0->view_src ? src0->view_src : src0;
+    bool is_repacked = false;
+    if (src0_base->buffer && ggml_backend_buffer_is_hexagon(src0_base->buffer) && src0_base->extra) {
+        const auto * extra = (const ggml_hexagon_tensor_extra *) src0_base->extra;
+        is_repacked = (extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0;
+        if (is_repacked && src0->type != GGML_TYPE_Q4_0 && src0->type != GGML_TYPE_Q8_0) {
             return false;
         }
     }
+    is_repacked = is_repacked || sess->needs_repack.count(src0_base) || sess->needs_repack.count(src0);
 
-    if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_I32 && src0->ne[0] < 32) {
+    // View offsets use the raw quantized layout and cannot address a tiled allocation.
+    if (src0->view_src && is_repacked) {
+        return false;
+    }
+
+    if (src0->type == GGML_TYPE_Q4_0 && src0->buffer && !is_repacked) {
+        return false;
+    }
+
+    if (src0->type != dst->type && src0->ne[0] < 32) {
         return false;
     }
 
     if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16 &&
-        src0->type != GGML_TYPE_Q8_0 && src0->type != GGML_TYPE_I32) {
+        src0->type != GGML_TYPE_Q4_0 && src0->type != GGML_TYPE_Q8_0 && src0->type != GGML_TYPE_I32) {
         return false;
     }
 
@@ -6040,13 +6215,30 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session *
         return false;
     }
 
-    if (src0->type == GGML_TYPE_I32) {
-        if (dst->type != GGML_TYPE_I32) {
+    if (src0->type == dst->type) {
+        if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_I32 && src0->type != GGML_TYPE_F16) {
             return false;
         }
-    }
-    else if (dst->type != GGML_TYPE_F32) {
+    } else if (src0->type == GGML_TYPE_I32) {
         return false;
+    } else if (dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+
+    // Empty recurrent-state gathers are skipped at execution; do not split the graph for them.
+    if (ggml_is_empty(op)) {
+        return true;
+    }
+
+    struct htp_get_rows_kernel_params kparams;
+    ggml_hexagon_precompute_get_rows_params(sess, src0, src1, dst, &kparams);
+    if (kparams.n_threads == 0 || (size_t) kparams.vtcm_size > sess->vtcm_size) {
+        return false;
+    }
+
+    // Q4_0 has no raw fallback. Mark only accepted tensors for repacking.
+    if (src0->type == GGML_TYPE_Q4_0 && !src0->buffer) {
+        sess->needs_repack.insert(src0);
     }
 
     return true;
@@ -6058,42 +6250,36 @@ static bool ggml_hexagon_supported_argsort(const struct ggml_hexagon_session * s
     const struct ggml_tensor * src0 = op->src[0]; // values
     const struct ggml_tensor * dst  = op;         // indices
 
-    if (src0->type != GGML_TYPE_F32) {
+    if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_I32) {
         return false;
     }
 
-    if (dst->type != GGML_TYPE_I32) {
-        return false;
-    }
+    const uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3];
+    const uint32_t n_threads_max = sess->n_threads > 0 ? sess->n_threads : 4;
+    const size_t vtcm_budget = sess->vtcm_size > 0 ? sess->vtcm_size : (8 * 1024 * 1024);
 
-    if (src0->ne[0] > (16*1024)) {
-        // reject tensors with huge rows for now
+    struct htp_sort_vtcm_layout layout;
+    if (!htp_sort_solve_layout(&layout, src0->ne[0], total_rows, dst->ne[0], n_threads_max, vtcm_budget, false)) {
         return false;
     }
 
     return true;
-
-    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_top_k(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
     const struct ggml_tensor * src0 = op->src[0]; // values
     const struct ggml_tensor * dst  = op;         // indices
 
-    if (src0->type != GGML_TYPE_F32) {
+    if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_I32) {
         return false;
     }
 
-    if (dst->type != GGML_TYPE_I32) {
-        return false;
-    }
+    const uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3];
+    const uint32_t n_threads_max = sess->n_threads > 0 ? sess->n_threads : 4;
+    const size_t vtcm_budget = sess->vtcm_size > 0 ? sess->vtcm_size : (8 * 1024 * 1024);
 
-    // Single row uses the threaded chunk+merge path. Multi-row uses one full
-    // buffer per thread, so it keeps the tighter 64K cap.
-    const bool single_row = (src0->ne[1] == 1 && src0->ne[2] == 1 && src0->ne[3] == 1);
-    const int64_t max_ne00 = single_row ? (256*1024) : (64*1024);
-
-    if (src0->ne[0] > max_ne00) {
+    struct htp_sort_vtcm_layout layout;
+    if (!htp_sort_solve_layout(&layout, src0->ne[0], total_rows, dst->ne[0], n_threads_max, vtcm_budget, true)) {
         return false;
     }
 
@@ -6401,9 +6587,11 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) {
         case GGML_OP_CONT:            return HTP_OP_CPY;
         case GGML_OP_GET_ROWS:        return HTP_OP_GET_ROWS;
         case GGML_OP_SET_ROWS:        return HTP_OP_SET_ROWS;
+        case GGML_OP_SUM:             return HTP_OP_SUM;
         case GGML_OP_SUM_ROWS:        return HTP_OP_SUM_ROWS;
         case GGML_OP_ARGSORT:         return HTP_OP_ARGSORT;
         case GGML_OP_TOP_K:           return HTP_OP_TOP_K;
+        case GGML_OP_ARGMAX:          return HTP_OP_ARGMAX;
         case GGML_OP_NORM:            return HTP_OP_NORM;
         case GGML_OP_L2_NORM:         return HTP_OP_L2_NORM;
         case GGML_OP_RMS_NORM:        return HTP_OP_RMS_NORM;
@@ -6440,6 +6628,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) {
                 case GGML_UNARY_OP_TANH:       return HTP_OP_UNARY_TANH;
                 case GGML_UNARY_OP_ABS:        return HTP_OP_UNARY_ABS;
                 case GGML_UNARY_OP_RELU:       return HTP_OP_UNARY_RELU;
+                case GGML_UNARY_OP_STEP:       return HTP_OP_UNARY_STEP;
             default:
                 break;
             }
@@ -6670,6 +6859,12 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg
                     node.node,
                     (struct htp_gdn_kernel_params *)node.kernel_params
                 );
+            } else if (node.opcode == HTP_OP_ARGSORT || node.opcode == HTP_OP_TOP_K) {
+                ggml_hexagon_precompute_sort_params(sess,
+                    node.node,
+                    node.opcode == HTP_OP_TOP_K,
+                    (struct htp_sort_kernel_params *) node.kernel_params
+                );
             }
             computed_nodes.push_back(std::move(node));
         }
@@ -7308,16 +7503,25 @@ static bool ggml_hexagon_supported_cpy(const struct ggml_hexagon_session * sess,
     if (dst->type != GGML_TYPE_F32 && dst->type != GGML_TYPE_F16 &&
         dst->type != GGML_TYPE_I32) return false;
 
+    const bool is_scalar  = (ggml_nelements(src0) == 1 && ggml_nelements(dst) == 1);
     const bool sametype   = (src0->type == dst->type);
-    const bool transposed = ggml_is_transposed(src0) || ggml_is_transposed(dst);
-    const bool sameshape  = !transposed && ggml_are_same_shape(src0, dst);
+    const bool transposed = !is_scalar && (ggml_is_transposed(src0) || ggml_is_transposed(dst));
+    const bool sameshape  = is_scalar || (!transposed && ggml_are_same_shape(src0, dst));
 
-    // Same-type copies also support I32.
+    if (src0->type == GGML_TYPE_I32 || dst->type == GGML_TYPE_I32) {
+        if (!sameshape) return false;
+        if (sametype) return true;
+        if ((src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_I32) ||
+            (src0->type == GGML_TYPE_I32 && dst->type == GGML_TYPE_F32)) {
+            return true;
+        }
+        return false;
+    }
+
+    // can handle any shape and any same-type (pretty slow if reshaping is required)
     if (sametype) return true;
 
-    // Type conversion is only supported between F32 and F16.
-    if (src0->type == GGML_TYPE_I32 || dst->type == GGML_TYPE_I32) return false;
-
+    // cannot handle re-shaping and type conversion at the same time
     if (!sameshape) return false;
 
     return true;
@@ -7466,10 +7670,18 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
             supp = ggml_hexagon_supported_unary(sess, op);
             break;
 
+        case GGML_OP_SUM:
+            supp = ggml_hexagon_supported_sum(sess, op);
+            break;
+
         case GGML_OP_SUM_ROWS:
             supp = ggml_hexagon_supported_sum_rows(sess, op);
             break;
 
+        case GGML_OP_ARGMAX:
+            supp = ggml_hexagon_supported_argmax(sess, op);
+            break;
+
         case GGML_OP_SOFT_MAX:
             supp = ggml_hexagon_supported_softmax(sess, op);
             break;
@@ -7486,6 +7698,7 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
                 case GGML_UNARY_OP_GELU:
                 case GGML_UNARY_OP_GELU_QUICK:
                 case GGML_UNARY_OP_RELU:
+                case GGML_UNARY_OP_STEP:
                     supp = ggml_hexagon_supported_unary(sess, op);
                     break;
                 default:

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

@@ -19,6 +19,8 @@
 #include "htp/ssm-conv.h"
 #include "htp/gated-delta-net-ops.h"
 #include "htp/softmax-ops.h"
+#include "htp/argsort-ops.h"
+#include "htp/get-rows-ops.h"
 
 struct htp_opnode {
     ggml_tensor * node   { nullptr };
@@ -359,17 +361,31 @@ struct htp_opformat {
         } else if (node.opcode == HTP_OP_GATED_DELTA_NET) {
             const auto * kparams = (const struct htp_gdn_kernel_params *) node.kernel_params;
             const char * path = (kparams->kernel_type == HTP_GDN_KERNEL_HMX_CHUNKED) ? "hmx-chunked" : "hvx-recurrent";
-            snprintf(str, max_size, "%s-%s vtcm %u",
-                     path,
-                     kparams->kda ? "kda" : "scalar",
+            snprintf(str, max_size, "%s-%s vtcm %u", path, kparams->kda ? "kda" : "scalar",
                      (unsigned int) (kparams->vtcm_size ? kparams->vtcm_size : kparams->vtcm_per_thread * kparams->n_threads));
         } else if (node.opcode == HTP_OP_MUL || node.opcode == HTP_OP_ADD || node.opcode == HTP_OP_ADD_ID ||
                    node.opcode == HTP_OP_SUB || node.opcode == HTP_OP_DIV) {
             const auto * kparams = (const struct htp_binary_kernel_params *) node.kernel_params;
             snprintf(str, max_size, "vtcm %u", (unsigned int) kparams->vtcm_size);
+        } else if (node.opcode == HTP_OP_ARGSORT || node.opcode == HTP_OP_TOP_K) {
+            const auto * kparams = (const struct htp_sort_kernel_params *) node.kernel_params;
+            snprintf(str, max_size, "%s nth %d nchk %d chk %d vtcm %d",
+                     node.opcode == HTP_OP_TOP_K ? "top_k" : "argsort",
+                     (int) kparams->n_threads, (int) kparams->n_chunks,
+                     (int) kparams->chunk_elems, (int) kparams->vtcm_size);
+        } else if (node.opcode == HTP_OP_GET_ROWS) {
+            const auto * kparams = (const struct htp_get_rows_kernel_params *) node.kernel_params;
+            const char * ktype_str = "unknown";
+            switch (kparams->kernel_type) {
+                case HTP_GET_ROWS_KERNEL_SAMETYPE: ktype_str = "sametype"; break;
+                case HTP_GET_ROWS_KERNEL_TILED:    ktype_str = "tiled"; break;
+                case HTP_GET_ROWS_KERNEL_FLAT:     ktype_str = "flat"; break;
+            }
+            snprintf(str, max_size, "%s%s vtcm %u", ktype_str, kparams->n_threads > 1 ? "-multi" : "", (unsigned int) kparams->vtcm_size);
         } else {
             snprintf(str, max_size, "----");
         }
+
     }
 
     void format(const htp_opnode & node) {

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

@@ -0,0 +1,178 @@
+#ifndef HTP_ARGSORT_OPS_H
+#define HTP_ARGSORT_OPS_H
+
+#include <stdint.h>
+#include <stddef.h>
+#include <stdbool.h>
+#include <string.h>
+
+#include "hex-fastdiv.h"
+
+struct htp_sort_kernel_params {
+    int32_t  n_threads;
+    int32_t  total_rows;
+    int32_t  row_start;
+    int32_t  row_end;
+    int32_t  ne00;
+    int32_t  k;
+    int32_t  order;             // GGML_SORT_ORDER_ASC (0) or GGML_SORT_ORDER_DESC (1)
+    int32_t  is_top_k;          // 1 if TOP_K, 0 if ARGSORT
+    int32_t  use_dma;           // 1 if DMA enabled
+    int32_t  chunk_elems;
+    int32_t  n_chunks;
+    int32_t  vtcm_size;
+    int32_t  phase1_slot_size;
+    int32_t  merge_values_off;
+    int32_t  merge_indices_off;
+    int32_t  merge_elems;
+    int32_t  n_slots;
+    int32_t  pad[15];
+};
+
+struct htp_sort_vtcm_layout {
+    size_t   total_bytes;
+    size_t   phase1_slot_size;
+    size_t   merge_values_off;
+    size_t   merge_indices_off;
+    uint32_t chunk_elems;
+    uint32_t n_chunks;
+    uint32_t merge_elems;
+    uint32_t n_threads;
+    uint32_t n_slots;
+};
+
+static inline bool htp_sort_solve_layout(
+    struct htp_sort_vtcm_layout * layout,
+    uint32_t ne00,
+    uint32_t total_rows,
+    uint32_t k,
+    uint32_t n_threads_max,
+    size_t vtcm_budget,
+    bool is_top_k) {
+
+    memset(layout, 0, sizeof(*layout));
+
+    if (total_rows > 1) {
+        uint32_t n_threads = total_rows < n_threads_max ? total_rows : n_threads_max;
+
+        uint32_t n_vec = (ne00 + 31) / 32;
+        uint32_t n_vec_pow2 = 1;
+        while (n_vec_pow2 < n_vec) n_vec_pow2 <<= 1;
+        uint32_t ne00_padded = n_vec_pow2 * 32;
+
+        size_t values_size  = ((ne00_padded * sizeof(float)) + 127) & ~127;
+        size_t indices_size = ((ne00_padded * sizeof(int32_t)) + 127) & ~127;
+        size_t spad_per_slot = ((values_size + indices_size) + 255) & ~255;
+
+        uint32_t n_slots = 2;
+        if (spad_per_slot * 2 > vtcm_budget) {
+            n_slots = 1;
+        }
+
+        size_t spad_per_thread = spad_per_slot * n_slots;
+
+        while (n_threads > 1 && (spad_per_thread * n_threads) > vtcm_budget) {
+            n_threads--;
+        }
+
+        size_t total_bytes = spad_per_thread * n_threads;
+        if (total_bytes > vtcm_budget) {
+            return false;
+        }
+
+        layout->total_bytes      = total_bytes;
+        layout->phase1_slot_size = spad_per_slot;
+        layout->chunk_elems      = ne00_padded;
+        layout->n_chunks         = 1;
+        layout->n_threads        = n_threads;
+        layout->n_slots          = n_slots;
+        return true;
+    }
+
+    uint32_t n_vec = (ne00 + 31) / 32;
+    uint32_t n_vec_pow2 = 1;
+    while (n_vec_pow2 < n_vec) n_vec_pow2 <<= 1;
+
+    uint32_t n_chunks = 1;
+    if (ne00 > 1024) {
+        while (n_chunks * 2 <= n_threads_max && n_chunks * 2 <= n_vec_pow2) {
+            n_chunks *= 2;
+        }
+    }
+
+    uint32_t chunk_n_vec = n_vec_pow2 / n_chunks;
+    uint32_t chunk_elems = chunk_n_vec * 32;
+
+    size_t phase1_values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
+    size_t phase1_indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
+    size_t phase1_slot_size    = ((phase1_values_size + phase1_indices_size) + 255) & ~255;
+    size_t phase1_total_size   = phase1_slot_size * n_chunks;
+
+    size_t merge_values_size  = 0;
+    size_t merge_indices_size = 0;
+    size_t merge_values_off   = phase1_total_size;
+    size_t merge_indices_off  = 0;
+    uint32_t merge_elems      = 0;
+
+    if (n_chunks > 1) {
+        if (is_top_k) {
+            uint32_t local_k = k < chunk_elems ? k : chunk_elems;
+            uint32_t total_candidates = n_chunks * local_k;
+            uint32_t merge_n_vec = (total_candidates + 31) / 32;
+            uint32_t merge_n_vec_pow2 = 1;
+            while (merge_n_vec_pow2 < merge_n_vec) merge_n_vec_pow2 <<= 1;
+            merge_elems = merge_n_vec_pow2 * 32;
+        } else {
+            merge_elems = n_vec_pow2 * 32;
+        }
+        merge_values_size  = ((merge_elems * sizeof(float)) + 127) & ~127;
+        merge_indices_size = ((merge_elems * sizeof(int32_t)) + 127) & ~127;
+        merge_indices_off  = merge_values_off + merge_values_size;
+    }
+
+    size_t total_bytes = phase1_total_size + merge_values_size + merge_indices_size;
+
+    if (total_bytes > vtcm_budget && n_chunks > 1) {
+        n_chunks = 1;
+        chunk_elems = n_vec_pow2 * 32;
+        phase1_values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
+        phase1_indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
+        phase1_slot_size    = ((phase1_values_size + phase1_indices_size) + 255) & ~255;
+        phase1_total_size   = phase1_slot_size;
+        merge_values_size   = 0;
+        merge_indices_size  = 0;
+        merge_values_off    = phase1_total_size;
+        merge_indices_off   = 0;
+        merge_elems         = 0;
+        total_bytes         = phase1_total_size;
+    }
+
+    if (total_bytes > vtcm_budget) {
+        return false;
+    }
+
+    uint32_t n_slots = 1;
+    if (n_chunks == 1 && phase1_slot_size * 2 <= vtcm_budget) {
+        n_slots = 2;
+        total_bytes = phase1_slot_size * 2;
+    }
+
+    layout->total_bytes       = total_bytes;
+    layout->phase1_slot_size  = phase1_slot_size;
+    layout->merge_values_off  = merge_values_off;
+    layout->merge_indices_off = merge_indices_off;
+    layout->chunk_elems       = chunk_elems;
+    layout->n_chunks          = n_chunks;
+    layout->merge_elems       = merge_elems;
+    layout->n_threads         = n_chunks;
+    layout->n_slots           = n_slots;
+    return true;
+}
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_sort_kernel_params) <= 128, "htp_sort_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_sort_kernel_params) <= 128, "htp_sort_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_ARGSORT_OPS_H

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

@@ -917,6 +917,154 @@ static void binary_thread_add_id_f32(unsigned int nth, unsigned int ith, void *
     dma_queue_flush(dma_q);
 }
 
+static inline void hvx_div_scalar_f32(uint8_t * restrict dst, const uint8_t * restrict src, const float val, const uint32_t num_elems) {
+    hvx_mul_scalar_f32(dst, src, 1.0f / val, num_elems);
+}
+
+static inline void hvx_div_scalar_f16(uint8_t * restrict dst, const uint8_t * restrict src, const _Float16 val, const uint32_t num_elems) {
+    hvx_div_scalar_f16_aa(dst, src, val, num_elems);
+}
+
+typedef void (*compute_binary_chunked_t)(
+    uint8_t * restrict dst,
+    const uint8_t * restrict src0,
+    const uint8_t * restrict src1,
+    const uint32_t num_elems
+);
+
+typedef void (*compute_binary_scalar_chunked_f32_t)(
+    uint8_t * restrict dst,
+    const uint8_t * restrict src0,
+    const float val,
+    const uint32_t num_elems
+);
+
+typedef void (*compute_binary_scalar_chunked_f16_t)(
+    uint8_t * restrict dst,
+    const uint8_t * restrict src0,
+    const _Float16 val,
+    const uint32_t num_elems
+);
+
+struct binary_chunked_context {
+    struct htp_ops_context * octx;
+    struct htp_binary_vtcm_layout vtcm_layout;
+    uint8_t *                vtcm_base;
+    uint32_t                 elem_start;
+    uint32_t                 nelem;
+    uint32_t                 chunk_size;
+    uint32_t                 chunks_per_thread;
+    uint32_t                 total_chunks;
+    bool                     is_scalar;
+    float                    scalar_f32;
+    _Float16                 scalar_f16;
+    compute_binary_chunked_t compute;
+    compute_binary_scalar_chunked_f32_t compute_scalar_f32;
+    compute_binary_scalar_chunked_f16_t compute_scalar_f16;
+};
+
+static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * data) {
+    (void) nth;
+    struct binary_chunked_context * ctx = (struct binary_chunked_context *) data;
+    struct htp_ops_context * octx = ctx->octx;
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * src1 = octx->src[1];
+    const struct htp_tensor * dst  = octx->dst;
+
+    const uint32_t start_chunk = ctx->chunks_per_thread * ith;
+    const uint32_t end_chunk   = MIN(start_chunk + ctx->chunks_per_thread, ctx->total_chunks);
+    if (start_chunk >= end_chunk) {
+        return;
+    }
+
+    const uint32_t src0_type   = src0->type;
+    const size_t   elem_size   = (src0_type == HTP_TYPE_F32) ? sizeof(float) : sizeof(_Float16);
+    const uint32_t chunk_size  = ctx->chunk_size;
+    const size_t   chunk_bytes = ctx->vtcm_layout.src0_spad_half_size;
+
+    FARF(HIGH, "binary-chunked: %d/%d (%u:%u) chunks %u elems %u",
+         ith, nth, start_chunk, end_chunk, ctx->total_chunks, ctx->nelem);
+
+    const struct htp_binary_vtcm_layout * layout = &ctx->vtcm_layout;
+    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread);
+    uint8_t * src1_spad_base = ctx->is_scalar ? NULL : (VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_src1) + (ith * layout->src1_bytes_per_thread));
+    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);
+
+    dma_queue * dma_q = octx->ctx->dma[ith];
+    uint32_t prefetch_chunk = start_chunk;
+    int spad_idx = 0;
+
+    for (int k = 0; k < 2 && prefetch_chunk < end_chunk; k++) {
+        const uint32_t c_start   = ctx->elem_start + prefetch_chunk * chunk_size;
+        const uint32_t c_end     = MIN(c_start + chunk_size, ctx->elem_start + ctx->nelem);
+        const uint32_t cur_elems = c_end - c_start;
+        const uint32_t cur_bytes = cur_elems * elem_size;
+
+        dma_addr_t s0_curr = src0->data + (size_t) c_start * elem_size;
+        dma_addr_t d_curr  = dst->data  + (size_t) c_start * elem_size;
+
+        uint8_t * s0_spad = src0_spad_base + spad_idx * chunk_bytes;
+        uint8_t * d_spad  = dst_spad_base  + spad_idx * chunk_bytes;
+
+        dma_queue_push(dma_q, dma_make_data(d_curr, d_spad), chunk_bytes, chunk_bytes, cur_bytes, 0);
+        dma_queue_push(dma_q, dma_make_data(s0_spad, s0_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
+        if (!ctx->is_scalar) {
+            dma_addr_t s1_curr = src1->data + (size_t) c_start * elem_size;
+            uint8_t * s1_spad = src1_spad_base + spad_idx * chunk_bytes;
+            dma_queue_push(dma_q, dma_make_data(s1_spad, s1_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
+        }
+
+        prefetch_chunk++;
+        spad_idx ^= 1;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
+    for (uint32_t c = start_chunk; c < end_chunk; c++) {
+        const uint32_t c_start   = ctx->elem_start + c * chunk_size;
+        const uint32_t c_end     = MIN(c_start + chunk_size, ctx->elem_start + ctx->nelem);
+        const uint32_t cur_elems = c_end - c_start;
+        const uint32_t cur_bytes = cur_elems * elem_size;
+
+        uint8_t * d_spad  = (uint8_t *) dma_queue_pop(dma_q).src;
+        uint8_t * s0_spad = (uint8_t *) dma_queue_pop(dma_q).dst;
+        uint8_t * s1_spad = ctx->is_scalar ? NULL : (uint8_t *) dma_queue_pop(dma_q).dst;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) c);
+        if (ctx->is_scalar) {
+            if (src0_type == HTP_TYPE_F32) {
+                ctx->compute_scalar_f32(d_spad, s0_spad, ctx->scalar_f32, cur_elems);
+            } else {
+                ctx->compute_scalar_f16(d_spad, s0_spad, ctx->scalar_f16, cur_elems);
+            }
+        } else {
+            ctx->compute(d_spad, s0_spad, s1_spad, cur_elems);
+        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) c);
+
+        dma_addr_t dst_curr = dst->data + (size_t) c_start * elem_size;
+        dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), chunk_bytes, chunk_bytes, cur_bytes, 1);
+
+        if (prefetch_chunk < end_chunk) {
+            const uint32_t pc_start  = ctx->elem_start + prefetch_chunk * chunk_size;
+            const uint32_t pc_end    = MIN(pc_start + chunk_size, ctx->elem_start + ctx->nelem);
+            const uint32_t p_elems   = pc_end - pc_start;
+            const uint32_t p_bytes   = p_elems * elem_size;
+
+            dma_addr_t s0_next = src0->data + (size_t) pc_start * elem_size;
+            dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), chunk_bytes, chunk_bytes, p_bytes, 1);
+            if (!ctx->is_scalar) {
+                dma_addr_t s1_next = src1->data + (size_t) pc_start * elem_size;
+                dma_queue_push(dma_q, dma_make_data(s1_spad, s1_next), chunk_bytes, chunk_bytes, p_bytes, 1);
+            }
+
+            prefetch_chunk++;
+        }
+    }
+
+    dma_queue_flush(dma_q);
+}
+
 static int execute_op_binary(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * src1 = octx->src[1];
@@ -932,6 +1080,135 @@ static int execute_op_binary(struct htp_ops_context * octx) {
     const size_t src1_row_size = src1->ne[0] * elem_size;
     const size_t dst_row_size  = dst->ne[0]  * elem_size;
 
+    if (kparams->kernel_type == HTP_BINARY_KERNEL_CHUNKED) {
+        const bool is_scalar = kparams->is_scalar || (src1->ne[0] == 1 && src1->ne[1] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1);
+
+        compute_binary_chunked_t compute = NULL;
+        compute_binary_scalar_chunked_f32_t compute_scalar_f32 = NULL;
+        compute_binary_scalar_chunked_f16_t compute_scalar_f16 = NULL;
+
+        if (is_scalar) {
+            if (src0_type == HTP_TYPE_F32) {
+                switch (octx->op) {
+                    case HTP_OP_ADD: compute_scalar_f32 = hvx_add_scalar_f32; break;
+                    case HTP_OP_SUB: compute_scalar_f32 = hvx_sub_scalar_f32; break;
+                    case HTP_OP_MUL: compute_scalar_f32 = hvx_mul_scalar_f32; break;
+                    case HTP_OP_DIV: compute_scalar_f32 = hvx_div_scalar_f32; break;
+                    default: break;
+                }
+            } else if (src0_type == HTP_TYPE_F16) {
+                switch (octx->op) {
+                    case HTP_OP_ADD: compute_scalar_f16 = hvx_add_scalar_f16; break;
+                    case HTP_OP_SUB: compute_scalar_f16 = hvx_sub_scalar_f16; break;
+                    case HTP_OP_MUL: compute_scalar_f16 = hvx_mul_scalar_f16; break;
+                    case HTP_OP_DIV: compute_scalar_f16 = hvx_div_scalar_f16; break;
+                    default: break;
+                }
+            }
+            if (!compute_scalar_f32 && !compute_scalar_f16) {
+                return HTP_STATUS_NO_SUPPORT;
+            }
+        } else {
+            if (src0_type == HTP_TYPE_F32) {
+                switch (octx->op) {
+                    case HTP_OP_ADD: compute = hvx_add_f32; break;
+                    case HTP_OP_SUB: compute = hvx_sub_f32; break;
+                    case HTP_OP_MUL: compute = hvx_mul_f32; break;
+                    case HTP_OP_DIV: compute = hvx_div_f32; break;
+                    default: break;
+                }
+            } else if (src0_type == HTP_TYPE_F16) {
+                switch (octx->op) {
+                    case HTP_OP_ADD: compute = hvx_add_f16; break;
+                    case HTP_OP_SUB: compute = hvx_sub_f16; break;
+                    case HTP_OP_MUL: compute = hvx_mul_f16; break;
+                    case HTP_OP_DIV: compute = hvx_div_f16; break;
+                    default: break;
+                }
+            }
+            if (!compute) {
+                return HTP_STATUS_NO_SUPPORT;
+            }
+        }
+
+        const uint32_t total_elems = (uint32_t) (src0->ne[0] * src0->ne[1] * src0->ne[2] * src0->ne[3]);
+        if (total_elems == 0) {
+            return HTP_STATUS_OK;
+        }
+
+        uint32_t elem_start = 0;
+        uint32_t nelem      = total_elems;
+        const uint32_t elems_per_line = (src0_type == HTP_TYPE_F32) ? 32 : 64;
+
+        if (octx->ctx->mdev.count > 1) {
+            const bool can_split = htp_tensor_mdev_data_aligned(dst) &&
+                                   htp_tensor_mdev_data_aligned(src0) &&
+                                   (is_scalar || htp_tensor_mdev_data_aligned(src1)) &&
+                                   htp_tensor_is_contiguous(dst, elem_size) &&
+                                   htp_tensor_is_contiguous(src0, elem_size) &&
+                                   (is_scalar || htp_tensor_is_contiguous(src1, elem_size));
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
+                total_elems, can_split ? elems_per_line : 0,
+                octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            elem_start = range.start;
+            nelem      = range.count;
+        }
+
+        if (nelem == 0) {
+            return HTP_STATUS_OK;
+        }
+
+        if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+            return HTP_STATUS_INVAL_PARAMS;
+        }
+
+        struct htp_binary_vtcm_layout vtcm_layout;
+        htp_binary_vtcm_layout_build(&vtcm_layout, kparams, octx->ctx->vtcm_size);
+        if (vtcm_layout.total_bytes == 0 || vtcm_layout.total_bytes > octx->ctx->vtcm_size) {
+            return HTP_STATUS_VTCM_TOO_SMALL;
+        }
+
+        const uint32_t chunk_size = kparams->chunk_size > 0 ? kparams->chunk_size : (32768 / elem_size);
+        const uint32_t total_chunks = (nelem + chunk_size - 1) / chunk_size;
+        const uint32_t n_threads = (total_chunks >= 2) ? MIN(octx->n_threads, total_chunks) : 1;
+        const uint32_t chunks_per_thread = (total_chunks + n_threads - 1) / n_threads;
+
+        float scalar_f32 = 0.0f;
+        _Float16 scalar_f16 = 0;
+        if (is_scalar) {
+            uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, vtcm_layout.off_src1);
+            dma_queue * dma_q = octx->ctx->dma[0];
+            dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), 128, 0, elem_size, 1);
+            dma_queue_pop(dma_q);
+
+            if (src0_type == HTP_TYPE_F32) {
+                scalar_f32 = ((const float *) vtcm_src1)[0];
+            } else {
+                scalar_f16 = ((const _Float16 *) vtcm_src1)[0];
+            }
+        }
+
+        struct binary_chunked_context cctx = {
+            .octx                = octx,
+            .vtcm_layout         = vtcm_layout,
+            .vtcm_base           = (uint8_t *) octx->ctx->vtcm_base,
+            .elem_start          = elem_start,
+            .nelem               = nelem,
+            .chunk_size          = chunk_size,
+            .chunks_per_thread   = chunks_per_thread,
+            .total_chunks        = total_chunks,
+            .is_scalar           = is_scalar,
+            .scalar_f32          = scalar_f32,
+            .scalar_f16          = scalar_f16,
+            .compute             = compute,
+            .compute_scalar_f32  = compute_scalar_f32,
+            .compute_scalar_f16  = compute_scalar_f16,
+        };
+
+        work_queue_run(octx->ctx->work_queue, binary_thread_chunked, &cctx, n_threads);
+        return HTP_STATUS_OK;
+    }
+
     uint32_t row_start = 0;
     uint32_t nrows     = src0_nrows;
 

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

@@ -17,6 +17,7 @@ enum htp_binary_kernel_type {
     HTP_BINARY_KERNEL_ADD_ID,
     HTP_BINARY_KERNEL_COMPLEX,
     HTP_BINARY_KERNEL_REPEAT,
+    HTP_BINARY_KERNEL_CHUNKED,
 };
 
 struct htp_binary_kernel_params {
@@ -30,6 +31,10 @@ struct htp_binary_kernel_params {
 
     uint32_t src1_size;
     uint32_t vtcm_size;
+
+    uint32_t chunk_size;
+    uint32_t chunk_bytes;
+    uint32_t is_scalar;
 };
 
 #if defined(__cplusplus)
@@ -68,6 +73,40 @@ static inline void htp_binary_vtcm_layout_build(
         return;
     }
 
+    if (kparams->kernel_type == HTP_BINARY_KERNEL_CHUNKED) {
+        const size_t chunk_bytes = kparams->chunk_bytes;
+        if (chunk_bytes == 0) {
+            return;
+        }
+
+        L->src0_bytes_per_thread = 2 * chunk_bytes;
+        L->src1_bytes_per_thread = kparams->is_scalar ? 0 : (2 * chunk_bytes);
+        L->dst_bytes_per_thread  = 2 * chunk_bytes;
+
+        L->src0_spad_half_size = chunk_bytes;
+        L->src1_spad_half_size = kparams->is_scalar ? 0 : chunk_bytes;
+        L->dst_spad_half_size  = chunk_bytes;
+
+        L->rows_per_buffer = 1;
+        L->src1_size = 0;
+
+        const size_t src0_total = n_threads * L->src0_bytes_per_thread;
+        const size_t src1_total = kparams->is_scalar ? 128 : (n_threads * L->src1_bytes_per_thread);
+        const size_t dst_total  = n_threads * L->dst_bytes_per_thread;
+
+        size_t off = 0;
+        VTCM_LAYOUT_ALLOC(off, off_src0, src0_total);
+        VTCM_LAYOUT_ALLOC(off, off_src1, src1_total);
+        VTCM_LAYOUT_ALLOC(off, off_dst,  dst_total);
+
+        if (off > vtcm_size) {
+            return;
+        }
+
+        L->total_bytes = off;
+        return;
+    }
+
     const size_t spad_row_total = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE)
         ? 2 * (kparams->src0_row_size_aligned + kparams->src1_row_size_aligned + kparams->dst_row_size_aligned)
         : 2 * (kparams->src0_row_size_aligned + kparams->dst_row_size_aligned);

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

@@ -301,6 +301,86 @@ static void cpy_thread_f32_f16_sameshape(unsigned int nth, unsigned int ith, voi
     }
 }
 
+static void cpy_thread_i32_f32_sameshape(unsigned int nth, unsigned int ith, void * data) {
+    struct htp_copy_context * ct = (struct htp_copy_context *) data;
+    struct htp_ops_context * octx = ct->octx;
+    cpy_preamble;
+
+    const uint32_t dr  = ct->src0_nrows_per_thread;
+    const uint32_t ir0 = ct->row_start + dr * ith;
+    const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows);
+    if (ir0 >= ir1) return;
+
+    const uint32_t ne02_ne01 = ne02 * ne01;
+    uint32_t i03 = fastdiv(ir0, &ct->div_ne02_ne01);
+    uint32_t rem = ir0 - i03 * ne02_ne01;
+    uint32_t i02 = fastdiv(rem, &ct->div_ne01);
+    uint32_t i01 = rem - i02 * ne01;
+
+    uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;
+    uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
+
+    for (uint32_t r = ir0; r < ir1; r++) {
+        hex_l2fetch(src0_ptr, ne00 * sizeof(float), nb01, 2);
+        const float * restrict src_row = (const float *) src0_ptr;
+        int32_t * restrict dst_row = (int32_t *) dst_ptr;
+        for (uint32_t i = 0; i < ne00; i++) {
+            dst_row[i] = (int32_t) src_row[i];
+        }
+        dst_ptr  += nb1;
+        src0_ptr += nb01;
+        if (++i01 == ne01) {
+            i01 = 0;
+            if (++i02 == ne02) {
+                i02 = 0;
+                i03++;
+            }
+            dst_ptr  = (uint8_t*) dst->data  + i02*nb2  + i03*nb3;
+            src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03;
+        }
+    }
+}
+
+static void cpy_thread_f32_i32_sameshape(unsigned int nth, unsigned int ith, void * data) {
+    struct htp_copy_context * ct = (struct htp_copy_context *) data;
+    struct htp_ops_context * octx = ct->octx;
+    cpy_preamble;
+
+    const uint32_t dr  = ct->src0_nrows_per_thread;
+    const uint32_t ir0 = ct->row_start + dr * ith;
+    const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows);
+    if (ir0 >= ir1) return;
+
+    const uint32_t ne02_ne01 = ne02 * ne01;
+    uint32_t i03 = fastdiv(ir0, &ct->div_ne02_ne01);
+    uint32_t rem = ir0 - i03 * ne02_ne01;
+    uint32_t i02 = fastdiv(rem, &ct->div_ne01);
+    uint32_t i01 = rem - i02 * ne01;
+
+    uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;
+    uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
+
+    for (uint32_t r = ir0; r < ir1; r++) {
+        hex_l2fetch(src0_ptr, ne00 * sizeof(int32_t), nb01, 2);
+        const int32_t * restrict src_row = (const int32_t *) src0_ptr;
+        float * restrict dst_row = (float *) dst_ptr;
+        for (uint32_t i = 0; i < ne00; i++) {
+            dst_row[i] = (float) src_row[i];
+        }
+        dst_ptr  += nb1;
+        src0_ptr += nb01;
+        if (++i01 == ne01) {
+            i01 = 0;
+            if (++i02 == ne02) {
+                i02 = 0;
+                i03++;
+            }
+            dst_ptr  = (uint8_t*) dst->data  + i02*nb2  + i03*nb3;
+            src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03;
+        }
+    }
+}
+
 static inline void cpy_dma_push_2d_chunked(
     dma_queue * dma_q,
     dma_addr_t  dst,
@@ -366,6 +446,42 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) {
     cpy_preamble;
     *use_dma = false;
 
+    const uint32_t total_elems_src = ne00 * ne01 * ne02 * ne03;
+    const uint32_t total_elems_dst = ne0 * ne1 * ne2 * ne3;
+    if (total_elems_src == 1 && total_elems_dst == 1) {
+        if (octx->ctx->mdev.count > 1 && octx->ctx->mdev.idx > 0) {
+            return HTP_STATUS_OK;
+        }
+        if (src0->type == HTP_TYPE_F32 && dst->type == HTP_TYPE_I32) {
+            ((int32_t *) dst->data)[0] = (int32_t) (((const float *) src0->data)[0]);
+            return HTP_STATUS_OK;
+        }
+        if (src0->type == HTP_TYPE_I32 && dst->type == HTP_TYPE_F32) {
+            ((float *) dst->data)[0] = (float) (((const int32_t *) src0->data)[0]);
+            return HTP_STATUS_OK;
+        }
+        if (src0->type == HTP_TYPE_I32 && dst->type == HTP_TYPE_I32) {
+            ((int32_t *) dst->data)[0] = ((const int32_t *) src0->data)[0];
+            return HTP_STATUS_OK;
+        }
+        if (src0->type == HTP_TYPE_F32 && dst->type == HTP_TYPE_F32) {
+            ((float *) dst->data)[0] = ((const float *) src0->data)[0];
+            return HTP_STATUS_OK;
+        }
+        if (src0->type == HTP_TYPE_F16 && dst->type == HTP_TYPE_F16) {
+            ((__fp16 *) dst->data)[0] = ((const __fp16 *) src0->data)[0];
+            return HTP_STATUS_OK;
+        }
+        if (src0->type == HTP_TYPE_F32 && dst->type == HTP_TYPE_F16) {
+            ((__fp16 *) dst->data)[0] = (__fp16) (((const float *) src0->data)[0]);
+            return HTP_STATUS_OK;
+        }
+        if (src0->type == HTP_TYPE_F16 && dst->type == HTP_TYPE_F32) {
+            ((float *) dst->data)[0] = (float) (((const __fp16 *) src0->data)[0]);
+            return HTP_STATUS_OK;
+        }
+    }
+
     struct htp_copy_context ct;
     ct.octx = octx;
 
@@ -450,6 +566,10 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) {
                 copy_fun = cpy_thread_f16_f32_sameshape;
             } else if (dst->type == HTP_TYPE_F32 && src0->type == HTP_TYPE_F16) {
                 copy_fun = cpy_thread_f32_f16_sameshape;
+            } else if (dst->type == HTP_TYPE_I32 && src0->type == HTP_TYPE_F32) {
+                copy_fun = cpy_thread_i32_f32_sameshape;
+            } else if (dst->type == HTP_TYPE_F32 && src0->type == HTP_TYPE_I32) {
+                copy_fun = cpy_thread_f32_i32_sameshape;
             } else {
                 return HTP_STATUS_NO_SUPPORT;
             }

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

@@ -426,6 +426,22 @@ static inline bool dma_queue_push(dma_queue *q, dma_data ddata, size_t dst_strid
 
 #endif
 
+static inline void dma_sync_read(dma_queue * dma_q, void * dst, dma_addr_t src, size_t bytes) {
+    const uint32_t b = (uint32_t) bytes;
+    if (b > 0) {
+        dma_queue_push(dma_q, dma_make_data(dst, src), b, b, b, 1);
+        dma_queue_pop(dma_q);
+    }
+}
+
+static inline void dma_sync_write(dma_queue * dma_q, dma_addr_t dst, const void * src, size_t bytes) {
+    const uint32_t b = (uint32_t) bytes;
+    if (b > 0) {
+        dma_queue_push(dma_q, dma_make_data(dst, src), b, b, b, 1);
+        dma_queue_pop(dma_q);
+    }
+}
+
 #define DMA_CACHE_MAX_SIZE 256U
 
 // Fully assoc LRU cache

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

@@ -17,6 +17,7 @@
 #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"
 
@@ -28,6 +29,9 @@ struct get_rows_context {
     uint32_t task_start;
     uint32_t tasks;
     uint32_t tasks_per_thread;
+    uint32_t tile_size;
+    uint32_t tile_stride;
+    bool index_i32;
 };
 
 #define get_rows_preamble                      \
@@ -195,32 +199,191 @@ static void get_rows_thread_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned
     dma_queue_flush(dma_q);                                                                                                     \
 }
 
-#define F32_BYTES(n)  ((n) * sizeof(float))
 #define F16_BYTES(n)  ((n) * sizeof(__fp16))
 #define Q8_0_BYTES(n) (((n) / 32) * sizeof(block_q8_0))
 
-GET_ROWS_THREAD_DT_FN(f32,  F32_BYTES,  int32_t, { if (cur_elems > 0) hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, cur_elems); })
-GET_ROWS_THREAD_DT_FN(f32,  F32_BYTES,  int64_t, { if (cur_elems > 0) hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, cur_elems); })
+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);
+}
 
-GET_ROWS_THREAD_DT_FN(f16,  F16_BYTES,  int32_t, { hvx_dequantize_row_f16_f32((float *)dst_spad, src_spad, ne00); })
-GET_ROWS_THREAD_DT_FN(f16,  F16_BYTES,  int64_t, { hvx_dequantize_row_f16_f32((float *)dst_spad, src_spad, ne00); })
+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);
+}
 
-GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int32_t, { hvx_dequantize_row_q8_0_f32((float *)dst_spad, src_spad, ne00); })
-GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int64_t, { hvx_dequantize_row_q8_0_f32((float *)dst_spad, src_spad, ne00); })
+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); })
+
+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); })
+
+
+static __attribute__((noinline)) void compute_get_rows_tiled(float * dst, const uint8_t * tile, uint32_t row, bool q4) {
+    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));
+        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));
+        }
+    }
+    const HVX_Vector scales = hvx_vmem(tile + (q4 ? 512 : 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));
+    const HVX_Vector values = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale);
+    *(HVX_Vector *) dst = values;
+}
+
+struct get_rows_tiled_task {
+    dma_addr_t tile_src_base;
+    dma_addr_t dst_data;
+    uint32_t   row;
+};
+
+static inline struct get_rows_tiled_task get_rows_tiled_calc_task(
+    const struct htp_ops_context * octx,
+    const struct get_rows_context * grctx,
+    uint32_t i,
+    uint32_t n_k_tiles,
+    uint32_t tile_size
+) {
+    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;
+    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;
+    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 dma_addr_t matrix = octx->src[0]->data + i02*nb02 + i03*nb03;
+
+    struct get_rows_tiled_task task;
+    task.tile_src_base = matrix + (column_tile * n_k_tiles) * tile_size;
+    task.dst_data      = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3;
+    task.row           = row;
+    return 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;
+    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 n_k_tiles = ne00 / HTP_MM_HMX_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;
+    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);
+
+    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_tiled_task task = get_rows_tiled_calc_task(octx, grctx, i, n_k_tiles, tile_size);
+
+        // Dummy writeback to prime the queue with dst descriptor
+        dma_queue_push(dma_q,
+                       dma_make_data(task.dst_data, dst_spad_base + spad_idx * vtcm_layout->dst_spad_half_size),
+                       dst_bytes, vtcm_layout->dst_spad_half_size, dst_bytes, 0);
+
+        // Prefetch row tiles
+        dma_queue_push(dma_q,
+                       dma_make_data(src_spad_base + spad_idx * vtcm_layout->src0_spad_half_size, task.tile_src_base),
+                       tile_stride, tile_size, tile_size, n_k_tiles);
+    }
+
+    for (uint32_t step = 0; step < ir1 - ir0; ++step) {
+        const uint32_t i = ir0 + step;
+        float * dst_spad   = (float *) dma_queue_pop(dma_q).src;
+        uint8_t * src_spad = (uint8_t *) dma_queue_pop(dma_q).dst;
+
+        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;
+            compute_get_rows_tiled(dst_block, tile, task.row, is_q4);
+        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
+
+        // Real writeback of dst_spad
+        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 ni = ir0 + next_step;
+            struct get_rows_tiled_task next_task = get_rows_tiled_calc_task(octx, grctx, ni, n_k_tiles, tile_size);
+            dma_queue_push(dma_q,
+                           dma_make_data(src_spad, next_task.tile_src_base),
+                           tile_stride, tile_size, tile_size, n_k_tiles);
+        }
+    }
+
+    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_Q8_0 &&
-        octx->src[0]->type != HTP_TYPE_I32) {
+         octx->src[0]->type != HTP_TYPE_F16 &&
+         octx->src[0]->type != HTP_TYPE_Q4_0 &&
+         octx->src[0]->type != HTP_TYPE_Q8_0 &&
+         octx->src[0]->type != HTP_TYPE_I32) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    if ((octx->src[0]->type == HTP_TYPE_I32 && octx->dst->type != HTP_TYPE_I32) ||
-        (octx->src[0]->type != HTP_TYPE_I32 && octx->dst->type != HTP_TYPE_F32)) {
-        return HTP_STATUS_NO_SUPPORT;
+    if (kparams->kernel_type == HTP_GET_ROWS_KERNEL_SAMETYPE) {
+        if (octx->src[0]->type != octx->dst->type) {
+            return HTP_STATUS_NO_SUPPORT;
+        }
+    } else {
+        if (octx->dst->type != HTP_TYPE_F32) {
+            return HTP_STATUS_NO_SUPPORT;
+        }
     }
 
     if (octx->src[1]->type != HTP_TYPE_I32 && octx->src[1]->type != HTP_TYPE_I64) {
@@ -262,33 +425,48 @@ int op_get_rows(struct htp_ops_context * octx) {
     grctx.vtcm_base = (uint8_t *)octx->ctx->vtcm_base;
     grctx.task_start = task_start;
     grctx.tasks = tasks;
-    grctx.tasks_per_thread = fastdiv(tasks + n_threads - 1, &octx->n_threads_div);
+    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 = octx->src[0]->type == HTP_TYPE_Q4_0 ? HTP_MM_WEIGHT_TILE_SIZE_Q4_0 : HTP_MM_WEIGHT_TILE_SIZE_Q8_0;
+    grctx.tile_stride = (grctx.tile_size + 127) & ~127;
+    grctx.index_i32 = octx->src[1]->type == HTP_TYPE_I32;
 
     const uint32_t ne00 = octx->src[0]->ne[0];
-    htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, octx->src[0]->type, ne00, n_threads);
+    htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, kparams->kernel_type, octx->src[0]->type, ne00, 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",
+             octx->ctx->vtcm_size, grctx.vtcm_layout.total_bytes);
+        return HTP_STATUS_INVAL_PARAMS;
+    }
 
     const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32);
 
     work_queue_func_t q_func = NULL;
-    if (kparams->use_dma) {
-        q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_st_int32_t : get_rows_thread_st_int64_t);
-    } else {
-        switch (octx->src[0]->type) {
-            case HTP_TYPE_F32:  q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_f32_int32_t  : get_rows_thread_f32_int64_t);  break;
-            case HTP_TYPE_F16:  q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_f16_int32_t  : get_rows_thread_f16_int64_t);  break;
-            case HTP_TYPE_Q8_0: q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_q8_0_int32_t : get_rows_thread_q8_0_int64_t); break;
-            case HTP_TYPE_I32:  q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_st_int32_t   : get_rows_thread_st_int64_t);   break;
-            default:            return HTP_STATUS_NO_SUPPORT;
-        }
+    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);
+            break;
+        case HTP_GET_ROWS_KERNEL_TILED:
+            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;
+                default:            return HTP_STATUS_NO_SUPPORT;
+            }
+            break;
+        default:
+            return HTP_STATUS_NO_SUPPORT;
     }
 
-    FARF(HIGH, "get-rows: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-vtcm-size %zu dst-vtcm-size %zu use-dma %d n-threads %d\n",
+    FARF(HIGH, "get-rows: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-vtcm-size %zu dst-vtcm-size %zu 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],
          octx->dst->ne[0], octx->dst->ne[1], octx->dst->ne[2], octx->dst->ne[3],
          grctx.vtcm_layout.src0_bytes_per_thread * n_threads,
          grctx.vtcm_layout.dst_bytes_per_thread  * n_threads,
-         kparams->use_dma, n_threads);
+         kparams->kernel_type, n_threads);
 
     work_queue_run(octx->ctx->work_queue, q_func, &grctx, n_threads);
     return HTP_STATUS_OK;

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

@@ -1,11 +1,21 @@
 #ifndef HTP_GET_ROWS_OPS_H
 #define HTP_GET_ROWS_OPS_H
 
+#include <stdbool.h>
+#include <string.h>
+
 #include "hex-fastdiv.h"
+#include "matmul-ops.h"
+
+enum htp_get_rows_kernel_type {
+    HTP_GET_ROWS_KERNEL_SAMETYPE = 0,
+    HTP_GET_ROWS_KERNEL_TILED,
+    HTP_GET_ROWS_KERNEL_FLAT,
+};
 
 struct htp_get_rows_kernel_params {
     int32_t  n_threads;
-    int32_t  use_dma;
+    int32_t  kernel_type;
     int32_t  chunks_per_row;
     int32_t  chunk_size;
     int32_t  total_tasks;
@@ -34,19 +44,37 @@ struct htp_get_rows_vtcm_layout {
 
 static inline void htp_get_rows_vtcm_layout_build(
     struct htp_get_rows_vtcm_layout * vtcm_layout,
+    int kernel_type,
     int type,
     uint32_t ne00,
     uint32_t n_threads) {
 
+    if (kernel_type == HTP_GET_ROWS_KERNEL_SAMETYPE) {
+        memset(vtcm_layout, 0, sizeof(*vtcm_layout));
+        return;
+    }
+
+    if (kernel_type == HTP_GET_ROWS_KERNEL_TILED) {
+        const size_t tile_size   = type == HTP_TYPE_Q4_0 ? HTP_MM_WEIGHT_TILE_SIZE_Q4_0 : HTP_MM_WEIGHT_TILE_SIZE_Q8_0;
+        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 0: // HTP_TYPE_F32
-            src0_row_size = ne00 * 4;
-            break;
-        case 1: // HTP_TYPE_F16
+        case HTP_TYPE_F16:
             src0_row_size = ne00 * 2;
             break;
-        case 8: // HTP_TYPE_Q8_0
+        case HTP_TYPE_Q8_0:
             src0_row_size = (ne00 / 32) * 34;
             break;
         default:

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

@@ -150,7 +150,9 @@ int op_matmul_nx(struct htp_ops_context * octx);
 int op_matmul_id_nx(struct htp_ops_context * octx);
 int op_binary(struct htp_ops_context * octx);
 int op_unary(struct htp_ops_context * octx);
+int op_sum(struct htp_ops_context * octx);
 int op_sum_rows(struct htp_ops_context * octx);
+int op_argmax(struct htp_ops_context * octx);
 int op_activations(struct htp_ops_context * octx);
 int op_softmax(struct htp_ops_context * octx);
 int op_add_id(struct htp_ops_context * octx);

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

@@ -69,6 +69,7 @@ enum htp_op_code {
     HTP_OP_UNARY_ABS,
     HTP_OP_UNARY_LOG,
     HTP_OP_UNARY_RELU,
+    HTP_OP_UNARY_STEP,
     HTP_OP_GLU_SWIGLU,
     HTP_OP_GLU_SWIGLU_OAI,
     HTP_OP_GLU_GEGLU,
@@ -86,6 +87,7 @@ enum htp_op_code {
     HTP_OP_TOP_K,
     HTP_OP_SQR,
     HTP_OP_SQRT,
+    HTP_OP_SUM,
     HTP_OP_SUM_ROWS,
     HTP_OP_SSM_CONV,
     HTP_OP_REPEAT,
@@ -108,6 +110,7 @@ enum htp_op_code {
     HTP_OP_GLU_SWIGLU_CLAMP,
     HTP_OP_MDEV_GROUP,
     HTP_OP_ROLL,
+    HTP_OP_ARGMAX,
 
     HTP_OP_INVALID
 };

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

@@ -579,6 +579,58 @@ static inline void hvx_abs_f16(uint8_t * restrict dst, const uint8_t * restrict
     }
 }
 
+//
+// Step
+//
+
+static inline void hvx_step_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) dst % 128 == 0);
+    assert((unsigned long) src % 128 == 0);
+
+    HVX_Vector * restrict vdst = (HVX_Vector *) dst;
+    HVX_Vector * restrict vsrc = (HVX_Vector *) src;
+
+    const uint32_t elem_size = sizeof(float);
+    const uint32_t epv       = 128 / elem_size;
+    const uint32_t nvec      = n / epv;
+    const uint32_t nloe      = n % epv;
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        vdst[i] = hvx_vec_step_f32(vsrc[i]);
+    }
+    if (nloe) {
+        HVX_Vector v = hvx_vec_step_f32(vsrc[i]);
+        hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v);
+    }
+}
+
+static inline void hvx_step_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) dst % 128 == 0);
+    assert((unsigned long) src % 128 == 0);
+
+    HVX_Vector * restrict vdst = (HVX_Vector *) dst;
+    HVX_Vector * restrict vsrc = (HVX_Vector *) src;
+
+    const uint32_t elem_size = sizeof(_Float16);
+    const uint32_t epv       = 128 / elem_size;
+    const uint32_t nvec      = n / epv;
+    const uint32_t nloe      = n % epv;
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        vdst[i] = hvx_vec_step_f16(vsrc[i]);
+    }
+    if (nloe) {
+        HVX_Vector v = hvx_vec_step_f16(vsrc[i]);
+        hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v);
+    }
+}
+
 //
 // Square
 //

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

@@ -111,6 +111,20 @@ static inline HVX_Vector hvx_vec_neg_f32(HVX_Vector v) {
 #endif  // __HVX_ARCH__ > 75
 }
 
+static inline HVX_Vector hvx_vec_step_f32(HVX_Vector v) {
+    const HVX_Vector zero = Q6_V_vzero();
+    const HVX_Vector one  = hvx_vec_splat_f32(1.0f);
+    HVX_VectorPred q = Q6_Q_vcmp_gt_VsfVsf(v, zero);
+    return Q6_V_vmux_QVV(q, one, zero);
+}
+
+static inline HVX_Vector hvx_vec_step_f16(HVX_Vector v) {
+    const HVX_Vector zero = Q6_V_vzero();
+    const HVX_Vector one  = hvx_vec_splat_f16((_Float16) 1.0f);
+    HVX_VectorPred q = Q6_Q_vcmp_gt_VhfVhf(v, zero);
+    return Q6_V_vmux_QVV(q, one, zero);
+}
+
 static inline HVX_VectorPred hvx_vec_is_nan_f16(HVX_Vector v) {
     const HVX_Vector vnan_exp  = Q6_Vh_vsplat_R(0x7C00);
     const HVX_Vector vnan_frac = Q6_Vh_vsplat_R(0x7FFF);

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

@@ -326,6 +326,79 @@ static inline int32_t hvx_reduce_max_i32(const uint8_t * restrict src, const int
     }
 }
 
+static inline void hvx_argmax_f32(
+    const float * restrict src,
+    uint32_t n,
+    uint32_t offset,
+    float * out_val,
+    int32_t * out_idx
+) {
+    if (n == 0) {
+        *out_val = -INFINITY;
+        *out_idx = (int32_t) offset;
+        return;
+    }
+
+    if (n < 32 || !hex_is_aligned((void *) src, 128)) {
+        float best_val = src[0];
+        int32_t best_idx = (int32_t) offset;
+        for (uint32_t i = 1; i < n; i++) {
+            if (src[i] > best_val) {
+                best_val = src[i];
+                best_idx = (int32_t) (offset + i);
+            }
+        }
+        *out_val = best_val;
+        *out_idx = best_idx;
+        return;
+    }
+
+    static const int32_t c_lane_idx[32] __attribute__((aligned(128))) = {
+         0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15,
+        16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31
+    };
+
+    const HVX_Vector v_init_idx = *(const HVX_Vector *) c_lane_idx;
+    const HVX_Vector v_step     = Q6_V_vsplat_R(32);
+    HVX_Vector v_cur_idx        = Q6_Vw_vadd_VwVw(v_init_idx, Q6_V_vsplat_R((int32_t) offset));
+    HVX_Vector v_max_val        = hvx_vec_splat_f32(-INFINITY);
+    HVX_Vector v_max_idx        = v_cur_idx;
+
+    const HVX_Vector * vsrc = (const HVX_Vector *) src;
+    const uint32_t nvec = n / 32;
+
+    for (uint32_t vi = 0; vi < nvec; vi++) {
+        HVX_Vector v = vsrc[vi];
+        HVX_VectorPred pred = Q6_Q_vcmp_gt_VsfVsf(v, v_max_val);
+        v_max_val = Q6_V_vmux_QVV(pred, v, v_max_val);
+        v_max_idx = Q6_V_vmux_QVV(pred, v_cur_idx, v_max_idx);
+        v_cur_idx = Q6_Vw_vadd_VwVw(v_cur_idx, v_step);
+    }
+
+    HVX_VectorAlias u_val, u_idx;
+    u_val.v = v_max_val;
+    u_idx.v = v_max_idx;
+
+    float best_val = u_val.fp32[0];
+    int32_t best_idx = (int32_t) u_idx.w[0];
+    for (int i = 1; i < 32; i++) {
+        if (u_val.fp32[i] > best_val) {
+            best_val = u_val.fp32[i];
+            best_idx = (int32_t) u_idx.w[i];
+        }
+    }
+
+    for (uint32_t i = nvec * 32; i < n; i++) {
+        if (src[i] > best_val) {
+            best_val = src[i];
+            best_idx = (int32_t) (offset + i);
+        }
+    }
+
+    *out_val = best_val;
+    *out_idx = best_idx;
+}
+
 #undef hvx_reduce_loop_body
 #undef HVX_REDUCE_MAX_OP
 #undef HVX_REDUCE_SUM_OP

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

@@ -857,6 +857,7 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_UNARY_ABS:
         case HTP_OP_UNARY_LOG:
         case HTP_OP_UNARY_RELU:
+        case HTP_OP_UNARY_STEP:
         case HTP_OP_L2_NORM:
             return op_unary(octx);
 
@@ -882,6 +883,9 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_GET_ROWS:
             return op_get_rows(octx);
 
+        case HTP_OP_SUM:
+            return op_sum(octx);
+
         case HTP_OP_SUM_ROWS:
             return op_sum_rows(octx);
 
@@ -898,6 +902,9 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_TOP_K:
             return op_top_k(octx);
 
+        case HTP_OP_ARGMAX:
+            return op_argmax(octx);
+
         case HTP_OP_SSM_CONV:
             return op_ssm_conv(octx);
 

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

@@ -83,16 +83,17 @@ static void sum_rows_thread_f32(unsigned int nth, unsigned int ith, void *data)
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start_row);
 
     for (uint32_t ir = 0; ir < n_rows; ir++) {
-        const float * restrict src_local = src_th + (ir * (src_stride / sizeof(float)));
+        const float * restrict src_local = (const float *) ((const uint8_t *) src_th + ir * src_stride);
+        float       * restrict dst_local = (float *)       ((uint8_t *)       dst_th + ir * dst_stride);
 
         if (ir + 1 < n_rows) {
-            hex_l2fetch(src_local + (src_stride / sizeof(float)), src_stride, src_stride, 1);
+            hex_l2fetch((const uint8_t *) src_local + src_stride, src_stride, src_stride, 1);
         }
 
         if (opt_path) {
-            dst_th[ir] = hvx_reduce_sum_f32_a((const uint8_t *) src_local, ne00);
+            *dst_local = hvx_reduce_sum_f32_a((const uint8_t *) src_local, ne00);
         } else {
-            dst_th[ir] = hvx_reduce_sum_f32((const uint8_t *) src_local, ne00);
+            *dst_local = hvx_reduce_sum_f32((const uint8_t *) src_local, ne00);
         }
     }
 
@@ -152,3 +153,258 @@ int op_sum_rows(struct htp_ops_context * octx) {
 
     return HTP_STATUS_OK;
 }
+
+struct sum_context {
+    struct htp_ops_context * octx;
+    const float            * src_data;
+    float                    partial_sums[HTP_MAX_NTHREADS];
+    uint32_t                 total_elems;
+    uint32_t                 elems_per_thread;
+};
+
+static void sum_thread_f32(unsigned int nth, unsigned int ith, void * data) {
+    struct sum_context * sctx = (struct sum_context *) data;
+    const uint32_t start = sctx->elems_per_thread * ith;
+    const uint32_t end   = MIN(start + sctx->elems_per_thread, sctx->total_elems);
+
+    if (start >= end) {
+        sctx->partial_sums[ith] = 0.0f;
+        return;
+    }
+
+    const uint32_t n = end - start;
+    const float * src = sctx->src_data + start;
+
+    struct htp_thread_trace * tr = &sctx->octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start);
+
+    hex_l2fetch_block((const void *) src, n * sizeof(float));
+
+    sctx->partial_sums[ith] = hvx_reduce_sum_f32((const uint8_t *) src, n);
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start);
+}
+
+int op_sum(struct htp_ops_context * octx) {
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * dst  = octx->dst;
+
+    if (src0->type != HTP_TYPE_F32) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    if (octx->ctx->mdev.count > 1 && octx->ctx->mdev.idx > 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t total_elems = (uint32_t) (src0->ne[0] * src0->ne[1] * src0->ne[2] * src0->ne[3]);
+    if (total_elems == 0) {
+        ((float *) dst->data)[0] = 0.0f;
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = (total_elems >= 1024) ? MIN(octx->n_threads, HTP_MAX_NTHREADS) : 1;
+    const uint32_t raw_chunk = (total_elems + n_threads - 1) / n_threads;
+    const uint32_t elems_per_thread = hex_round_up(raw_chunk, 32);
+
+    struct sum_context sctx = {
+        .octx             = octx,
+        .src_data         = (const float *) src0->data,
+        .total_elems      = total_elems,
+        .elems_per_thread = elems_per_thread,
+    };
+
+    work_queue_run(octx->ctx->work_queue, sum_thread_f32, &sctx, n_threads);
+
+    float sum = 0.0f;
+    for (uint32_t i = 0; i < n_threads; i++) {
+        sum += sctx.partial_sums[i];
+    }
+    ((float *) dst->data)[0] = sum;
+
+    return HTP_STATUS_OK;
+}
+
+static inline void argmax_slice_f32(
+    const float * restrict src,
+    uint32_t n,
+    uint32_t offset,
+    float * out_val,
+    int32_t * out_idx
+) {
+    hvx_argmax_f32(src, n, offset, out_val, out_idx);
+}
+
+struct argmax_context {
+    struct htp_ops_context * octx;
+    const float            * src_data;
+    int32_t                * dst_data;
+    uint32_t                 ne00;
+    uint32_t                 src_stride;
+    uint32_t                 dst_stride;
+    uint32_t                 row_start;
+    uint32_t                 nrows;
+    uint32_t                 rows_per_thread;
+    uint32_t                 elems_per_thread;
+
+    float                    partial_max[HTP_MAX_NTHREADS];
+    int32_t                  partial_idx[HTP_MAX_NTHREADS];
+};
+
+static void argmax_thread_single_row(unsigned int nth, unsigned int ith, void * data) {
+    struct argmax_context * actx = (struct argmax_context *) data;
+    const uint32_t start = actx->elems_per_thread * ith;
+    const uint32_t end   = MIN(start + actx->elems_per_thread, actx->ne00);
+
+    if (start >= end) {
+        actx->partial_max[ith] = -INFINITY;
+        actx->partial_idx[ith] = 0;
+        return;
+    }
+
+    const uint32_t n = end - start;
+    const float * src = actx->src_data + start;
+
+    struct htp_thread_trace * tr = &actx->octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start);
+
+    hex_l2fetch_block((const void *) src, n * sizeof(float));
+
+    float max_val;
+    int32_t max_idx;
+    argmax_slice_f32(src, n, start, &max_val, &max_idx);
+
+    actx->partial_max[ith] = max_val;
+    actx->partial_idx[ith] = max_idx;
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start);
+}
+
+static void argmax_thread_multi_row(unsigned int nth, unsigned int ith, void * data) {
+    struct argmax_context * actx = (struct argmax_context *) data;
+    const uint32_t r0 = actx->row_start + actx->rows_per_thread * ith;
+    const uint32_t r1 = MIN(r0 + actx->rows_per_thread, actx->row_start + actx->nrows);
+
+    if (r0 >= r1) {
+        return;
+    }
+
+    struct htp_thread_trace * tr = &actx->octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r0);
+
+    for (uint32_t r = r0; r < r1; r++) {
+        const float * src_row = (const float *) ((const uint8_t *) actx->src_data + r * actx->src_stride);
+        int32_t * dst_val     = (int32_t *) ((uint8_t *) actx->dst_data + r * actx->dst_stride);
+
+        hex_l2fetch_block((const void *) src_row, actx->ne00 * sizeof(float));
+
+        float max_val;
+        int32_t max_idx;
+        argmax_slice_f32(src_row, actx->ne00, 0, &max_val, &max_idx);
+
+        *dst_val = max_idx;
+    }
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r0);
+}
+
+int op_argmax(struct htp_ops_context * octx) {
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * dst  = octx->dst;
+
+    if (src0->type != HTP_TYPE_F32 || dst->type != HTP_TYPE_I32) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    const uint32_t ne00 = src0->ne[0];
+    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
+
+    if (ne00 == 0 || src0_nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    if (src0_nrows == 1) {
+        if (octx->ctx->mdev.count > 1 && octx->ctx->mdev.idx > 0) {
+            return HTP_STATUS_OK;
+        }
+
+        if (ne00 == 1) {
+            ((int32_t *) dst->data)[0] = 0;
+            return HTP_STATUS_OK;
+        }
+
+        const uint32_t n_threads = (ne00 >= 1024) ? MIN(octx->n_threads, HTP_MAX_NTHREADS) : 1;
+        const uint32_t raw_chunk = (ne00 + n_threads - 1) / n_threads;
+        const uint32_t elems_per_thread = hex_round_up(raw_chunk, 32);
+
+        struct argmax_context actx = {
+            .octx             = octx,
+            .src_data         = (const float *) src0->data,
+            .dst_data         = (int32_t *) dst->data,
+            .ne00             = ne00,
+            .src_stride       = src0->nb[1] > 0 ? (uint32_t) src0->nb[1] : (uint32_t) (ne00 * sizeof(float)),
+            .dst_stride       = dst->nb[0] > 0 ? (uint32_t) dst->nb[0] : (uint32_t) sizeof(int32_t),
+            .row_start        = 0,
+            .nrows            = 1,
+            .rows_per_thread  = 1,
+            .elems_per_thread = elems_per_thread,
+        };
+
+        work_queue_run(octx->ctx->work_queue, argmax_thread_single_row, &actx, n_threads);
+
+        float best_val = actx.partial_max[0];
+        int32_t best_idx = actx.partial_idx[0];
+        for (uint32_t i = 1; i < n_threads; i++) {
+            if (actx.partial_max[i] > best_val) {
+                best_val = actx.partial_max[i];
+                best_idx = actx.partial_idx[i];
+            }
+        }
+        ((int32_t *) dst->data)[0] = best_idx;
+        return HTP_STATUS_OK;
+    }
+
+    uint32_t row_start = 0;
+    uint32_t nrows     = src0_nrows;
+
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_mdev_data_aligned(dst) &&
+                               htp_tensor_is_contiguous(dst, sizeof(int32_t));
+        const uint32_t elems_per_chunk = can_split ? 32 : 0;
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
+            src0_nrows, elems_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = MIN(octx->n_threads, nrows);
+    const uint32_t rows_per_thread = (nrows + n_threads - 1) / n_threads;
+
+    struct argmax_context actx = {
+        .octx             = octx,
+        .src_data         = (const float *) src0->data,
+        .dst_data         = (int32_t *) dst->data,
+        .ne00             = ne00,
+        .src_stride       = src0->nb[1] > 0 ? (uint32_t) src0->nb[1] : (uint32_t) (ne00 * sizeof(float)),
+        .dst_stride       = dst->nb[0] > 0 ? (uint32_t) dst->nb[0] : (uint32_t) sizeof(int32_t),
+        .row_start        = row_start,
+        .nrows            = nrows,
+        .rows_per_thread  = rows_per_thread,
+        .elems_per_thread = 0,
+    };
+
+    work_queue_run(octx->ctx->work_queue, argmax_thread_multi_row, &actx, n_threads);
+    return HTP_STATUS_OK;
+}

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

@@ -409,6 +409,20 @@ static void log_f16(const void * restrict src,
     }
 }
 
+static void step_f16(const void * restrict src,
+                     void * restrict dst,
+                     const uint32_t num_rows,
+                     const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_step_f16_aa((uint8_t *) dst_local, (const uint8_t *) src_local, ne0);
+    }
+}
+
 static void l2_norm_f16(const void * restrict src,
                         void * restrict dst,
                         const uint32_t num_rows,
@@ -662,6 +676,20 @@ static void relu_f32(const void * restrict src,
     }
 }
 
+static void step_f32(const void * restrict src,
+                     void * restrict dst,
+                     const uint32_t num_rows,
+                     const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_step_f32_aa(dst_local, src_local, ne0);
+    }
+}
+
 static void log_f32(const void * restrict src,
                     void * restrict dst,
                     const uint32_t num_rows,
@@ -767,6 +795,11 @@ static void tile_relu_f32(void * restrict dst, const void * restrict src, uint32
     hvx_max_scalar_f32((uint8_t *) dst, (const uint8_t *) src, 0.0f, tw);
 }
 
+static void tile_step_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) {
+    (void) uctx;
+    hvx_step_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw);
+}
+
 static void tri_apply_tile_f32(const void * restrict src, void * restrict dst,
                                uint32_t tile_elems, uint32_t col_start, uint32_t i01,
                                uint32_t ne0, int32_t ttype) {
@@ -1486,6 +1519,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
         case HTP_OP_UNARY_ABS:       op_type = is_f16 ? "abs-f16"      : "abs-f32";          break;
         case HTP_OP_UNARY_LOG:       op_type = is_f16 ? "log-f16"      : "log-f32";          break;
         case HTP_OP_UNARY_RELU:      op_type = "relu-f32";                                   break;
+        case HTP_OP_UNARY_STEP:      op_type = is_f16 ? "step-f16"     : "step-f32";         break;
         case HTP_OP_L2_NORM:         op_type = is_f16 ? "l2norm-f16"   : "l2norm-f32";       break;
         case HTP_OP_TRI:             op_type = "tri-f32";                                    break;
         default:
@@ -1506,6 +1540,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
             case HTP_OP_L2_NORM:
             case HTP_OP_UNARY_ABS:
             case HTP_OP_UNARY_LOG:
+            case HTP_OP_UNARY_STEP:
                 break;
             default:
                 FARF(ERROR, "unary-%s: not supported for F16\n", op_type);
@@ -1634,6 +1669,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
             case HTP_OP_UNARY_ABS:       compute_func = (void *) tile_abs_f32;            break;
             case HTP_OP_UNARY_LOG:       compute_func = (void *) tile_log_f32;            break;
             case HTP_OP_UNARY_RELU:      compute_func = (void *) tile_relu_f32;           break;
+            case HTP_OP_UNARY_STEP:      compute_func = (void *) tile_step_f32;           break;
             case HTP_OP_TRI:
                 task_func    = unary_thread_tiled_tri_f32;
                 compute_func = (void *) tri_apply_tile_f32;
@@ -1652,6 +1688,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
             case HTP_OP_L2_NORM:         compute_func = (void *) l2_norm_f16;             break;
             case HTP_OP_UNARY_ABS:       compute_func = (void *) abs_f16;                 break;
             case HTP_OP_UNARY_LOG:       compute_func = (void *) log_f16;                 break;
+            case HTP_OP_UNARY_STEP:      compute_func = (void *) step_f16;                break;
             default:                     break;
         }
     } else {
@@ -1678,6 +1715,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
             case HTP_OP_UNARY_ABS:       compute_func = (void *) abs_f32;                 break;
             case HTP_OP_UNARY_LOG:       compute_func = (void *) log_f32;                 break;
             case HTP_OP_UNARY_RELU:      compute_func = (void *) relu_f32;                break;
+            case HTP_OP_UNARY_STEP:      compute_func = (void *) step_f32;                break;
             case HTP_OP_L2_NORM:         compute_func = (void *) l2_norm_f32;             break;
             case HTP_OP_TRI:
                 task_func    = unary_thread_tri_f32;

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

@@ -59,6 +59,7 @@ static inline bool htp_op_is_unary(uint32_t opcode) {
         case HTP_OP_UNARY_ABS:
         case HTP_OP_UNARY_LOG:
         case HTP_OP_UNARY_RELU:
+        case HTP_OP_UNARY_STEP:
         case HTP_OP_L2_NORM:
         case HTP_OP_TRI:
             return true;

ggml/src/ggml-opencl/ggml-opencl.cpp

@@ -1186,7 +1186,7 @@ struct ggml_backend_opencl_context {
             return nullptr;
         }
 
-        size_t sz;
+        size_t sz = 0;
         const void * kernel_bin = get_adreno_bin_kernel_func(
             kernel_name.c_str(), device_name.c_str(), driver_version.c_str(), &sz);
         if (bin_size) {
@@ -3881,7 +3881,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     backend_ctx->kernel_gemv_noshuffle_q4_0_f32_32b_trans = nullptr;
     backend_ctx->kernel_gemm_noshuffle_q4_0_f32_32b_trans_ila_a8_bin = nullptr;
     backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a_ila_a8_bin = nullptr;
-    if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E) {
+    {
         {
             std::string opts = std::string("-cl-std=") + opencl_c_std +
                                            " -cl-mad-enable "
@@ -4381,8 +4381,8 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     backend_ctx->kernel_gemv_noshuffle_q4_k_f32_32b_trans = nullptr;
     backend_ctx->kernel_gemm_noshuffle_q4_k_f32_32b_trans_ila_a8_bin = nullptr;
     backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_ila_a8_bin = nullptr;
-    if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E) {
-        {
+    {
+        if (backend_ctx->has_vector_subgroup_broadcast) {
             std::string opts = std::string("-cl-std=") + opencl_c_std +
                                            " -cl-mad-enable "
                                            " -DSIMDGROUP_WIDTH=" +
@@ -4430,8 +4430,8 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     backend_ctx->kernel_gemv_noshuffle_q6_k_f32_32b_trans = nullptr;
     backend_ctx->kernel_gemm_noshuffle_q6_k_f32_32b_trans_ila_a8_bin = nullptr;
     backend_ctx->kernel_gemm_noshuffle_q6_k_q8_1_dp4a_ila_a8_bin = nullptr;
-    if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E) {
-        {
+    {
+        if (backend_ctx->has_vector_subgroup_broadcast) {
             std::string opts = std::string("-cl-std=") + opencl_c_std +
                                            " -cl-mad-enable "
                                            " -DSIMDGROUP_WIDTH=" +
@@ -4479,8 +4479,8 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     backend_ctx->kernel_gemv_noshuffle_q5_k_f32_32b_trans = nullptr;
     backend_ctx->kernel_gemm_noshuffle_q5_k_f32_32b_trans_ila_a8_bin = nullptr;
     backend_ctx->kernel_gemm_noshuffle_q5_k_q8_1_dp4a_ila_a8_bin = nullptr;
-    if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E) {
-        {
+    {
+        if (backend_ctx->has_vector_subgroup_broadcast) {
             std::string opts = std::string("-cl-std=") + opencl_c_std +
                                            " -cl-mad-enable "
                                            " -DSIMDGROUP_WIDTH=" +

ggml/src/ggml-rpc/transport-apple.cpp

@@ -26,6 +26,7 @@
 // so every SEND posts a whole 128KiB stride over the wire, even when partially filled.
 // (In testing 128KiB was the best performing among 32, 64, 128, 256)
 
+//TODO: add mechanism similar to https://github.com/ggml-org/llama.cpp/pull/29440 to prevent idle CPU from spinning
 static constexpr uint32_t RDMA_SEG_MAGIC   = 0x52534547u; // "RSEG"
 static constexpr int      RDMA_NBUF        = 16;          // ring depth (frames per direction)
 static constexpr size_t   RDMA_FRAME       = 4096;        // Thunderbolt frame (fixed on Apple)

ggml/src/ggml-rpc/transport.cpp

@@ -25,6 +25,8 @@
 #ifdef GGML_RPC_RDMA
 #  include <infiniband/verbs.h>
 #  include <array>
+#  include <cerrno>
+#  include <chrono>
 #  include <time.h>
 #  ifndef _WIN32
 #    include <poll.h>
@@ -54,6 +56,8 @@ using rdma_gid_t = std::array<uint8_t, RDMA_GID_SIZE>;
 #if defined(GGML_RPC_RDMA) && !defined(GGML_RPC_RDMA_APPLE)
 static constexpr size_t RDMA_CHUNK    = 256 * 1024;   // 256 KiB per send/recv (fits default 8 MiB memlock)
 static constexpr int    RDMA_RX_DEPTH = 24;            // pre-posted recv ring: 24 × 256 KiB = 6 MiB
+// keep polling the CQ for this long after the last activity, then sleep until the next completion
+static constexpr auto   RDMA_SPIN_TIME = std::chrono::milliseconds(100);
 
 struct rdma_conn {
     struct ibv_context * ctx = nullptr;
@@ -61,6 +65,9 @@ struct rdma_conn {
     struct ibv_cq * scq = nullptr;   // send completions
     struct ibv_cq * rcq = nullptr;   // recv completions
     struct ibv_qp * qp  = nullptr;
+    struct ibv_comp_channel * ch = nullptr; // CQ events, so an idle connection can sleep instead of spinning
+
+    std::chrono::steady_clock::time_point last_active; // last completion or posted send
 
     void          * tx_buf = nullptr;
     struct ibv_mr * tx_mr  = nullptr;
@@ -95,6 +102,7 @@ struct rdma_conn {
         if (qp)  ibv_destroy_qp(qp);
         if (scq) ibv_destroy_cq(scq);
         if (rcq) ibv_destroy_cq(rcq);
+        if (ch)  ibv_destroy_comp_channel(ch);
         if (pd)  ibv_dealloc_pd(pd);
         if (ctx) ibv_close_device(ctx);
     }
@@ -142,6 +150,7 @@ struct socket_t::impl {
     bool tcp_peer_closed();
     bool rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, const uint8_t * remote_gid);
     bool rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc);
+    bool rdma_wait_event();
 
     std::unique_ptr<rdma_conn> rdma;
     rdma_local_info            rdma_local = {};
@@ -291,8 +300,10 @@ bool socket_t::impl::rdma_probe() {
     rdma->pd = ibv_alloc_pd(ibctx);
     if (!rdma->pd) return false;
 
-    rdma->scq = ibv_create_cq(ibctx, 16, nullptr, nullptr, 0);
-    rdma->rcq = ibv_create_cq(ibctx, RDMA_RX_DEPTH + 4, nullptr, nullptr, 0);
+    // without a completion channel rdma_poll() spins all the time, as before
+    rdma->ch  = ibv_create_comp_channel(ibctx);
+    rdma->scq = ibv_create_cq(ibctx, 16, nullptr, rdma->ch, 0);
+    rdma->rcq = ibv_create_cq(ibctx, RDMA_RX_DEPTH + 4, nullptr, rdma->ch, 0);
     if (!rdma->scq || !rdma->rcq) return false;
 
     ibv_qp_init_attr qia = {};
@@ -394,15 +405,45 @@ bool socket_t::impl::rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, con
         }
     }
 
+    rdma->last_active = std::chrono::steady_clock::now();
+
     GGML_LOG_INFO("RDMA activated: qpn=%u->%u mtu=%d rx_depth=%d\n",
                   rdma_local.qpn, remote_qpn, 128 << rdma_local.path_mtu, RDMA_RX_DEPTH);
     return true;
 }
 
+// Sleep until the completion channel has an event or the TCP peer closes.
+bool socket_t::impl::rdma_wait_event() {
+    rdma_conn * c = rdma.get();
+    // POLLHUP and POLLERR are always reported, the TCP socket carries no data after the RDMA upgrade
+    struct pollfd pfds[2] = {
+        { c->ch->fd, POLLIN,    0 },
+        { fd,        POLLRDHUP, 0 },
+    };
+    if (poll(pfds, 2, -1) < 0) {
+        return errno == EINTR;
+    }
+    if (pfds[1].revents & (POLLHUP | POLLERR | POLLRDHUP)) {
+        return false;
+    }
+    if (pfds[0].revents & POLLIN) {
+        struct ibv_cq * ev_cq  = nullptr;
+        void          * ev_ctx = nullptr;
+        if (ibv_get_cq_event(c->ch, &ev_cq, &ev_ctx) != 0) {
+            return false;
+        }
+        ibv_ack_cq_events(ev_cq, 1);
+    }
+    return true;
+}
+
 bool socket_t::impl::rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc) {
-    for (uint64_t s = 0; ; s++) {
+    rdma_conn * c = rdma.get();
+    bool armed = false;
+    for (uint64_t s = 1; ; s++) {
         int n = ibv_poll_cq(cq, 1, wc);
         if (n > 0) {
+            c->last_active = std::chrono::steady_clock::now();
             if (wc->status != IBV_WC_SUCCESS) {
                 GGML_LOG_ERROR("RDMA CQ wc error: status=%d (%s) vendor_err=0x%x\n",
                     wc->status, ibv_wc_status_str(wc->status), wc->vendor_err);
@@ -410,7 +451,24 @@ bool socket_t::impl::rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc) {
             return wc->status == IBV_WC_SUCCESS;
         }
         if (n < 0) return false;
-        if ((s & 0xFFFFF) == 0 && s > 0) {
+        if (armed) {
+            // armed and still empty: sleep until the next completion
+            if (!rdma_wait_event()) {
+                return false;
+            }
+            armed = false;
+            continue;
+        }
+        // spin while the connection is busy, arm the CQ once it has been idle for RDMA_SPIN_TIME
+        // a completion that arrives before arming raises no event, so poll once more after arming
+        if (c->ch && (s & 0x3FF) == 0 && std::chrono::steady_clock::now() - c->last_active > RDMA_SPIN_TIME) {
+            if (ibv_req_notify_cq(cq, 0) != 0) {
+                return false;
+            }
+            armed = true;
+            continue;
+        }
+        if ((s & 0xFFFFF) == 0) {
             if (tcp_peer_closed()) {
                 return false;
             }
@@ -444,6 +502,7 @@ bool socket_t::impl::rdma_send(const void * data, size_t size) {
         }
 
         if (ibv_post_send(c->qp, &wr, &bad) != 0) return false;
+        c->last_active = std::chrono::steady_clock::now();
         struct ibv_wc wc;
         if (!rdma_poll(c->scq, &wc)) return false;
 

ggml/src/ggml-sycl/fwht.cpp

@@ -124,6 +124,107 @@ static void launch_fwht(const float * src, float * dst, const int64_t n_rows, co
                          });
 }
 
+// Wide blocks: one row per work-group instead of per sub-group, so each work-item
+// keeps N/NT values rather than N/WARP_SIZE. Butterflies below the sub-group width
+// still shuffle; those up to NT go through work-group local memory; the rest stay
+// in registers.
+template <int N, int NT>
+static void fwht_kernel_wide(const float * __restrict__ src,
+                             float * __restrict__ dst,
+                             const int64_t            n_rows,
+                             const float              scale,
+                             const sycl::nd_item<2> & item,
+                             float *                  smem) {
+    const int64_t r = item.get_global_id(0);
+    if (r >= n_rows) {
+        return;
+    }
+
+    src += r * N;
+    dst += r * N;
+
+    constexpr int el_w = N / NT;
+    static_assert(el_w >= 1 && N % NT == 0, "row must be a whole number of work-group widths");
+
+    const int tid = item.get_local_id(1);
+
+    float reg[el_w];
+#pragma unroll
+    for (int i = 0; i < el_w; ++i) {
+        reg[i] = src[i * NT + tid] * scale;
+    }
+
+    const sycl::sub_group sg   = item.get_sub_group();
+    const int             lane = sg.get_local_linear_id();
+
+    // Butterflies inside the sub-group, same pattern as the narrow kernel.
+#pragma unroll
+    for (int h = 1; h < WARP_SIZE; h *= 2) {
+#pragma unroll
+        for (int j = 0; j < el_w; ++j) {
+            const float val  = reg[j];
+            const float val2 = dpct::permute_sub_group_by_xor(sg, val, h, WARP_SIZE);
+
+            reg[j] = (lane & h) == 0 ? val + val2 : val2 - val;
+        }
+    }
+
+    // Butterflies from the sub-group width up to NT: the partner lane is outside
+    // this sub-group, so it goes through work-group local memory instead of a shuffle.
+    for (int h = WARP_SIZE; h < NT; h *= 2) {
+#pragma unroll
+        for (int j = 0; j < el_w; ++j) {
+            smem[j * NT + tid] = reg[j];
+        }
+        item.barrier(sycl::access::fence_space::local_space);
+#pragma unroll
+        for (int j = 0; j < el_w; ++j) {
+            const float val  = reg[j];
+            const float val2 = smem[j * NT + (tid ^ h)];
+            reg[j]           = (tid & h) == 0 ? val + val2 : val2 - val;
+        }
+        item.barrier(sycl::access::fence_space::local_space);
+    }
+
+    // Butterflies across registers: h is a multiple of NT, so the partner of element
+    // i*NT + tid lives in reg[i + h/NT] on the same work-item.
+    for (int h = NT; h < N; h *= 2) {
+        const int step = h / NT;
+        for (int j = 0; j < el_w; j += 2 * step) {
+            for (int k = 0; k < step; ++k) {
+                const float x = reg[j + k];
+                const float y = reg[j + k + step];
+
+                reg[j + k]        = x + y;
+                reg[j + k + step] = x - y;
+            }
+        }
+    }
+
+#pragma unroll
+    for (int i = 0; i < el_w; ++i) {
+        dst[i * NT + tid] = reg[i];
+    }
+}
+
+template <int N, int NT>
+static void launch_fwht_wide(const float *   src,
+                             float *         dst,
+                             const int64_t   n_rows,
+                             const float     scale,
+                             dpct::queue_ptr stream) {
+    const sycl::range<2> global(n_rows, NT);
+    const sycl::range<2> local(1, NT);
+
+    stream->submit([&](sycl::handler & cgh) {
+        sycl::local_accessor<float, 1> smem(sycl::range<1>(N), cgh);
+        cgh.parallel_for(sycl::nd_range<2>(global, local),
+                         [=](sycl::nd_item<2> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                             fwht_kernel_wide<N, NT>(src, dst, n_rows, scale, item, get_pointer(smem));
+                         });
+    });
+}
+
 template <int N, int m>
 static void kronecker_kernel(const float * __restrict__ src,
                              float * __restrict__ dst,
@@ -285,6 +386,18 @@ bool ggml_sycl_op_fwht(ggml_backend_sycl_context & ctx, const ggml_tensor * src,
         case 1280:
             launch_kronecker<1280, 20>(src_d, dst_d, rows, scale, stream);
             return true;
+        case 1024:
+            launch_fwht_wide<1024, 256>(src_d, dst_d, rows, scale, stream);
+            return true;
+        case 2048:
+            launch_fwht_wide<2048, 256>(src_d, dst_d, rows, scale, stream);
+            return true;
+        case 4096:
+            launch_fwht_wide<4096, 256>(src_d, dst_d, rows, scale, stream);
+            return true;
+        case 8192:
+            launch_fwht_wide<8192, 256>(src_d, dst_d, rows, scale, stream);
+            return true;
         default:
             return false;
     }

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

@@ -11172,8 +11172,10 @@ void ggml_vk_argsort(ggml_backend_vk_context * ctx, vk_context& subctx, const gg
     // Pick the largest workgroup size <= ncolsp2
     uint32_t pipeline_idx = std::min(ncols_pad_log2, num_argsort_pipelines - 1);
 
+    uint32_t max_wg_log2 = std::min(ctx->device->max_workgroup_size_log2, num_argsort_pipelines - 1);
+
     // Use the "small" argsort shader if the whole sort can be done by a single workgroup.
-    bool use_small = ncols_pad_log2 <= ctx->device->max_workgroup_size_log2 &&
+    bool use_small = ncols_pad_log2 <= max_wg_log2 &&
                      ctx->device->pipeline_argsort_f32[pipeline_idx] != nullptr;
 
     vk_pipeline pipeline = use_small ? ctx->device->pipeline_argsort_f32[pipeline_idx]
@@ -11208,7 +11210,7 @@ void ggml_vk_argsort(ggml_backend_vk_context * ctx, vk_context& subctx, const gg
     {
         vk_op_argsort_push_constants pc2 = pc;
         pc2.outer_start = 0;
-        pc2.outer_end = std::min(ncols_pad_log2, ctx->device->max_workgroup_size_log2);
+        pc2.outer_end = std::min(ncols_pad_log2, max_wg_log2);
         pc2.inner_start = 0;
         pc2.inner_end = 100;
         ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
@@ -11217,7 +11219,7 @@ void ggml_vk_argsort(ggml_backend_vk_context * ctx, vk_context& subctx, const gg
     if (!use_small) {
         ggml_vk_sync_buffers(ctx, subctx);
         // Loop over outer/inner passes, synchronizing between each pass.
-        for (uint32_t outer = ctx->device->max_workgroup_size_log2; outer < ncols_pad_log2; ++outer) {
+        for (uint32_t outer = max_wg_log2; outer < ncols_pad_log2; ++outer) {
             for (uint32_t inner = 0; inner < outer + 1; ++inner) {
                 vk_op_argsort_push_constants pc2 = pc;
                 pc2.outer_start = outer;
@@ -15612,7 +15614,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
                 if (device->vulkan_memory_model) {
                     return true;
                 } else {
-                    return op->ne[0] <= (1 << device->max_workgroup_size_log2);
+                    return op->ne[0] <= (1 << std::min(device->max_workgroup_size_log2, num_argsort_pipelines - 1));
                 }
             }
         case GGML_OP_TOP_K:

ggml/src/ggml-zdnn/ggml-zdnn.cpp

@@ -438,15 +438,13 @@ static ggml_backend_i ggml_backend_zdnn_i = {
 };
 
 static ggml_guid_t ggml_backend_zdnn_guid(void) {
-    static const char * guid_str = "IBM-ZDNN-ACCELER";
-    return reinterpret_cast<ggml_guid_t>((void *)guid_str);
+    static char guid_str[] = "IBM-ZDNN-ACCELER";
+    return reinterpret_cast<ggml_guid_t>(guid_str);
 }
 
 bool ggml_backend_is_zdnn(ggml_backend_t backend) {
     return backend != NULL &&
            ggml_guid_matches(backend->guid, ggml_backend_zdnn_guid());
-
-    GGML_UNUSED(backend);
 }
 
 //
@@ -483,7 +481,7 @@ static void ggml_backend_zdnn_device_get_props(ggml_backend_dev_t dev, ggml_back
     props->description = ggml_backend_zdnn_device_get_description(dev);
     props->type        = ggml_backend_zdnn_device_get_type(dev);
     ggml_backend_zdnn_device_get_memory(dev, &props->memory_free, &props->memory_total);
-    props->caps = (ggml_backend_dev_caps) {
+    props->caps = {
         /* .async                = */ false,
         /* .host_buffer          = */ false,
         /* .buffer_from_host_ptr = */ false,
@@ -500,7 +498,7 @@ static ggml_backend_t ggml_backend_zdnn_device_init(ggml_backend_dev_t dev, cons
     }
 
     ggml_backend_t backend = (ggml_backend *)malloc(sizeof(ggml_backend));
-    *backend = (ggml_backend) {
+    *backend = {
         /* .guid       = */ ggml_backend_zdnn_guid(),
         /* .iface      = */ ggml_backend_zdnn_i,
         /* .device     = */ dev,
@@ -619,13 +617,13 @@ ggml_backend_reg_t ggml_backend_zdnn_reg(void) {
     atexit(ggml_zdnn_cleanup);
 
     {
-        g_ggml_backend_zdnn_reg = (ggml_backend_reg) {
+        g_ggml_backend_zdnn_reg = {
             /* .api_version = */ GGML_ZDNN_VERSION,
             /* .iface       = */ ggml_backend_zdnn_reg_i,
             /* .context     = */ NULL
         };
 
-        g_ggml_backend_zdnn_device = (ggml_backend_device) {
+        g_ggml_backend_zdnn_device = {
             /* .iface       = */ ggml_backend_zdnn_device_i,
             /* .reg         = */ &g_ggml_backend_zdnn_reg,
             /* .context     = */ &g_ggml_ctx_dev_main

scripts/hip/gcn-cdna-vgpr-check.py

@@ -64,7 +64,15 @@ def main():
         '_ZL12rwkv_wkv_f32ILi128EEviiiiPKfS1_S1_S1_S1_S1_Pf',
         '_ZL9mul_mat_qIL9ggml_type10ELi64ELb1EEvPKcPKiS4_S4_PfS5_PKf15HIP_vector_typeIjLj3EEiiiiiS9_S9_iiiS9_S9_iiiS9_',
         '_ZL9mul_mat_qIL9ggml_type42ELi128ELb1EEvPKcPKiS4_S4_PfS5_PKf15HIP_vector_typeIjLj3EEiiiiiS9_S9_iiiS9_S9_iiiS9_',
-        '_ZL18flash_attn_ext_f16ILi576ELi512ELi2ELi32ELb0ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil'
+        '_ZL18flash_attn_ext_f16ILi576ELi512ELi2ELi32ELb0ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
+        '_ZL18flash_attn_ext_f16ILi512ELi512ELi16ELi4ELb0ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
+        '_ZL18flash_attn_ext_f16ILi512ELi512ELi16ELi4ELb1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
+        '_ZL18flash_attn_ext_f16ILi512ELi512ELi32ELi2ELb0ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
+        '_ZL18flash_attn_ext_f16ILi512ELi512ELi32ELi2ELb1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
+        '_ZL18flash_attn_ext_f16ILi512ELi512ELi8ELi8ELb0ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
+        '_ZL18flash_attn_ext_f16ILi512ELi512ELi8ELi8ELb1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
+        '_ZL18flash_attn_ext_f16ILi576ELi512ELi16ELi4ELb0ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
+        '_ZL18flash_attn_ext_f16ILi576ELi512ELi4ELi16ELb0ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
     }
 
     functions = parse_log_file(log_file)

src/llama-arch.cpp

@@ -722,7 +722,7 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
     {LLM_TENSOR_TOKEN_EMBD,                 {LLM_TENSOR_LAYER_INPUT,     GGML_OP_GET_ROWS}},
     {LLM_TENSOR_POS_EMBD,                   {LLM_TENSOR_LAYER_INPUT,     GGML_OP_GET_ROWS}},
     {LLM_TENSOR_TOKEN_TYPES,                {LLM_TENSOR_LAYER_INPUT,     GGML_OP_GET_ROWS}},
-    {LLM_TENSOR_HRM_Z_L_INIT,               {LLM_TENSOR_LAYER_INPUT,     GGML_OP_ADD}},
+    {LLM_TENSOR_HRM_Z_L_INIT,               {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}},
     {LLM_TENSOR_TOKEN_EMBD_NORM,            {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},  // do the norms on the first layer (not the input layer)
     {LLM_TENSOR_OUTPUT,                     {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL_MAT}},
     {LLM_TENSOR_CLS,                        {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL_MAT}},

src/llama-context.cpp

@@ -2736,6 +2736,10 @@ public:
         buf_size -= size;
     }
 
+    void discard() override {
+        rinfos.clear();
+    }
+
     size_t n_bytes() override {
         return size_read;
     }
@@ -3086,6 +3090,11 @@ public:
         rinfos.push_back({tensor, ptr, size, offset});
     }
 
+    void discard() override {
+        rinfos.clear();
+        buf_size = 0;
+    }
+
     size_t n_bytes() override {
         return size_read;
     }
@@ -3132,6 +3141,7 @@ size_t llama_context::state_set_data(const uint8_t * src, size_t size) {
         return state_read_data(io);
     } catch (const std::exception & err) {
         LLAMA_LOG_ERROR("%s: error loading state: %s\n", __func__, err.what());
+        io.discard();
         return 0;
     }
 }
@@ -3205,6 +3215,7 @@ size_t llama_context::state_seq_set_data(llama_seq_id seq_id, const uint8_t * sr
         return state_seq_read_data(*io, seq_id, flags);
     } catch (const std::exception & err) {
         LLAMA_LOG_ERROR("%s: error loading state: %s\n", __func__, err.what());
+        io->discard();
         return 0;
     }
 }

src/llama-impl.cpp

@@ -1,8 +1,10 @@
 #include "llama-impl.h"
 
+#include "ggml-backend.h"
 #include "gguf.h"
 #include "llama.h"
 
+#include <algorithm>
 #include <cinttypes>
 #include <climits>
 #include <cstdarg>
@@ -66,6 +68,16 @@ void llama_log_callback_default(ggml_log_level level, const char * text, void *
     fflush(stderr);
 }
 
+void llama_clear_tensor_data(ggml_tensor * t, size_t offset, size_t size) {
+    static const std::vector<uint8_t> zeros(1024*1024, 0);
+
+    // not all backend buffers implement ggml_backend_tensor_memset(), so write zeros instead
+    // TODO: make this a generic fallback in `ggml_backend_tensor_memset` when `set_tensor` is available
+    for (size_t ofs = 0; ofs < size; ofs += zeros.size()) {
+        ggml_backend_tensor_set(t, zeros.data(), offset + ofs, std::min(size - ofs, zeros.size()));
+    }
+}
+
 void replace_all(std::string & s, const std::string & search, const std::string & replace) {
     if (search.empty()) {
         return;

src/llama-impl.h

@@ -93,6 +93,8 @@ struct buffer_view {
     }
 };
 
+void llama_clear_tensor_data(ggml_tensor * t, size_t offset, size_t size);
+
 void replace_all(std::string & s, const std::string & search, const std::string & replace);
 
 // TODO: rename to llama_format ?

src/llama-io.h

@@ -28,6 +28,9 @@ public:
     virtual void read(void * dst, size_t size) = 0;
     virtual void read_tensor(ggml_tensor * tensor, size_t offset, size_t size) = 0;
 
+    // drop tensor data that has been read but not yet applied (e.g. when a restore fails)
+    virtual void discard() {}
+
     // bytes read so far
     virtual size_t n_bytes() = 0;
 

src/llama-kv-cache-dsa.cpp

@@ -168,7 +168,15 @@ void llama_kv_cache_dsa::state_write(llama_io_write_i & io, llama_seq_id seq_id,
 
 void llama_kv_cache_dsa::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
     kv_mla->state_read(io, seq_id, flags);
-    kv_lid->state_read(io, seq_id, flags);
+
+    try {
+        kv_lid->state_read(io, seq_id, flags);
+    } catch (...) {
+        // the MLA part is already restored - undo it, so that a failed restore leaves nothing behind
+        kv_mla->state_clear(seq_id);
+
+        throw;
+    }
 }
 
 llama_kv_cache * llama_kv_cache_dsa::get_mla() const {

src/llama-kv-cache.cpp

@@ -2188,11 +2188,7 @@ const slot_info_vec_t *   sinfos_in) {
         }
 
         if (!res) {
-            if (seq_id == -1) {
-                clear(true);
-            } else {
-                seq_rm(seq_id, -1, -1);
-            }
+            state_clear(seq_id, strm, sinfo);
             throw std::runtime_error("failed to restore kv cache");
         }
 
@@ -2340,6 +2336,11 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
 
     if (dest_seq_id != -1) {
         // single sequence
+        if (cell_count > cells.size()) {
+            LLAMA_LOG_ERROR("%s: not enough cells in kv cache\n", __func__);
+            return false;
+        }
+
         seq_rm(dest_seq_id, -1, -1);
 
         llama_batch_allocr balloc(hparams.n_pos_per_embd());
@@ -2663,6 +2664,112 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
     return true;
 }
 
+void llama_kv_cache::state_clear(llama_seq_id seq_id) {
+    if (seq_id == -1) {
+        clear(true);
+        return;
+    }
+
+    GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
+
+    const uint32_t strm = seq_to_stream[seq_id];
+
+    const auto & cells = v_cells[strm];
+
+    slot_info sinfo;
+    sinfo.s0 = strm;
+    sinfo.s1 = strm;
+    sinfo.resize(1);
+    sinfo.strm[0] = strm;
+
+    // a cell that another sequence still uses keeps its data
+    for (uint32_t i = 0; i < cells.size(); ++i) {
+        if (cells.seq_has(i, seq_id) && cells.seq_count(i) == 1) {
+            sinfo.idxs[0].push_back(i);
+        }
+    }
+
+    state_clear(seq_id, strm, sinfo);
+}
+
+// the cleared ranges mirror the write pattern of state_read_data() - keep both in sync
+void llama_kv_cache::state_clear(llama_seq_id seq_id, uint32_t strm, const slot_info & sinfo) {
+    if (seq_id == -1) {
+        clear(true);
+        return;
+    }
+
+    seq_rm(seq_id, -1, -1);
+
+    // zero the K/V data of the failed restore attempt - the attention can still read the data of free cells
+    if (sinfo.empty() || sinfo.size() == 0) {
+        return;
+    }
+
+    const auto & cells = v_cells[strm];
+
+    const uint32_t cell_count = sinfo.size();
+
+    const bool is_contiguous = sinfo.is_contiguous();
+
+    for (const auto & layer : layers) {
+        const uint32_t il = layer.il;
+
+        const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);
+
+        auto * k = layer.k_stream[strm];
+
+        const size_t k_size_row = ggml_row_size(k->type, n_embd_k_gqa);
+
+        if (is_contiguous) {
+            llama_clear_tensor_data(k, sinfo.head() * k_size_row, cell_count * k_size_row);
+        } else {
+            for (uint32_t i = 0; i < cell_count; ++i) {
+                llama_clear_tensor_data(k, sinfo.idxs[0][i] * k_size_row, k_size_row);
+            }
+        }
+    }
+
+    for (const auto & layer : layers) {
+        const uint32_t il = layer.il;
+
+        const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(il);
+
+        auto * v = layer.v_stream[strm];
+        if (!v) {
+            continue;
+        }
+
+        if (!v_trans) {
+            const size_t v_size_row = ggml_row_size(v->type, n_embd_v_gqa);
+
+            if (is_contiguous) {
+                llama_clear_tensor_data(v, sinfo.head() * v_size_row, cell_count * v_size_row);
+            } else {
+                for (uint32_t i = 0; i < cell_count; ++i) {
+                    llama_clear_tensor_data(v, sinfo.idxs[0][i] * v_size_row, v_size_row);
+                }
+            }
+        } else {
+            const size_t v_size_el = ggml_type_size(v->type);
+
+            if (is_contiguous) {
+                const uint32_t h = sinfo.head();
+
+                for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
+                    llama_clear_tensor_data(v, (h + j * cells.size()) * v_size_el, cell_count * v_size_el);
+                }
+            } else {
+                for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
+                    for (uint32_t i = 0; i < cell_count; ++i) {
+                        llama_clear_tensor_data(v, (sinfo.idxs[0][i] + j * cells.size()) * v_size_el, v_size_el);
+                    }
+                }
+            }
+        }
+    }
+}
+
 //
 // llama_kv_cache_context
 //

src/llama-kv-cache.h

@@ -179,6 +179,9 @@ public:
           slot_info_vec_t *   sinfos_out,
     const slot_info_vec_t *   sinfos_in);
 
+    // undo a state_read() of seq_id (-1 for the whole cache) that another memory module failed to complete
+    void state_clear(llama_seq_id seq_id);
+
     //
     // graph_build API
     //
@@ -345,6 +348,8 @@ private:
     // sinfo_in, when set, replaces the find_slot call: the cells are given by the caller
     bool state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count,       slot_info & sinfo, llama_seq_id dest_seq_id = -1, const slot_info * sinfo_in = nullptr);
     bool state_read_data(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, const slot_info & sinfo);
+
+    void state_clear(llama_seq_id seq_id, uint32_t strm, const slot_info & sinfo);
 };
 
 class llama_kv_cache_context : public llama_memory_context_i {

src/llama-memory-hybrid.cpp

@@ -195,10 +195,22 @@ void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id
 }
 
 void llama_memory_hybrid::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
-    if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
+    const bool read_attn = (flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0;
+
+    if (read_attn) {
         mem_attn->state_read(io, seq_id, flags);
     }
-    mem_recr->state_read(io, seq_id, flags);
+
+    try {
+        mem_recr->state_read(io, seq_id, flags);
+    } catch (...) {
+        // the attention part is already restored - undo it
+        if (read_attn) {
+            mem_attn->state_clear(seq_id);
+        }
+
+        throw;
+    }
 }
 
 llama_kv_cache * llama_memory_hybrid::get_mem_attn() const {

src/llama-memory-recurrent.cpp

@@ -852,7 +852,12 @@ void llama_memory_recurrent::state_read(llama_io_read_i & io, llama_seq_id seq_i
 
     bool res = true;
 
-    res = res && state_read_meta(io, cell_count, seq_id);
+    // save the head of the restored cells - could be needed to clear the state
+    // the head is valid only when state_read_meta() succeeded
+    const bool meta_read = state_read_meta(io, cell_count, seq_id);
+    const uint32_t cell_head = head;
+
+    res = res && meta_read;
 
     try {
         res = res && state_read_data(io, cell_count);
@@ -861,12 +866,7 @@ void llama_memory_recurrent::state_read(llama_io_read_i & io, llama_seq_id seq_i
     }
 
     if (!res) {
-        // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG]
-        if (seq_id == -1) {
-            clear(true);
-        } else {
-            seq_rm(seq_id, -1, -1);
-        }
+        state_clear(seq_id, cell_head, meta_read ? cell_count : 0);
         throw std::runtime_error("failed to restore kv cache");
     }
 
@@ -992,6 +992,11 @@ void llama_memory_recurrent::state_write_data(llama_io_write_i & io, const std::
 bool llama_memory_recurrent::state_read_meta(llama_io_read_i & io, uint32_t cell_count, llama_seq_id dest_seq_id) {
     if (dest_seq_id != -1) {
         // single sequence
+        if (cell_count > size) {
+            LLAMA_LOG_ERROR("%s: not enough cells in kv cache\n", __func__);
+            return false;
+        }
+
         seq_rm(dest_seq_id, -1, -1);
 
         if (cell_count == 0) {
@@ -1223,6 +1228,41 @@ bool llama_memory_recurrent::state_read_data(llama_io_read_i & io, uint32_t cell
     return true;
 }
 
+// the cleared ranges mirror the write pattern of state_read_data() - keep both in sync
+// the transposed s layout is not handled - state_read_data() rejects it before any write
+void llama_memory_recurrent::state_clear(llama_seq_id seq_id, uint32_t cell_head, uint32_t cell_count) {
+    // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG]
+    if (seq_id == -1) {
+        clear(true);
+        return;
+    }
+
+    seq_rm(seq_id, -1, -1);
+
+    if (cell_count == 0) {
+        return;
+    }
+
+    const uint32_t n_layer = hparams.n_layer();
+
+    for (uint32_t il = 0; il < n_layer; ++il) {
+        if (r_l[il] != nullptr) {
+            const size_t r_size_row = ggml_row_size(r_l[il]->type, hparams.n_embd_r());
+            llama_clear_tensor_data(r_l[il], cell_head * r_size_row, cell_count * r_size_row);
+        }
+
+        if (s_l[il] != nullptr) {
+            const size_t s_size_row = ggml_row_size(s_l[il]->type, hparams.n_embd_s());
+            llama_clear_tensor_data(s_l[il], cell_head * s_size_row, cell_count * s_size_row);
+        }
+
+        if (p_l[il] != nullptr) {
+            const size_t p_size_row = ggml_row_size(p_l[il]->type, hparams.ple_conv_state());
+            llama_clear_tensor_data(p_l[il], cell_head * p_size_row, cell_count * p_size_row);
+        }
+    }
+}
+
 //
 // llama_memory_recurrent_context
 //

src/llama-memory-recurrent.h

@@ -134,6 +134,8 @@ private:
 
     bool state_read_meta(llama_io_read_i & io, uint32_t cell_count, llama_seq_id dest_seq_id = -1);
     bool state_read_data(llama_io_read_i & io, uint32_t cell_count);
+
+    void state_clear(llama_seq_id seq_id, uint32_t cell_head, uint32_t cell_count);
 };
 
 class llama_memory_recurrent_context : public llama_memory_context_i {

src/models/hrm-text.cpp

@@ -43,7 +43,7 @@ void llama_model_hrm_text::load_arch_tensors(llama_model_loader &) {
         output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);
     }
 
-    hrm_z_l_init = create_tensor(tn(LLM_TENSOR_HRM_Z_L_INIT), { n_embd }, 0);
+    hrm_z_l_init = create_tensor(tn(LLM_TENSOR_HRM_Z_L_INIT, 0), { n_embd }, 0);
 
     const int lps = hparams.n_hrm_layers_per_stack;
 

tests/CMakeLists.txt

@@ -309,6 +309,7 @@ endif()
 # TODO: make this test (and others) not link `libllama` as it is not needed [TAG_TESTS_LLAMA_LINK]
 llama_build(test-backend-ops.cpp)
 
+llama_build_and_test(test-tiled-mulmat.cpp)
 llama_build_and_test(test-model-load-cancel.cpp LABEL "model")
 llama_build_and_test(test-autorelease.cpp       LABEL "model")
 llama_build_and_test(test-backend-sampler.cpp   LABEL "model")

tests/test-backend-ops.cpp

@@ -9830,6 +9830,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         add_test_bin_bcast(type, {5120, 1, 1, 1}, {1, 256, 1, 1});
         add_test_bin_bcast(type, {640, 1, 1, 1}, {1, 1, 1, 1});
         add_test_bin_bcast(type, {64, 262144, 1, 1}, {1, 1, 1, 1});
+        add_test_bin_bcast(type, {128, 1, 8, 1}, {1, 4, 1, 1});
         //add_test_bin_bcast(type, {3, 3, 2560, 1280}, {1, 1, 1, 1});
         //add_test_bin_bcast(type, {3, 3, 2560, 1280}, {2, 1, 1, 1});
     }
@@ -9979,6 +9980,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     }
 
     test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 16, 1, 1024, 1, 32, 4)); // Mamba-1
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 96, 64, 128, 8, 1, 1)); // Nemotron-3-Puzzle decode (scan path, unused warp in last block)
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 96, 64, 16, 8, 300, 2, false, /*K=*/1, /*weak_decay=*/true)); // d_state=96 SSD multi-chunk (partial 2nd chunk, 2 seqs)
+    test_cases.emplace_back(new test_ssm_scan_rollback(GGML_TYPE_F32, 96, 64, 16, 8, 8, 2, /*K=*/3)); // d_state=96 rollback snapshots match prefix states
     test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 32, 4)); // Mamba-2
     test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 256, 64,  8, 2, 32, 4)); // Falcon-H1
     test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 128, 4, 4, 16, 2, true)); // x/B/C overlap
@@ -10017,6 +10021,12 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 128, 32, 128));
     test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 128, 4, 128, {2, 3}));
     test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 256, 512, 256)); // many rows
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 64, 1, 64));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 128, 1, 128));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 256, 1, 256));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 512, 1, 512));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 128, 32, 128));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 256, 512, 256)); // many rows
     test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 32, 1, 32)); // too small (N<64)
     test_cases.emplace_back(
         new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 16384, 1, 16384));               // too big (N>8192)

tests/test-jinja.cpp

@@ -1153,7 +1153,7 @@ static void test_tests(testing & t) {
     );
 
     test_template(t, "is not equalto",
-        "{{ 'yes' if 3 is not equalto(4) }}",
+        "{{ 'yes' if 3 is not equalto 4 }}",
         json::object(),
         "yes"
     );
@@ -1165,7 +1165,7 @@ static void test_tests(testing & t) {
     );
 
     test_template(t, "is gt",
-        "{{ 'yes' if 3 is gt(2) }}",
+        "{{ 'yes' if 3 is gt 2 }}",
         json::object(),
         "yes"
     );
@@ -1177,7 +1177,7 @@ static void test_tests(testing & t) {
     );
 
     test_template(t, "is lt",
-        "{{ 'yes' if 2 is lt(3) }}",
+        "{{ 'yes' if 2 is lt 3 }}",
         json::object(),
         "yes"
     );
@@ -1194,6 +1194,12 @@ static void test_tests(testing & t) {
         "yes"
     );
 
+    test_template(t, "is lt and gt",
+        "{{ 'yes' if x is lt 3 and x is gt 1 }}",
+        {{"x", 2}},
+        "yes"
+    );
+
     test_template(t, "is lower",
         "{{ 'yes' if 'lowercase' is lower }}",
         json::object(),
@@ -1206,12 +1212,36 @@ static void test_tests(testing & t) {
         "yes"
     );
 
-    test_template(t, "is sameas",
+    test_template(t, "is sameas boolean",
         "{{ 'yes' if x is sameas(false) }}",
         {{"x", false}},
         "yes"
     );
 
+    test_template(t, "is sameas integer",
+        "{{ 'yes' if x is sameas(1) }}",
+        {{"x", 1}},
+        "yes"
+    );
+
+    test_template(t, "is sameas object",
+        "{{ 'yes' if x is sameas(x) }}",
+        {{"x", {{"y", false}}}},
+        "yes"
+    );
+
+    test_template(t, "is sameas ref object",
+        "{% set y = x.y %}{{ 'yes' if x.y is sameas(y) and x.y is not sameas(x.z) }}",
+        {{"x", {{"y", {{"z", 1}}}, {"z", {{"z", 1}}}}}},
+        "yes"
+    );
+
+    test_template(t, "is sameas undefined",
+        "{{ 'yes' if x is sameas(x) else 'no' }}",
+        json::object(),
+        "no"
+    );
+
     test_template(t, "is boolean",
         "{{ 'yes' if x is boolean }}",
         {{"x", true}},
@@ -1248,6 +1278,12 @@ static void test_tests(testing & t) {
         "yes"
     );
 
+    test_template(t, "is integer or float",
+        "{{ 'yes' if x.y is integer or x.y is float else 'no' }}",
+        {{"x", {{"y", 1.1}}}},
+        "yes"
+    );
+
     test_template(t, "is sequence",
         "{{ 'yes' if x is sequence }}",
         {{"x", json::array({1, 2, 3})}},
@@ -1940,6 +1976,30 @@ static void test_object_methods(testing & t) {
         json::object(),
         ""
     );
+
+    test_template(t, "dict from dict",
+        "{% set o = dict({'a': 3, 'b': 1, 'c': 2}) %}{{ o|tojson }}",
+        json::object(),
+        "{\"a\": 3, \"b\": 1, \"c\": 2}"
+    );
+
+    test_template(t, "dict from kwargs",
+        "{% set o = dict(a=3, b=1, c=2) %}{{ o|tojson }}",
+        json::object(),
+        "{\"a\": 3, \"b\": 1, \"c\": 2}"
+    );
+
+    test_template(t, "dict from tuples",
+        "{% set o = dict((obj | items | list)) %}{{ o|tojson }}",
+        {{"obj", {{"a", 3}, {"b", 1}, {"c", 2}}}},
+        "{\"a\": 3, \"b\": 1, \"c\": 2}"
+    );
+
+    test_template(t, "dict from tuples and kwargs",
+        "{% set o = dict((obj | items | list), c=2) %}{{ o|tojson }}",
+        {{"obj", {{"a", 3}, {"b", 1}}}},
+        "{\"a\": 3, \"b\": 1, \"c\": 2}"
+    );
 }
 
 static void test_hasher(testing & t) {

tests/test-save-load-state.cpp

@@ -5,10 +5,16 @@
 
 #include <algorithm>
 #include <clocale>
+#include <cmath>
+#include <cstdio>
 #include <cstring>
 #include <filesystem>
+#include <fstream>
+#include <functional>
+#include <iterator>
 #include <random>
 #include <string>
+#include <utility>
 #include <vector>
 
 constexpr double NMSE_THRESHOLD = 1e-5;
@@ -599,6 +605,161 @@ static bool test_state_roundtrip(struct llama_model * model, const struct common
 }
 
 
+// overwrite the tensor data with 0xff bytes (NaN when read as f16/f32), so that the restore fails
+static bool corrupt_state(std::vector<uint8_t> & data) {
+    if (data.size() < 3*4096) {
+        LOG_ERR("%s: state of %zu bytes is too small to corrupt\n", __func__, data.size());
+        return false;
+    }
+
+    std::fill(data.begin() + 4096, data.end() - data.size()/4, 0xff);
+    return true;
+}
+
+
+// Test 9: state restore failure
+// a failed restore must leave the sequence empty and must not change the logits of other sequences
+static bool test_state_restore_failure(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) {
+    auto params_ctx = common_context_params_to_llama(params);
+    params_ctx.n_ctx      = 256;
+    params_ctx.n_seq_max  = 4;
+    params_ctx.kv_unified = true;
+
+    // without flash attention, corrupted data left behind by the restore shows up as NaN logits on the other sequences
+    params_ctx.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
+
+    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
+    if (!ctx) {
+        LOG_ERR("%s: failed to create context\n", __func__);
+        return false;
+    }
+
+    LOGV(LOG_LEVEL_INFO, "\n=== Test 9: state restore failure ===\n");
+
+    llama_memory_t mem = llama_get_memory(ctx.get());
+    if (mem == nullptr) {
+        LOGV(LOG_LEVEL_INFO, "PASS (model has no memory)\n");
+        return true;
+    }
+
+    const auto decode = [&](const llama_tokens & inp, llama_seq_id seq_id, std::vector<float> * logits_out) {
+        llama_batch_ptr batch(inp.size(), 0, 1);
+        for (size_t i = 0; i < inp.size(); ++i) {
+            common_batch_add(batch.get(), inp[i], i, { seq_id }, i == inp.size() - 1);
+        }
+
+        if (llama_decode(ctx.get(), batch.get())) {
+            LOG_ERR("%s: failed to decode on sequence %d\n", __func__, seq_id);
+            return false;
+        }
+
+        if (logits_out && !get_current_logits(ctx.get(), *logits_out)) {
+            LOG_ERR("%s: failed to get logits\n", __func__);
+            return false;
+        }
+
+        return true;
+    };
+
+    const llama_tokens tokens_save  (tokens.begin(), tokens.begin() + std::min<size_t>(24, tokens.size()));
+    const llama_tokens tokens_verify(tokens.end() - std::min<size_t>(8, tokens.size()), tokens.end());
+
+    // the registered tests share a working directory, so the state file is named after the model
+    const std::string path = "state-restore-failure." + std::filesystem::path(params.model.path).filename().string() + ".tmp.bin";
+
+    llama_memory_clear(mem, true);
+
+    std::vector<float> baseline;
+    if (!decode(tokens_verify, 1, &baseline)) {
+        return false;
+    }
+
+    const std::vector<std::pair<const char *, std::function<bool()>>> cases = {
+        { "buffer", [&]() {
+            std::vector<uint8_t> state(llama_state_seq_get_size(ctx.get(), 0));
+            GGML_ASSERT(llama_state_seq_get_data(ctx.get(), state.data(), state.size(), 0) == state.size());
+            llama_memory_seq_rm(mem, 0, -1, -1);
+
+            if (!corrupt_state(state)) {
+                return false;
+            }
+
+            return llama_state_seq_set_data(ctx.get(), state.data(), state.size(), 0) == 0;
+        }},
+        { "file", [&]() {
+            GGML_ASSERT(llama_state_seq_save_file(ctx.get(), path.c_str(), 0, tokens_save.data(), tokens_save.size()) > 0);
+            llama_memory_seq_rm(mem, 0, -1, -1);
+
+            std::vector<uint8_t> data;
+            {
+                std::ifstream f(path, std::ios::binary);
+                data.assign(std::istreambuf_iterator<char>(f), std::istreambuf_iterator<char>());
+            }
+
+            if (!corrupt_state(data)) {
+                std::remove(path.c_str());
+                return false;
+            }
+
+            {
+                std::ofstream f(path, std::ios::binary);
+                f.write((const char *) data.data(), data.size());
+            }
+
+            llama_tokens tokens_out(tokens_save.size());
+            size_t n_token_count = 0;
+            const size_t nread = llama_state_seq_load_file(ctx.get(), path.c_str(), 0, tokens_out.data(), tokens_out.size(), &n_token_count);
+            std::remove(path.c_str());
+
+            return nread == 0;
+        }},
+    };
+
+    for (const auto & [name, restore_failed] : cases) {
+        llama_memory_clear(mem, true);
+
+        if (!decode(tokens_save, 0, nullptr)) {
+            return false;
+        }
+
+        if (!restore_failed()) {
+            LOG_ERR("%s: %s: restoring a corrupted state did not fail\n", __func__, name);
+            return false;
+        }
+
+        if (llama_memory_seq_pos_max(mem, 0) != -1) {
+            LOG_ERR("%s: %s: sequence not empty after failed restore\n", __func__, name);
+            return false;
+        }
+
+        std::vector<float> logits;
+        if (!decode(tokens_verify, 1, &logits)) {
+            return false;
+        }
+
+        float  diff_max = 0.0f;
+        size_t n_nan    = 0;
+        for (size_t i = 0; i < logits.size(); ++i) {
+            if (std::isnan(logits[i]) || std::isnan(baseline[i])) {
+                n_nan++;
+            } else {
+                diff_max = std::max(diff_max, std::fabs(logits[i] - baseline[i]));
+            }
+        }
+
+        if (n_nan > 0 || diff_max > 1e-6f) {
+            LOG_ERR("%s: %s: logits changed after failed restore (max diff = %g, nan = %zu)\n", __func__, name, diff_max, n_nan);
+            return false;
+        }
+
+        LOG_TRC("%s: %s: logits match (max diff = %g)\n", __func__, name, diff_max);
+    }
+
+    LOGV(LOG_LEVEL_INFO, "\nPASS\n");
+    return true;
+}
+
+
 struct test_suite {
     std::vector<test_status> results;
 
@@ -609,10 +770,10 @@ struct test_suite {
 
 // column headers for the --models table, one per test, in the order they are run
 static const std::vector<const char *> test_names = {
-    "baseline", "seq_rm", "state_load", "cp_h", "cp_d", "cp_h_s", "cp_d_s", "rt",
+    "baseline", "seq_rm", "state_load", "cp_h", "cp_d", "cp_h_s", "cp_d_s", "rt", "rf",
 };
 
-// Run the full save/load test suite (tests 1-8) for a single model.
+// Run the full save/load test suite (tests 1-9) for a single model.
 // Returns the per-test results.
 static test_suite run_save_load_tests_for_model(const std::string & model_path, const struct common_params & base_params) {
     test_suite suite;
@@ -688,6 +849,9 @@ static test_suite run_save_load_tests_for_model(const std::string & model_path,
     // Test 8: state blob round-trip
     suite.results.push_back(test_state_roundtrip(model, params, tokens) ? test_status::PASS : test_status::FAIL);
 
+    // Test 9: state restore failure
+    suite.results.push_back(test_state_restore_failure(model, params, tokens) ? test_status::PASS : test_status::FAIL);
+
     return suite;
 }
 

tests/test-tiled-mulmat.cpp

@@ -0,0 +1,852 @@
+#include "ggml-alloc.h"
+#include "ggml-backend.h"
+#include "ggml.h"
+
+#include <time.h>
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#if defined(_WIN32)
+#include <windows.h>
+#endif
+
+// failed checks; main exits non-zero on any failures
+static int n_failed = 0;
+
+static float * gen_rand_f32(int64_t n) {
+    float * data = (float *) malloc(n * sizeof(float));
+    for (int64_t i = 0; i < n; ++i) {
+        data[i] = (float)rand() / (float)RAND_MAX - 0.5f;
+        data[i] *= 5.0f;
+    }
+    return data;
+}
+
+// updates pointers to max_err and rms_err
+static void compare_f32(const float * ref, const float * out, int64_t n, float * max_err, float * rms_err) {
+    *max_err = 0.0f;
+    double sum_sq_err = 0.0;
+    for (int64_t i = 0; i < n; ++i) {
+        float err = fabsf(ref[i] - out[i]);
+        if (err > *max_err) {
+            *max_err = err;
+        }
+        sum_sq_err += (double)err * err;
+    }
+    *rms_err = sqrt(sum_sq_err / n);
+}
+
+// Fill from a flat row-major f32 source: rows * cols floats.
+static void fill_tensor(struct ggml_tensor * t, const float * src, int64_t rows, int64_t cols, ggml_type qtype) {
+    GGML_ASSERT(t->ne[0] == cols);
+    GGML_ASSERT(rows * cols == ggml_nelements(t));
+    if (qtype == GGML_TYPE_F32) {
+        ggml_backend_tensor_set(t, src, 0, ggml_nbytes(t));
+        return;
+    }
+    void * q = malloc(ggml_nbytes(t));
+    // some iq types (iq2_xxs, iq2_xs, iq1_s) steer their code choice with an imatrix; the
+    // quantized bytes are shared by the reference and tiled paths, so a dummy imatrix is enough
+    const float * imatrix = NULL;
+    float * im = NULL;
+    if (ggml_quantize_requires_imatrix(qtype)) {
+        im = (float *) malloc(rows * cols * sizeof(float));
+        for (int64_t i = 0; i < rows * cols; i++) im[i] = 1.0f;
+        imatrix = im;
+    }
+    ggml_quantize_chunk(qtype, src, q, 0, rows, cols, imatrix);
+    ggml_backend_tensor_set(t, q, 0, ggml_nbytes(t));
+    free(q);
+    free(im);
+}
+
+// The CPU-specific control API is resolved through the backend registry: with
+// GGML_BACKEND_DL the CPU backend is a runtime-loaded module and its symbols
+// are not available at link time
+static void cpu_set_use_ref(ggml_backend_t backend, bool use_ref) {
+    ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
+    void (* set_use_ref)(ggml_backend_t, bool) =
+        (void (*)(ggml_backend_t, bool)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cpu_set_use_ref");
+    if (!set_use_ref) {
+        fprintf(stderr, "ggml_backend_cpu_set_use_ref not available\n");
+        exit(1);
+    }
+    set_use_ref(backend, use_ref);
+}
+
+static void cpu_set_n_threads(ggml_backend_t backend, int n_threads) {
+    ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
+    void (* set_n_threads)(ggml_backend_t, int) =
+        (void (*)(ggml_backend_t, int)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads");
+    if (!set_n_threads) {
+        fprintf(stderr, "ggml_backend_set_n_threads not available\n");
+        exit(1);
+    }
+    set_n_threads(backend, n_threads);
+}
+
+static void test_matmul(ggml_backend_t backend, int64_t M, int64_t N, int64_t K, ggml_type quant_type) {
+    srand(0xBEEF);
+
+    float * src1_ref = gen_rand_f32(M * N);
+    float * src0_ref = gen_rand_f32(N * K);
+    float * dst_out = (float *) malloc(M * K * sizeof(float));
+    float * dst_tiled = (float *) malloc(M * K * sizeof(float));
+
+    struct ggml_init_params ip = { 1024*1024*1024, nullptr, true };
+    struct ggml_context * ctx = ggml_init(ip);
+
+    // ggml_mul_mat(t0, t1) computes t1 * t0^T: t0 (src0, quant) is N wide x
+    // K high, t1 (src1, f32) is N wide x M high, dst = (K, M).
+    struct ggml_tensor * src1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, N, M);
+    struct ggml_tensor * src0 = ggml_new_tensor_2d(ctx, quant_type,   N, K);
+
+    struct ggml_cgraph * gf  = ggml_new_graph(ctx);
+    struct ggml_tensor * dst   = ggml_mul_mat(ctx, src0, src1);
+    ggml_build_forward_expand(gf, dst);
+
+    ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend);
+
+    fill_tensor(src1, src1_ref, M, N, GGML_TYPE_F32);
+    fill_tensor(src0, src0_ref, K, N, quant_type);
+
+    cpu_set_use_ref(backend, true);
+    ggml_backend_graph_compute(backend, gf);
+    ggml_backend_tensor_get(dst, dst_out, 0, ggml_nbytes(dst));
+    cpu_set_use_ref(backend, false);
+    ggml_backend_graph_compute(backend, gf);
+    ggml_backend_tensor_get(dst, dst_tiled, 0, ggml_nbytes(dst));
+
+    // std vs tiled: identical quantized inputs, so any large difference here is a bug in the tiled kernel
+    float max_err, rms_err;
+    compare_f32(dst_out, dst_tiled, M*K, &max_err, &rms_err);
+    float tol = 1e-3f;
+
+    printf("TEST %lldx%lld * %lldx%lld (%s): %s (max_err: %f, rms: %f, tolerance: %f)\n",
+           (long long)M, (long long)N, (long long)N, (long long)K,
+           ggml_type_name(quant_type), (max_err <= tol) ? "PASS" : "FAIL", max_err, rms_err, tol);
+
+    // if the tiled kernel deviates from std beyond quantization tolerance,
+    // dump a few offenders
+    if (max_err > tol) {
+        int64_t shown = 0;
+        for (int64_t i = 0; i < M*K && shown < 8; ++i) {
+            float err = fabsf(dst_out[i] - dst_tiled[i]);
+            if (err > tol) {
+                printf("  tiled vs std: i=%lld (m=%lld k=%lld) std=%f tiled=%f err=%f\n",
+                       (long long)i, (long long)(i/K), (long long)(i%K), dst_out[i], dst_tiled[i], err);
+                ++shown;
+            }
+        }
+        ++n_failed;
+    }
+
+    ggml_backend_buffer_free(buf);
+    ggml_free(ctx);
+    free(src1_ref); free(src0_ref); free(dst_out); free(dst_tiled);
+}
+
+// Higher-dim (ne[2], ne[3] > 1) check, std (ggml_mul_mat) is trusted as the reference
+//
+// src1 (F32)  : [N, M, src1_2, src1_3]   ne0=N (reduction), ne1=M (out0)
+// src0 (quant): [N, K, src0_2, src0_3]   ne0=N (reduction), ne1=K (out1)
+// dst = ggml_mul_mat(src0, src1) : [K, M, src1_2, src1_3]
+static void test_matmul_highdim(ggml_backend_t backend, int64_t M, int64_t N, int64_t K,
+                         int64_t src1_2, int64_t src1_3,
+                         int64_t src0_2, int64_t src0_3,
+                         ggml_type quant_type) {
+    srand(0xBEEF);
+
+    const int64_t n_src1 = N*M*src1_2*src1_3;
+    const int64_t n_src0 = N*K*src0_2*src0_3;
+    const int64_t n_dst = K*M*src1_2*src1_3;
+
+    float * src1_ref   = gen_rand_f32(n_src1);
+    float * src0_ref   = gen_rand_f32(n_src0);
+    float * dst_std   = (float *) malloc(n_dst * sizeof(float));
+    float * dst_tiled = (float *) malloc(n_dst * sizeof(float));
+
+    struct ggml_init_params ip = { 1024*1024*1024, nullptr, true };
+    struct ggml_context * ctx = ggml_init(ip);
+
+    int64_t ne_src1[4] = { N, M, src1_2, src1_3 };
+    int64_t ne_src0[4] = { N, K, src0_2, src0_3 };
+    struct ggml_tensor * src1  = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_src1);
+    struct ggml_tensor * src0 = ggml_new_tensor(ctx, quant_type,  4, ne_src0);
+
+    struct ggml_cgraph * gf  = ggml_new_graph(ctx);
+    struct ggml_tensor * dst   = ggml_mul_mat(ctx, src0, src1);
+    ggml_build_forward_expand(gf, dst);
+
+    ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend);
+
+    // the batch dims fold into the row count (rows are contiguous, ne0 fastest)
+    fill_tensor(src1, src1_ref, M * src1_2 * src1_3, N, GGML_TYPE_F32);
+    fill_tensor(src0, src0_ref, K * src0_2 * src0_3, N, quant_type);
+
+    // reference = stock path (use_ref keeps the op off the tiled hook), tiled
+    // = the gated path; same weights, same op, run twice into separate buffers
+    cpu_set_use_ref(backend, true);
+    ggml_backend_graph_compute(backend, gf);
+    ggml_backend_tensor_get(dst, dst_std, 0, ggml_nbytes(dst));
+    cpu_set_use_ref(backend, false);
+    ggml_backend_graph_compute(backend, gf);
+    ggml_backend_tensor_get(dst, dst_tiled, 0, ggml_nbytes(dst));
+
+    // max |dst|: scale for the quantization tolerance
+    float scale = 0.0f;
+    for (int64_t i = 0; i < n_dst; ++i) {
+        scale = fmaxf(scale, fabsf(dst_std[i]));
+    }
+
+    // std vs tiled: identical quantized inputs, so a large difference is a bug
+    float max_err, rms_err;
+    compare_f32(dst_std, dst_tiled, n_dst, &max_err, &rms_err);
+    float tol = (quant_type == GGML_TYPE_F32) ? 1e-4f : fmaxf(1e-3f, 1e-3f*scale);
+
+    // which path did the tiled op take (mirrors the hard-constraint gate; the
+    // harness runs with force on, so the profitability check is bypassed)
+    bool tiled_kernel = (N % 256 == 0) && (src0_2 > 0) && (src0_3 > 0)
+        && (src1_2 % src0_2 == 0) && (src1_3 % src0_3 == 0);
+
+    printf("TEST %lldx%lldx%lldx%lld * %lldx%lldx%lldx%lld (%s, %s): %s (max_err: %f, rms: %f, scale: %f)\n",
+           (long long)M, (long long)N, (long long)src1_2, (long long)src1_3,
+           (long long)K, (long long)N, (long long)src0_2, (long long)src0_3,
+           ggml_type_name(quant_type), tiled_kernel ? "tiled-kernel" : "stock",
+           (max_err <= tol) ? "PASS" : "FAIL", max_err, rms_err, scale);
+
+    if (max_err > tol) {
+        int64_t shown = 0;
+        for (int64_t i = 0; i < n_dst && shown < 8; ++i) {
+            float err = fabsf(dst_std[i] - dst_tiled[i]);
+            if (err > tol) {
+                printf("  tiled vs std: i=%lld std=%f tiled=%f err=%f\n",
+                       (long long)i, dst_std[i], dst_tiled[i], err);
+                ++shown;
+            }
+        }
+        ++n_failed;
+    }
+
+    ggml_backend_buffer_free(buf);
+    ggml_free(ctx);
+    free(src1_ref); free(src0_ref); free(dst_std); free(dst_tiled);
+}
+
+// MUL_MAT_ID (MoE) check, std (vec_dot) is trusted as the reference
+//
+// src0 (as, quant): [K, R, n_experts]   ne0 = K (reduction), ne1 = R rows per expert, ne2 = experts
+// src1 (b, F32)   : [K, b_slots, batch]  ne0 = K (reduction), ne1 = b_slots (b rows, broadcast over k columns)
+// ids             : [k, batch]           expert picked per (slot, batch row); k % b_slots == 0
+// dst = ggml_mul_mat_id(as, b, ids) : [R, k, batch]; column (id, t) = GEMV of src1 row (id % b_slots + t*b_slots)
+// against expert ids[t*k+id]. b_slots = 1 is the common MoE case (all top-k experts see one input row).
+static void test_mul_mat_id(ggml_backend_t backend, int64_t K, int64_t R, int64_t n_experts,
+                            int64_t k, int64_t b_slots, int64_t batch, ggml_type quant_type,
+                            bool src1_strided = false) {
+    srand(0xBEEF);
+
+    const int64_t n_as  = K * R * n_experts;
+    const int64_t n_b   = K * b_slots * batch;
+    const int64_t n_dst = R * k * batch;
+
+    float * as_ref  = gen_rand_f32(n_as);
+    float * b_ref   = gen_rand_f32(n_b);
+    int32_t * ids   = (int32_t *) malloc(k * batch * sizeof(int32_t));
+    float * dst_ref = (float *) malloc(n_dst * sizeof(float));
+    float * dst_tiled = (float *) malloc(n_dst * sizeof(float));
+
+    struct ggml_init_params ip = { 1024*1024*1024, nullptr, true };
+    struct ggml_context * ctx = ggml_init(ip);
+
+    int64_t ne_as[4]  = { K, R, n_experts, 1 };
+    int64_t ne_b[4]   = { K, b_slots, batch, 1 };
+    int64_t ne_ids[4] = { k, batch, 1, 1 };
+    struct ggml_tensor * src0  = ggml_new_tensor(ctx, quant_type,  4, ne_as);
+    struct ggml_tensor * ids_t = ggml_new_tensor(ctx, GGML_TYPE_I32, 4, ne_ids);
+    struct ggml_tensor * src1;
+    struct ggml_tensor * src1_base = NULL;
+    if (src1_strided) {
+        // strided view: nb[1] = 4 bytes, the tiled path reads src1 through strides into wdata
+        src1_base = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K, b_slots * batch);
+        src1 = ggml_view_3d(ctx, src1_base, K, b_slots, batch, 4, (size_t) b_slots * K * 4, 0);
+    } else {
+        src1 = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_b);
+    }
+
+    struct ggml_cgraph * gf  = ggml_new_graph(ctx);
+    struct ggml_tensor * dst = ggml_mul_mat_id(ctx, src0, src1, ids_t);
+    ggml_build_forward_expand(gf, dst);
+
+    ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend);
+
+    // deterministic balanced routing: each expert gets ~ k*batch/n_experts rows
+    for (int64_t t = 0; t < batch; t++) {
+        for (int64_t e = 0; e < k; e++) {
+            ids[t * k + e] = (int32_t) ((t * k + e) % n_experts);
+        }
+    }
+
+    fill_tensor(src1, b_ref, b_slots * batch, K, GGML_TYPE_F32);
+    fill_tensor(src0, as_ref, R * n_experts, K, quant_type);
+    ggml_backend_tensor_set(ids_t, ids, 0, ggml_nbytes(ids_t));
+
+    cpu_set_use_ref(backend, true);
+    ggml_backend_graph_compute(backend, gf);
+    ggml_backend_tensor_get(dst, dst_ref, 0, ggml_nbytes(dst));
+    cpu_set_use_ref(backend, false);
+    ggml_backend_graph_compute(backend, gf);
+    ggml_backend_tensor_get(dst, dst_tiled, 0, ggml_nbytes(dst));
+
+    // same quantized inputs, so a large difference is a bug in the tiled path
+    float max_err, rms_err;
+    compare_f32(dst_ref, dst_tiled, n_dst, &max_err, &rms_err);
+    float tol = 1e-3f;
+
+    printf("TEST mul_mat_id %lldx%lldx%lld * %lldx%lldx%lld ids[%lld,%lld] (%s%s): %s (max_err: %f, rms: %f, tolerance: %f)\n",
+           (long long)K, (long long)R, (long long)n_experts,
+           (long long)K, (long long)b_slots, (long long)batch,
+           (long long)k, (long long)batch,
+           ggml_type_name(quant_type), src1_strided ? ", strided-b" : "",
+           (max_err <= tol) ? "PASS" : "FAIL", max_err, rms_err, tol);
+
+    if (max_err > tol) {
+        int64_t shown = 0;
+        for (int64_t i = 0; i < n_dst && shown < 8; ++i) {
+            float err = fabsf(dst_ref[i] - dst_tiled[i]);
+            if (err > tol) {
+                printf("  tiled vs std: i=%lld (row=%lld slot=%lld batch=%lld) std=%f tiled=%f err=%f\n",
+                       (long long)i, (long long)(i % R), (long long)((i / R) % k), (long long)(i / (R * k)),
+                       dst_ref[i], dst_tiled[i], err);
+                ++shown;
+            }
+        }
+        ++n_failed;
+    }
+
+    ggml_backend_buffer_free(buf);
+    ggml_free(ctx);
+    free(as_ref); free(b_ref); free(ids); free(dst_ref); free(dst_tiled);
+}
+
+static double time_graph_compute(ggml_backend_t backend, struct ggml_cgraph * gf) {
+#if defined(_WIN32)
+    // high-res timer; clock_gettime is not portable across Windows toolchains
+    LARGE_INTEGER freq, t0, t1;
+    QueryPerformanceFrequency(&freq);
+    QueryPerformanceCounter(&t0);
+    ggml_backend_graph_compute(backend, gf);
+    QueryPerformanceCounter(&t1);
+    return (double) (t1.QuadPart - t0.QuadPart) / (double) freq.QuadPart;
+#else
+    struct timespec t0, t1;
+    clock_gettime(CLOCK_MONOTONIC, &t0);
+    ggml_backend_graph_compute(backend, gf);
+    clock_gettime(CLOCK_MONOTONIC, &t1);
+    return (t1.tv_sec - t0.tv_sec) + (t1.tv_nsec - t0.tv_nsec) / 1e9;
+#endif
+}
+
+// Warm up, then time n runs and return the best
+static double time_graph_compute_best(ggml_backend_t backend, struct ggml_cgraph * gf, int n,
+                                      void * flush_buf = NULL, size_t flush_size = 0) {
+
+    time_graph_compute(backend, gf); // warmup
+    double best = 1e30;
+    for (int i = 0; i < n; ++i) {
+        // Evict L3 cache first
+        if (flush_buf) { memset(flush_buf, 0xAB, flush_size); }
+        const double t = time_graph_compute(backend, gf);
+        if (t < best) best = t;
+    }
+    return best;
+}
+
+// Fetch the repack extra buffer type through the public proc-address API
+static ggml_backend_buffer_type_t get_cpu_repack_buft(void) {
+    ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
+    if (!cpu_dev) {
+        return NULL;
+    }
+    ggml_backend_reg_t cpu_reg = ggml_backend_dev_backend_reg(cpu_dev);
+    ggml_backend_dev_get_extra_bufts_t get_extra =
+        (ggml_backend_dev_get_extra_bufts_t) ggml_backend_reg_get_proc_address(cpu_reg, "ggml_backend_dev_get_extra_bufts");
+    if (!get_extra) {
+        return NULL;
+    }
+    ggml_backend_buffer_type_t * bufts = get_extra(cpu_dev);
+    // the only extra buffer type the CPU backend exposes is CPU_REPACK
+    return (bufts && *bufts) ? *bufts : NULL;
+}
+
+struct bench_row {
+    const char * name;
+    double time_std, time_repack, time_tiled;
+    float max_err_repack, rmse_repack;
+    float max_err_tiled, rmse_tiled;
+    bool have_repack;
+};
+
+// MUL_MAT three-way bench: std (default optimized GEMM) vs repack (CPU_REPACK buffer) vs
+// tiled kernel. std is timed in-process with use_ref=true (bypasses the tiled gate); tiled
+// and repack with use_ref=false (repack's tiled gate declines on the buffer's extra). repack
+// is n/a where no repack kernel exists for the type. Errors are vs the vec-only reference (use_ref).
+static bench_row bench_three_way(ggml_backend_t backend, int64_t M, int64_t N, int64_t K, ggml_type quant_type) {
+    bench_row row;
+    row.name = ggml_type_name(quant_type);
+    row.time_std = row.time_repack = row.time_tiled = 0.0;
+    row.max_err_repack = row.rmse_repack = row.max_err_tiled = row.rmse_tiled = 0.0f;
+    row.have_repack = false;
+
+    struct ggml_init_params ip = { 1024*1024*1024, nullptr, true };
+    struct ggml_context * ctx = ggml_init(ip);
+
+    struct ggml_tensor * src1     = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, N, M);
+    struct ggml_tensor * src0_std = ggml_new_tensor_2d(ctx, quant_type,  N, K);
+    struct ggml_tensor * src0_rep = ggml_new_tensor_2d(ctx, quant_type,  N, K);
+
+    struct ggml_cgraph * gf = ggml_new_graph(ctx);
+    struct ggml_tensor * dst = ggml_mul_mat(ctx, src0_std, src1);
+    ggml_build_forward_expand(gf, dst);
+
+    struct ggml_cgraph * gf_repack = ggml_new_graph(ctx);
+    struct ggml_tensor * dst_repack = ggml_mul_mat(ctx, src0_rep, src1);
+    ggml_build_forward_expand(gf_repack, dst_repack);
+
+    ggml_backend_buffer_type_t repack_buft = get_cpu_repack_buft();
+    ggml_backend_buffer_t buf_rep = NULL;
+    if (repack_buft && N % 8 == 0 && K % 8 == 0) {
+        buf_rep = ggml_backend_buft_alloc_buffer(repack_buft, ggml_nbytes(src0_rep));
+        src0_rep->buffer = buf_rep;
+        src0_rep->data   = ggml_backend_buffer_get_base(buf_rep);
+        ggml_backend_buffer_init_tensor(buf_rep, src0_rep);
+        row.have_repack = (src0_rep->extra != NULL);
+    }
+
+    ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend);
+
+    srand(0xBEEF);
+    float * src1_data = gen_rand_f32(M * N);
+    float * src0_data = gen_rand_f32(N * K);
+    fill_tensor(src1, src1_data, M, N, GGML_TYPE_F32);
+    fill_tensor(src0_std, src0_data, K, N, quant_type);
+    if (row.have_repack) {
+        fill_tensor(src0_rep, src0_data, K, N, quant_type);
+    }
+    free(src1_data); free(src0_data);
+
+    // one warmup + N timings; best (min) wins. The flush buffer is larger than the
+    // whole-chip cache, so its memset evicts L3 and each timed iteration streams DRAM.
+    const size_t flush_size = 256 * 1024 * 1024;
+    void * flush_buf = malloc(flush_size);
+    const int n_reps = 5;
+
+    // std in-process (use_ref bypasses the tiled gate; the vec_dot is the same optimized one)
+    cpu_set_use_ref(backend, true);
+    row.time_std = time_graph_compute_best(backend, gf, n_reps, flush_buf, flush_size);
+
+    // tiled in-process (the parent holds TILED_MM=1 forced on)
+    cpu_set_use_ref(backend, false);
+    row.time_tiled = time_graph_compute_best(backend, gf, n_reps, flush_buf, flush_size);
+
+    // repack in-process (the tiled gate declines on the buffer's extra)
+    if (row.have_repack) {
+        row.time_repack = time_graph_compute_best(backend, gf_repack, n_reps, flush_buf, flush_size);
+    }
+    free(flush_buf);
+
+    // errors vs the vec-only reference (use_ref)
+    float * out_ref    = (float *) malloc(M * K * sizeof(float));
+    float * out_tiled  = (float *) malloc(M * K * sizeof(float));
+    float * out_repack = (float *) malloc(M * K * sizeof(float));
+    cpu_set_use_ref(backend, true);
+    ggml_backend_graph_compute(backend, gf);
+    ggml_backend_tensor_get(dst, out_ref, 0, ggml_nbytes(dst));
+    cpu_set_use_ref(backend, false);
+    ggml_backend_graph_compute(backend, gf);
+    ggml_backend_tensor_get(dst, out_tiled, 0, ggml_nbytes(dst));
+    compare_f32(out_ref, out_tiled, M * K, &row.max_err_tiled, &row.rmse_tiled);
+    if (row.have_repack) {
+        ggml_backend_graph_compute(backend, gf_repack);
+        ggml_backend_tensor_get(dst_repack, out_repack, 0, ggml_nbytes(dst_repack));
+        compare_f32(out_ref, out_repack, M * K, &row.max_err_repack, &row.rmse_repack);
+    }
+    free(out_ref); free(out_tiled); free(out_repack);
+
+    ggml_backend_buffer_free(buf_rep);
+    ggml_backend_buffer_free(buf);
+    ggml_free(ctx);
+    return row;
+}
+
+static void print_bench_table(int64_t M, int64_t N, int64_t K, const bench_row * rows, size_t n_types) {
+    const double flops = 2.0 * M * N * K;
+
+    printf("\nBENCH dense %lldx%lld * %lldx%lld (std vs repack vs tiled), min of 5 timings, 8 threads\n",
+           (long long)M, (long long)N, (long long)N, (long long)K);
+    printf("%-8s %10s %12s %12s %11s %11s %17s %17s %17s %17s\n",
+           "type", "std TF", "repack TF", "tiled TF", "repack/std", "tiled/std",
+           "max_err(repack)", "rmse(repack)", "max_err(tiled)", "rmse(tiled)");
+    for (size_t i = 0; i < n_types; ++i) {
+        const bench_row * r = &rows[i];
+        if (r->have_repack) {
+            printf("%-8s %10.3f %12.3f %12.3f %11.2f %11.2f %17.5e %17.5e %17.5e %17.5e\n",
+                   r->name,
+                   flops / (r->time_std * 1e12),
+                   flops / (r->time_repack * 1e12), flops / (r->time_tiled * 1e12),
+                   r->time_std / r->time_repack, r->time_std / r->time_tiled,
+                   r->max_err_repack, r->rmse_repack, r->max_err_tiled, r->rmse_tiled);
+        } else {
+            printf("%-8s %10.3f %12s %12.3f %11s %11.2f %17s %17s %17.5e %17.5e\n",
+                   r->name,
+                   flops / (r->time_std * 1e12),
+                   "n/a", flops / (r->time_tiled * 1e12), "n/a",
+                   r->time_std / r->time_tiled,
+                   "n/a", "n/a", r->max_err_tiled, r->rmse_tiled);
+        }
+    }
+}
+
+struct bench_row_mmid {
+    const char * name;
+    double time_std, time_repack, time_tiled;
+    float max_err_repack, rmse_repack;
+    float max_err_tiled, rmse_tiled;
+    bool have_repack;
+};
+
+// MUL_MAT_ID (MoE) three-way bench: std (default GEMM, use_ref) vs repack (CPU_REPACK buffer)
+// vs tiled kernel (default; TILED_MM_FORCE is set in main so ragged experts take it too).
+// Column (id, t) = GEMV of src1 row (id % b_slots + t * b_slots) against expert ids[t*k+id];
+// cne1 = k * batch / n_experts rows per expert on average. repack is n/a where no repack
+// kernel exists. Errors are vs the vec-only reference (use_ref).
+static bench_row_mmid bench_mul_mat_id(ggml_backend_t backend, int64_t K, int64_t R, int64_t n_experts,
+                                       int64_t k, int64_t b_slots, int64_t batch, ggml_type quant_type) {
+    bench_row_mmid row;
+    row.name = ggml_type_name(quant_type);
+    row.time_std = row.time_repack = row.time_tiled = 0.0;
+    row.max_err_repack = row.rmse_repack = row.max_err_tiled = row.rmse_tiled = 0.0f;
+    row.have_repack = false;
+
+    struct ggml_init_params ip = { 1024*1024*1024, nullptr, true };
+    struct ggml_context * ctx = ggml_init(ip);
+
+    int64_t ne_as[4]  = { K, R, n_experts, 1 };
+    int64_t ne_b[4]   = { K, b_slots, batch, 1 };
+    int64_t ne_ids[4] = { k, batch, 1, 1 };
+    struct ggml_tensor * src1     = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_b);
+    struct ggml_tensor * src0_std = ggml_new_tensor(ctx, quant_type,  4, ne_as);
+    struct ggml_tensor * src0_rep = ggml_new_tensor(ctx, quant_type,  4, ne_as);
+    struct ggml_tensor * ids_t    = ggml_new_tensor(ctx, GGML_TYPE_I32, 4, ne_ids);
+
+    struct ggml_cgraph * gf = ggml_new_graph(ctx);
+    struct ggml_tensor * dst = ggml_mul_mat_id(ctx, src0_std, src1, ids_t);
+    ggml_build_forward_expand(gf, dst);
+
+    struct ggml_cgraph * gf_repack = ggml_new_graph(ctx);
+    struct ggml_tensor * dst_repack = ggml_mul_mat_id(ctx, src0_rep, src1, ids_t);
+    ggml_build_forward_expand(gf_repack, dst_repack);
+
+    ggml_backend_buffer_type_t repack_buft = get_cpu_repack_buft();
+    ggml_backend_buffer_t buf_rep = NULL;
+    if (repack_buft && R % 8 == 0) {
+        buf_rep = ggml_backend_buft_alloc_buffer(repack_buft, ggml_nbytes(src0_rep));
+        src0_rep->buffer = buf_rep;
+        src0_rep->data   = ggml_backend_buffer_get_base(buf_rep);
+        ggml_backend_buffer_init_tensor(buf_rep, src0_rep);
+        row.have_repack = (src0_rep->extra != NULL);
+    }
+
+    ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend);
+
+    // deterministic balanced routing: each expert gets ~ k*batch/n_experts rows
+    int32_t * ids = (int32_t *) malloc(k * batch * sizeof(int32_t));
+    for (int64_t t = 0; t < batch; t++) {
+        for (int64_t e = 0; e < k; e++) {
+            ids[t * k + e] = (int32_t) ((t * k + e) % n_experts);
+        }
+    }
+    srand(0xBEEF);
+    float * b_data  = gen_rand_f32(K * b_slots * batch);
+    float * as_data = gen_rand_f32(K * R * n_experts);
+    fill_tensor(src1, b_data, b_slots * batch, K, GGML_TYPE_F32);
+    fill_tensor(src0_std, as_data, R * n_experts, K, quant_type);
+    if (row.have_repack) {
+        fill_tensor(src0_rep, as_data, R * n_experts, K, quant_type);
+    }
+    ggml_backend_tensor_set(ids_t, ids, 0, ggml_nbytes(ids_t));
+    free(b_data); free(as_data);
+
+    const size_t flush_size = 256 * 1024 * 1024;
+    void * flush_buf = malloc(flush_size);
+    const int n_reps = 5;
+
+    // std in-process (use_ref bypasses the tiled gate; the vec_dot is the same optimized one)
+    cpu_set_use_ref(backend, true);
+    row.time_std = time_graph_compute_best(backend, gf, n_reps, flush_buf, flush_size);
+
+    // tiled in-process (the parent holds TILED_MM=1 forced on)
+    cpu_set_use_ref(backend, false);
+    row.time_tiled = time_graph_compute_best(backend, gf, n_reps, flush_buf, flush_size);
+
+    // repack in-process (the tiled gate declines on the buffer's extra)
+    if (row.have_repack) {
+        row.time_repack = time_graph_compute_best(backend, gf_repack, n_reps, flush_buf, flush_size);
+    }
+    free(flush_buf);
+
+    // errors vs the vec-only reference (use_ref)
+    float * out_ref    = (float *) malloc(R * k * batch * sizeof(float));
+    float * out_tiled  = (float *) malloc(R * k * batch * sizeof(float));
+    float * out_repack = (float *) malloc(R * k * batch * sizeof(float));
+    cpu_set_use_ref(backend, true);
+    ggml_backend_graph_compute(backend, gf);
+    ggml_backend_tensor_get(dst, out_ref, 0, ggml_nbytes(dst));
+    cpu_set_use_ref(backend, false);
+    ggml_backend_graph_compute(backend, gf);
+    ggml_backend_tensor_get(dst, out_tiled, 0, ggml_nbytes(dst));
+    compare_f32(out_ref, out_tiled, R * k * batch, &row.max_err_tiled, &row.rmse_tiled);
+    if (row.have_repack) {
+        ggml_backend_graph_compute(backend, gf_repack);
+        ggml_backend_tensor_get(dst_repack, out_repack, 0, ggml_nbytes(dst_repack));
+        compare_f32(out_ref, out_repack, R * k * batch, &row.max_err_repack, &row.rmse_repack);
+    }
+    free(out_ref); free(out_tiled); free(out_repack); free(ids);
+
+    ggml_backend_buffer_free(buf_rep);
+    ggml_backend_buffer_free(buf);
+    ggml_free(ctx);
+    return row;
+}
+
+static void print_mmid_table(int64_t K, int64_t R, int64_t n_experts, int64_t k,
+                             int64_t b_slots, int64_t batch, const bench_row_mmid * rows, size_t n_types) {
+    const double flops = 2.0 * (double) K * R * k * batch;
+
+    printf("\nBENCH mmid %lldx%lldx%lld k=%lld b_slots=%lld batch=%lld (cne1=%lld, std vs repack vs tiled), min of 5 timings, 8 threads\n",
+           (long long)K, (long long)R, (long long)n_experts, (long long)k,
+           (long long)b_slots, (long long)batch, (long long)(k * batch / n_experts));
+    printf("%-8s %10s %12s %12s %11s %11s %17s %17s %17s %17s\n",
+           "type", "std TF", "repack TF", "tiled TF", "repack/std", "tiled/std",
+           "max_err(repack)", "rmse(repack)", "max_err(tiled)", "rmse(tiled)");
+    for (size_t i = 0; i < n_types; ++i) {
+        const bench_row_mmid * r = &rows[i];
+        if (r->have_repack) {
+            printf("%-8s %10.3f %12.3f %12.3f %11.2f %11.2f %17.5e %17.5e %17.5e %17.5e\n",
+                   r->name,
+                   flops / (r->time_std * 1e12),
+                   flops / (r->time_repack * 1e12), flops / (r->time_tiled * 1e12),
+                   r->time_std / r->time_repack, r->time_std / r->time_tiled,
+                   r->max_err_repack, r->rmse_repack, r->max_err_tiled, r->rmse_tiled);
+        } else {
+            printf("%-8s %10.3f %12s %12.3f %11s %11.2f %17s %17s %17.5e %17.5e\n",
+                   r->name,
+                   flops / (r->time_std * 1e12),
+                   "n/a", flops / (r->time_tiled * 1e12), "n/a",
+                   r->time_std / r->time_tiled,
+                   "n/a", "n/a", r->max_err_tiled, r->rmse_tiled);
+        }
+    }
+}
+
+int main(int argc, char ** argv) {
+    bool run_bench = false;
+    bool run_fuzz = false;
+    for (int i = 1; i < argc; ++i) {
+        if (strcmp(argv[i], "--bench") == 0) {
+            run_bench = true;
+        } else if (strcmp(argv[i], "--fuzz") == 0) {
+            run_fuzz = true;
+        } else {
+            fprintf(stderr, "error: unknown argument: %s\n", argv[i]);
+            return 1;
+        }
+    }
+
+    // Enable tiled MM, also force tiled MM even when unprofitable for benchmarks
+#if defined(_WIN32)
+    _putenv("GGML_CPU_TILED_MM=1");
+    _putenv("GGML_CPU_TILED_MM_FORCE=1");
+#else
+    setenv("GGML_CPU_TILED_MM", "1", 1);
+    setenv("GGML_CPU_TILED_MM_FORCE", "1", 1);
+#endif
+
+    ggml_backend_load_all();
+    ggml_backend_t backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, NULL);
+    if (!backend) {
+        fprintf(stderr, "failed to initialize the CPU backend\n");
+        return 1;
+    }
+    cpu_set_n_threads(backend, 8);
+    // base case, smoke each quant type
+    test_matmul(backend, 256, 512, 256, GGML_TYPE_Q6_K);
+    test_matmul(backend, 256, 512, 256, GGML_TYPE_Q5_K);
+    test_matmul(backend, 256, 512, 256, GGML_TYPE_Q4_K);
+    test_matmul(backend, 256, 512, 256, GGML_TYPE_Q3_K);
+    test_matmul(backend, 256, 512, 256, GGML_TYPE_Q2_K);
+    test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ4_XS);
+    test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ2_XXS);
+    test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ2_XS);
+    test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ2_S);
+    test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ3_XXS);
+    test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ3_S);
+    test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ1_S);
+    test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ1_M);
+    test_mul_mat_id(backend,  256,  256,   2, 8, 1,  64, GGML_TYPE_Q4_K);  // single block, cne1 = 256
+    test_mul_mat_id(backend,  512,  300,   4, 2, 1, 150, GGML_TYPE_IQ2_XXS); // iq type, imatrix quant
+
+    if (run_fuzz) {
+        printf("Running fuzz tests.\n");
+        // ragged edges, both subblock lengths (Q5_K = 32, Q6_K = 16)
+        test_matmul(backend, 256, 1024, 8192, GGML_TYPE_Q6_K);   // long K, int32 accumulation
+        test_matmul(backend, 357, 1024, 137, GGML_TYPE_Q6_K);    // ragged M and K
+        test_matmul(backend, 16, 1024, 16, GGML_TYPE_Q5_K);
+        test_matmul(backend, 18, 1024, 7, GGML_TYPE_Q5_K);       // ragged M tail, ragged K in first subblock
+        test_matmul(backend, 8, 256, 8, GGML_TYPE_Q5_K);
+        test_matmul(backend, 18, 1024, 256, GGML_TYPE_Q5_K);
+        test_matmul(backend, 9, 256, 256, GGML_TYPE_Q5_K);
+        test_matmul(backend, 8, 256, 8, GGML_TYPE_Q6_K);         // tiny M, single QK_K block
+        test_matmul(backend, 256, 1024, 8192, GGML_TYPE_Q3_K);
+        test_matmul(backend, 357, 1024, 137, GGML_TYPE_Q3_K);
+        test_matmul(backend, 8, 256, 8, GGML_TYPE_Q3_K);
+        test_matmul(backend, 17, 1024, 257, GGML_TYPE_Q3_K);     // K past a full 256 tile
+        test_matmul(backend, 257, 1024, 17, GGML_TYPE_Q3_K);     // K past a full microtile
+        test_matmul(backend, 256, 1024, 8192, GGML_TYPE_Q2_K);
+        test_matmul(backend, 357, 1024, 137, GGML_TYPE_Q2_K);
+        test_matmul(backend, 8, 256, 8, GGML_TYPE_Q2_K);
+        test_matmul(backend, 17, 1024, 257, GGML_TYPE_Q2_K);
+        test_matmul(backend, 257, 1024, 17, GGML_TYPE_Q2_K);
+        test_matmul(backend, 357, 1024, 137, GGML_TYPE_IQ4_XS);    // ragged M and K
+        test_matmul(backend,   8,  256,   8, GGML_TYPE_IQ4_XS);    // tiny M, single QK_K block
+        test_matmul(backend,  17, 1024, 257, GGML_TYPE_IQ4_XS);    // K past a full 256 tile
+
+        // fuzz: M/K around the microtile (16) and tile (256) boundaries
+        // Q4_K = subblock 32, Q6_K = subblock 16
+        test_matmul(backend, 1, 1024, 1, GGML_TYPE_Q4_K);
+        test_matmul(backend, 2, 1024, 3, GGML_TYPE_Q4_K);
+        test_matmul(backend, 15, 1024, 15, GGML_TYPE_Q4_K);      // M and K in the first microtile
+        test_matmul(backend, 17, 1024, 17, GGML_TYPE_Q4_K);      // M and K past the first microtile
+        test_matmul(backend, 255, 1024, 255, GGML_TYPE_Q4_K);    // M and K one short of a tile
+        test_matmul(backend, 257, 1024, 257, GGML_TYPE_Q4_K);    // M and K one past a tile
+        test_matmul(backend, 271, 1024, 271, GGML_TYPE_Q4_K);    // 1 tile + 15
+        test_matmul(backend, 272, 1024, 272, GGML_TYPE_Q4_K);    // 1 tile + 16
+        test_matmul(backend, 511, 1024, 511, GGML_TYPE_Q4_K);    // 2 tiles - 1
+        test_matmul(backend, 513, 1024, 513, GGML_TYPE_Q4_K);    // 2 tiles + 1
+        test_matmul(backend, 17, 512, 257, GGML_TYPE_Q4_K);      // ragged M and K, N = 2 blocks
+
+        test_matmul(backend, 1, 1024, 1, GGML_TYPE_Q6_K);
+        test_matmul(backend, 2, 1024, 3, GGML_TYPE_Q6_K);
+        test_matmul(backend, 15, 1024, 15, GGML_TYPE_Q6_K);
+        test_matmul(backend, 17, 1024, 17, GGML_TYPE_Q6_K);
+        test_matmul(backend, 255, 1024, 255, GGML_TYPE_Q6_K);
+        test_matmul(backend, 257, 1024, 257, GGML_TYPE_Q6_K);
+        test_matmul(backend, 271, 1024, 271, GGML_TYPE_Q6_K);
+        test_matmul(backend, 272, 1024, 272, GGML_TYPE_Q6_K);
+        test_matmul(backend, 511, 1024, 511, GGML_TYPE_Q6_K);
+        test_matmul(backend, 513, 1024, 513, GGML_TYPE_Q6_K);
+        test_matmul(backend, 17, 512, 257, GGML_TYPE_Q6_K);
+
+        // MUL_MAT_ID (MoE): K = reduction (tiled gate needs K % 256 == 0), R = output rows per expert,
+        // k = top-k slots, b_slots = b rows (1 = broadcast MoE, k = i11-diverse gather), cne1 = k*batch/E
+        test_mul_mat_id(backend,  512,  300,   4, 2, 1, 150, GGML_TYPE_Q4_K);  // ragged R, broadcast, cne1 = 75
+        test_mul_mat_id(backend,  512,  256,   4, 2, 2, 150, GGML_TYPE_Q5_K);  // i11-diverse gather, HAS_MIN
+        test_mul_mat_id(backend,  512,  256,   4, 2, 1, 100, GGML_TYPE_Q6_K);  // ragged cne1 = 50
+        test_mul_mat_id(backend,  768,  256,   4, 2, 1, 150, GGML_TYPE_Q3_K);  // K = 3 slabs
+        test_mul_mat_id(backend, 1024, 2048, 128, 8, 1, 2048, GGML_TYPE_Q4_K); // large, cne1 = 128
+        test_mul_mat_id(backend,  512,  256,  32, 2, 1,   8, GGML_TYPE_Q4_K);  // tiny cne1 (~0.5), forced only
+
+        // cne1 > 256: the k-outer ring sweeps multiple 256-row windows per expert (the multi-window path)
+        test_mul_mat_id(backend,  256,  256,   8, 8, 1,  512, GGML_TYPE_Q4_K);  // cne1 = 512, two full windows
+        test_mul_mat_id(backend,  512,  256,   8, 8, 1,  300, GGML_TYPE_Q4_K);  // cne1 = 300, 2nd window ragged (44 rows)
+        test_mul_mat_id(backend,  256,   64,   8, 8, 1,  257, GGML_TYPE_Q4_K);  // cne1 = 257, 1-row tail window, R = 1 group
+        test_mul_mat_id(backend,  512,  128,   8, 8, 1,  511, GGML_TYPE_Q4_K);  // cne1 = 511, 15-row ragged tail
+        test_mul_mat_id(backend,  768,  300,   8, 8, 1,  300, GGML_TYPE_Q4_K);  // 3 K slabs x 2 windows, ragged R and cne1
+        test_mul_mat_id(backend,  768,  300,   8, 2, 2,  300, GGML_TYPE_Q5_K);  // cne1 = 600, 3 slabs, 3 windows (last ragged), i11-diverse, HAS_MIN
+
+        // cne1 = 255: one short of a full ring, last ring row zero-padded by the unpack
+        test_mul_mat_id(backend,  512,  256,   4, 4, 1,  255, GGML_TYPE_Q4_K);
+
+        // long K: 8 slabs per ring row
+        test_mul_mat_id(backend, 2048,  300,   8, 2, 1,  400, GGML_TYPE_Q4_K);  // cne1 = 100
+
+        // R group edges: 64 = TILED_MMID_GROUP (ceil div, ragged group tail, single-group early return)
+        test_mul_mat_id(backend,  512,   64,   4, 2, 1,  128, GGML_TYPE_Q4_K);  // ngroups = 1, most threads idle
+        test_mul_mat_id(backend,  512,   65,   4, 2, 1,  128, GGML_TYPE_Q4_K);  // ceil-div boundary, 1-row tail group
+        test_mul_mat_id(backend,  512,   63,   4, 2, 1,  128, GGML_TYPE_Q4_K);  // ragged group tail
+        test_mul_mat_id(backend,  512,    1,   4, 1, 1,  256, GGML_TYPE_Q4_K);  // R = 1 row, cne1 = 256
+
+        // strided (non-contiguous) F32 src1: nb[1] = 4 bytes, the tiled path reads through strides
+        test_mul_mat_id(backend,  512,  256,   8, 8, 1,  128, GGML_TYPE_Q4_K, true);  // cne1 = 128
+
+        // narrow path (cne1 in [8,16], K a multiple of k_extent): the weight is read as long per-row K streams
+        test_mul_mat_id(backend, 1024,  256,   2, 2, 1,  16, GGML_TYPE_Q4_K);  // cne1 = 16
+        test_mul_mat_id(backend, 1024,  256,   2, 1, 1,  16, GGML_TYPE_Q4_K);  // cne1 = 8
+        test_mul_mat_id(backend,  512,  128,   2, 3, 1,   8, GGML_TYPE_Q5_K);  // cne1 = 12
+
+        // higher-dim (ne[2], ne[3] > 1)
+        // equal-batch and broadcast shapes (src1_2=2*src0_2, src1_3=2*src0_3)
+        test_matmul_highdim(backend, 1024, 1024, 1024, 2, 1, 2, 1, GGML_TYPE_Q4_K); // 3D, tiled
+        test_matmul_highdim(backend, 1024, 1024, 1024, 1, 2, 1, 2, GGML_TYPE_Q4_K); // 3D, tiled
+        test_matmul_highdim(backend,  512, 1024,  512, 2, 2, 2, 2, GGML_TYPE_Q4_K); // 4D, tiled
+        test_matmul_highdim(backend, 1024, 1024, 1024, 2, 1, 1, 1, GGML_TYPE_Q4_K); // broadcast r2=2 (src1_2>src0_2)
+        test_matmul_highdim(backend, 1024, 1024, 1024, 2, 2, 2, 1, GGML_TYPE_Q4_K); // broadcast r3=2 (src1_3>src0_3)
+        test_matmul_highdim(backend, 1024, 1024, 1024, 2, 1, 2, 1, GGML_TYPE_Q6_K); // 3D, tiled, q6_K
+        test_matmul_highdim(backend, 1024, 1024, 1024, 2, 2, 2, 2, GGML_TYPE_Q3_K); // 4D, tiled, q3_K
+        test_matmul_highdim(backend, 1024, 1024, 1024, 2, 2, 2, 2, GGML_TYPE_Q2_K); // 4D, tiled, q2_K
+    }
+
+
+    // benchmark: std (use_ref) vs repack (CPU_REPACK buffer) vs tiled kernel, over all tiled
+    // types. std is timed with use_ref=true (bypasses the tiled gate); tiled and repack with
+    // use_ref=false. Errors are vs the vec-only reference (use_ref).
+    if (run_bench) {
+        printf("Starting benchmark runs.\n");
+        const ggml_type bench_types[] = { GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K,
+                                          GGML_TYPE_IQ4_XS, GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XS, GGML_TYPE_IQ2_S,
+                                          GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ3_S, GGML_TYPE_IQ1_S, GGML_TYPE_IQ1_M };
+        const size_t n_types = sizeof(bench_types) / sizeof(bench_types[0]);
+
+        struct { int64_t M, N, K; } shapes[] = {
+            { 4096, 4096, 4096 },
+            { 4096, 4096,   64 },
+            { 4096, 4096,   32 },
+            { 4096, 4096,   24 },
+            { 4096, 4096,   16 },
+            { 4096, 4096,   8 },
+            { 4096, 4096,   1 },
+        };
+        for (size_t s = 0; s < sizeof(shapes) / sizeof(shapes[0]); ++s) {
+            bench_row rows[n_types];
+            for (size_t i = 0; i < n_types; ++i) {
+                rows[i] = bench_three_way(backend, shapes[s].M, shapes[s].N, shapes[s].K, bench_types[i]);
+            }
+            print_bench_table(shapes[s].M, shapes[s].N, shapes[s].K, rows, n_types);
+        }
+
+        // MUL_MAT_ID (MoE): K must be a multiple of 256 (the tiled slab). cne1 = k*batch/E.
+        struct { int64_t K, R, E, k, b_slots, batch; } mmid_shapes[] = {
+            { 1024, 1024,  8,   8, 1,    1 },  // cne1 = 1, single-token decode
+            { 1024, 1024,  4,   2, 1,   16 },  // cne1 = 8, narrow path
+            { 1024, 1024,  4,   2, 1,   32 },  // cne1 = 16, narrow path
+            {  512,  512,  8,   2, 1,  128 },  // cne1 = 32
+            { 1024, 1024, 16,   8, 1,   64 },  // cne1 = 32
+            { 1024, 1024, 16,   8, 1,  256 },  // cne1 = 128
+            { 1024, 2048, 32,   8, 1,  512 },  // cne1 = 128, wide experts
+            { 2048, 1024, 16,   8, 1, 1024 },  // cne1 = 512, long dot
+        };
+        for (size_t s = 0; s < sizeof(mmid_shapes) / sizeof(mmid_shapes[0]); ++s) {
+            bench_row_mmid rows[n_types];
+            for (size_t i = 0; i < n_types; ++i) {
+                rows[i] = bench_mul_mat_id(backend, mmid_shapes[s].K, mmid_shapes[s].R, mmid_shapes[s].E,
+                                           mmid_shapes[s].k, mmid_shapes[s].b_slots, mmid_shapes[s].batch, bench_types[i]);
+            }
+            print_mmid_table(mmid_shapes[s].K, mmid_shapes[s].R, mmid_shapes[s].E, mmid_shapes[s].k,
+                             mmid_shapes[s].b_slots, mmid_shapes[s].batch, rows, n_types);
+        }
+    }
+
+    ggml_backend_free(backend);
+    return n_failed ? 1 : 0;
+}

tools/llama-bench/README.md

@@ -76,6 +76,7 @@ test parameters:
                                             (default: disabled)
   -nopo, --no-op-offload <0|1>              (default: 0)
   --no-host <0|1>                           (default: 0)
+  --repack <0|1>                            (default: 1)
 
 Multiple values can be given for each parameter by separating them with ','
 or by specifying the parameter multiple times. Ranges can be given as
@@ -210,7 +211,7 @@ $ ./llama-bench -o csv
 ```
 
 ```csv
-build_commit,build_number,cpu_info,gpu_info,backends,model_filename,model_type,model_size,model_n_params,n_batch,n_ubatch,n_threads,cpu_mask,cpu_strict,poll,type_k,type_v,n_gpu_layers,n_cpu_moe,split_mode,main_gpu,no_kv_offload,flash_attn,devices,tensor_split,tensor_buft_overrides,use_mmap,use_direct_io,embeddings,no_op_offload,no_host,fit_target,fit_min_ctx,n_prompt,n_gen,n_depth,test_time,avg_ns,stddev_ns,avg_ts,stddev_ts
+build_commit,build_number,cpu_info,gpu_info,backends,model_filename,model_type,model_size,model_n_params,n_batch,n_ubatch,n_threads,cpu_mask,cpu_strict,poll,type_k,type_v,n_gpu_layers,n_cpu_moe,split_mode,main_gpu,no_kv_offload,flash_attn,devices,tensor_split,tensor_buft_overrides,use_mmap,use_direct_io,embeddings,no_op_offload,no_host,repack,fit_target,fit_min_ctx,n_prompt,n_gen,n_depth,test_time,avg_ns,stddev_ns,avg_ts,stddev_ts
 "8cf427ff","5163","AMD Ryzen 7 7800X3D 8-Core Processor","NVIDIA GeForce RTX 4080","CUDA","models/Qwen2.5-7B-Instruct-Q4_K_M.gguf","qwen2 7B Q4_K - Medium","4677120000","7615616512","2048","512","8","0x0","0","50","f16","f16","-1","0","layer","0","0","-1","auto","0.00","none","1","0","0","0","0","0","0","512","0","0","2025-04-24T11:57:09Z","70285660","982040","7285.676949","100.064434"
 "8cf427ff","5163","AMD Ryzen 7 7800X3D 8-Core Processor","NVIDIA GeForce RTX 4080","CUDA","models/Qwen2.5-7B-Instruct-Q4_K_M.gguf","qwen2 7B Q4_K - Medium","4677120000","7615616512","2048","512","8","0x0","0","50","f16","f16","-1","0","layer","0","0","-1","auto","0.00","none","1","0","0","0","0","0","0","0","128","0","2025-04-24T11:57:10Z","1067431600","3834831","119.915244","0.430617"
 ```
@@ -255,6 +256,7 @@ $ ./llama-bench -o json
     "embeddings": false,
     "no_op_offload": 0,
     "no_host": false,
+    "repack": true,
     "fit_target": 0,
     "fit_min_ctx": 0,
     "n_prompt": 512,
@@ -300,6 +302,7 @@ $ ./llama-bench -o json
     "embeddings": false,
     "no_op_offload": 0,
     "no_host": false,
+    "repack": true,
     "fit_target": 0,
     "fit_min_ctx": 0,
     "n_prompt": 0,
@@ -324,8 +327,8 @@ $ ./llama-bench -o jsonl
 ```
 
 ```json lines
-{"build_commit": "8cf427ff", "build_number": 5163, "cpu_info": "AMD Ryzen 7 7800X3D 8-Core Processor", "gpu_info": "NVIDIA GeForce RTX 4080", "backends": "CUDA", "model_filename": "models/Qwen2.5-7B-Instruct-Q4_K_M.gguf", "model_type": "qwen2 7B Q4_K - Medium", "model_size": 4677120000, "model_n_params": 7615616512, "n_batch": 2048, "n_ubatch": 512, "n_threads": 8, "cpu_mask": "0x0", "cpu_strict": false, "poll": 50, "type_k": "f16", "type_v": "f16", "n_gpu_layers": -1, "n_cpu_moe": 0, "split_mode": "layer", "main_gpu": 0, "no_kv_offload": false, "flash_attn": -1, "devices": "auto", "tensor_split": "0.00", "tensor_buft_overrides": "none", "use_mmap": true, "use_direct_io": false, "embeddings": false, "no_op_offload": 0, "no_host": false, "fit_target": 0, "fit_min_ctx": 0, "n_prompt": 512, "n_gen": 0, "n_depth": 0, "test_time": "2025-04-24T11:59:33Z", "avg_ns": 70497220, "stddev_ns": 883196, "avg_ts": 7263.609157, "stddev_ts": 90.940578, "samples_ns": [ 71551000, 71222800, 70364100, 69439100, 69909100 ],"samples_ts": [ 7155.74, 7188.71, 7276.44, 7373.37, 7323.8 ]}
-{"build_commit": "8cf427ff", "build_number": 5163, "cpu_info": "AMD Ryzen 7 7800X3D 8-Core Processor", "gpu_info": "NVIDIA GeForce RTX 4080", "backends": "CUDA", "model_filename": "models/Qwen2.5-7B-Instruct-Q4_K_M.gguf", "model_type": "qwen2 7B Q4_K - Medium", "model_size": 4677120000, "model_n_params": 7615616512, "n_batch": 2048, "n_ubatch": 512, "n_threads": 8, "cpu_mask": "0x0", "cpu_strict": false, "poll": 50, "type_k": "f16", "type_v": "f16", "n_gpu_layers": -1, "n_cpu_moe": 0, "split_mode": "layer", "main_gpu": 0, "no_kv_offload": false, "flash_attn": -1, "devices": "auto", "tensor_split": "0.00", "tensor_buft_overrides": "none", "use_mmap": true, "use_direct_io": false, "embeddings": false, "no_op_offload": 0, "no_host": false, "fit_target": 0, "fit_min_ctx": 0, "n_prompt": 0, "n_gen": 128, "n_depth": 0, "test_time": "2025-04-24T11:59:33Z", "avg_ns": 1068078400, "stddev_ns": 6279455, "avg_ts": 119.844681, "stddev_ts": 0.699739, "samples_ns": [ 1066331700, 1064864900, 1079042600, 1063328400, 1066824400 ],"samples_ts": [ 120.038, 120.203, 118.624, 120.377, 119.982 ]}
+{"build_commit": "8cf427ff", "build_number": 5163, "cpu_info": "AMD Ryzen 7 7800X3D 8-Core Processor", "gpu_info": "NVIDIA GeForce RTX 4080", "backends": "CUDA", "model_filename": "models/Qwen2.5-7B-Instruct-Q4_K_M.gguf", "model_type": "qwen2 7B Q4_K - Medium", "model_size": 4677120000, "model_n_params": 7615616512, "n_batch": 2048, "n_ubatch": 512, "n_threads": 8, "cpu_mask": "0x0", "cpu_strict": false, "poll": 50, "type_k": "f16", "type_v": "f16", "n_gpu_layers": -1, "n_cpu_moe": 0, "split_mode": "layer", "main_gpu": 0, "no_kv_offload": false, "flash_attn": -1, "devices": "auto", "tensor_split": "0.00", "tensor_buft_overrides": "none", "use_mmap": true, "use_direct_io": false, "embeddings": false, "no_op_offload": 0, "no_host": false, "repack": true, "fit_target": 0, "fit_min_ctx": 0, "n_prompt": 512, "n_gen": 0, "n_depth": 0, "test_time": "2025-04-24T11:59:33Z", "avg_ns": 70497220, "stddev_ns": 883196, "avg_ts": 7263.609157, "stddev_ts": 90.940578, "samples_ns": [ 71551000, 71222800, 70364100, 69439100, 69909100 ],"samples_ts": [ 7155.74, 7188.71, 7276.44, 7373.37, 7323.8 ]}
+{"build_commit": "8cf427ff", "build_number": 5163, "cpu_info": "AMD Ryzen 7 7800X3D 8-Core Processor", "gpu_info": "NVIDIA GeForce RTX 4080", "backends": "CUDA", "model_filename": "models/Qwen2.5-7B-Instruct-Q4_K_M.gguf", "model_type": "qwen2 7B Q4_K - Medium", "model_size": 4677120000, "model_n_params": 7615616512, "n_batch": 2048, "n_ubatch": 512, "n_threads": 8, "cpu_mask": "0x0", "cpu_strict": false, "poll": 50, "type_k": "f16", "type_v": "f16", "n_gpu_layers": -1, "n_cpu_moe": 0, "split_mode": "layer", "main_gpu": 0, "no_kv_offload": false, "flash_attn": -1, "devices": "auto", "tensor_split": "0.00", "tensor_buft_overrides": "none", "use_mmap": true, "use_direct_io": false, "embeddings": false, "no_op_offload": 0, "no_host": false, "repack": true, "fit_target": 0, "fit_min_ctx": 0, "n_prompt": 0, "n_gen": 128, "n_depth": 0, "test_time": "2025-04-24T11:59:33Z", "avg_ns": 1068078400, "stddev_ns": 6279455, "avg_ts": 119.844681, "stddev_ts": 0.699739, "samples_ns": [ 1066331700, 1064864900, 1079042600, 1063328400, 1066824400 ],"samples_ts": [ 120.038, 120.203, 118.624, 120.377, 119.982 ]}
 ```
 
 
@@ -370,6 +373,7 @@ CREATE TABLE IF NOT EXISTS llama_bench (
   embeddings INTEGER,
   no_op_offload INTEGER,
   no_host INTEGER,
+  repack INTEGER,
   fit_target INTEGER,
   fit_min_ctx INTEGER,
   n_prompt INTEGER,
@@ -382,6 +386,6 @@ CREATE TABLE IF NOT EXISTS llama_bench (
   stddev_ts REAL
 );
 
-INSERT INTO llama_bench (build_commit, build_number, cpu_info, gpu_info, backends, model_filename, model_type, model_size, model_n_params, n_batch, n_ubatch, n_threads, cpu_mask, cpu_strict, poll, type_k, type_v, n_gpu_layers, n_cpu_moe, split_mode, main_gpu, no_kv_offload, flash_attn, devices, tensor_split, tensor_buft_overrides, use_mmap, use_direct_io, embeddings, no_op_offload, no_host, fit_target, fit_min_ctx, n_prompt, n_gen, n_depth, test_time, avg_ns, stddev_ns, avg_ts, stddev_ts) VALUES ('8cf427ff', '5163', 'AMD Ryzen 7 7800X3D 8-Core Processor', 'NVIDIA GeForce RTX 4080', 'CUDA', 'models/Qwen2.5-7B-Instruct-Q4_K_M.gguf', 'qwen2 7B Q4_K - Medium', '4677120000', '7615616512', '2048', '512', '8', '0x0', '0', '50', 'f16', 'f16', '-1', '0', 'layer', '0', '0', '-1', 'auto', '0.00', 'none', '1', '0', '0', '0', '0', '0', '0', '512', '0', '0', '2025-04-24T12:00:08Z', '69905000', '519516', '7324.546977', '54.032613');
-INSERT INTO llama_bench (build_commit, build_number, cpu_info, gpu_info, backends, model_filename, model_type, model_size, model_n_params, n_batch, n_ubatch, n_threads, cpu_mask, cpu_strict, poll, type_k, type_v, n_gpu_layers, n_cpu_moe, split_mode, main_gpu, no_kv_offload, flash_attn, devices, tensor_split, tensor_buft_overrides, use_mmap, use_direct_io, embeddings, no_op_offload, no_host, fit_target, fit_min_ctx, n_prompt, n_gen, n_depth, test_time, avg_ns, stddev_ns, avg_ts, stddev_ts) VALUES ('8cf427ff', '5163', 'AMD Ryzen 7 7800X3D 8-Core Processor', 'NVIDIA GeForce RTX 4080', 'CUDA', 'models/Qwen2.5-7B-Instruct-Q4_K_M.gguf', 'qwen2 7B Q4_K - Medium', '4677120000', '7615616512', '2048', '512', '8', '0x0', '0', '50', 'f16', 'f16', '-1', '0', 'layer', '0', '0', '-1', 'auto', '0.00', 'none', '1', '0', '0', '0', '0', '0', '0', '0', '128', '0', '2025-04-24T12:00:09Z', '1063608780', '4464130', '120.346696', '0.504647');
+INSERT INTO llama_bench (build_commit, build_number, cpu_info, gpu_info, backends, model_filename, model_type, model_size, model_n_params, n_batch, n_ubatch, n_threads, cpu_mask, cpu_strict, poll, type_k, type_v, n_gpu_layers, n_cpu_moe, split_mode, main_gpu, no_kv_offload, flash_attn, devices, tensor_split, tensor_buft_overrides, use_mmap, use_direct_io, embeddings, no_op_offload, no_host, repack, fit_target, fit_min_ctx, n_prompt, n_gen, n_depth, test_time, avg_ns, stddev_ns, avg_ts, stddev_ts) VALUES ('8cf427ff', '5163', 'AMD Ryzen 7 7800X3D 8-Core Processor', 'NVIDIA GeForce RTX 4080', 'CUDA', 'models/Qwen2.5-7B-Instruct-Q4_K_M.gguf', 'qwen2 7B Q4_K - Medium', '4677120000', '7615616512', '2048', '512', '8', '0x0', '0', '50', 'f16', 'f16', '-1', '0', 'layer', '0', '0', '-1', 'auto', '0.00', 'none', '1', '0', '0', '0', '0', '1', '0', '0', '512', '0', '0', '2025-04-24T12:00:08Z', '69905000', '519516', '7324.546977', '54.032613');
+INSERT INTO llama_bench (build_commit, build_number, cpu_info, gpu_info, backends, model_filename, model_type, model_size, model_n_params, n_batch, n_ubatch, n_threads, cpu_mask, cpu_strict, poll, type_k, type_v, n_gpu_layers, n_cpu_moe, split_mode, main_gpu, no_kv_offload, flash_attn, devices, tensor_split, tensor_buft_overrides, use_mmap, use_direct_io, embeddings, no_op_offload, no_host, repack, fit_target, fit_min_ctx, n_prompt, n_gen, n_depth, test_time, avg_ns, stddev_ns, avg_ts, stddev_ts) VALUES ('8cf427ff', '5163', 'AMD Ryzen 7 7800X3D 8-Core Processor', 'NVIDIA GeForce RTX 4080', 'CUDA', 'models/Qwen2.5-7B-Instruct-Q4_K_M.gguf', 'qwen2 7B Q4_K - Medium', '4677120000', '7615616512', '2048', '512', '8', '0x0', '0', '50', 'f16', 'f16', '-1', '0', 'layer', '0', '0', '-1', 'auto', '0.00', 'none', '1', '0', '0', '0', '0', '1', '0', '0', '0', '128', '0', '2025-04-24T12:00:09Z', '1063608780', '4464130', '120.346696', '0.504647');
 ```

tools/llama-bench/llama-bench.cpp

@@ -364,6 +364,7 @@ struct cmd_params {
     std::vector<bool>                embeddings;
     std::vector<bool>                no_op_offload;
     std::vector<bool>                no_host;
+    std::vector<bool>                repack;
     std::vector<size_t>              fit_params_target;
     std::vector<uint32_t>            fit_params_min_ctx;
     ggml_numa_strategy               numa;
@@ -409,6 +410,7 @@ static const cmd_params cmd_params_defaults = {
     /* embeddings           */ { false },
     /* no_op_offload        */ { false },
     /* no_host              */ { false },
+    /* repack               */ { llama_model_default_params().use_extra_bufts },
     /* fit_params_target    */ { 0 },
     /* fit_params_min_ctx   */ { 0 },
     /* numa                 */ GGML_NUMA_STRATEGY_DISABLED,
@@ -483,6 +485,7 @@ static void print_usage(int /* argc */, char ** argv) {
     printf("                                                    (default: disabled)\n");
     printf("  -nopo, --no-op-offload <0|1>                      (default: 0)\n");
     printf("  --no-host <0|1>                                   (default: %s)\n", join(cmd_params_defaults.no_host, ",").c_str());
+    printf("  --repack <0|1>                                    (default: %s)\n", join(cmd_params_defaults.repack, ",").c_str());
     printf("\n");
     printf(
         "Multiple values can be given for each parameter by separating them with ','\n"
@@ -906,6 +909,13 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
                 }
                 auto p = string_split<bool>(argv[i], split_delim);
                 params.no_host.insert(params.no_host.end(), p.begin(), p.end());
+            } else if (arg == "--repack") {
+                if (++i >= argc) {
+                    invalid_param = true;
+                    break;
+                }
+                auto p = string_split<bool>(argv[i], split_delim);
+                params.repack.insert(params.repack.end(), p.begin(), p.end());
             } else if (arg == "-ts" || arg == "--tensor-split") {
                 if (++i >= argc) {
                     invalid_param = true;
@@ -1087,7 +1097,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
             p.hf_token      = params.hf_token;
             p.offline       = params.offline;
             p.model.hf_repo = params.hf_repo[i];
-            if (!params.hf_file.empty() && !params.hf_file[i].empty()) {
+            if (i < params.hf_file.size() && !params.hf_file[i].empty()) {
                 p.model.hf_file = params.hf_file[i];
             }
 
@@ -1173,6 +1183,9 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
     if (params.no_host.empty()) {
         params.no_host = cmd_params_defaults.no_host;
     }
+    if (params.repack.empty()) {
+        params.repack = cmd_params_defaults.repack;
+    }
     if (params.n_threads.empty()) {
         params.n_threads = cmd_params_defaults.n_threads;
     }
@@ -1222,6 +1235,7 @@ struct cmd_params_instance {
     bool               embeddings;
     bool               no_op_offload;
     bool               no_host;
+    bool               repack;
     size_t             fit_target;
     uint32_t           fit_min_ctx;
 
@@ -1238,6 +1252,7 @@ struct cmd_params_instance {
         mparams.main_gpu      = main_gpu;
         mparams.tensor_split  = tensor_split.data();
         mparams.no_host       = no_host;
+        mparams.use_extra_bufts = repack;
 
         if (n_cpu_moe <= 0) {
             if (tensor_buft_overrides.empty()) {
@@ -1283,7 +1298,7 @@ struct cmd_params_instance {
                split_mode == other.split_mode &&
                main_gpu == other.main_gpu && tensor_split == other.tensor_split &&
                load_mode == other.load_mode && lazy_mode == other.lazy_mode &&
-               devices == other.devices && no_host == other.no_host &&
+               devices == other.devices && no_host == other.no_host && repack == other.repack &&
                vec_tensor_buft_override_equal(tensor_buft_overrides, other.tensor_buft_overrides);
     }
 
@@ -1323,6 +1338,7 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
     for (const auto & ts : params.tensor_split)
     for (const auto & ot : params.tensor_buft_overrides)
     for (const auto & noh : params.no_host)
+    for (const auto & rpk : params.repack)
     for (const auto & embd : params.embeddings)
     for (const auto & nopo : params.no_op_offload)
     for (const auto & nb : params.n_batch)
@@ -1367,6 +1383,7 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
                 /* .embeddings            = */ embd,
                 /* .no_op_offload         = */ nopo,
                 /* .no_host               = */ noh,
+                /* .repack                = */ rpk,
                 /* .fit_target            = */ fpt,
                 /* .fit_min_ctx           = */ fpc,
             };
@@ -1404,6 +1421,7 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
                 /* .embeddings            = */ embd,
                 /* .no_op_offload         = */ nopo,
                 /* .no_host               = */ noh,
+                /* .repack                = */ rpk,
                 /* .fit_target            = */ fpt,
                 /* .fit_min_ctx           = */ fpc,
             };
@@ -1441,6 +1459,7 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
                 /* .embeddings            = */ embd,
                 /* .no_op_offload         = */ nopo,
                 /* .no_host               = */ noh,
+                /* .repack                = */ rpk,
                 /* .fit_target            = */ fpt,
                 /* .fit_min_ctx           = */ fpc,
             };
@@ -1483,6 +1502,7 @@ struct test {
     bool                     embeddings;
     bool                     no_op_offload;
     bool                     no_host;
+    bool                     repack;
     size_t                   fit_target;
     uint32_t                 fit_min_ctx;
     int                      n_prompt;
@@ -1523,6 +1543,7 @@ struct test {
         embeddings     = inst.embeddings;
         no_op_offload  = inst.no_op_offload;
         no_host        = inst.no_host;
+        repack         = inst.repack;
         fit_target     = inst.fit_target;
         fit_min_ctx    = inst.fit_min_ctx;
         n_prompt       = inst.n_prompt;
@@ -1583,7 +1604,7 @@ struct test {
             "main_gpu",       "no_kv_offload",  "flash_attn",    "devices",        "tensor_split",
             "tensor_buft_overrides",            "load_mode",     "lazy_mode",
             "embeddings",
-            "no_op_offload",  "no_host",        "fit_target",    "fit_min_ctx",
+            "no_op_offload",  "no_host",        "repack",        "fit_target",    "fit_min_ctx",
             "n_prompt",       "n_gen",          "n_depth",
             "test_time",      "avg_ns",         "stddev_ns",     "avg_ts",         "stddev_ts"
         };
@@ -1601,7 +1622,7 @@ struct test {
             return INT;
         }
         if (field == "f16_kv" || field == "no_kv_offload" || field == "cpu_strict" ||
-            field == "embeddings" || field == "no_host") {
+            field == "embeddings" || field == "no_host" || field == "repack") {
             return BOOL;
         }
         if (field == "avg_ts" || field == "stddev_ts") {
@@ -1681,6 +1702,7 @@ struct test {
                                             std::to_string(embeddings),
                                             std::to_string(no_op_offload),
                                             std::to_string(no_host),
+                                            std::to_string(repack),
                                             std::to_string(fit_target),
                                             std::to_string(fit_min_ctx),
                                             std::to_string(n_prompt),
@@ -1874,6 +1896,9 @@ struct markdown_printer : public printer {
         if (field == "no_host") {
             return 4;
         }
+        if (field == "repack") {
+            return 3;
+        }
 
         int width = std::max((int) field.length(), 10);
 
@@ -1911,6 +1936,9 @@ struct markdown_printer : public printer {
         if (field == "no_host") {
             return "noh";
         }
+        if (field == "repack") {
+            return "rpk";
+        }
         if (field == "devices") {
             return "dev";
         }
@@ -2004,6 +2032,9 @@ struct markdown_printer : public printer {
         if (params.no_host.size() > 1 || params.no_host != cmd_params_defaults.no_host) {
             fields.emplace_back("no_host");
         }
+        if (params.repack.size() > 1 || params.repack != cmd_params_defaults.repack) {
+            fields.emplace_back("repack");
+        }
         if (params.fit_params_target.size() > 1 || params.fit_params_target != cmd_params_defaults.fit_params_target) {
             fields.emplace_back("fit_target");
         }

tools/server/server-common.h

@@ -643,7 +643,9 @@ struct server_subproc {
         void wait(const std::vector<server_subproc *> & procs, std::vector<bool> & ready, int64_t timeout_ms);
 
     private:
+#ifndef _WIN32
         intptr_t wake_fd[2] = { -1, -1 }; // POSIX self-pipe
+#endif
     };
 
 private:

tools/server/server.cpp

@@ -102,12 +102,12 @@ int llama_server(int argc, char ** argv) {
     // touch it. lifecycle is symmetric, stop_gc() runs in clean_up() before backend free
     server_stream_session_manager_start();
 
-    SRV_INF("%s", "initializing ...\n");
-
     if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SERVER)) {
         return 1;
     }
 
+    SRV_INF("%s", "initializing ...\n");
+
     llama_backend_init();
     llama_numa_init(params.numa);
 

vendor/stb/CMakeLists.txt

@@ -4,3 +4,7 @@ add_library(stb INTERFACE)
 add_library(vendor::stb ALIAS stb)
 
 target_include_directories(stb INTERFACE ..)
+
+if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
+    target_compile_options(stb INTERFACE -Wno-maybe-uninitialized)
+endif()