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

* spec : add probabilistic sampling for DFlash2 drafts

* spec : apply review suggestions

---------

Co-authored-by: Pranesh Gonegandla <pgonegandla@nvidia.com>

ci/run.sh

@@ -447,10 +447,14 @@ function gg_run_qwen3_0_6b {
 
     (time ./bin/llama-imatrix --model ${model_f16} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-imatrix.log
 
-    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa off --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
-    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa on  --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
-    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa off                ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
-    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa on                 ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+    # TODO: these tests are extremely slow on WebGPU - re-enable when fixed
+    #       https://github.com/ggml-org/llama.cpp/actions/runs/39639985732
+    if [ -z ${GG_BUILD_WEBGPU} ]; then
+        (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa off --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+        (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa on  --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+        (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa off                ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+        (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa on                 ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+    fi
 
     function check_ppl {
         qnt="$1"

common/jinja/lexer.cpp

@@ -1,7 +1,9 @@
 #include "lexer.h"
 #include "runtime.h"
+#include "unicode.h"
 
 #include <cctype>
+#include <charconv>
 #include <functional>
 #include <map>
 #include <string>
@@ -75,6 +77,24 @@ lexer_result lexer::tokenize(const std::string & source) {
                 }
                 // add escaped char
                 char escaped_char = src[pos++];
+                if (escaped_char == 'x') {
+                    if (pos + 1 >= src.size()) {
+                        throw lexer_exception("incomplete hexadecimal escape", source, pos);
+                    }
+
+                    unsigned int value = 0;
+                    const char * first = &src[pos];
+                    const char * last  = first + 2;
+
+                    auto [ptr, ec] = std::from_chars(first, last, value, 16);
+                    if (ec != std::errc{} || ptr != last) {
+                        throw lexer_exception("invalid hexadecimal escape", source, pos);
+                    }
+
+                    str += common_unicode_cpt_to_utf8(value);
+                    pos += 2;
+                    continue;
+                }
                 if (escape_chars.find(escaped_char) == escape_chars.end()) {
                     throw lexer_exception(std::string("unknown escape character \\") + escaped_char, source, pos);
                 }

common/jinja/value.cpp

@@ -425,7 +425,12 @@ const func_builtins & global_builtins() {
             throw raised_exception("Jinja Exception: " + msg);
         }},
         {"dict", toobject},
-        {"namespace", toobject},
+        {"namespace", [](const func_args & args) -> value {
+            // unlike a dict, a namespace has no methods: ns.items is the attribute, as in jinja2
+            value out = toobject(args);
+            cast_val<value_object>(out)->has_builtins = false;
+            return out;
+        }},
         {"strftime_now", [](const func_args & args) -> value {
             args.ensure_vals<value_string>();
             std::string format = args.get_pos(0)->as_string().str();

common/speculative.cpp

@@ -19,6 +19,7 @@
 #include <fstream>
 #include <iomanip>
 #include <map>
+#include <random>
 #include <cinttypes>
 
 #define SPC_DBG(fmt, ...) LOG_DBG("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)
@@ -1010,6 +1011,10 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
     bool    is_mrope       = false;
     int32_t selector_top_k = 0;
 
+    // DFlash2 probabilistic drafting: one generator per sequence and the seed it was last seeded with (-1: reseed)
+    std::vector<std::mt19937> dflash2_rng;
+    std::vector<int64_t>      dflash2_seed;
+
     // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout
     const bool is_dspark;
 
@@ -1060,6 +1065,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
 
         selector_top_k = llama_model_dflash_selector_top_k(model_dft);
         is_dflash2     = selector_top_k > 0;
+        dflash2_rng.resize(n_seq);
+        dflash2_seed.assign(n_seq, -1);
         char lilicorr_meta[32] = {};
         lilicorr_top_k = llama_model_meta_val_str(model_dft, "lilicorr.candidate_topk", lilicorr_meta, sizeof(lilicorr_meta)) >= 0
                 ? std::atoi(lilicorr_meta) : 0;
@@ -1158,6 +1165,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
             return;
         }
 
+        // reseed per request rather than per round, or two identical requests differ
+        dflash2_seed[seq_id] = -1;
+
         const int32_t N = (int32_t) prompt.size();
         if (N <= 0) {
             return;
@@ -1338,6 +1348,21 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
                 const float * lattice = llama_get_embeddings_nextn(ctx_dft);
                 GGML_ASSERT(lattice && "DFlash2 selector produced no lattice");
 
+                // DFlash2 can instead sample each step and hand q to the verifier for rejection sampling
+                const bool sample = is_dflash2 && params.probabilistic && dp.result_q;
+
+                // greedy keeps temperature 1, sampling uses the target's
+                const float temp = sample ? dp.temp : 1.0f;
+
+                if (sample && dflash2_seed[seq_id] != dp.seed) {
+                    dflash2_seed[seq_id] = dp.seed;
+                    // the default seed means random; mixed so it differs from the target's, as in spec_retune
+                    dflash2_rng[seq_id].seed(dp.seed == LLAMA_DEFAULT_SEED ? std::random_device{}() : dp.seed ^ 0x85ebca6bu);
+                }
+
+                std::vector<float>   probs;
+                std::vector<int32_t> sorted_idx;
+
                 int32_t predecessor = 0;
                 for (int32_t i = 1; i < (is_dflash2 ? n_block_tokens : params.n_max + 1); ++i) {
                     const float * row = lattice + (size_t) (beg + i) * n_embd_dec;
@@ -1346,21 +1371,53 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
 
                     predecessor = (int32_t) std::distance(scores,
                             std::max_element(scores, scores + top_k));
-                    if (params.p_min > 0.0f) {
-                        // softmax(scores) at the argmax, i.e. 1 / sum(exp(s_k - s_max))
+                    if (params.p_min > 0.0f || sample) {
+                        // softmax(scores / temp) at the argmax, i.e. 1 / sum(exp((s_k - s_max) / temp))
+                        probs.resize(top_k);
                         float sum = 0.0f;
                         for (int32_t k = 0; k < top_k; ++k) {
-                            sum += std::exp(scores[k] - scores[predecessor]);
+                            probs[k] = std::exp((scores[k] - scores[predecessor]) / temp);
+                            sum += probs[k];
                         }
                         if (1.0f / sum < params.p_min) {
                             break;
                         }
+                        for (int32_t k = 0; k < top_k; ++k) {
+                            probs[k] /= sum;
+                        }
+                    }
+                    if (sample) {
+                        // ggml_top_k leaves candidates unordered, so walk them by token id to keep a seed's draft stable
+                        sorted_idx.resize(top_k);
+                        for (int32_t k = 0; k < top_k; ++k) {
+                            sorted_idx[k] = k;
+                        }
+                        std::sort(sorted_idx.begin(), sorted_idx.end(), [row](int32_t a, int32_t b) { return row[a] < row[b]; });
+
+                        // a round-off overshoot keeps the argmax, which has q > 0
+                        double u = std::uniform_real_distribution<double>(0.0, 1.0)(dflash2_rng[seq_id]);
+                        for (const int32_t k : sorted_idx) {
+                            u -= probs[k];
+                            if (u < 0.0) {
+                                predecessor = k;
+                                break;
+                            }
+                        }
+
+                        auto & q = dp.result_q->emplace_back();
+                        q.reserve(top_k);
+                        for (int32_t k = 0; k < top_k; ++k) {
+                            q.push_back({ (llama_token) row[k], scores[k], probs[k] });
+                        }
                     }
                     result.push_back((llama_token) row[predecessor]);
                 }
 
                 if (result.size() < (size_t) params.n_min) {
                     result.clear();
+                    if (sample) {
+                        dp.result_q->clear();
+                    }
                 }
                 continue;
             }

common/speculative.h

@@ -73,7 +73,7 @@ struct common_speculative_draft_params {
     // the generated draft from the last _draft() call
     llama_tokens * result;
 
-    // candidate distribution per drafted token; set it to make draft-simple and draft-mtp sample
+    // candidate distribution per drafted token; set it to make draft-simple, draft-mtp and draft-dflash (DFlash2 only) sample
     std::vector<std::vector<llama_token_data>> * result_q = nullptr;
 
     // the target's temp and seed, read only when the drafter samples probabilistically

conversion/__init__.py

@@ -17,8 +17,10 @@ __all__ = [
 
 
 TEXT_MODEL_MAP: dict[str, str] = {
+    "NEOChatModel": "sensenova_u1",
     "AfmoeForCausalLM": "afmoe",
     "LagunaForCausalLM": "laguna",
+    "FlintForCausalLM": "flint",
     "ApertusForCausalLM": "llama",
     "ArceeForCausalLM": "llama",
     "ArcticForCausalLM": "arctic",
@@ -194,6 +196,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "MiniCPMForCausalLM": "minicpm",
     "MiniCPMV4_6ForConditionalGeneration": "minicpm",
     "MiniCPMV4_7ForConditionalGeneration": "minicpm",
+    "MiniCPMOcrForConditionalGeneration": "minicpm",
     "MiniMaxText01ForCausalLM": "minimax",
     "MiniMaxM1ForCausalLM": "minimax",
     "MiniMaxM2ForCausalLM": "minimax",
@@ -359,6 +362,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
     "MiniMaxM3SparseForConditionalGeneration": "minimax",
     "MiniCPMV4_6ForConditionalGeneration": "minicpm",
     "MiniCPMV4_7ForConditionalGeneration": "minicpm",
+    "MiniCPMOcrForConditionalGeneration": "minicpm",
     "Mistral3ForConditionalGeneration": "llava",
     "Mistral3Model": "llava",
     "NemotronH_Omni_Reasoning_V3": "nemotron",
@@ -367,6 +371,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
     "MuseSparkForConditionalGeneration": "muse_spark",
     "PaddleOCRVisionModel": "ernie",
     "Phi4ForCausalLMV": "phi",
+    "NEOVisionModel": "sensenova_u1",
     "Qwen2AudioForConditionalGeneration": "ultravox",
     "Qwen2VLForConditionalGeneration": "qwenvl",
     "Qwen2VLModel": "qwenvl",

conversion/base.py

@@ -1701,7 +1701,7 @@ class TextModel(ModelBase):
             self.gguf_writer.add_expert_group_used_count(n_group_used)
             logger.info(f"gguf: expert groups used count = {n_group_used}")
 
-        if (score_func := self.find_hparam(["score_function", "scoring_func", "score_func", "moe_router_activation", "moe_router_activation_func", "expert_selection_fn", "router_score_func"], optional=True)) is not None:
+        if (score_func := self.find_hparam(["score_function", "scoring_func", "score_func", "moe_router_activation", "moe_router_activation_func", "expert_selection_fn", "router_score_func", "router_score_function"], optional=True)) is not None:
             if score_func == "sigmoid":
                 self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)
             elif score_func == "softmax":
@@ -1837,6 +1837,9 @@ class TextModel(ModelBase):
         if chkhsh == "1431a23e583c97432bc230bff598d103ddb5a1f89960c8f1d1051aaa944d0b35":
             # ref: https://huggingface.co/sapienzanlp/Minerva-7B-base-v1.0
             res = "minerva-7b"
+        if chkhsh == "6e4d72c6e5260b90606e26a0931fa15e4b42ec6e97719d12919f2b1f93478262":
+            # ref: https://huggingface.co/domyn/Domyn-Small-v1.0
+            res = "domyn-small"
         if chkhsh == "7e57df22b1fe23a7b1e1c7f3dc4e3f96d43a4eb0836d0c6bdc3436d7b2f1c664":
             # ref: https://huggingface.co/tencent/Hunyuan-A13B-Instruct
             res = "hunyuan"
@@ -1939,6 +1942,9 @@ class TextModel(ModelBase):
         if chkhsh == "52df12b4c8d4176e7481aab4b6e8454d1fd0a210a04a574f6d4e067d10e23c3e":
             # ref: https://huggingface.co/CohereLabs/North-Mini-Code-1.0
             res = "cohere2moe"
+        if chkhsh == "d5420e601e979078c670bc503d9ac411c311fa76b76a76267c31e0d6ba816637":
+            # ref: https://huggingface.co/RecursiveSI/flint
+            res = "flint"
         if chkhsh == "e636dc30a262dcc0d8c323492e32ae2b70728f4df7dfe9737d9f920a282b8aea":
             # ref: https://huggingface.co/Qwen/Qwen1.5-7B
             res = "qwen2"
@@ -3143,7 +3149,7 @@ def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> st
     # Kimi-K3's text_config reports "KimiLinearForCausalLM", which is the older
     # Kimi-Linear-48B architecture and cannot load K3 (no attention residuals,
     # latent MoE, situ, ...). Route on the top-level architecture instead.
-    if model_type == ModelType.TEXT and arch in ("StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration"):
+    if model_type == ModelType.TEXT and arch in ("NEOChatModel", "StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration"):
         return arch
 
     # if "architectures" is found in the sub-config, use that instead

conversion/flint.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from .base import ModelBase, gguf
+from .qwen import Qwen2MoeModel
+
+
+@ModelBase.register("FlintForCausalLM")
+@ModelBase.example("RecursiveSI/flint")
+class FlintModel(Qwen2MoeModel):
+    model_arch = gguf.MODEL_ARCH.FLINT
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        dense = self.hparams["mlp_only_layers"]
+        if dense != list(range(len(dense))):
+            raise ValueError("Flint dense layers must precede the MoE layers")
+        self.gguf_writer.add_leading_dense_block_count(len(dense))
+        self.gguf_writer.add_expert_shared_count(
+            self.hparams["shared_expert_intermediate_size"] // self.hparams["moe_intermediate_size"]
+        )
+        self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])
+        self.gguf_writer.add_expert_weights_scale(1.0)
+        self.gguf_writer.add_sliding_window(self.hparams["sliding_window"])
+        self.gguf_writer.add_sliding_window_pattern([
+            kind == "sliding_attention" for kind in self.hparams["layer_types"]
+        ])
+        self.gguf_writer.add_rope_pattern([
+            i not in self.hparams["no_rope_layer_indices"] for i in range(self.block_count)
+        ])
+
+    def modify_tensors(self, data_torch, name, bid):
+        if name.endswith(".softmax_offset"):
+            name += ".weight"
+        if name.endswith(".gate.expert_bias"):
+            if not self.hparams["router_expert_bias"]:
+                return
+            name += ".bias"
+        yield from super().modify_tensors(data_torch, name, bid)

conversion/minicpm.py

@@ -272,3 +272,18 @@ class MiniCPMV4_7VisionModel(MiniCPMV4_6VisionModel):
             yield self.format_tensor_name(gguf.MODEL_TENSOR.V_TOK_EMBD_SEP, suffix=""), data_torch[ids]
             return
         yield from super().modify_tensors(data_torch, name, bid)
+
+
+# MiniCPM-OCR 1.0: the 4.6 stack with model_type minicpm_ocr.
+
+@ModelBase.register("MiniCPMOcrForConditionalGeneration")
+@ModelBase.example("openbmb/MiniCPM-OCR-1.0")
+class MiniCPMOCR1_0TextModel(MiniCPMV4_6TextModel):
+    model_arch = gguf.MODEL_ARCH.QWEN35
+    no_mtp = True
+
+
+@ModelBase.register("MiniCPMOcrForConditionalGeneration")
+@ModelBase.example("openbmb/MiniCPM-OCR-1.0")
+class MiniCPMOCR1_0VisionModel(MiniCPMV4_6VisionModel):
+    projector_type = gguf.VisionProjectorType.MINICPM_OCR_1_0

conversion/mistral3.py

@@ -39,6 +39,11 @@ class Mistral3Model(TextModel):
             self.gguf_writer.add_add_bos_token(False)
             self.gguf_writer.add_add_eos_token(False)
 
+        def get_vocab_base_pre(self, tokenizer) -> str:
+            if self.origin_hf_arch == "Mistral3Model":
+                return "nemotron-3.5-embed"
+            return super().get_vocab_base_pre(tokenizer)
+
         def set_gguf_parameters(self):
             super().set_gguf_parameters()
             self._try_set_pooling_type()

conversion/nemotron.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import re
 from typing import Any, Callable, Iterable, Sequence, TYPE_CHECKING
 
 import numpy as np
@@ -367,9 +368,21 @@ class NemotronModel(TextModel):
     model_arch = gguf.MODEL_ARCH.NEMOTRON
 
     def set_vocab(self):
-        self._set_vocab_sentencepiece()
-        self.gguf_writer.add_pad_token_id(0)
-        self.gguf_writer.add_unk_token_id(1)
+        if (self.dir_model / "tokenizer.model").exists():
+            self._set_vocab_sentencepiece()
+            self.gguf_writer.add_pad_token_id(0)
+            self.gguf_writer.add_unk_token_id(1)
+        else:
+            # SPM-style BPE (Metaspace + byte fallback): mark <0xXX> tokens as BYTE so they decode to raw bytes
+            tokens, toktypes, tokpre = self.get_vocab_base()
+            toktypes = [gguf.TokenType.BYTE if re.fullmatch(r"<0x[0-9A-F]{2}>", tok) else typ for tok, typ in zip(tokens, toktypes)]
+            self.gguf_writer.add_tokenizer_model("gpt2")
+            self.gguf_writer.add_tokenizer_pre(tokpre)
+            self.gguf_writer.add_token_list(tokens)
+            self.gguf_writer.add_token_types(toktypes)
+
+            special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
+            special_vocab.add_to_gguf(self.gguf_writer)
 
     def set_gguf_parameters(self):
         super().set_gguf_parameters()
@@ -383,7 +396,10 @@ class NemotronModel(TextModel):
         rot_pct = self.rope_parameters["partial_rotary_factor"]
         n_embd = self.find_hparam(["hidden_size", "n_embd"])
         n_head = self.find_hparam(["num_attention_heads", "n_head"])
-        self.gguf_writer.add_rope_dimension_count(int(rot_pct * n_embd) // n_head)
+        # Use explicit head_dim when available; fallback keeps standard Nemotron working.
+        # Without this, models where n_head * head_dim != n_embd (e.g. 48 heads, head_dim=128) produce the wrong rope dimension count.
+        head_dim = self.hparams.get("head_dim") or n_embd // n_head
+        self.gguf_writer.add_rope_dimension_count(int(rot_pct * head_dim))
 
         # * RopeScaling for Nemotron
         factor = self.hparams.get("factor") or self.rope_parameters.get("factor")

conversion/sensenova_u1.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+from .base import MmprojModel, ModelBase, gguf
+from .qwen import Qwen2Model
+
+
+@ModelBase.example("sensenova/SenseNova-U1.5-8B-MoT")
+@ModelBase.register("NEOChatModel")
+class SenseNovaU1Model(Qwen2Model):
+    model_arch = gguf.MODEL_ARCH.SENSENOVA_U1
+
+    @classmethod
+    def filter_tensors(cls, item):
+        name, _ = item
+        if not name.startswith("language_model.") or "_mot_gen" in name:
+            return None
+        return super().filter_tensors(item)
+
+    def set_gguf_parameters(self):
+        if int(self.hparams.get("num_experts") or 0) > 1:
+            raise ValueError("SenseNova U1 conversion currently supports the dense understanding branch")
+        if self.rope_parameters.get("rope_type", "default") not in ("default", None):
+            raise ValueError("SenseNova U1 scaled RoPE is not supported")
+        if self.hparams.get("attention_bias", False):
+            raise ValueError("SenseNova U1 attention bias is not supported")
+        if self.hparams.get("use_sliding_window", False):
+            raise ValueError("SenseNova U1 sliding attention is not supported")
+        super().set_gguf_parameters()
+        self.gguf_writer.add_rope_freq_base_spatial(self.hparams["rope_theta_hw"])
+
+    def get_tensors(self):
+        for name, gen in self.model_tensors.items():
+            if name.endswith((".q_norm_hw.weight", ".k_norm_hw.weight")):
+                continue
+            data = gen()
+            if name.endswith((".q_norm.weight", ".k_norm.weight")):
+                spatial = self.model_tensors[name.replace("_norm.weight", "_norm_hw.weight")]()
+                # Keep one weight vector; the graph normalizes its two halves independently.
+                data = torch.cat((data, spatial), dim=0)
+            yield name, data
+
+
+@ModelBase.register("NEOVisionModel")
+@ModelBase.example("sensenova/SenseNova-U1.5-8B-MoT")
+class SenseNovaU1VisionModel(MmprojModel):
+    def __init__(self, dir_model: Path, *args, **kwargs):
+        hparams = kwargs.pop("hparams", None)
+        if hparams is None:
+            hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False)
+        else:
+            hparams = dict(hparams)
+
+        vision_config = dict(hparams["vision_config"])
+        # NEOVisionModel has no transformer blocks
+        # these defaults satisfy the generic mmproj metadata contract without replacing checkpoint values
+        vision_config.setdefault("image_size", 0)
+        vision_config.setdefault("intermediate_size", 0)
+        vision_config.setdefault("num_hidden_layers", 0)
+        vision_config.setdefault("num_attention_heads", 1)
+        hparams["vision_config"] = vision_config
+        super().__init__(dir_model, *args, hparams=hparams, **kwargs)
+
+        self.preprocessor_config.setdefault("image_mean", [0.485, 0.456, 0.406])
+        self.preprocessor_config.setdefault("image_std", [0.229, 0.224, 0.225])
+
+    @classmethod
+    def filter_tensors(cls, item):
+        name, gen = item
+        if not name.startswith("vision_model."):
+            return None
+        return super().filter_tensors((name, gen))
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_clip_projector_type("sensenova_u1")
+        vision_config = self.hparams["vision_config"]
+        # These defaults match the known U1.5 checkpoint when older configs omit fields.
+        hidden_size = int(vision_config.get("hidden_size", 1024))
+        min_pixels = int(vision_config.get("min_pixels", 65536))
+        max_pixels = int(vision_config.get("max_pixels", 16777216))
+        downsample_ratio = float(vision_config.get("downsample_ratio", 0.5))
+        if downsample_ratio <= 0:
+            raise ValueError(f"SenseNova U1 downsample_ratio must be positive, got {downsample_ratio}")
+        merge_size = round(1.0 / downsample_ratio)
+        if abs((1.0 / downsample_ratio) - merge_size) > 1e-6:
+            raise ValueError(f"SenseNova U1 downsample_ratio must have an integer reciprocal, got {downsample_ratio}")
+        norm_eps = float(vision_config.get("layer_norm_eps", vision_config.get("rms_norm_eps", 1e-6)))
+        self.gguf_writer.add_vision_head_dim(hidden_size)
+        self.gguf_writer.add_vision_attention_layernorm_eps(norm_eps)
+        self.gguf_writer.add_vision_min_pixels(min_pixels)
+        self.gguf_writer.add_vision_max_pixels(max_pixels)
+        self.gguf_writer.add_vision_spatial_merge_size(merge_size)
+
+    def modify_tensors(self, data_torch, name, bid):
+        tensor_names = {
+            "vision_model.embeddings.patch_embedding.weight": "v.patch_embd.weight",
+            "vision_model.embeddings.patch_embedding.bias":   "v.patch_embd.bias",
+            "vision_model.embeddings.dense_embedding.weight": self.format_tensor_name(gguf.MODEL_TENSOR.V_MMPROJ, 0),
+            "vision_model.embeddings.dense_embedding.bias": self.format_tensor_name(
+                gguf.MODEL_TENSOR.V_MMPROJ, 0, suffix=".bias"
+            ),
+        }
+        if name not in tensor_names:
+            raise ValueError(f"Unexpected SenseNova U1 vision tensor: {name!r}")
+        return [(tensor_names[name], data_torch)]
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if new_name == self.format_tensor_name(gguf.MODEL_TENSOR.V_MMPROJ, 0):
+            return gguf.GGMLQuantizationType.F16 if self.ftype == gguf.LlamaFileType.MOSTLY_F16 else gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)

convert_hf_to_gguf_update.py

@@ -101,6 +101,7 @@ models = [
     {"name": "command-r",        "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/CohereForAI/c4ai-command-r-v01", },
     {"name": "tiny_aya",         "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/CohereLabs/tiny-aya-base", },
     {"name": "cohere2moe",       "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/CohereLabs/North-Mini-Code-1.0", },
+    {"name": "flint",            "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/RecursiveSI/flint", },
     {"name": "qwen2",            "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/Qwen/Qwen1.5-7B", },
     {"name": "olmo",             "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/allenai/OLMo-1.7-7B-hf", },
     {"name": "dbrx",             "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/databricks/dbrx-base", },
@@ -177,6 +178,7 @@ pre_computed_hashes = [
     {"name": "glm4", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/zai-org/GLM-4.5-Air", "chkhsh": "9ca2dd618e8afaf09731a7cf6e2105b373ba6a1821559f258b272fe83e6eb902"},
     {"name": "glm4", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/zai-org/GLM-4.7-Flash", "chkhsh": "cdf5f35325780597efd76153d4d1c16778f766173908894c04afc20108536267"},
     {"name": "minerva-7b", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/sapienzanlp/Minerva-7B-base-v1.0", "chkhsh": "1431a23e583c97432bc230bff598d103ddb5a1f89960c8f1d1051aaa944d0b35"},
+    {"name": "domyn-small", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/domyn/Domyn-Small-v1.0", "chkhsh": "6e4d72c6e5260b90606e26a0931fa15e4b42ec6e97719d12919f2b1f93478262"},
     {"name": "hunyuan", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/tencent/Hunyuan-A13B-Instruct", "chkhsh": "7e57df22b1fe23a7b1e1c7f3dc4e3f96d43a4eb0836d0c6bdc3436d7b2f1c664"},
     {"name": "hunyuan-dense", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/tencent/Hunyuan-4B-Instruct", "chkhsh": "bba3b3366b646dbdded5dbc42d59598b849371afc42f7beafa914afaa5b70aa6"},
     {"name": "hy_v4", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/tencent/Hy4-preview", "chkhsh": "e6ddf9c6686791c12d698d34c31ab9be1fea9af5a3d9a6909783ab382198ae1c"},

docs/backend/SYCL.md

@@ -783,7 +783,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL          | ON (mandatory)                        | Enable build with SYCL code path.           |
 | GGML_SYCL_TARGET   | INTEL *(default)*                     | Set the SYCL target device type.            |
 | GGML_SYCL_DEVICE_ARCH | Optional                           | Set the SYCL device architecture. Setting the device architecture can improve the performance. See the table [--offload-arch](https://github.com/intel/llvm/blob/sycl/sycl/doc/design/OffloadDesign.md#--offload-arch) for a list of valid architectures. |
-| GGML_SYCL_F16      | OFF *(default)* \|ON *(optional)*     | Enable FP16 build with SYCL code path. (1.) |
+| GGML_SYCL_F16      | OFF *(default)* \|ON *(optional)*     | Enable FP16 build with SYCL code path. [1.](#sycl-f16-footnote) |
 | GGML_SYCL_GRAPH    | ON *(default)* \|OFF *(Optional)*     | Enable build with [SYCL Graph extension](https://github.com/intel/llvm/blob/sycl/sycl/doc/extensions/experimental/sycl_ext_oneapi_graph.asciidoc). |
 | GGML_SYCL_DNN      | ON *(default)* \|OFF *(Optional)*     | Enable build with oneDNN.                   |
 | GGML_SYCL_HOST_MEM_FALLBACK | ON *(default)* \|OFF *(Optional)* | Allow host memory fallback when device memory is full during quantized weight reorder. Enables inference to continue at reduced speed (reading over PCIe) instead of failing. Requires Linux kernel 6.8+. |
@@ -791,7 +791,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | CMAKE_C_COMPILER   | `icx` *(Linux)*, `icx/cl` *(Windows)* | Set `icx` compiler for SYCL code path.      |
 | CMAKE_CXX_COMPILER | `icpx` *(Linux)*, `icx` *(Windows)*   | Set `icpx/icx` compiler for SYCL code path. |
 
-1. FP32 or FP16 have different performance impact to LLM. Recommended to test them for better prompt processing performance on your models. You need to rebuild the code after change `GGML_SYCL_F16=OFF/ON`.
+1. <a id="sycl-f16-footnote"></a>FP32 or FP16 have different performance impact to LLM. Recommended to test them for better prompt processing performance on your models. You need to rebuild the code after change `GGML_SYCL_F16=OFF/ON`.
 
 ### Runtime
 
@@ -802,7 +802,8 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_DEV2DEV_MEMCPY | 0 (default), 1, 2 | Choose the method of dev2dev memory copy.<br>Value: <br>*  0: SYCL API (default), only support dGPUs.<br>* 1: L0 API -- Better performance, only support dGPUs, found to lead to abnormal crash in some case. <br>* 2: Host Forward -- Most stable method for all cases (including iGPU + dGPU*N), but with lower performance (-2% to -5%).<br>SYCL & L0 API are easy to be impacted by Intel GPU driver issue. When you meet the garbled output or crash issues in multiple GPUs case, try with this debug flag to work around or check the issue.|
 | GGML_SYCL_ENABLE_FLASH_ATTN | 1 (default) or 0| Enable Flash-Attention. It can reduce memory usage. The performance impact depends on the LLM.|
 | GGML_SYCL_ENABLE_OPT | 0 or 1 (default)| Enable optimize features for Intel GPUs. (Recommended to 0 for Intel devices older than Gen 10) |
-| GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still on development, no better performance. |
+| GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still in development, performance depends on model and device. [1.](#sycl-graph-footnote) |
+| GGML_SYCL_GRAPH_EVICTION_TIMEOUT | 10 (default) or positive integer | Seconds a cached SYCL graph can go unused before it is evicted. Checked on a sweep that runs every half of this timeout. |
 | GGML_SYCL_ENABLE_HOST_PINNED_MEM | 0 or 1 (default) | Enable host pinned memory to speed up copy data from host to device. When disable it, host memory will common malloc() on CPU. Disable it when use `--load-model mlock`.|
 | GGML_SYCL_HOST_PINNED_MEM_2G | 0 (default) or 1 | Limit the max memory allocation to be no more than 2GB when enable host pinned memory. USM allocations above 2 GiB take the relaxed/large-allocation path, which serializes H2D copies with compute and prevents copy/compute overlap. It will impact the startup time. Need more test. Depend on `GGML_SYCL_ENABLE_HOST_PINNED_MEM=1`.|
 | GGML_SYCL_UPLOAD_STAGING_SLOTS | 4 (default) or non-negative integer | Number of 8 MiB pinned host slots used to stage tensor uploads (model loading), so the host copy of one slot overlaps the transfer of the previous one. Set to 0 to use the old path: a malloc'd bounce buffer and a blocking copy per tensor. |
@@ -819,6 +820,9 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_MKL_FA_DIAG | 0 (default) or 1 | Enable output fingerprinting for MKL flash attention. Dumps the first 64 float output values for the first 6 FA calls with n_kv ≥ 1024, labeled with kernel type (MKL/TILE/VEC) for cross-kernel comparison. |
 | GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute. Unsupported types and layouts fall back to the standalone op kernels. See `ggml_sycl_can_fuse()`. |
 | GGML_SYCL_ENABLE_ESIMD | 0 or 1 (default)| Enable ESIMD kernels when available. |
+| GGML_SYCL_ENABLE_XMX_ESIMD | 0 or 1 (default) | Use the XMX (DPAS) multi-column mat-vec kernels for reordered Q2_K to Q6_K and Q8_0 weights on Xe2 GPUs. Speeds up speculative decoding and short prompt batches. |
+| GGML_SYCL_XMX_MAX_COLS | non-negative integer, 80 (default) | Most activation columns (batch size) the `GGML_SYCL_ENABLE_XMX_ESIMD` kernels take; wider batches use the other mat-mul kernels. A larger value keeps XMX for bigger batches (as repeated launches), `0` never uses it. |
+| GGML_SYCL_XMX_GLU_MAX_COLS | 0 to 16, 16 (default) | Most columns the fused gate and up XMX kernel takes; wider batches run gate and up as separate mat-vecs. 16 is the most the kernel is built for, larger values are clamped to 16 and `0` turns the fused kernel off. |
 | GGML_SYCL_XMX_GATHER_TYPES | decimal bitmask, all bits set (default) | Weight formats that may use the XMX dequant-GEMM paths, which dequantize weights straight into the XMX tiles. This speeds up prompt processing of MoE models on GPUs with XMX units (Arc A- and B-series, Arc Pro, Data Center GPU Max), for example pp512 of Qwen3-30B-A3B UD-IQ3_XXS by about 50% on an Arc Pro B60. Bits:<br>* 1: IQ4_NL, 2: IQ3_S, 4: IQ4_XS, 8: IQ3_XXS, 16: IQ2_XXS, 32: IQ2_XS, 64: IQ2_S, 128: IQ1_S, 256: IQ1_M<br>* 512: Q8_0, 1024: Q4_K, 2048: Q5_K, 4096: Q6_K (MoE `MUL_MAT_ID` only)<br>Add values to combine them, for example `3` for IQ4_NL and IQ3_S; `0` disables the paths. A set bit does not force the path: batches of more than 64 tokens per expert or row lengths that are not a multiple of 256 (32 for IQ4_NL and Q8_0) use the library GEMM. |
 | GGML_SYCL_XMX_GATHER_SHAPES | decimal bitmask, 255 (default) | XMX `joint_matrix` combinations the paths of `GGML_SYCL_XMX_GATHER_TYPES` may use; the operand type comes from `GGML_SYCL_DYNAMIC_PRECISION` and the best supported combination is picked automatically (logged as `fg_pick_combo`). Bits:<br>* Xe2, Xe3, Xe-HPC: 1: f16 8x16x16, 2: f16 16x16x16, 4: f16 32x64x16, 8: f16 32x64x32, 32: tf32 8x16x8, 64: bf16 8x16x16<br>* Xe-HPG (Arc A770, ARL-H): 16: f16 8x8x16, 128: bf16 8x8x16<br>Clear a bit to exclude a combination, or set a single bit to force one for testing. |
 | GGML_SYCL_DYNAMIC_PRECISION | `F16` (default with `GGML_SYCL_F16=ON`), `BF16`, `TF32` or `F32` (default otherwise) | Operand type of the XMX dequant-GEMM paths (`GGML_SYCL_XMX_GATHER_TYPES`); accumulation is always f32. `F16` is the fastest, but activations above 65504 overflow. `BF16` keeps the f32 range at a 7-bit mantissa, `TF32` keeps the range and the f16 mantissa but is about 30% slower and needs Xe2, Xe3 or Xe-HPC, and `F32` turns the XMX paths off. Ops that request a higher src1 precision ([TAG_GGML_PREC]) get it regardless of this setting. |
@@ -831,6 +835,17 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS | 0 (default) or 1 | Allow SYCL/Unified Runtime Level Zero device allocations larger than 4 GiB. llama.cpp's direct Level Zero allocation path requests the relaxed maximum-size limit itself when GGML_SYCL_ENABLE_LEVEL_ZERO=1. |
 | GGML_SYCL_USM_SYSTEM | 0 (default) or 1 | Enable experimental support for [USM system allocations](https://github.khronos.org/SYCL_Reference/iface/usm_basic_concept.html#system-allocations) for large GPU buffers. This requires enough host memory for model weights and caches, an Intel Xe2+ GPU such as BMG or newer and supported on Linux only, with CONFIG_DRM_XE_GPUSVM enabled. |
 
+1. <a id="sycl-graph-footnote"></a>SYCL Graph support (`GGML_SYCL_ENABLE_GRAPH=1`) is limited to single-GPU workloads and may also depend on the model architecture and kernels in use.
+   Performance depends on how often a recorded graph can be reused, which will vary by workload.
+   For this reason, it is recommended to measure performance both with and without graphs enabled on your workloads.
+
+   To check whether graphs are supported for your workload, enable INFO logging (`-lv 4`) and look for `enabling SYCL graphs` or `disabling SYCL graphs`.
+   Support may differ between prefill and decode, so check the status for the phases that matter to you.
+   DEBUG logging (`-lv 5`) reports why graphs are disabled.
+
+   Then benchmark with your preferred tool, such as `llama-bench`.
+   For graph activity details (including reuse), set `GGML_SYCL_DEBUG=1` and look for `[SYCL-GRAPH]` messages.
+
 ## Compile-time Flags
 
 Pass these via `CXXFLAGS` or add a one-off `#define` to enable a flag on the spot.

docs/ops.md

@@ -27,7 +27,7 @@ Legend:
 |                        COL2IM_1D | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           CONCAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                             CONT | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
-|                       CONV_1D_DW | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ |
+|                       CONV_1D_DW | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ |
 |                          CONV_2D | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
 |                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          CONV_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -76,10 +76,10 @@ Legend:
 |                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
 |                 MUL_MAT_HADAMARD | ❌ | ❌ | ✅ | ❌ | ✅ | 🟡 | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 |
 |                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | 🟡 | ❌ |
-|                  MUL_MAT_ID_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-|                  MUL_MAT_ID_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-|                     MUL_MAT_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-|                     MUL_MAT_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |

ggml/include/ggml-rpc.h

@@ -8,10 +8,10 @@ extern "C" {
 
 #define RPC_PROTO_MAJOR_VERSION    8
 #define RPC_PROTO_MINOR_VERSION    0
-#define RPC_PROTO_PATCH_VERSION    0
+#define RPC_PROTO_PATCH_VERSION    1
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif
 
 #define GGML_RPC_MAX_SERVERS       16

ggml/include/ggml.h

@@ -585,6 +585,7 @@ extern "C" {
         GGML_OP_DSV4_HC_COMB,
         GGML_OP_DSV4_HC_PRE,
         GGML_OP_DSV4_HC_POST,
+        GGML_OP_TOP_K_MASK,
 
         GGML_OP_UNARY,
 
@@ -2732,6 +2733,26 @@ extern "C" {
             struct ggml_tensor  * post,
             struct ggml_tensor  * comb);
 
+    // top-k sparse attention mask: keep the mask cells of the live picks and the extra cells, set the rest to -inf
+    //
+    // mask:   [n_kv,    m1,  m2,  m3]  f16 or f32, broadcast to the rows of top_k
+    // top_k:  [n_top,   ne1, ne2, ne3] i32, picks of each row
+    // score:  [n_pick,  ne1, ne2, ne3] f32, can be NULL: a pick is live when its score is above -inf, else every pick is live
+    // blocks: [n_blk,   n_pick]        i32, can be NULL: the member cells of each pick, else a pick is a cell
+    // extra:  [n_extra, ne1, ne2, ne3] i32, can be NULL: cells of each row kept without a pick
+    // res:    [n_kv,    ne1, ne2, ne3] same type as mask
+    //
+    // res[c, i] = mask[c, i % m] when c is a cell of a live pick of row i or an extra cell of row i, else -inf
+    // picks and cells out of range are skipped
+    //
+    GGML_API struct ggml_tensor * ggml_top_k_mask(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * mask,
+            struct ggml_tensor  * top_k,
+            struct ggml_tensor  * score,
+            struct ggml_tensor  * blocks,
+            struct ggml_tensor  * extra);
+
     // custom operators
 
     typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata);

ggml/src/ggml-backend-meta.cpp

@@ -1055,7 +1055,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             } break;
             case GGML_OP_DSV4_HC_COMB:
             case GGML_OP_DSV4_HC_PRE:
-            case GGML_OP_DSV4_HC_POST: {
+            case GGML_OP_DSV4_HC_POST:
+            case GGML_OP_TOP_K_MASK: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ true);
             } break;
             case GGML_OP_UNARY: {

ggml/src/ggml-cann/ggml-cann.cpp

@@ -920,7 +920,7 @@ static void * ggml_backend_cann_buffer_get_base(ggml_backend_buffer_t buffer) {
  * @param dst Pointer to the destination buffer where transformed data will be
  * stored.
  */
-static void ggml_backend_cann_transform_q4_0(ggml_tensor * tensor, const void * src, void * dst) {
+static void ggml_backend_cann_transform_q4_0(const ggml_tensor * tensor, const void * src, void * dst) {
     int64_t n_elems     = ggml_nelements(tensor);
     int64_t groups      = n_elems / QK4_0;
     size_t  quant_bytes = n_elems * sizeof(uint8_t) / 2;
@@ -1014,7 +1014,7 @@ static void ggml_backend_cann_transform_back_q4_0(const ggml_tensor * tensor, vo
  * @param dst Pointer to the destination buffer where transformed data will be
  * stored.
  */
-static void ggml_backend_cann_transform_q8_0(ggml_tensor * tensor, const void * src, void * dst) {
+static void ggml_backend_cann_transform_q8_0(const ggml_tensor * tensor, const void * src, void * dst) {
     int64_t n_elems     = ggml_nelements(tensor);
     int64_t groups      = n_elems / QK8_0;
     size_t  quant_bytes = n_elems * sizeof(uint8_t);
@@ -1075,7 +1075,7 @@ static void ggml_backend_cann_transform_back_q8_0(const ggml_tensor * tensor, co
  * @param dst Pointer to the destination buffer where transformed data will be
  * stored.
  */
-static void ggml_backend_cann_transform(ggml_tensor * tensor, const void * src, void * dst) {
+static void ggml_backend_cann_transform(const ggml_tensor * tensor, const void * src, void * dst) {
     switch (tensor->type) {
         case GGML_TYPE_Q4_0:
             ggml_backend_cann_transform_q4_0(tensor, src, dst);

ggml/src/ggml-cpu/CMakeLists.txt

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

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

@@ -2105,6 +2105,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_dsv4_hc_post(params, tensor);
             } break;
+        case GGML_OP_TOP_K_MASK:
+            {
+                ggml_compute_forward_top_k_mask(params, tensor);
+            } break;
         case GGML_OP_MAP_CUSTOM1:
             {
                 ggml_compute_forward_map_custom1(params, tensor);
@@ -2288,6 +2292,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_DSV4_HC_COMB:
         case GGML_OP_DSV4_HC_PRE:
         case GGML_OP_DSV4_HC_POST:
+        case GGML_OP_TOP_K_MASK:
             {
                 n_tasks = n_threads;
             } break;

ggml/src/ggml-cpu/kleidiai/kernels.cpp

@@ -201,85 +201,117 @@ static struct kai_matmul_uker_config q4_k_matmul_config() {
     return config;
 }
 
-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_m_step() {
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    return api.get_step(&config).m;
+static struct kai_matmul_uker_config default_matmul_config() {
+    return {};
 }
 
-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_n_step() {
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    return api.get_step(&config).n;
-}
+template<
+        struct kai_matmul_uker_api (*GetApi)(void),
+        struct kai_matmul_uker_config (*GetConfig)(void),
+        size_t MR, size_t NR, size_t KR, size_t SR>
+struct matmul_uker_adapter {
+    static size_t get_m_step() {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        return api.get_step(&config).m;
+    }
 
-template<size_t Value>
-static size_t q4_k_get_constant() {
-    return Value;
-}
+    static size_t get_n_step() {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        return api.get_step(&config).n;
+    }
 
-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_dst_offset(size_t m_idx, size_t n_idx, size_t dst_stride) {
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_dst_dim_args index = { m_idx, n_idx };
-    const struct kai_matmul_uker_dst_stride_args stride = { dst_stride };
-    return api.get_dst_offset(&config, &index, &stride);
-}
+    static size_t get_mr() {
+        if constexpr (MR == 0) {
+            return get_m_step();
+        }
+        return MR;
+    }
 
-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_lhs_packed_offset(size_t m_idx, size_t k, size_t bl) {
-    GGML_ASSERT(bl == KLEIDIAI_Q4_K_BLOCK_LENGTH);
+    static size_t get_nr() {
+        if constexpr (NR == 0) {
+            return get_n_step();
+        }
+        return NR;
+    }
 
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_lhs_dim_args shape = { m_idx, k };
-    const struct kai_matmul_uker_lhs_dim_args index = { m_idx, 0 };
-    const struct kai_matmul_uker_lhs_stride_args stride = api.get_lhs_stride(&config, &shape);
-    return api.get_lhs_offset(&config, &index, &stride);
-}
+    static size_t get_kr() {
+        return KR;
+    }
 
-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_rhs_packed_offset(size_t n_idx, size_t k, size_t bl) {
-    GGML_ASSERT(bl == KLEIDIAI_Q4_K_BLOCK_LENGTH);
+    static size_t get_sr() {
+        return SR;
+    }
 
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_rhs_dim_args shape = { n_idx, k };
-    const struct kai_matmul_uker_rhs_dim_args index = { n_idx, 0 };
-    const struct kai_matmul_uker_rhs_stride_args stride = api.get_rhs_stride(&config, &shape);
-    return api.get_rhs_offset(&config, &index, &stride);
-}
+    static size_t get_dst_offset(size_t m_idx, size_t n_idx, size_t dst_stride) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        const struct kai_matmul_uker_dst_dim_args index = { m_idx, n_idx };
+        const struct kai_matmul_uker_dst_stride_args stride = { dst_stride };
+        return api.get_dst_offset(&config, &index, &stride);
+    }
 
-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static void q4_k_run_matmul(
-        size_t m, size_t n, size_t k, size_t bl,
-        const void * lhs_packed, const void * rhs_packed,
-        void * dst, size_t dst_stride_row, size_t dst_stride_col,
-        float clamp_min, float clamp_max) {
-    GGML_ASSERT(bl == KLEIDIAI_Q4_K_BLOCK_LENGTH);
-    GGML_ASSERT(dst_stride_col == sizeof(float));
+    static size_t get_lhs_packed_offset(size_t m_idx, size_t k, size_t bl) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        GGML_ASSERT(bl == config.format.bl);
+        const struct kai_matmul_uker_lhs_dim_args shape = { m_idx, k };
+        const struct kai_matmul_uker_lhs_dim_args index = { m_idx, 0 };
+        const struct kai_matmul_uker_lhs_stride_args stride = api.get_lhs_stride(&config, &shape);
+        return api.get_lhs_offset(&config, &index, &stride);
+    }
 
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_lhs_dim_args lhs_shape = { m, k };
-    const struct kai_matmul_uker_rhs_dim_args rhs_shape = { n, k };
+    static size_t get_rhs_packed_offset(size_t n_idx, size_t k, size_t bl) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        GGML_ASSERT(bl == config.format.bl);
+        const struct kai_matmul_uker_rhs_dim_args shape = { n_idx, k };
+        const struct kai_matmul_uker_rhs_dim_args index = { n_idx, 0 };
+        const struct kai_matmul_uker_rhs_stride_args stride = api.get_rhs_stride(&config, &shape);
+        return api.get_rhs_offset(&config, &index, &stride);
+    }
 
-    struct kai_matmul_uker_args args = {};
-    args.flags = KAI_MATMUL_UKER_FLAGS_ARGS_CLAMP;
-    args.shape = { m, n, k };
-    args.operand.lhs.ptr = lhs_packed;
-    args.operand.lhs.stride = api.get_lhs_stride(&config, &lhs_shape);
-    args.operand.rhs.ptr = rhs_packed;
-    args.operand.rhs.stride = api.get_rhs_stride(&config, &rhs_shape);
-    args.operand.dst.ptr = dst;
-    args.operand.dst.stride.m = dst_stride_row;
-    args.activation.clamp.min_ptr = &clamp_min;
-    args.activation.clamp.max_ptr = &clamp_max;
-    api.run(&config, &args);
-}
+    static void run(
+            size_t m, size_t n, size_t k, size_t bl,
+            const void * lhs_packed, const void * rhs_packed,
+            void * dst, size_t dst_stride_row, size_t dst_stride_col,
+            float clamp_min, float clamp_max) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        GGML_ASSERT(bl == config.format.bl);
+        GGML_ASSERT(dst_stride_col == sizeof(float));
+        const struct kai_matmul_uker_lhs_dim_args lhs_shape = { m, k };
+        const struct kai_matmul_uker_rhs_dim_args rhs_shape = { n, k };
+
+        struct kai_matmul_uker_args args = {};
+        args.flags = KAI_MATMUL_UKER_FLAGS_ARGS_CLAMP;
+        args.shape = { m, n, k };
+        args.operand.lhs.ptr = lhs_packed;
+        args.operand.lhs.stride = api.get_lhs_stride(&config, &lhs_shape);
+        args.operand.rhs.ptr = rhs_packed;
+        args.operand.rhs.stride = api.get_rhs_stride(&config, &rhs_shape);
+        args.operand.dst.ptr = dst;
+        args.operand.dst.stride.m = dst_stride_row;
+        args.activation.clamp.min_ptr = &clamp_min;
+        args.activation.clamp.max_ptr = &clamp_max;
+        api.run(&config, &args);
+    }
+};
+
+using q4_k_gemm_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa,
+    q4_k_matmul_config, 0, 0, 4, 2>;
+using q4_k_gemv_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot,
+    q4_k_matmul_config, 0, 0, 4, 2>;
+// GEMM and GEMV share the RHS layout only at SVE256 (NR=8, KR=8, SR=1).
+using sve_q8_gemm_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_qai8dxp4x8sf32_qsi8cxp4vsx8sf32bf32_16x4vs_sve_i8mm,
+    default_matmul_config, 4, 0, 8, 1>;
+using sve_q8_gemv_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_qai8dxp1x8_qsi8cxp8x8_1x8_sve_dot,
+    default_matmul_config, 1, 8, 8, 1>;
 
 static struct kai_matmul_pack_lhs_uker_config q4_k_lhs_pack_config() {
     struct kai_matmul_pack_lhs_uker_config config = {};
@@ -628,7 +660,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
             /* .packed_size_ex        = */ &rhs_ps_fn5<kai_get_rhs_packed_size_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
             /* .pack_func_ex          = */ &rhs_pack_fn12<kai_run_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
         },
-        /* .required_cpu       = */ CPU_FEATURE_SVE | CPU_FEATURE_I8MM | CPU_FEATURE_DOTPROD,
+        /* .required_cpu       = */ CPU_FEATURE_SVE | CPU_FEATURE_SVEI8MM | CPU_FEATURE_DOTPROD,
         /* .lhs_type           = */ GGML_TYPE_F32,
         /* .rhs_type           = */ GGML_TYPE_Q4_0,
         /* .op_type            = */ GGML_TYPE_F32,
@@ -737,16 +769,16 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q4_k[] = {
     {
         /* SME2 GEMM */
         {
-            /* .get_m_step            = */ &q4_k_get_m_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_n_step            = */ &q4_k_get_n_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_mr                = */ &q4_k_get_m_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_nr                = */ &q4_k_get_n_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_kr                = */ &q4_k_get_constant<4>,
-            /* .get_sr                = */ &q4_k_get_constant<2>,
-            /* .get_dst_offset        = */ &q4_k_get_dst_offset<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_lhs_offset_ex     = */ &q4_k_get_lhs_packed_offset<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_rhs_packed_offset_ex = */ &q4_k_get_rhs_packed_offset<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .run_kernel_ex         = */ &q4_k_run_matmul<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
+            /* .get_m_step            = */ q4_k_gemm_adapter::get_m_step,
+            /* .get_n_step            = */ q4_k_gemm_adapter::get_n_step,
+            /* .get_mr                = */ q4_k_gemm_adapter::get_mr,
+            /* .get_nr                = */ q4_k_gemm_adapter::get_nr,
+            /* .get_kr                = */ q4_k_gemm_adapter::get_kr,
+            /* .get_sr                = */ q4_k_gemm_adapter::get_sr,
+            /* .get_dst_offset        = */ q4_k_gemm_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ q4_k_gemm_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ q4_k_gemm_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ q4_k_gemm_adapter::run,
         },
         {
             /* .get_offset            = */ kai_get_lhs_offset_lhs_pack_f16pmrx2_f32_neon,
@@ -756,16 +788,16 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q4_k[] = {
         },
         /* SME2 GEMV */
         {
-            /* .get_m_step            = */ &q4_k_get_m_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_n_step            = */ &q4_k_get_n_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_mr                = */ &q4_k_get_m_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_nr                = */ &q4_k_get_n_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_kr                = */ &q4_k_get_constant<4>,
-            /* .get_sr                = */ &q4_k_get_constant<2>,
-            /* .get_dst_offset        = */ &q4_k_get_dst_offset<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_lhs_offset_ex     = */ &q4_k_get_lhs_packed_offset<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_rhs_packed_offset_ex = */ &q4_k_get_rhs_packed_offset<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .run_kernel_ex         = */ &q4_k_run_matmul<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
+            /* .get_m_step            = */ q4_k_gemv_adapter::get_m_step,
+            /* .get_n_step            = */ q4_k_gemv_adapter::get_n_step,
+            /* .get_mr                = */ q4_k_gemv_adapter::get_mr,
+            /* .get_nr                = */ q4_k_gemv_adapter::get_nr,
+            /* .get_kr                = */ q4_k_gemv_adapter::get_kr,
+            /* .get_sr                = */ q4_k_gemv_adapter::get_sr,
+            /* .get_dst_offset        = */ q4_k_gemv_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ q4_k_gemv_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ q4_k_gemv_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ q4_k_gemv_adapter::run,
         },
         {
             /* .get_offset            = */ q4_k_lhs_get_offset,
@@ -883,6 +915,54 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
         /* .rhs_type           = */ GGML_TYPE_Q8_0,
         /* .op_type            = */ GGML_TYPE_F32,
     },
+    {
+        /* SVE I8MM GEMM */
+        {
+            /* .get_m_step            = */ sve_q8_gemm_adapter::get_m_step,
+            /* .get_n_step            = */ sve_q8_gemm_adapter::get_n_step,
+            /* .get_mr                = */ sve_q8_gemm_adapter::get_mr,
+            /* .get_nr                = */ sve_q8_gemm_adapter::get_nr,
+            /* .get_kr                = */ sve_q8_gemm_adapter::get_kr,
+            /* .get_sr                = */ sve_q8_gemm_adapter::get_sr,
+            /* .get_dst_offset        = */ sve_q8_gemm_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ sve_q8_gemm_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ sve_q8_gemm_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ sve_q8_gemm_adapter::run,
+        },
+        /* .gemm_lhs_info = */ {
+            /* .get_offset            = */ kai_get_lhs_offset_lhs_quant_pack_qai8dxp_f32,
+            /* .get_packed_offset_ex  = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_quant_pack_qai8dxp_f32>,
+            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32>,
+            /* .pack_func_ex          = */ &lhs_pack_float_fn9_no_bl<kai_run_lhs_quant_pack_qai8dxp_f32>,
+        },
+        /* SVE DOT GEMV */
+        {
+            /* .get_m_step            = */ sve_q8_gemv_adapter::get_m_step,
+            /* .get_n_step            = */ sve_q8_gemv_adapter::get_n_step,
+            /* .get_mr                = */ sve_q8_gemv_adapter::get_mr,
+            /* .get_nr                = */ sve_q8_gemv_adapter::get_nr,
+            /* .get_kr                = */ sve_q8_gemv_adapter::get_kr,
+            /* .get_sr                = */ sve_q8_gemv_adapter::get_sr,
+            /* .get_dst_offset        = */ sve_q8_gemv_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ sve_q8_gemv_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ sve_q8_gemv_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ sve_q8_gemv_adapter::run,
+        },
+        /* .gemv_lhs_info = */ {
+            /* .get_offset            = */ kai_get_lhs_offset_lhs_quant_pack_qai8dxp_f32,
+            /* .get_packed_offset_ex  = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_quant_pack_qai8dxp_f32>,
+            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32>,
+            /* .pack_func_ex          = */ &lhs_pack_float_fn9_no_bl<kai_run_lhs_quant_pack_qai8dxp_f32>,
+        },
+        /* .rhs_info = */ {
+            /* .packed_size_ex        = */ &rhs_ps_fn5<kai_get_rhs_packed_size_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
+            /* .pack_func_ex          = */ &rhs_pack_scale_fn12<kai_run_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
+        },
+        /* .required_cpu       = */ CPU_FEATURE_SVE | CPU_FEATURE_SVEI8MM | CPU_FEATURE_DOTPROD,
+        /* .lhs_type           = */ GGML_TYPE_F32,
+        /* .rhs_type           = */ GGML_TYPE_Q8_0,
+        /* .op_type            = */ GGML_TYPE_F32,
+    },
     {
         /* I8MM GEMM */
         {

ggml/src/ggml-cpu/kleidiai/kernels.h

@@ -13,7 +13,8 @@ enum cpu_feature {
     CPU_FEATURE_SVE     = 4,
     CPU_FEATURE_SME     = 8,
     CPU_FEATURE_SME2    = 16,
-    CPU_FEATURE_FP16    = 32
+    CPU_FEATURE_FP16    = 32,
+    CPU_FEATURE_SVEI8MM = 64
 };
 
 inline cpu_feature& operator|=(cpu_feature& lhs, cpu_feature rhs) {

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

@@ -319,8 +319,10 @@ static void init_kleidiai_context(void) {
 
         size_t detected_smcus = 0;
 
+        // The fixed-width SVE kernels require 32-byte vectors, including the Q8 shared RHS layout.
         ctx.features  = (runtime_feat.has_dotprod  ? CPU_FEATURE_DOTPROD : CPU_FEATURE_NONE) |
                         (runtime_feat.has_i8mm     ? CPU_FEATURE_I8MM    : CPU_FEATURE_NONE) |
+                        (runtime_feat.has_svei8mm  ? CPU_FEATURE_SVEI8MM : CPU_FEATURE_NONE) |
                         (runtime_feat.has_fp16     ? CPU_FEATURE_FP16    : CPU_FEATURE_NONE) |
                         (runtime_feat.sve_cnt == QK8_0 ? CPU_FEATURE_SVE : CPU_FEATURE_NONE);
 
@@ -1616,6 +1618,11 @@ public:
             const size_t kr = kernel->get_kr();
             const size_t sr = kernel->get_sr();
             const ggml_type rhs_type = kernels->rhs_type;
+            if (rhs_type == GGML_TYPE_Q8_0 && (kernels->required_cpu & CPU_FEATURE_SVE)) {
+                GGML_ASSERT(nr == kernels->gemv.get_nr());
+                GGML_ASSERT(kr == kernels->gemv.get_kr());
+                GGML_ASSERT(sr == kernels->gemv.get_sr());
+            }
             const size_t block_len = rhs_type == GGML_TYPE_Q8_0 ? QK8_0 :
                                      rhs_type == GGML_TYPE_Q4_0 ? QK4_0 :
                                      rhs_type == GGML_TYPE_Q4_K ? GGML_KLEIDIAI_Q4_K_BL :

ggml/src/ggml-cpu/ops.cpp

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

ggml/src/ggml-cpu/ops.h

@@ -108,6 +108,7 @@ void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * p
 void ggml_compute_forward_dsv4_hc_comb(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_dsv4_hc_pre(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_dsv4_hc_post(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_top_k_mask(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst);

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

@@ -70,6 +70,7 @@
 #include "ggml-cuda/cumsum.cuh"
 #include "ggml-cuda/fill.cuh"
 #include "ggml-cuda/lightning-indexer.cuh"
+#include "ggml-cuda/top-k-mask.cuh"
 #include "ggml.h"
 
 #include <algorithm>
@@ -2570,6 +2571,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_LIGHTNING_INDEXER:
             ggml_cuda_lightning_indexer(ctx, dst);
             break;
+        case GGML_OP_TOP_K_MASK:
+            ggml_cuda_op_top_k_mask(ctx, dst);
+            break;
         default:
             return false;
     }
@@ -6003,6 +6007,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
             return true;
         case GGML_OP_LIGHTNING_INDEXER:
             return ggml_cuda_lightning_indexer_supported(dev_ctx->device, op);
+        case GGML_OP_TOP_K_MASK:
+            return op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32;
 
         default:
             return false;

ggml/src/ggml-cuda/top-k-mask.cu

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

ggml/src/ggml-cuda/top-k-mask.cuh

@@ -0,0 +1,3 @@
+#include "common.cuh"
+
+void ggml_cuda_op_top_k_mask(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

ggml/src/ggml-feats.h

@@ -22,6 +22,10 @@
 #define HWCAP2_I8MM (1ULL << 13)
 #endif
 
+#if !defined(HWCAP2_SVEI8MM)
+#define HWCAP2_SVEI8MM (1ULL << 9)
+#endif
+
 #if !defined(HWCAP_ASIMDDP)
 #define HWCAP_ASIMDDP (1 << 20)
 #endif
@@ -63,6 +67,10 @@
 #define PF_ARM_SVE2_INSTRUCTIONS_AVAILABLE 47
 #endif
 
+#if !defined(PF_ARM_SVE_I8MM_INSTRUCTIONS_AVAILABLE)
+#define PF_ARM_SVE_I8MM_INSTRUCTIONS_AVAILABLE 57
+#endif
+
 #if !defined(PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE)
 #define PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE 66
 #endif
@@ -87,6 +95,7 @@ typedef struct ggml_feats_arch64_runtime {
     bool has_sve;
     bool has_sve2;
     bool has_i8mm;
+    bool has_svei8mm;
     bool has_sme;
     bool has_sme2;
     int sve_cnt;
@@ -104,6 +113,7 @@ static inline ggml_feats_arch64_runtime_t ggml_feats_get_arch64_runtime(void) {
     runtime_feat.has_sve     = !!(hwcap & HWCAP_SVE);
     runtime_feat.has_sve2    = !!(hwcap2 & HWCAP2_SVE2);
     runtime_feat.has_i8mm    = !!(hwcap2 & HWCAP2_I8MM);
+    runtime_feat.has_svei8mm = !!(hwcap2 & HWCAP2_SVEI8MM);
     runtime_feat.has_sme     = !!(hwcap2 & HWCAP2_SME);
     runtime_feat.has_sme2    = !!(hwcap2 & HWCAP2_SME2);
 
@@ -153,6 +163,7 @@ static inline ggml_feats_arch64_runtime_t ggml_feats_get_arch64_runtime(void) {
     runtime_feat.has_sve     = IsProcessorFeaturePresent(PF_ARM_SVE_INSTRUCTIONS_AVAILABLE) != 0;
     runtime_feat.has_sve2    = IsProcessorFeaturePresent(PF_ARM_SVE2_INSTRUCTIONS_AVAILABLE) != 0;
     runtime_feat.has_i8mm    = IsProcessorFeaturePresent(PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE) != 0;
+    runtime_feat.has_svei8mm = IsProcessorFeaturePresent(PF_ARM_SVE_I8MM_INSTRUCTIONS_AVAILABLE) != 0;
     runtime_feat.has_sme     = IsProcessorFeaturePresent(PF_ARM_SME_INSTRUCTIONS_AVAILABLE) != 0;
     runtime_feat.has_sme2    = IsProcessorFeaturePresent(PF_ARM_SME2_INSTRUCTIONS_AVAILABLE) != 0;
 

ggml/src/ggml-hexagon/CMakeLists.txt

@@ -51,6 +51,13 @@ endif()
 set(TARGET_NAME ggml-hexagon)
 ggml_add_backend_library(${TARGET_NAME}
     ggml-hexagon.cpp
+    htp-common.h
+    htp-kparams.cpp
+    htp-kparams.h
+    htp-opnode.cpp
+    htp-opnode.h
+    htp-fusion.cpp
+    htp-fusion.h
     htp-repack.cpp
     htp-repack.h
     htp-drv.cpp

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

@@ -0,0 +1,106 @@
+#ifndef HTP_COMMON_H
+#define HTP_COMMON_H
+
+#include "htp-opnode.h"
+#include "htp-repack.h"
+
+struct htp_tensor_extra {
+    std::vector<uint8_t> shadow_buf;
+    size_t               shadow_size { 0 };
+    uint32_t             flags { 0 };
+};
+
+extern int    opt_arch;
+extern size_t opt_ndev;
+extern size_t opt_nhvx;
+extern int    opt_nhmx;
+extern size_t opt_vmem;
+extern int    opt_etm;
+extern int    opt_verbose;
+extern int    opt_profile;
+extern bool   opt_hostbuf;
+extern bool   opt_dma64;
+extern int    opt_power_mode;
+
+extern size_t opt_mbuf_dyn;
+extern size_t opt_mbuf_static;
+extern size_t opt_mbuf_total;
+
+extern int    opt_mm_select;
+extern int    opt_fa_select;
+extern int    opt_fa_head_split;
+extern int    opt_gdn_select;
+extern int    opt_ar_select;
+extern int    opt_ar_scatter;
+
+extern int    opt_opbatch;
+extern int    opt_opqueue;
+extern int    opt_optrace;
+extern int    opt_oppoll;
+extern int    opt_opfusion;
+extern int    opt_rt_repack;
+
+#ifndef HEX_VERBOSE
+#define HEX_VERBOSE(...) \
+    if (opt_verbose) GGML_LOG_DEBUG(__VA_ARGS__)
+#endif
+
+static inline bool htp_tensor_is_fuseable(const struct ggml_tensor * t) {
+    if (!t || !t->extra) return false;
+    const auto * extra = (const struct htp_tensor_extra *) t->extra;
+    return (extra->flags & HTP_TENSOR_FUSEABLE) != 0;
+}
+
+static inline bool htp_can_rt_repack(const struct ggml_tensor * t) {
+    if (opt_rt_repack == 0 || !htp_is_tiled_weight_tensor(t)) {
+        return false;
+    }
+    if (t->buffer && ggml_backend_buft_is_host(t->buffer->buft)) {
+        return false;
+    }
+    return opt_rt_repack > 1 || ggml_is_quantized(t->type);
+}
+
+static inline bool htp_tensor_is_tiled(const struct ggml_tensor * t) {
+    const auto * extra = (const struct htp_tensor_extra *) t->extra;
+    if (extra && (extra->flags & HTP_TENSOR_TILED)) {
+        return true;
+    }
+    return htp_can_rt_repack(t);
+}
+
+static inline bool htp_needs_repack(const struct ggml_tensor * t) {
+    const auto * extra = (const struct htp_tensor_extra *) t->extra;
+    if (extra && (extra->flags & HTP_TENSOR_TILED)) {
+        return false;
+    }
+    return htp_can_rt_repack(t);
+}
+
+static inline bool htp_tensors_overlap(const struct ggml_tensor * a, const struct ggml_tensor * b) {
+    const uintptr_t a0 = (uintptr_t) a->data;
+    const uintptr_t b0 = (uintptr_t) b->data;
+    const uintptr_t a1 = a0 + ggml_nbytes(a);
+    const uintptr_t b1 = b0 + ggml_nbytes(b);
+
+    return a0 < b1 && b0 < a1;
+}
+
+static inline bool htp_can_row_partition(const struct ggml_tensor * t) {
+    if (t->ne[1] > 1 && (t->nb[1] & 127) != 0) return false;
+    if (t->ne[2] > 1 && (t->nb[2] & 127) != 0) return false;
+    if (t->ne[3] > 1 && (t->nb[3] & 127) != 0) return false;
+    return true;
+}
+
+struct ggml_hexagon_session;
+
+const char * htp_session_name(const struct ggml_hexagon_session * sess);
+size_t       htp_session_vtcm_size(const struct ggml_hexagon_session * sess);
+uint32_t     htp_session_n_threads(const struct ggml_hexagon_session * sess);
+uint32_t     htp_session_n_hmx(const struct ggml_hexagon_session * sess);
+uint32_t     htp_session_mdev_count(const struct ggml_hexagon_session * sess);
+
+htp_op_code  htp_op_remap(const struct ggml_tensor * t);
+
+#endif // HTP_COMMON_H

ggml/src/ggml-hexagon/htp-fusion.cpp

@@ -0,0 +1,650 @@
+#include "htp-fusion.h"
+#include "htp-kparams.h"
+#include "ggml-impl.h"
+
+#include <algorithm>
+#include <cstring>
+
+static inline bool htp_same_tensor_data(const ggml_tensor * a, const ggml_tensor * b) {
+    return a && b && (a == b || a->data == b->data);
+}
+
+static inline bool htp_same_shape_and_contiguity(const ggml_tensor * a, const ggml_tensor * b) {
+    return ggml_are_same_shape(a, b) && (ggml_is_contiguous(a) == ggml_is_contiguous(b));
+}
+
+static inline bool htp_inplace_overlap_unsafe(const ggml_tensor * dst, const ggml_tensor * in) {
+    return htp_tensors_overlap(dst, in) && (dst->data != in->data || !ggml_are_same_stride(dst, in));
+}
+
+static inline bool htp_inputs_overlap_unsafe(const htp_opnode & prev, const ggml_tensor * dst, const ggml_tensor * extra = nullptr) {
+    if (extra && htp_inplace_overlap_unsafe(dst, extra)) return true;
+    for (const auto * in : prev.inputs) {
+        if (htp_inplace_overlap_unsafe(dst, in)) return true;
+    }
+    return false;
+}
+
+static inline bool htp_match_consumer(
+    const htp_opnode & prev,
+    const htp_opnode & cur,
+    bool commutative,
+    const ggml_tensor ** other_src = nullptr
+) {
+    if (prev.outputs.size() != 1) {
+        return false;
+    }
+    const ggml_tensor * prev_out = prev.dst();
+    if (!htp_tensor_is_fuseable(prev_out)) {
+        return false;
+    }
+    if (htp_same_tensor_data(cur.src0(), prev_out)) {
+        if (other_src) *other_src = cur.src1();
+        return true;
+    }
+    if (commutative && htp_same_tensor_data(cur.src1(), prev_out)) {
+        if (other_src) *other_src = cur.src0();
+        return true;
+    }
+    return false;
+}
+
+static bool is_supported_mul_mat_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams) {
+    if (kparams->n_hmx) {
+        return kparams->kernel_type == HTP_MM_KERNEL_HMX_2D;
+    }
+    return htp_tensor_is_tiled(src0) && (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW);
+}
+
+static bool is_supported_mul_mat_id_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams) {
+    if (kparams->n_hmx) {
+        return kparams->kernel_type == HTP_MM_KERNEL_HMX_2D;
+    }
+    return htp_tensor_is_tiled(src0) &&
+           (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK);
+}
+
+// --- Fusion rules ---
+
+static bool fuse_allreduce_add(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (opt_ar_select != 2 || prev.opcode != HTP_OP_ALLREDUCE) return false;
+
+    const auto * ar_kparams = prev.kparams<htp_allreduce_kernel_params>();
+    const uint32_t rank    = (uint32_t) ar_kparams->rank;
+    const uint32_t n_ranks = (uint32_t) ar_kparams->n_ranks;
+
+    const ggml_tensor * ar_local = prev.inputs[rank];
+    if (!htp_tensor_is_fuseable(ar_local)) return false;
+
+    const ggml_tensor * add_dst    = cur.dst();
+    const ggml_tensor * res_tensor = nullptr;
+    if (htp_same_tensor_data(cur.src0(), ar_local)) {
+        res_tensor = cur.src1();
+    } else if (htp_same_tensor_data(cur.src1(), ar_local)) {
+        res_tensor = cur.src0();
+    } else {
+        return false;
+    }
+
+    if (ar_local->type != res_tensor->type) return false;
+
+    const bool is_same_shape = ggml_are_same_shape(ar_local, res_tensor);
+    const bool is_row_bcast  = !is_same_shape && (ar_local->ne[0] == res_tensor->ne[0] &&
+                                                  res_tensor->ne[1] == 1 && res_tensor->ne[2] == 1 && res_tensor->ne[3] == 1);
+    if (!is_same_shape && !is_row_bcast) return false;
+
+    if (is_same_shape && !ggml_are_same_stride(ar_local, res_tensor)) {
+        return false;
+    }
+    if (ggml_is_contiguous(ar_local) != ggml_is_contiguous(add_dst)) return false;
+
+    const bool is_shard_ok = n_ranks <= HTP_OP_MAX_OUTPUTS && add_dst->data == ar_local->data;
+
+    struct htp_allreduce_kernel_params kparams;
+    if (!htp_kparams_solve_allreduce(sess, add_dst, rank, n_ranks, true, is_row_bcast, is_shard_ok, &kparams)) {
+        HEX_VERBOSE("ggml-hex: %s skip ALLREDUCE_ADD fusion: solver failed\n", htp_session_name(sess));
+        return false;
+    }
+
+    const bool scatter_ok = kparams.n_dsts > 1;
+    kparams.mode = scatter_ok ? HTP_ALLREDUCE_SHARDED_FANOUT : HTP_ALLREDUCE_FULL;
+
+    for (uint32_t r = 0; r < n_ranks; r++) {
+        const ggml_tensor * ar_src = prev.inputs[r];
+        if (!htp_tensors_overlap(add_dst, ar_src)) continue;
+        if (scatter_ok && r == rank) continue;
+        HEX_VERBOSE("ggml-hex: %s skip ALLREDUCE_ADD fusion: dst overlaps allreduce src %u\n", htp_session_name(sess), r);
+        return false;
+    }
+
+    prev.opcode = HTP_OP_ALLREDUCE_ADD;
+    prev.inputs.push_back(res_tensor);
+    if (kparams.mode != HTP_ALLREDUCE_SHARDED_FANOUT) {
+        prev.outputs = { add_dst };
+    }
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_rms_norm_mul(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_RMS_NORM) return false;
+
+    const ggml_tensor * weight = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &weight)) return false;
+
+    const ggml_tensor * src0 = prev.src0();
+    const ggml_tensor * dst  = cur.dst();
+    if (src0->type != GGML_TYPE_F32 || src0->ne[0] != weight->ne[0] || !htp_same_shape_and_contiguity(prev.dst(), dst)) return false;
+    if (htp_inputs_overlap_unsafe(prev, dst, weight)) return false;
+
+    const bool is_row_bcast  = (weight->ne[1] == 1 && weight->ne[2] == 1 && weight->ne[3] == 1);
+    const bool is_same_shape = ggml_are_same_shape(src0, weight);
+    if (!is_row_bcast && !is_same_shape) return false;
+
+    struct htp_unary_kernel_params kparams;
+    if (!htp_kparams_solve_rms_norm_mul(sess, src0, weight, dst, &kparams)) {
+        HEX_VERBOSE("ggml-hex: %s skip RMS_NORM_MUL fusion: VTCM needed (%d) > budget (%zu)\n",
+                    htp_session_name(sess), kparams.vtcm_size, htp_session_vtcm_size(sess));
+        return false;
+    }
+
+    prev.opcode  = HTP_OP_RMS_NORM_MUL;
+    prev.inputs  = { src0, weight };
+    prev.outputs = { dst };
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_rms_norm_scale(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    GGML_UNUSED(sess);
+    if (prev.opcode != HTP_OP_RMS_NORM && prev.opcode != HTP_OP_RMS_NORM_MUL) return false;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false)) return false;
+
+    const ggml_tensor * src0 = prev.src0();
+    const ggml_tensor * dst  = cur.dst();
+    if (src0->type != GGML_TYPE_F32 || !htp_same_shape_and_contiguity(prev.dst(), dst)) return false;
+    if (htp_inputs_overlap_unsafe(prev, dst)) return false;
+
+    float scale = 0.0f;
+    float bias  = 0.0f;
+    memcpy(&scale, (const float *) cur.params + 0, sizeof(float));
+    memcpy(&bias,  (const float *) cur.params + 1, sizeof(float));
+    if (bias != 0.0f || scale == 0.0f) return false;
+
+    float prev_scale = 0.0f;
+    memcpy(&prev_scale, &prev.params[1], sizeof(float));
+    if (prev_scale != 0.0f) {
+        scale *= prev_scale;
+    }
+    memcpy(&prev.params[1], &scale, sizeof(float));
+
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_mul_mat_add(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL_MAT) return false;
+
+    const ggml_tensor * src2 = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &src2)) return false;
+    if (src2->type != GGML_TYPE_F32) return false;
+
+    const ggml_tensor * src0 = prev.src0();
+    const ggml_tensor * src1 = prev.src1();
+    const ggml_tensor * dst  = cur.dst();
+    if (!ggml_are_same_shape(prev.dst(), dst)) return false;
+
+    const auto * orig_kparams = prev.kparams<htp_mm_kernel_params>();
+    struct htp_mm_kernel_params kparams;
+    if (!htp_kparams_solve_matmul_add(sess, src0, src1, src2, dst, &kparams) ||
+        orig_kparams->collapse != kparams.collapse) {
+        return false;
+    }
+
+    const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3];
+    if (kparams.n_hmx == 0 && src1_nrows != 1) return false;
+
+    if (kparams.n_hmx > 0 && orig_kparams->n_hmx > 0) {
+        if (kparams.m_chunk < orig_kparams->m_chunk ||
+            kparams.n_chunk < orig_kparams->n_chunk ||
+            kparams.n_act_threads < orig_kparams->n_act_threads) {
+            HEX_VERBOSE("ggml-hex: %s skip MUL_MAT_ADD fusion: HMX efficiency reduced (m %d->%d, n %d->%d, th %d->%d)\n",
+                        htp_session_name(sess), orig_kparams->m_chunk, kparams.m_chunk,
+                        orig_kparams->n_chunk, kparams.n_chunk,
+                        orig_kparams->n_act_threads, kparams.n_act_threads);
+            return false;
+        }
+    }
+
+    prev.opcode  = HTP_OP_MUL_MAT_ADD;
+    prev.inputs  = { src0, src1, src2 };
+    prev.outputs = { dst };
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_mul_mat_nx(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL_MAT && prev.opcode != HTP_OP_MUL_MAT_NX) return false;
+
+    const uint32_t curr_n = (uint32_t) prev.outputs.size();
+    if (curr_n >= HTP_OP_MAX_OUTPUTS || curr_n + 1 >= HTP_OP_MAX_INPUTS) return false;
+
+    const ggml_tensor * w0   = prev.inputs[0];
+    const ggml_tensor * x    = prev.inputs[curr_n];
+    const ggml_tensor * w_in = cur.src0();
+    const ggml_tensor * x_in = cur.src1();
+    const ggml_tensor * d_in = cur.dst();
+
+    auto * orig_kparams = prev.kparams<htp_mm_kernel_params>();
+    const auto * cur_kparams = cur.kparams<htp_mm_kernel_params>();
+
+    if (x_in != x || w_in->type != w0->type || w_in->ne[0] != w0->ne[0] ||
+        w0->ne[2] != 1 || w0->ne[3] != 1 || w_in->ne[2] != 1 || w_in->ne[3] != 1 ||
+        htp_tensor_is_tiled(w_in) != htp_tensor_is_tiled(w0) ||
+        orig_kparams->n_hmx != cur_kparams->n_hmx) {
+        return false;
+    }
+
+    const bool collapse = orig_kparams->collapse && htp_matmul_can_collapse(w_in, x, d_in);
+    if ((x->ne[2] > 1 || x->ne[3] > 1) && !collapse) return false;
+
+    if (prev.opcode == HTP_OP_MUL_MAT_NX) {
+        orig_kparams->n_weights = curr_n + 1;
+    } else {
+        struct htp_mm_kernel_params kparams;
+        if (collapse) {
+            const ggml_tensor x_collapsed = htp_tensor_collapse_rows(x);
+            if (!htp_kparams_solve_mmnx(sess, w0, &x_collapsed, curr_n + 1, &kparams)) return false;
+            kparams.collapse = 1;
+        } else if (!htp_kparams_solve_mmnx(sess, w0, x, curr_n + 1, &kparams)) {
+            return false;
+        }
+        if (!is_supported_mul_mat_nx_kernel(w0, &kparams)) return false;
+
+        prev.set_kparams(kparams);
+        prev.opcode = HTP_OP_MUL_MAT_NX;
+    }
+
+    prev.inputs.insert(prev.inputs.begin() + curr_n, w_in);
+    prev.outputs.push_back(d_in);
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_mul_mat_id_nx(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL_MAT_ID && prev.opcode != HTP_OP_MUL_MAT_ID_NX) return false;
+
+    const uint32_t curr_n = (uint32_t) prev.outputs.size();
+    if (curr_n >= HTP_OP_MAX_OUTPUTS || curr_n + 2 >= HTP_OP_MAX_INPUTS) return false;
+
+    const ggml_tensor * w0     = prev.inputs[0];
+    const ggml_tensor * x      = prev.inputs[curr_n];
+    const ggml_tensor * ids    = prev.inputs[curr_n + 1];
+    const ggml_tensor * w_in   = cur.src0();
+    const ggml_tensor * x_in   = cur.src1();
+    const ggml_tensor * ids_in = cur.node->src[2];
+    const ggml_tensor * d_in   = cur.dst();
+
+    auto * orig_kparams = prev.kparams<htp_mm_kernel_params>();
+    const auto * cur_kparams = cur.kparams<htp_mm_kernel_params>();
+
+    if (x_in != x || ids_in != ids || w_in->type != w0->type ||
+        w_in->ne[0] != w0->ne[0] || w_in->ne[2] != w0->ne[2] ||
+        htp_tensor_is_tiled(w_in) != htp_tensor_is_tiled(w0) ||
+        orig_kparams->n_hmx != cur_kparams->n_hmx) {
+        return false;
+    }
+
+    struct htp_mm_kernel_params kparams;
+    if (d_in->nb[1] == prev.node->nb[1]) {
+        kparams = *orig_kparams;
+        kparams.n_weights = curr_n + 1;
+    } else if (!htp_kparams_solve_mmidnx(sess, w0, x, d_in, curr_n + 1, &kparams)) {
+        return false;
+    }
+    if (!is_supported_mul_mat_id_nx_kernel(w0, &kparams)) return false;
+
+    prev.opcode = HTP_OP_MUL_MAT_ID_NX;
+    prev.inputs.insert(prev.inputs.begin() + curr_n, w_in);
+    prev.outputs.push_back(d_in);
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_gdn_cpy(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    GGML_UNUSED(sess);
+    if (prev.opcode != HTP_OP_GATED_DELTA_NET || prev.outputs.size() != 1) return false;
+
+    const ggml_tensor * gdn_out  = prev.dst();
+    const ggml_tensor * cpy_node = cur.node;
+    const ggml_tensor * cpy_src  = cur.src0();
+    const ggml_tensor * cpy_dst  = cur.dst();
+
+    if (gdn_out->type != GGML_TYPE_F32 || cpy_src->type != GGML_TYPE_F32 || cpy_dst->type != GGML_TYPE_F32) return false;
+    if ((gdn_out->flags & GGML_TENSOR_FLAG_OUTPUT) || (cpy_node->flags & GGML_TENSOR_FLAG_OUTPUT)) return false;
+
+    const ggml_tensor * v = prev.node->src[2];
+    if (!v) return false;
+
+    const int64_t S_v       = v->ne[0];
+    const int64_t H         = v->ne[1];
+    const int64_t n_tokens  = v->ne[2];
+    const int64_t n_seqs    = v->ne[3];
+    const int64_t K         = ggml_get_op_params_i32(prev.node, 0);
+    const size_t  tail_off  = (size_t) S_v * H * n_tokens * n_seqs * sizeof(float);
+    const int64_t D         = S_v * S_v * H;
+    const int64_t n_written = std::min<int64_t>(n_tokens, K);
+
+    if (cpy_src->op != GGML_OP_VIEW || !htp_same_tensor_data(cpy_src->view_src, gdn_out) ||
+        cpy_src->view_offs != tail_off || !ggml_is_contiguous(cpy_src)) {
+        return false;
+    }
+
+    if (cpy_dst->ne[0] != D || cpy_dst->ne[1] != n_seqs || cpy_dst->ne[2] != n_written || !ggml_is_contiguous(cpy_dst)) {
+        return false;
+    }
+
+    prev.outputs.push_back(cpy_dst);
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_ssm_conv_silu(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    GGML_UNUSED(sess);
+    if (prev.opcode != HTP_OP_SSM_CONV) return false;
+
+    auto * kparams = prev.kparams<htp_ssm_conv_kernel_params>();
+    if (kparams->silu) return false;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false)) return false;
+
+    const ggml_tensor * conv_dst = prev.dst();
+    const ggml_tensor * dst      = cur.dst();
+    if (!ggml_are_same_layout(conv_dst, dst)) return false;
+    if (htp_tensors_overlap(dst, prev.src0()) || htp_tensors_overlap(dst, prev.src1())) return false;
+
+    kparams->silu   = 1;
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_repeat_get_rows(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_REPEAT) return false;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false)) return false;
+
+    const ggml_tensor * repeat_out  = prev.dst();
+    const ggml_tensor * repeat_src0 = prev.src0();
+    const ggml_tensor * gr_src1     = cur.src1();
+    const ggml_tensor * gr_dst      = cur.dst();
+
+    if (repeat_src0->ne[0] != repeat_out->ne[0] || repeat_src0->ne[1] != repeat_out->ne[1]) return false;
+    if (repeat_src0->type != repeat_out->type) return false;
+
+    struct htp_get_rows_kernel_params kparams;
+    if (!htp_kparams_solve_repeat_get_rows(sess, repeat_src0, gr_src1, gr_dst, &kparams) || kparams.vtcm_size == 0) return false;
+
+    prev.opcode  = HTP_OP_GET_ROWS;
+    prev.inputs  = { repeat_src0, gr_src1 };
+    prev.outputs = { gr_dst };
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_unary_mul_glu(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    htp_op_code glu_op = HTP_OP_INVALID;
+    switch (prev.opcode) {
+        case HTP_OP_UNARY_SILU:       glu_op = HTP_OP_GLU_SWIGLU;      break;
+        case HTP_OP_UNARY_GELU:       glu_op = HTP_OP_GLU_GEGLU;       break;
+        case HTP_OP_UNARY_GELU_QUICK: glu_op = HTP_OP_GLU_GEGLU_QUICK; break;
+        case HTP_OP_UNARY_GELU_ERF:   glu_op = HTP_OP_GLU_GEGLU_ERF;   break;
+        case HTP_OP_UNARY_RELU:       glu_op = HTP_OP_GLU_REGLU;       break;
+        case HTP_OP_UNARY_SIGMOID:    glu_op = HTP_OP_GLU_SIGMOID;     break;
+        default: return false;
+    }
+
+    const ggml_tensor * up = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &up)) return false;
+
+    const ggml_tensor * gate = prev.src0();
+    const ggml_tensor * dst  = cur.dst();
+    if (htp_inputs_overlap_unsafe(prev, dst, up)) return false;
+
+    struct htp_glu_kernel_params kparams;
+    if (!htp_kparams_solve_unary_mul_glu(sess, glu_op, gate, up, dst, &kparams)) return false;
+
+    prev.opcode  = glu_op;
+    prev.inputs  = { gate, up };
+    prev.outputs = { dst };
+    prev.fused.push_back(cur.node);
+    memset(prev.params, 0, sizeof(prev.params));
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_mul_mul_scalar(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL || prev.inputs.size() != 2) return false;
+
+    const auto * k0 = prev.kparams<htp_binary_kernel_params>();
+    const auto * k1 = cur.kparams<htp_binary_kernel_params>();
+    if (k0->kernel_type != HTP_BINARY_KERNEL_SCALAR || k1->kernel_type != HTP_BINARY_KERNEL_SCALAR) return false;
+    if (k0->src1_strided || k1->src1_strided) return false;
+
+    const ggml_tensor * src2 = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false, &src2)) return false;
+
+    const ggml_tensor * src1 = prev.src1();
+    const ggml_tensor * dst  = cur.dst();
+    if (!ggml_are_same_shape(src1, src2) || !htp_same_shape_and_contiguity(prev.dst(), dst)) return false;
+    if (htp_inputs_overlap_unsafe(prev, dst, src2)) return false;
+
+    const size_t vtcm_size = htp_session_vtcm_size(sess);
+    struct htp_binary_kernel_params kparams = *k0;
+    kparams.src2_size       = kparams.src1_size;
+    kparams.rows_per_buffer = 0;
+    if (kparams.src1_size > 65536 || kparams.src1_size + kparams.src2_size > vtcm_size) return false;
+
+    struct htp_binary_vtcm_layout L;
+    htp_binary_vtcm_layout_build(&L, &kparams, vtcm_size);
+    if (L.rows_per_buffer == 0 || L.total_bytes > vtcm_size) return false;
+
+    kparams.rows_per_buffer = L.rows_per_buffer;
+    kparams.vtcm_size       = L.total_bytes;
+
+    prev.inputs.push_back(src2);
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_binary_nx(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != cur.opcode) return false;
+
+    const uint32_t curr_n = (uint32_t) prev.inputs.size();
+    if (curr_n < 2 || curr_n >= HTP_OP_MAX_INPUTS) return false;
+
+    const auto * k0 = prev.kparams<htp_binary_kernel_params>();
+    const auto * k1 = cur.kparams<htp_binary_kernel_params>();
+    const bool k0_ok = (k0->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) ||
+                       (k0->kernel_type == HTP_BINARY_KERNEL_CHUNKED && !k0->is_scalar);
+    const bool k1_ok = (k1->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) ||
+                       (k1->kernel_type == HTP_BINARY_KERNEL_CHUNKED && !k1->is_scalar);
+    if (!k0_ok || !k1_ok) return false;
+
+    const bool is_comm = (cur.opcode == HTP_OP_ADD || cur.opcode == HTP_OP_MUL);
+    const ggml_tensor * new_src = nullptr;
+    if (!htp_match_consumer(prev, cur, is_comm, &new_src)) return false;
+
+    const ggml_tensor * src0 = prev.inputs[0];
+    const ggml_tensor * dst  = cur.dst();
+    if (!ggml_are_same_shape(src0, prev.inputs[1]) ||
+        !ggml_are_same_shape(src0, new_src) ||
+        !htp_same_shape_and_contiguity(prev.dst(), dst)) {
+        return false;
+    }
+    if (htp_inputs_overlap_unsafe(prev, dst, new_src)) return false;
+
+    const size_t vtcm_size = htp_session_vtcm_size(sess);
+    struct htp_binary_kernel_params kparams = *k0;
+    struct htp_binary_vtcm_layout L;
+    memset(&L, 0, sizeof(L));
+
+    const uint32_t n_srcs = curr_n + 1;
+    if (k0->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE || k1->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) {
+        if (k0->kernel_type != HTP_BINARY_KERNEL_SAME_SHAPE) {
+            kparams = *k1;
+        }
+        kparams.n_srcs          = (uint8_t) n_srcs;
+        kparams.rows_per_buffer = 0;
+
+        auto check_split = [&](const ggml_tensor * t) {
+            if (src0->ne[2] > 1 && t->nb[2] != t->ne[1] * t->nb[1]) kparams.split_at_ne01 = 1;
+            if (src0->ne[3] > 1 && t->nb[3] != t->ne[2] * t->nb[2]) kparams.split_at_ne02 = 1;
+        };
+        for (const auto * in : prev.inputs) check_split(in);
+        check_split(new_src);
+        check_split(dst);
+
+        htp_binary_vtcm_layout_build(&L, &kparams, vtcm_size);
+    }
+
+    if (L.rows_per_buffer == 0 || L.total_bytes == 0 || L.total_bytes > vtcm_size) {
+        bool all_contig = ggml_is_contiguous(new_src) && ggml_is_contiguous(dst);
+        for (const auto * in : prev.inputs) {
+            all_contig = all_contig && ggml_is_contiguous(in);
+        }
+        if (!all_contig) return false;
+
+        const uint32_t total_elems = (uint32_t) ggml_nelements(src0);
+        const uint32_t elem_size   = (uint32_t) ggml_type_size(src0->type);
+        const uint32_t max_threads = htp_session_n_threads(sess);
+        if (!htp_binary_solve_chunked(&L, &kparams, total_elems, elem_size, max_threads, n_srcs, false, vtcm_size)) {
+            return false;
+        }
+    }
+
+    kparams.rows_per_buffer = L.rows_per_buffer;
+    kparams.vtcm_size       = (uint32_t) L.total_bytes;
+
+    prev.inputs.push_back(new_src);
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+bool htp_fusion_can_stack(const struct ggml_tensor * a) {
+    if (!a || !a->src[0] || !a->src[1]) return false;
+    switch (a->op) {
+        case GGML_OP_MUL_MAT:
+            if (!htp_fusion_enabled(HTP_FUSE_MUL_MAT_NX)) return false;
+            if (htp_is_hadamard(a)) return false;
+            if (a->src[0]->ne[2] != 1 || a->src[0]->ne[3] != 1) return false;
+            return ggml_is_quantized(a->src[0]->type);
+        case GGML_OP_MUL_MAT_ID:
+            if (!htp_fusion_enabled(HTP_FUSE_MUL_MAT_ID_NX)) return false;
+            return a->src[2] != nullptr && ggml_is_quantized(a->src[0]->type);
+        default:
+            return false;
+    }
+}
+
+bool htp_fusion_can_stack(const struct ggml_tensor * a, const struct ggml_tensor * b) {
+    if (!htp_fusion_can_stack(b) || a->op != b->op || a->src[1] != b->src[1]) {
+        return false;
+    }
+    if (a->op == GGML_OP_MUL_MAT_ID && a->src[2] != b->src[2]) {
+        return false;
+    }
+    const ggml_tensor * w0 = a->src[0];
+    const ggml_tensor * w1 = b->src[0];
+    return w0->type == w1->type && w0->ne[0] == w1->ne[0] && w0->ne[2] == w1->ne[2];
+}
+
+int htp_fusion_max(const struct ggml_cgraph * graph, int idx) {
+    const int n = graph->n_nodes;
+    const ggml_op op0 = graph->nodes[idx]->op;
+    if (op0 != GGML_OP_ADD && op0 != GGML_OP_NORM && op0 != GGML_OP_RMS_NORM) {
+        return 1;
+    }
+
+    enum ggml_op ops[HTP_OP_MAX_FUSED];
+    ops[0] = op0;
+
+    int f = idx + 1;
+    while (f < n && f < idx + HTP_OP_MAX_FUSED) {
+        const ggml_op op = graph->nodes[f]->op;
+        if (op != GGML_OP_ADD && op != GGML_OP_MUL && op != GGML_OP_SCALE &&
+            op != GGML_OP_NORM && op != GGML_OP_RMS_NORM) {
+            break;
+        }
+        ops[f - idx] = op;
+        f++;
+    }
+
+    f -= idx;
+    for (; f > 1; f--) {
+        if (ggml_can_fuse(graph, idx, ops, f)) {
+            break;
+        }
+    }
+    return f;
+}
+
+void htp_fusion_tag_graph(struct ggml_cgraph * graph) {
+    for (int i = 0; i < graph->n_nodes; i++) {
+        auto * extra = (htp_tensor_extra *) graph->nodes[i]->extra;
+        if (!extra) continue;
+
+        extra->flags &= ~HTP_TENSOR_FUSEABLE;
+
+        const ggml_op op = graph->nodes[i]->op;
+        if ((op == GGML_OP_RMS_NORM || op == GGML_OP_SCALE || op == GGML_OP_REPEAT ||
+             op == GGML_OP_UNARY    || op == GGML_OP_ADD   || op == GGML_OP_SUB    ||
+             op == GGML_OP_MUL      || op == GGML_OP_DIV   || op == GGML_OP_MUL_MAT_ID ||
+             op == GGML_OP_SSM_CONV ||
+             (op == GGML_OP_MUL_MAT && !htp_is_hadamard(graph->nodes[i]))) &&
+            ggml_node_has_n_uses(graph, i, 1)) {
+            extra->flags |= HTP_TENSOR_FUSEABLE;
+        }
+    }
+}
+
+bool htp_fusion_try_fuse(const struct ggml_hexagon_session * sess, struct htp_opnode & prev, const struct htp_opnode & cur) {
+    if (!opt_opfusion) return false;
+
+    switch (cur.opcode) {
+        case HTP_OP_ADD:
+            return (htp_fusion_enabled(HTP_FUSE_ALLREDUCE_ADD) && fuse_allreduce_add(sess, prev, cur)) ||
+                   (htp_fusion_enabled(HTP_FUSE_MUL_MAT_ADD)   && fuse_mul_mat_add(sess, prev, cur))   ||
+                   (htp_fusion_enabled(HTP_FUSE_BINARY_NX)     && fuse_binary_nx(sess, prev, cur));
+        case HTP_OP_MUL:
+            return (htp_fusion_enabled(HTP_FUSE_RMS_NORM_MUL)   && fuse_rms_norm_mul(sess, prev, cur))   ||
+                   (htp_fusion_enabled(HTP_FUSE_UNARY_MUL_GLU)  && fuse_unary_mul_glu(sess, prev, cur))  ||
+                   (htp_fusion_enabled(HTP_FUSE_MUL_MUL_SCALAR) && fuse_mul_mul_scalar(sess, prev, cur)) ||
+                   (htp_fusion_enabled(HTP_FUSE_BINARY_NX)      && fuse_binary_nx(sess, prev, cur));
+        case HTP_OP_SUB:
+        case HTP_OP_DIV:
+            return htp_fusion_enabled(HTP_FUSE_BINARY_NX) && fuse_binary_nx(sess, prev, cur);
+        case HTP_OP_SCALE:
+            return htp_fusion_enabled(HTP_FUSE_RMS_NORM_SCALE) && fuse_rms_norm_scale(sess, prev, cur);
+        case HTP_OP_UNARY_SILU:
+            return htp_fusion_enabled(HTP_FUSE_SSM_CONV_SILU) && fuse_ssm_conv_silu(sess, prev, cur);
+        case HTP_OP_GET_ROWS:
+            return htp_fusion_enabled(HTP_FUSE_REPEAT_GET_ROWS) && fuse_repeat_get_rows(sess, prev, cur);
+        case HTP_OP_CPY:
+            return htp_fusion_enabled(HTP_FUSE_GDN_CPY) && fuse_gdn_cpy(sess, prev, cur);
+        case HTP_OP_MUL_MAT:
+            return htp_fusion_enabled(HTP_FUSE_MUL_MAT_NX) && fuse_mul_mat_nx(sess, prev, cur);
+        case HTP_OP_MUL_MAT_ID:
+            return htp_fusion_enabled(HTP_FUSE_MUL_MAT_ID_NX) && fuse_mul_mat_id_nx(sess, prev, cur);
+        default:
+            return false;
+    }
+}

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

@@ -0,0 +1,33 @@
+#ifndef HTP_FUSION_H
+#define HTP_FUSION_H
+
+#include "htp-common.h"
+
+enum htp_fusion_flags {
+    HTP_FUSE_ALLREDUCE_ADD   = (1 << 1),  // 2
+    HTP_FUSE_RMS_NORM_MUL    = (1 << 2),  // 4
+    HTP_FUSE_MUL_MAT_ADD     = (1 << 3),  // 8
+    HTP_FUSE_MUL_MAT_NX      = (1 << 4),  // 16
+    HTP_FUSE_MUL_MAT_ID_NX   = (1 << 5),  // 32
+    HTP_FUSE_GDN_CPY         = (1 << 6),  // 64
+    HTP_FUSE_REPEAT_GET_ROWS = (1 << 7),  // 128
+    HTP_FUSE_UNARY_MUL_GLU   = (1 << 8),  // 256
+    HTP_FUSE_MUL_MUL_SCALAR  = (1 << 9),  // 512
+    HTP_FUSE_BINARY_NX       = (1 << 10), // 1024
+    HTP_FUSE_RMS_NORM_SCALE  = (1 << 11), // 2048
+    HTP_FUSE_SSM_CONV_SILU   = (1 << 12), // 4096
+};
+
+static inline bool htp_fusion_enabled(int flag) {
+    if (opt_opfusion <= 0) return false;
+    if (opt_opfusion == 1) return true; // 1 enables all
+    return (opt_opfusion & flag) != 0;
+}
+
+bool htp_fusion_can_stack(const struct ggml_tensor * a);
+bool htp_fusion_can_stack(const struct ggml_tensor * a, const struct ggml_tensor * b);
+int  htp_fusion_max(const struct ggml_cgraph * graph, int idx);
+void htp_fusion_tag_graph(struct ggml_cgraph * graph);
+bool htp_fusion_try_fuse(const struct ggml_hexagon_session * sess, struct htp_opnode & prev, const struct htp_opnode & cur);
+
+#endif // HTP_FUSION_H

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

@@ -0,0 +1,123 @@
+#ifndef HTP_KPARAMS_H
+#define HTP_KPARAMS_H
+
+#include "htp-common.h"
+
+#include "htp/allreduce-ops.h"
+#include "htp/argsort-ops.h"
+#include "htp/binary-ops.h"
+#include "htp/concat-ops.h"
+#include "htp/cpy-ops.h"
+#include "htp/cumsum-ops.h"
+#include "htp/diag-ops.h"
+#include "htp/fill-ops.h"
+#include "htp/flash-attn-ops.h"
+#include "htp/gated-delta-net-ops.h"
+#include "htp/get-rows-ops.h"
+#include "htp/glu-ops.h"
+#include "htp/im2col-ops.h"
+#include "htp/matmul-ops.h"
+#include "htp/pad-ops.h"
+#include "htp/pool-ops.h"
+#include "htp/repeat-ops.h"
+#include "htp/roll-ops.h"
+#include "htp/rope-ops.h"
+#include "htp/set-rows-ops.h"
+#include "htp/softmax-ops.h"
+#include "htp/solve-tri-ops.h"
+#include "htp/ssm-conv.h"
+#include "htp/sum-rows-ops.h"
+#include "htp/unary-ops.h"
+
+#include <type_traits>
+#include <utility>
+
+bool        htp_rows_stride(const int64_t * ne, const size_t * nb, size_t * stride);
+bool        htp_matmul_can_collapse(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst);
+ggml_tensor htp_tensor_collapse_rows(const struct ggml_tensor * t);
+
+template <typename... Args>
+auto htp_kparams_type(bool (*)(Args...)) -> std::remove_pointer_t<std::remove_reference_t<decltype((std::declval<Args>(), ...))>>;
+
+template <auto SolveFn, typename... Args>
+static inline bool htp_kparams_check(Args &&... args) {
+    decltype(htp_kparams_type(SolveFn)) kparams;
+    return SolveFn(std::forward<Args>(args)..., &kparams);
+}
+
+bool htp_kparams_solve(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    void * kparams = nullptr
+);
+
+static inline bool htp_kparams_check(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op
+) {
+    return htp_kparams_solve(sess, op, nullptr);
+}
+
+bool htp_kparams_solve_allreduce(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * dst,
+    uint32_t rank,
+    uint32_t n_ranks,
+    bool has_add,
+    bool is_row_bcast,
+    bool is_shard_ok,
+    struct htp_allreduce_kernel_params * kparams
+);
+
+bool htp_kparams_solve_rms_norm_mul(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * dst,
+    struct htp_unary_kernel_params * kparams
+);
+
+bool htp_kparams_solve_matmul_add(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * src2,
+    const struct ggml_tensor * dst,
+    struct htp_mm_kernel_params * kparams
+);
+
+bool htp_kparams_solve_mmnx(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * act,
+    int32_t n_weights,
+    struct htp_mm_kernel_params * kparams
+);
+
+bool htp_kparams_solve_mmidnx(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * act,
+    const struct ggml_tensor * dst,
+    int32_t n_weights,
+    struct htp_mm_kernel_params * kparams
+);
+
+bool htp_kparams_solve_repeat_get_rows(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * dst,
+    struct htp_get_rows_kernel_params * kparams
+);
+
+bool htp_kparams_solve_unary_mul_glu(
+    const struct ggml_hexagon_session * sess,
+    uint32_t op,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * dst,
+    struct htp_glu_kernel_params * kparams
+);
+
+#endif // HTP_KPARAMS_H

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

@@ -0,0 +1,260 @@
+#include "htp-opnode.h"
+
+#include "htp/allreduce-ops.h"
+#include "htp/argsort-ops.h"
+#include "htp/binary-ops.h"
+#include "htp/cumsum-ops.h"
+#include "htp/diag-ops.h"
+#include "htp/fill-ops.h"
+#include "htp/flash-attn-ops.h"
+#include "htp/gated-delta-net-ops.h"
+#include "htp/get-rows-ops.h"
+#include "htp/glu-ops.h"
+#include "htp/matmul-ops.h"
+#include "htp/pad-ops.h"
+#include "htp/repeat-ops.h"
+#include "htp/roll-ops.h"
+#include "htp/softmax-ops.h"
+#include "htp/solve-tri-ops.h"
+#include "htp/ssm-conv.h"
+#include "htp/sum-rows-ops.h"
+#include "htp/unary-ops.h"
+
+#include <stdio.h>
+
+std::string htp_opnode::name() const {
+    switch (opcode) {
+        case HTP_OP_MDEV_GROUP:    return "MDEV_GROUP";
+        case HTP_OP_CPY_FENCE:     return "CPY+FENCE";
+        case HTP_OP_FENCE:         return (node && node->op_params[1] == 0) ? "FENCE_WAIT" : "FENCE_SIGNAL";
+        case HTP_OP_ALLREDUCE:     return "ALLREDUCE";
+        case HTP_OP_ALLREDUCE_ADD: return "ALLREDUCE+ADD";
+        case HTP_OP_HADAMARD:      return "MUL_MAT_HADAMARD";
+        case HTP_OP_MUL_MAT_NX:    return "MUL_MAT_NX";
+        case HTP_OP_MUL_MAT_ID_NX: return "MUL_MAT_ID_NX";
+        case HTP_OP_ADD:
+        case HTP_OP_SUB:
+        case HTP_OP_MUL:
+        case HTP_OP_DIV:
+            if (!fused.empty()) {
+                const auto * kp = kparams<htp_binary_kernel_params>();
+                if (kp->kernel_type != HTP_BINARY_KERNEL_SCALAR) {
+                    switch (opcode) {
+                        case HTP_OP_ADD: return "ADD_NX";
+                        case HTP_OP_SUB: return "SUB_NX";
+                        case HTP_OP_MUL: return "MUL_NX";
+                        case HTP_OP_DIV: return "DIV_NX";
+                        default: break;
+                    }
+                }
+            }
+            break;
+        default:
+            break;
+    }
+
+    if (!node) {
+        return "";
+    }
+    if (htp_is_hadamard(node)) {
+        return "MUL_MAT_HADAMARD";
+    }
+
+    std::string res = ggml_op_desc(node);
+    for (const auto * t : fused) {
+        res += "+";
+        res += ggml_op_desc(t);
+    }
+    return res;
+}
+
+static const char * format_tensor_name(const struct ggml_tensor * t) {
+    return t ? t->name : "NONE";
+}
+
+static std::string format_tensor_dims(const struct ggml_tensor * t) {
+    if (!t) {
+        return "NONE";
+    }
+    char buf[64];
+    if (t->ne[2] == 1 && t->ne[3] == 1) {
+        snprintf(buf, sizeof(buf), "%d:%d", (int) t->ne[0], (int) t->ne[1]);
+    } else {
+        snprintf(buf, sizeof(buf), "%d:%d:%d:%d", (int) t->ne[0], (int) t->ne[1], (int) t->ne[2], (int) t->ne[3]);
+    }
+    return buf;
+}
+
+static const char * format_tensor_type(const struct ggml_tensor * t) {
+    return t ? ggml_type_name(t->type) : "NONE";
+}
+
+static std::string format_tensor_strides(const struct ggml_tensor * t) {
+    if (!t) {
+        return "NONE";
+    }
+    char buf[64];
+    const char * c = ggml_is_contiguous(t) ? "" : "!";
+    if (t->ne[2] == 1 && t->ne[3] == 1) {
+        snprintf(buf, sizeof(buf), "%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], c);
+    } else {
+        snprintf(buf, sizeof(buf), "%zu:%zu:%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], (size_t) t->nb[2], (size_t) t->nb[3], c);
+    }
+    return buf;
+}
+
+static const char * format_tensor_buff(const struct ggml_tensor * t) {
+    if (t && t->buffer) {
+        return ggml_backend_buffer_name(t->buffer);
+    }
+    return "NONE";
+}
+
+template <typename Fn>
+static std::string format_op_tensors(const htp_opnode & node, Fn fn) {
+    std::string res;
+    const auto & inputs  = node.get_inputs();
+    const auto & outputs = node.get_outputs();
+
+    if (!inputs.empty()) {
+        for (size_t i = 0; i < inputs.size(); i++) {
+            if (i > 0) res += " x ";
+            res += fn(inputs[i]);
+        }
+        res += " -> ";
+    }
+
+    if (outputs.empty()) {
+        res += fn(node.dst());
+    } else {
+        for (size_t i = 0; i < outputs.size(); i++) {
+            if (i > 0) res += " x ";
+            res += fn(outputs[i]);
+        }
+    }
+    return res;
+}
+
+static std::string format_kernel_params(const htp_opnode & node) {
+    char str[128];
+    const size_t max_size = sizeof(str);
+
+    if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID ||
+        node.opcode == HTP_OP_MUL_MAT_NX || node.opcode == HTP_OP_MUL_MAT_ID_NX ||
+        node.opcode == HTP_OP_MUL_MAT_ADD) {
+        const auto * kparams = node.kparams<htp_mm_kernel_params>();
+        const char * path = "unknown";
+        int32_t type = kparams->kernel_type;
+        if (type == HTP_MM_KERNEL_HMX_2D || type == HTP_MM_KERNEL_HMX_BATCHED) {
+            path = (kparams->tile_size > 0) ? "hmx-tiled" : "hmx-flat";
+        } else if (type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || type == HTP_MM_KERNEL_HVX_F32_F32_VTCM ||
+                   type == HTP_MM_KERNEL_HVX_BF16_F32_VTCM) {
+            path = "hvx-flat";
+        } else if (type == HTP_MM_KERNEL_HVX_QUANT_ROW || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
+            path = "hvx-tiled";
+        }
+        snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) {
+        const auto * kparams = node.kparams<htp_fa_kernel_params>();
+        const char * path = "unknown";
+        int32_t type = kparams->kernel_type;
+        if (type == HTP_FA_KERNEL_HMX) {
+            path = kparams->u.hmx.pipeline ? "hmx-pipe" : "hmx-seq";
+        } else if (type == HTP_FA_KERNEL_HVX) {
+            path = "hvx";
+        }
+        snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
+    } else if (htp_op_is_unary(node.opcode)) {
+        const auto * kparams = node.kparams<htp_unary_kernel_params>();
+        snprintf(str, max_size, "%s vtcm %d", kparams->col_tile ? "wide-row" : "row-block", (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_MDEV_GROUP && node.node) {
+        snprintf(str, max_size, "idx %d count %d", (int) node.node->op_params[0], (int) node.dst()->ne[1]);
+    } else if ((node.opcode == HTP_OP_FENCE || node.opcode == HTP_OP_CPY_FENCE) && node.node) {
+        snprintf(str, max_size, "seq 0x%x", (uint32_t) node.node->op_params[0]);
+    } else if (node.opcode == HTP_OP_ALLREDUCE && node.node) {
+        snprintf(str, max_size, "seq 0x%x -> 0x%x", (uint32_t) node.node->op_params[0], (uint32_t) node.node->op_params[1]);
+    } else if (node.opcode == HTP_OP_SSM_CONV) {
+        const auto * kparams = node.kparams<htp_ssm_conv_kernel_params>();
+        snprintf(str, max_size, "%s vtcm %d", kparams->n_t == 1 ? "decode" : "prefill", (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_SOFTMAX) {
+        const auto * kparams = node.kparams<htp_softmax_kernel_params>();
+        snprintf(str, max_size, "k%d nth %d vtcm %d", (int) kparams->kernel_type, (int) kparams->n_threads, (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_GATED_DELTA_NET) {
+        const auto * kparams = node.kparams<htp_gdn_kernel_params>();
+        const char * path = (kparams->kernel_type == HTP_GDN_KERNEL_HMX_CHUNKED) ? "hmx-chunked" : "hvx-recurrent";
+        snprintf(str, max_size, "%s-%s vtcm %u", path, kparams->kda ? "kda" : "scalar",
+                 (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_MUL || node.opcode == HTP_OP_ADD || node.opcode == HTP_OP_ADD_ID ||
+               node.opcode == HTP_OP_SUB || node.opcode == HTP_OP_DIV) {
+        const auto * kparams = node.kparams<htp_binary_kernel_params>();
+        snprintf(str, max_size, "vtcm %u", (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_ARGSORT || node.opcode == HTP_OP_TOP_K) {
+        const auto * kparams = node.kparams<htp_sort_kernel_params>();
+        snprintf(str, max_size, "%s nth %d nchk %d chk %d vtcm %d",
+                 node.opcode == HTP_OP_TOP_K ? "top_k" : "argsort",
+                 (int) kparams->n_threads, (int) kparams->n_chunks,
+                 (int) kparams->chunk_elems, (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_GET_ROWS) {
+        const auto * kparams = node.kparams<htp_get_rows_kernel_params>();
+        const char * ktype_str = "unknown";
+        switch (kparams->kernel_type) {
+            case HTP_GET_ROWS_KERNEL_SAMETYPE: ktype_str = "sametype"; break;
+            case HTP_GET_ROWS_KERNEL_TILED:    ktype_str = "tiled"; break;
+            case HTP_GET_ROWS_KERNEL_FLAT:     ktype_str = "flat"; break;
+        }
+        snprintf(str, max_size, "%s%s vtcm %u", ktype_str, kparams->n_threads > 1 ? "-multi" : "", (unsigned int) kparams->vtcm_size);
+    } else if (htp_op_is_glu(node.opcode)) {
+        const auto * kparams = node.kparams<htp_glu_kernel_params>();
+        snprintf(str, max_size, "nth %d vtcm %u", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_PAD) {
+        const auto * kparams = node.kparams<htp_pad_kernel_params>();
+        snprintf(str, max_size, "k%d nth %d vtcm %u", (int) kparams->kernel_type, (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_CUMSUM) {
+        const auto * kparams = node.kparams<htp_cumsum_kernel_params>();
+        snprintf(str, max_size, "%s nth %d vtcm %u", kparams->use_dma ? "dma" : "hvx", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_FILL) {
+        const auto * kparams = node.kparams<htp_fill_kernel_params>();
+        snprintf(str, max_size, "%s nth %d", kparams->kernel_type == HTP_FILL_KERNEL_CONTIGUOUS ? "contig" : "strided", (int) kparams->n_threads);
+    } else if (node.opcode == HTP_OP_ROLL) {
+        const auto * kparams = node.kparams<htp_roll_kernel_params>();
+        snprintf(str, max_size, "k%d nth %d", (int) kparams->kernel_type, (int) kparams->n_threads);
+    } else if (node.opcode == HTP_OP_DIAG) {
+        const auto * kparams = node.kparams<htp_diag_kernel_params>();
+        snprintf(str, max_size, "%s nth %d vtcm %u", kparams->use_dma ? "dma" : "hvx", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_SOLVE_TRI) {
+        const auto * kparams = node.kparams<htp_solve_tri_kernel_params>();
+        snprintf(str, max_size, "%s nth %d", kparams->batched ? "batched" : "chunked", (int) kparams->n_threads);
+    } else {
+        snprintf(str, max_size, "----");
+    }
+    return str;
+}
+
+static void append_section(std::string & res, const std::string & sec) {
+    if (!res.empty()) {
+        res += "|";
+    }
+    res += sec;
+}
+
+std::string htp_format_tensor(const ggml_tensor * t, uint32_t flags) {
+    std::string res;
+    if (flags & HTP_OPFMT_NAMES)   append_section(res, format_tensor_name(t));
+    if (flags & HTP_OPFMT_DIMS)    append_section(res, format_tensor_dims(t));
+    if (flags & HTP_OPFMT_TYPES)   append_section(res, format_tensor_type(t));
+    if (flags & HTP_OPFMT_STRIDES) append_section(res, format_tensor_strides(t));
+    if (flags & HTP_OPFMT_BUFFS)   append_section(res, format_tensor_buff(t));
+    return res;
+}
+
+std::string htp_format_op(const htp_opnode & node, uint32_t flags) {
+    std::string res;
+    if (flags & HTP_OPFMT_OP)      append_section(res, node.name());
+    if (flags & HTP_OPFMT_NAMES)   append_section(res, format_op_tensors(node, format_tensor_name));
+    if (flags & HTP_OPFMT_DIMS)    append_section(res, format_op_tensors(node, format_tensor_dims));
+    if (flags & HTP_OPFMT_TYPES)   append_section(res, format_op_tensors(node, format_tensor_type));
+    if (flags & HTP_OPFMT_STRIDES) append_section(res, format_op_tensors(node, format_tensor_strides));
+    if (flags & HTP_OPFMT_BUFFS)   append_section(res, format_op_tensors(node, format_tensor_buff));
+    if (flags & HTP_OPFMT_KPARAMS) append_section(res, format_kernel_params(node));
+    return res;
+}

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

@@ -1,38 +1,74 @@
 #ifndef HTP_OPNODE_H
 #define HTP_OPNODE_H
 
-#define GGML_COMMON_IMPL_CPP
 #include "ggml-backend-impl.h"
-#include "ggml-common.h"
+#include "ggml-impl.h"
+#include "htp/htp-ops.h"
 
-#include <algorithm>
+#include <cstring>
+#include <initializer_list>
+#include <memory>
 #include <string>
 #include <vector>
-#include <memory>
-#include <stdio.h>
-#include "htp-ops.h"
-#include "htp/matmul-ops.h"
-#include "htp/flash-attn-ops.h"
-#include "htp/unary-ops.h"
-#include "htp/binary-ops.h"
-#include "htp/allreduce-ops.h"
-#include "htp/ssm-conv.h"
-#include "htp/gated-delta-net-ops.h"
-#include "htp/softmax-ops.h"
-#include "htp/argsort-ops.h"
-#include "htp/get-rows-ops.h"
+
+#define HTP_OP_MAX_FUSED 16
+
+// local std::inplace_vector replacement
+template <typename T, size_t MAX_N>
+struct htp_opvec {
+    T      data[MAX_N] {};
+    size_t n = 0;
+
+    htp_opvec() = default;
+    htp_opvec(std::initializer_list<T> init) {
+        GGML_ASSERT(init.size() <= MAX_N);
+        for (const auto & v : init) {
+            data[n++] = v;
+        }
+    }
+
+    size_t    size()  const { return n; }
+    bool      empty() const { return n == 0; }
+    void      clear()       { n = 0; }
+    T *       begin()       { return data; }
+    T *       end()         { return data + n; }
+    const T * begin() const { return data; }
+    const T * end()   const { return data + n; }
+    T &       back()        { return data[n - 1]; }
+    const T & back()  const { return data[n - 1]; }
+    T &       operator[](size_t i)       { return data[i]; }
+    const T & operator[](size_t i) const { return data[i]; }
+
+    void push_back(T val) {
+        GGML_ASSERT(n < MAX_N);
+        data[n++] = val;
+    }
+
+    void insert(T * pos, T val) {
+        GGML_ASSERT(n < MAX_N);
+        for (T * p = data + n; p > pos; --p) {
+            *p = *(p - 1);
+        }
+        *pos = val;
+        n++;
+    }
+};
+
+static inline bool htp_is_hadamard(const struct ggml_tensor * t) {
+    return t && t->op == GGML_OP_MUL_MAT && ggml_get_op_params_i32(t, 1) == GGML_HINT_SRC0_IS_HADAMARD;
+}
 
 struct htp_opnode {
-    ggml_tensor * node   { nullptr };
-    htp_op_code   opcode { HTP_OP_INVALID };
-    int32_t       kernel_params[HTP_OP_MAX_KERN_PARAMS] {0};
+    const ggml_tensor * node   { nullptr };
+    htp_op_code         opcode { HTP_OP_INVALID };
+    int32_t             params[HTP_OP_MAX_PARAMS] {0};
+    int32_t             kernel_params[HTP_OP_MAX_KERN_PARAMS] {0};
 
-    std::vector<ggml_tensor *>                fused;
-    std::vector<std::shared_ptr<ggml_tensor>> dummy;
+    std::vector<std::shared_ptr<ggml_tensor>>          dummy;
 
-    std::vector<const ggml_tensor *> inputs;
-    std::vector<const ggml_tensor *> outputs;
-    std::string                      name;
+    htp_opvec<const ggml_tensor *, HTP_OP_MAX_FUSED>   fused;
+    htp_opvec<const ggml_tensor *, HTP_OP_MAX_INPUTS>  inputs;
+    htp_opvec<const ggml_tensor *, HTP_OP_MAX_OUTPUTS> outputs;
 
     int n_active_src(const ggml_tensor * t) const {
         if (!t) return 0;
@@ -44,34 +80,36 @@ struct htp_opnode {
         return 0;
     }
 
-    void init(ggml_tensor * node) {
+    void init(const ggml_tensor * node) {
+        static_assert(sizeof(this->params) == sizeof(this->node->op_params), "params size mismatch");
+
         this->node = node;
-        if (this->node) {
-            if (this->node->op == GGML_OP_MUL_MAT &&
-                ggml_get_op_params_i32(this->node, 1) == GGML_HINT_SRC0_IS_HADAMARD) {
-                this->name = "MUL_MAT_HADAMARD";
-                this->inputs.push_back(this->node->src[1]);
-            } else {
-                this->name = ggml_op_desc(this->node);
+        if (!node) return;
 
-                // Build inputs (preserving optional nullptrs)
-                int n_inputs = n_active_src(this->node);
-                this->inputs.resize(n_inputs, nullptr);
-                for (int i = 0; i < n_inputs; i++) {
-                    this->inputs[i] = this->node->src[i];
-                }
+        memcpy(this->params, this->node->op_params, sizeof(this->params));
+
+        // Build inputs (preserving optional nullptrs)
+        if (htp_is_hadamard(this->node)) {
+            this->inputs.push_back(this->node->src[1]);
+        } else {
+            int n_inputs = n_active_src(this->node);
+            for (int i = 0; i < n_inputs; i++) {
+                this->inputs.push_back(this->node->src[i]);
             }
-
-            // Build outputs
-            this->outputs.push_back(this->dst());
         }
+
+        // Build outputs
+        this->outputs.push_back(this->dst());
     }
 
-    htp_opnode(htp_op_code opcode = HTP_OP_INVALID, ggml_tensor * node = nullptr) : opcode(opcode) {
+    htp_opnode(htp_op_code opcode = HTP_OP_INVALID, const ggml_tensor * node = nullptr) : opcode(opcode) {
+        init(node);
+    }
+
+    htp_opnode(const ggml_tensor * node) : opcode(HTP_OP_INVALID) {
         init(node);
     }
 
-    ggml_op             op()   const { return node->op; }
     const ggml_tensor * src0() const { return node->src[0]; }
     const ggml_tensor * src1() const { return node->src[1]; }
     const ggml_tensor * dst()  const { return outputs.empty() ? node : outputs.back(); }
@@ -81,343 +119,54 @@ struct htp_opnode {
         return dummy.back().get();
     }
 
-    void add_fused(ggml_tensor * t, bool extra_dst = false) {
-        fused.push_back(t);
-
-        name += "+";
-        name += ggml_op_desc(t);
-
-        if (extra_dst) {
-            outputs.push_back(t);
-        } else {
-            outputs.clear();
-            outputs.push_back(t);
-        }
-
-        // Remove the newly fused intermediate output tensor t from inputs (if it was there)
-        inputs.erase(std::remove(inputs.begin(), inputs.end(), t), inputs.end());
-
-        // Append new inputs from t, preserving middle nullptrs
-        int n_inputs = n_active_src(t);
-        for (int i = 0; i < n_inputs; i++) {
-            const auto * src = t->src[i];
-            if (!src) {
-                inputs.push_back(nullptr);
-            } else if (src != node &&
-                       std::find(fused.begin(), fused.end(), src) == fused.end() &&
-                       std::find(inputs.begin(), inputs.end(), src) == inputs.end()) {
-                inputs.push_back(src);
-            }
-        }
-    }
-
-    const std::vector<const ggml_tensor *> & get_inputs() const {
+    const htp_opvec<const ggml_tensor *, HTP_OP_MAX_INPUTS> & get_inputs() const {
         return inputs;
     }
 
-    const std::vector<const ggml_tensor *> & get_outputs() const {
+    const htp_opvec<const ggml_tensor *, HTP_OP_MAX_OUTPUTS> & get_outputs() const {
         return outputs;
     }
 
-    std::string op_name() const {
-        return name;
+    template <typename T>
+    T * kparams() {
+        static_assert(sizeof(T) <= sizeof(kernel_params), "kparams size exceeds HTP_OP_MAX_KERN_PARAMS");
+        static_assert(alignof(T) <= alignof(int32_t), "kparams alignment exceeds int32_t");
+        return reinterpret_cast<T *>(kernel_params);
     }
 
-    bool is_empty() const {
-        return ggml_op_is_empty(node->op);
+    template <typename T>
+    const T * kparams() const {
+        static_assert(sizeof(T) <= sizeof(kernel_params), "kparams size exceeds HTP_OP_MAX_KERN_PARAMS");
+        static_assert(alignof(T) <= alignof(int32_t), "kparams alignment exceeds int32_t");
+        return reinterpret_cast<const T *>(kernel_params);
     }
 
-    bool stackable() const {
-        switch (this->op()) {
-            case GGML_OP_MUL_MAT:
-                if (ggml_get_op_params_i32(this->node, 1) == GGML_HINT_SRC0_IS_HADAMARD) {
-                    return false;
-                }
-                return ggml_is_quantized(this->src0()->type);
-            case GGML_OP_MUL_MAT_ID:
-                return ggml_is_quantized(this->src0()->type);
-            default:
-                return false;
-        }
+    template <typename T>
+    void set_kparams(const T & kp) {
+        static_assert(sizeof(T) <= sizeof(kernel_params), "kparams size exceeds HTP_OP_MAX_KERN_PARAMS");
+        static_assert(alignof(T) <= alignof(int32_t), "kparams alignment exceeds int32_t");
+        memcpy(kernel_params, &kp, sizeof(T));
     }
 
-    bool same_input(const htp_opnode& n) const {
-        return n.src1() == this->src1();
-    }
+    std::string name() const;
 };
 
-struct htp_opformat {
-    char strides[64 * GGML_MAX_SRC];
-    char dims[64 * GGML_MAX_SRC];
-    char types[16 * GGML_MAX_SRC];
-    char buffs[64 * GGML_MAX_SRC];
-    char names[64 * GGML_MAX_SRC];
-    char kparams[128];
+enum htp_opformat_flags {
+    HTP_OPFMT_OP      = (1 << 0),
+    HTP_OPFMT_NAMES   = (1 << 1),
+    HTP_OPFMT_DIMS    = (1 << 2),
+    HTP_OPFMT_TYPES   = (1 << 3),
+    HTP_OPFMT_STRIDES = (1 << 4),
+    HTP_OPFMT_BUFFS   = (1 << 5),
+    HTP_OPFMT_KPARAMS = (1 << 6),
 
-    int format_tensor_dims(char * str, size_t max_size, const struct ggml_tensor * t) {
-        if (!t) {
-            return snprintf(str, max_size, "NONE");
-        }
-        if (t->ne[2] == 1 && t->ne[3] == 1) {
-            return snprintf(str, max_size, "%d:%d", (int) t->ne[0], (int) t->ne[1]);
-        } else {
-            return snprintf(str, max_size, "%d:%d:%d:%d", (int) t->ne[0], (int) t->ne[1], (int) t->ne[2], (int) t->ne[3]);
-        }
-    }
-
-    void format_op_dims(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            p += std::min((size_t)format_tensor_dims(p, p_end - p, inputs[0]), (size_t)(p_end - p));
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)format_tensor_dims(p, p_end - p, inputs[i]), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        char self[64];
-        format_tensor_dims(self, sizeof(self), node.dst());
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", self), (size_t)(p_end - p));
-        }
-    }
-
-    int format_tensor_strides(char * str, size_t max_size, const struct ggml_tensor * t) {
-        if (!t) {
-            return snprintf(str, max_size, "NONE");
-        }
-        const char * c = ggml_is_contiguous(t) ? "" : "!";
-
-        if (t->ne[2] == 1 && t->ne[3] == 1) {
-            return snprintf(str, max_size, "%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], c);
-        } else {
-            return snprintf(str, max_size, "%zu:%zu:%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], (size_t) t->nb[2], (size_t) t->nb[3], c);
-        }
-    }
-
-    void format_op_strides(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            p += std::min((size_t)format_tensor_strides(p, p_end - p, inputs[0]), (size_t)(p_end - p));
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)format_tensor_strides(p, p_end - p, inputs[i]), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        char self[64];
-        format_tensor_strides(self, sizeof(self), node.dst());
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", self), (size_t)(p_end - p));
-        }
-    }
-
-    void format_op_types(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[0] ? ggml_type_name(inputs[0]->type) : "NONE"), (size_t)(p_end - p));
-            }
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[i] ? ggml_type_name(inputs[i]->type) : "NONE"), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", ggml_type_name(node.dst()->type)), (size_t)(p_end - p));
-        }
-    }
-
-    const char * tensor_buff_name(const struct ggml_tensor * t) {
-        if (t && t->buffer) {
-            return ggml_backend_buffer_name(t->buffer);
-        }
-        return "NONE";
-    }
-
-    void format_op_buffs(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, "%s", tensor_buff_name(inputs[0])), (size_t)(p_end - p));
-            }
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, "%s", tensor_buff_name(inputs[i])), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", tensor_buff_name(node.dst())), (size_t)(p_end - p));
-        }
-    }
-
-    void format_op_names(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[0] ? inputs[0]->name : "NONE"), (size_t)(p_end - p));
-            }
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[i] ? inputs[i]->name : "NONE"), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", node.dst()->name), (size_t)(p_end - p));
-        }
-    }
-    void format_kernel_params(char * str, size_t max_size, const htp_opnode & node) {
-        if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID ||
-            node.opcode == HTP_OP_MUL_MAT_NX || node.opcode == HTP_OP_MUL_MAT_ID_NX ||
-            node.opcode == HTP_OP_MUL_MAT_ADD) {
-            const auto * kparams = (const struct htp_mm_kernel_params *) node.kernel_params;
-            const char * path = "unknown";
-            int32_t type = kparams->kernel_type;
-            if (type == HTP_MM_KERNEL_HMX_2D || type == HTP_MM_KERNEL_HMX_BATCHED) {
-                path = (kparams->tile_size > 0) ? "hmx-tiled" : "hmx-flat";
-            } else if (type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || type == HTP_MM_KERNEL_HVX_F32_F32_VTCM ||
-                       type == HTP_MM_KERNEL_HVX_BF16_F32_VTCM) {
-                path = "hvx-flat";
-            } else if (type == HTP_MM_KERNEL_HVX_QUANT_ROW || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
-                path = "hvx-tiled";
-            }
-            snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) {
-            const auto * kparams = (const struct htp_fa_kernel_params *) node.kernel_params;
-            const char * path = "unknown";
-            int32_t type = kparams->kernel_type;
-            if (type == HTP_FA_KERNEL_HMX) {
-                path = kparams->u.hmx.pipeline ? "hmx-pipe" : "hmx-seq";
-            } else if (type == HTP_FA_KERNEL_HVX) {
-                path = "hvx";
-            }
-            snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
-        } else if (htp_op_is_unary(node.opcode)) {
-            const auto * kparams = (const struct htp_unary_kernel_params *) node.kernel_params;
-            snprintf(str, max_size, "%s vtcm %d", kparams->col_tile ? "wide-row" : "row-block", (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_MDEV_GROUP && node.node) {
-            snprintf(str, max_size, "idx %d count %d", (int) node.node->op_params[0], (int) node.dst()->ne[1]);
-        } else if ((node.opcode == HTP_OP_FENCE || node.opcode == HTP_OP_CPY_FENCE) && node.node) {
-            snprintf(str, max_size, "seq 0x%x", (uint32_t) node.node->op_params[0]);
-        } else if (node.opcode == HTP_OP_ALLREDUCE && node.node) {
-            snprintf(str, max_size, "seq 0x%x -> 0x%x", (uint32_t) node.node->op_params[0], (uint32_t) node.node->op_params[1]);
-        } else if (node.opcode == HTP_OP_SSM_CONV) {
-            const auto * kparams = (const struct htp_ssm_conv_kernel_params *) node.kernel_params;
-            snprintf(str, max_size, "%s vtcm %d", kparams->n_t == 1 ? "decode" : "prefill", (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_SOFTMAX) {
-            const auto * kparams = (const struct htp_softmax_kernel_params *) node.kernel_params;
-            snprintf(str, max_size, "k%d nth %d vtcm %d", (int) kparams->kernel_id, (int) kparams->n_threads, (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_GATED_DELTA_NET) {
-            const auto * kparams = (const struct htp_gdn_kernel_params *) node.kernel_params;
-            const char * path = (kparams->kernel_type == HTP_GDN_KERNEL_HMX_CHUNKED) ? "hmx-chunked" : "hvx-recurrent";
-            snprintf(str, max_size, "%s-%s vtcm %u", path, kparams->kda ? "kda" : "scalar",
-                     (unsigned int) (kparams->vtcm_size ? kparams->vtcm_size : kparams->vtcm_per_thread * kparams->n_threads));
-        } else if (node.opcode == HTP_OP_MUL || node.opcode == HTP_OP_ADD || node.opcode == HTP_OP_ADD_ID ||
-                   node.opcode == HTP_OP_SUB || node.opcode == HTP_OP_DIV) {
-            const auto * kparams = (const struct htp_binary_kernel_params *) node.kernel_params;
-            snprintf(str, max_size, "vtcm %u", (unsigned int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_ARGSORT || node.opcode == HTP_OP_TOP_K) {
-            const auto * kparams = (const struct htp_sort_kernel_params *) node.kernel_params;
-            snprintf(str, max_size, "%s nth %d nchk %d chk %d vtcm %d",
-                     node.opcode == HTP_OP_TOP_K ? "top_k" : "argsort",
-                     (int) kparams->n_threads, (int) kparams->n_chunks,
-                     (int) kparams->chunk_elems, (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_GET_ROWS) {
-            const auto * kparams = (const struct htp_get_rows_kernel_params *) node.kernel_params;
-            const char * ktype_str = "unknown";
-            switch (kparams->kernel_type) {
-                case HTP_GET_ROWS_KERNEL_SAMETYPE: ktype_str = "sametype"; break;
-                case HTP_GET_ROWS_KERNEL_TILED:    ktype_str = "tiled"; break;
-                case HTP_GET_ROWS_KERNEL_FLAT:     ktype_str = "flat"; break;
-            }
-            snprintf(str, max_size, "%s%s vtcm %u", ktype_str, kparams->n_threads > 1 ? "-multi" : "", (unsigned int) kparams->vtcm_size);
-        } else {
-            snprintf(str, max_size, "----");
-        }
-
-    }
-
-    void format(const htp_opnode & node) {
-        format_op_dims(dims, sizeof(dims), node);
-        format_op_strides(strides, sizeof(strides), node);
-        format_op_types(types, sizeof(types), node);
-        format_op_buffs(buffs, sizeof(buffs), node);
-        format_op_names(names, sizeof(names), node);
-        format_kernel_params(kparams, sizeof(kparams), node);
-    }
-
-    htp_opformat() {
-        strides[0] = '\0';
-        dims[0]    = '\0';
-        types[0]   = '\0';
-        buffs[0]   = '\0';
-        names[0]   = '\0';
-        kparams[0] = '\0';
-    }
-    htp_opformat(const htp_opnode & node) { format(node); }
+    HTP_OPFMT_EXEC    = HTP_OPFMT_OP | HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES | HTP_OPFMT_BUFFS | HTP_OPFMT_KPARAMS,
+    HTP_OPFMT_SUPP    = HTP_OPFMT_OP | HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES | HTP_OPFMT_BUFFS,
+    HTP_OPFMT_PROF    = HTP_OPFMT_OP | HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES | HTP_OPFMT_KPARAMS,
+    HTP_OPFMT_TENSOR  = HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES,
 };
 
+std::string htp_format_op(const htp_opnode & node, uint32_t flags = HTP_OPFMT_EXEC);
+std::string htp_format_tensor(const ggml_tensor * t, uint32_t flags = HTP_OPFMT_TENSOR);
+
 #endif // HTP_OPNODE_H

ggml/src/ggml-hexagon/htp-repack.cpp

@@ -9,6 +9,11 @@
 #include <cmath>
 #include <cstring>
 #include <initializer_list>
+#include <thread>
+
+#if defined(__ARM_NEON) || defined(__aarch64__)
+#    include <arm_neon.h>
+#endif
 
 // --- 4-bit block-32 nibble helpers (shared by Q4_0, IQ4_NL, Q4_1, MXFP4) ---
 
@@ -43,8 +48,309 @@ static inline void transpose_tile_to_nibbles(uint8_t tile_quants[32][32], const
     }
 }
 
+#if defined(__ARM_NEON) || defined(__aarch64__)
+static void repack_neon_transpose8x8(const uint8x8_t rows[8], uint8x8_t columns[8]) {
+    const uint8x8x2_t t0 = vtrn_u8(rows[0], rows[1]);
+    const uint8x8x2_t t1 = vtrn_u8(rows[2], rows[3]);
+    const uint8x8x2_t t2 = vtrn_u8(rows[4], rows[5]);
+    const uint8x8x2_t t3 = vtrn_u8(rows[6], rows[7]);
+    const uint16x4x2_t t4 = vtrn_u16(vreinterpret_u16_u8(t0.val[0]), vreinterpret_u16_u8(t1.val[0]));
+    const uint16x4x2_t t5 = vtrn_u16(vreinterpret_u16_u8(t0.val[1]), vreinterpret_u16_u8(t1.val[1]));
+    const uint16x4x2_t t6 = vtrn_u16(vreinterpret_u16_u8(t2.val[0]), vreinterpret_u16_u8(t3.val[0]));
+    const uint16x4x2_t t7 = vtrn_u16(vreinterpret_u16_u8(t2.val[1]), vreinterpret_u16_u8(t3.val[1]));
+    const uint32x2x2_t t8 = vtrn_u32(vreinterpret_u32_u16(t4.val[0]), vreinterpret_u32_u16(t6.val[0]));
+    const uint32x2x2_t t9 = vtrn_u32(vreinterpret_u32_u16(t4.val[1]), vreinterpret_u32_u16(t6.val[1]));
+    const uint32x2x2_t ta = vtrn_u32(vreinterpret_u32_u16(t5.val[0]), vreinterpret_u32_u16(t7.val[0]));
+    const uint32x2x2_t tb = vtrn_u32(vreinterpret_u32_u16(t5.val[1]), vreinterpret_u32_u16(t7.val[1]));
+    columns[0] = vreinterpret_u8_u32(t8.val[0]);
+    columns[1] = vreinterpret_u8_u32(ta.val[0]);
+    columns[2] = vreinterpret_u8_u32(t9.val[0]);
+    columns[3] = vreinterpret_u8_u32(tb.val[0]);
+    columns[4] = vreinterpret_u8_u32(t8.val[1]);
+    columns[5] = vreinterpret_u8_u32(ta.val[1]);
+    columns[6] = vreinterpret_u8_u32(t9.val[1]);
+    columns[7] = vreinterpret_u8_u32(tb.val[1]);
+}
+
+static uint8x8_t repack_neon_pack_4bit_pairs(uint8x8_t first, uint8x8_t second, bool high) {
+    const uint8x8x2_t separated = vuzp_u8(first, second);
+    const uint8x8_t mask = vdup_n_u8(0x0F);
+    uint8x8_t even = separated.val[0];
+    uint8x8_t odd  = separated.val[1];
+
+    if (high) {
+        even = vshr_n_u8(even, 4);
+        odd  = vshr_n_u8(odd, 4);
+    }
+
+    even = vand_u8(even, mask);
+    odd  = vand_u8(odd, mask);
+    return vorr_u8(vshl_n_u8(odd, 4), even);
+}
+
+static void repack_neon_4bit_rows(const uint8_t * qs[8], uint8_t * tile_dst, int row0) {
+    uint8x8_t low_rows[8];
+    uint8x8_t high_rows[8];
+    uint8x8_t columns[8];
+
+    for (int rr = 0; rr < 8; ++rr) {
+        const uint8x16_t packed = vld1q_u8(qs[rr]);
+        const uint8x8_t first = vget_low_u8(packed);
+        const uint8x8_t second = vget_high_u8(packed);
+        low_rows[rr]  = repack_neon_pack_4bit_pairs(first, second, false);
+        high_rows[rr] = repack_neon_pack_4bit_pairs(first, second, true);
+    }
+
+    repack_neon_transpose8x8(low_rows, columns);
+    for (int cp = 0; cp < 8; ++cp) {
+        vst1_u8(tile_dst + cp * 32 + row0, columns[cp]);
+    }
+
+    repack_neon_transpose8x8(high_rows, columns);
+    for (int cp = 0; cp < 8; ++cp) {
+        vst1_u8(tile_dst + (cp + 8) * 32 + row0, columns[cp]);
+    }
+}
+
+static void repack_neon_k4_rows(const uint8_t * qs[8], uint8_t * tile_dst, int row0, bool high) {
+    uint8x8_t first_rows[8];
+    uint8x8_t second_rows[8];
+    uint8x8_t columns[8];
+
+    for (int rr = 0; rr < 8; ++rr) {
+        first_rows[rr] = repack_neon_pack_4bit_pairs(vld1_u8(qs[rr]), vld1_u8(qs[rr] + 8), high);
+        second_rows[rr] = repack_neon_pack_4bit_pairs(vld1_u8(qs[rr] + 16), vld1_u8(qs[rr] + 24), high);
+    }
+
+    repack_neon_transpose8x8(first_rows, columns);
+    for (int cp = 0; cp < 8; ++cp) {
+        vst1_u8(tile_dst + cp * 32 + row0, columns[cp]);
+    }
+
+    repack_neon_transpose8x8(second_rows, columns);
+    for (int cp = 0; cp < 8; ++cp) {
+        vst1_u8(tile_dst + (cp + 8) * 32 + row0, columns[cp]);
+    }
+}
+
+static void repack_neon_q8_rows(const uint8_t * qs[8], uint8_t * tile_dst, int row0) {
+    uint8x8_t rows[8];
+    uint8x8_t columns[8];
+
+    for (int cp = 0; cp < 16; cp += 4) {
+        for (int rr = 0; rr < 8; ++rr) {
+            rows[rr] = vld1_u8(qs[rr] + 2 * cp);
+        }
+        repack_neon_transpose8x8(rows, columns);
+
+        for (int pair = 0; pair < 4; ++pair) {
+            const uint8x8x2_t interleaved = vzip_u8(columns[2 * pair], columns[2 * pair + 1]);
+            uint8_t * dst = tile_dst + (cp + pair) * 64 + 2 * row0;
+            vst1_u8(dst, interleaved.val[0]);
+            vst1_u8(dst + 8, interleaved.val[1]);
+        }
+    }
+}
+
+static inline void repack_store_u32(uint8_t * dst, uint32_t v) {
+    memcpy(dst, &v, sizeof(v));
+}
+
+static inline uint32x4_t repack_neon_transpose_2bit_4x4(uint32x4_t w) {
+    uint32x4_t t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 6)), vdupq_n_u32(0x00CC00CCu));
+    w = veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 6)));
+    t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 12)), vdupq_n_u32(0x0000F0F0u));
+    return veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 12)));
+}
+
+static inline void repack_neon_q2_0_block(const uint8_t * qs, uint8_t * tile, int row) {
+    const uint32x4_t w = repack_neon_transpose_2bit_4x4(vreinterpretq_u32_u8(vld1q_u8(qs)));
+    repack_store_u32(tile + 0 * 128 + 4 * row, vgetq_lane_u32(w, 0));
+    repack_store_u32(tile + 1 * 128 + 4 * row, vgetq_lane_u32(w, 1));
+    repack_store_u32(tile + 2 * 128 + 4 * row, vgetq_lane_u32(w, 2));
+    repack_store_u32(tile + 3 * 128 + 4 * row, vgetq_lane_u32(w, 3));
+}
+
+static inline void repack_neon_q1_0_block(const uint8_t * qs, uint8_t * tile, int row) {
+    uint32x4_t w = vreinterpretq_u32_u8(vld1q_u8(qs));
+    uint32x4_t t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 1)), vdupq_n_u32(0x22222222u));
+    w = veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 1)));
+    t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 3)), vdupq_n_u32(0x0A0A0A0Au));
+    w = veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 3)));
+    w = repack_neon_transpose_2bit_4x4(w);
+    repack_store_u32(tile + 0 * 128 + 4 * row, vgetq_lane_u32(w, 0));
+    repack_store_u32(tile + 1 * 128 + 4 * row, vgetq_lane_u32(w, 1));
+    repack_store_u32(tile + 2 * 128 + 4 * row, vgetq_lane_u32(w, 2));
+    repack_store_u32(tile + 3 * 128 + 4 * row, vgetq_lane_u32(w, 3));
+}
+
+static inline void repack_neon_pack_2bit_plane(uint8x16_t v0, uint8x16_t v1, uint8_t * dst, int row) {
+    const uint32x4x2_t z   = vuzpq_u32(vreinterpretq_u32_u8(v0), vreinterpretq_u32_u8(v1));
+    const uint32x4_t   p01 = vorrq_u32(z.val[0], vshlq_n_u32(z.val[1], 2));
+    const uint32x2x2_t u   = vuzp_u32(vget_low_u32(p01), vget_high_u32(p01));
+    const uint32x2_t   out = vorr_u32(u.val[0], vshl_n_u32(u.val[1], 4));
+    repack_store_u32(dst + 0 * 128 + row * 4, vget_lane_u32(out, 0));
+    repack_store_u32(dst + 1 * 128 + row * 4, vget_lane_u32(out, 1));
+}
+
+static inline void repack_neon_q2_3_k_low2(const uint8_t * qs, int kt_local, uint8_t * lo_pl, int row) {
+    const uint8_t *  qs_sub = qs + (kt_local / 4) * 32;
+    const int8x16_t  shift  = vdupq_n_s8(-2 * (kt_local % 4));
+    const uint8x16_t mask   = vdupq_n_u8(0x03);
+    const uint8x16_t v0     = vandq_u8(vshlq_u8(vld1q_u8(qs_sub +  0), shift), mask);
+    const uint8x16_t v1     = vandq_u8(vshlq_u8(vld1q_u8(qs_sub + 16), shift), mask);
+    repack_neon_pack_2bit_plane(v0, v1, lo_pl, row);
+}
+
+static inline void repack_neon_q3_k_hmask(const uint8_t * hmask, int kt_local, uint8_t * neg_pl, int row) {
+    static const uint32_t mask_words[4] = { 0x01010101u, 0x02020202u, 0x04040404u, 0x08080808u };
+    const uint8x16_t m0   = vreinterpretq_u8_u32(vld1q_u32(mask_words));
+    const uint8x16_t m1   = vshlq_n_u8(m0, 4);
+    const uint8x16_t hbit = vdupq_n_u8((uint8_t) (1u << kt_local));
+    const uint8x16_t b0   = vbicq_u8(m0, vtstq_u8(vld1q_u8(hmask +  0), hbit));
+    const uint8x16_t b1   = vbicq_u8(m1, vtstq_u8(vld1q_u8(hmask + 16), hbit));
+    const uint32x4_t b01  = vreinterpretq_u32_u8(vorrq_u8(b0, b1));
+    const uint32x2_t f0   = vorr_u32(vget_low_u32(b01), vget_high_u32(b01));
+    const uint32x2_t f1   = vpadd_u32(f0, f0);
+    repack_store_u32(neg_pl + row * 4, vget_lane_u32(f1, 0));
+}
+
+static inline void repack_neon_q6_k(const block_q6_K * b, int kt_local, uint8_t * lo_pl, uint8_t * hi_pl, int row) {
+    const int        c      = kt_local / 4;
+    const int        g_sb   = kt_local % 4;
+    const uint8_t *  ql     = b->ql + c * 64 + (g_sb & 1) * 32;
+    const uint8_t *  qh     = b->qh + c * 32;
+    const int8x16_t  lshift = vdupq_n_s8((g_sb >= 2) ? -4 : 0);
+    const int8x16_t  hshift = vdupq_n_s8(-2 * g_sb);
+    const uint8x16_t mask4  = vdupq_n_u8(0x0F);
+    const uint8x16_t mask2  = vdupq_n_u8(0x03);
+
+    const uint8x16_t lo0  = vandq_u8(vshlq_u8(vld1q_u8(ql +  0), lshift), mask4);
+    const uint8x16_t lo1  = vandq_u8(vshlq_u8(vld1q_u8(ql + 16), lshift), mask4);
+    const uint32x4x2_t lz = vuzpq_u32(vreinterpretq_u32_u8(lo0), vreinterpretq_u32_u8(lo1));
+    const uint32x4_t lout = vorrq_u32(lz.val[0], vshlq_n_u32(lz.val[1], 4));
+    repack_store_u32(lo_pl + 0 * 128 + row * 4, vgetq_lane_u32(lout, 0));
+    repack_store_u32(lo_pl + 1 * 128 + row * 4, vgetq_lane_u32(lout, 1));
+    repack_store_u32(lo_pl + 2 * 128 + row * 4, vgetq_lane_u32(lout, 2));
+    repack_store_u32(lo_pl + 3 * 128 + row * 4, vgetq_lane_u32(lout, 3));
+
+    const uint8x16_t hi0 = vandq_u8(vshlq_u8(vld1q_u8(qh +  0), hshift), mask2);
+    const uint8x16_t hi1 = vandq_u8(vshlq_u8(vld1q_u8(qh + 16), hshift), mask2);
+    repack_neon_pack_2bit_plane(hi0, hi1, hi_pl, row);
+}
+
+static inline void repack_neon_q5_k_qh(const uint8_t * qh, uint8_t hbit, uint8_t * plane, int row) {
+    static const uint8_t mask_bytes[16] = { 1, 2, 1, 2, 1, 2, 1, 2, 4, 8, 4, 8, 4, 8, 4, 8 };
+    const uint8x16_t m0   = vld1q_u8(mask_bytes);
+    const uint8x16_t m1   = vshlq_n_u8(m0, 4);
+    const uint8x16_t vbit = vdupq_n_u8(hbit);
+    const uint8x16_t b0   = vandq_u8(vtstq_u8(vld1q_u8(qh +  0), vbit), m0);
+    const uint8x16_t b1   = vandq_u8(vtstq_u8(vld1q_u8(qh + 16), vbit), m1);
+    const uint8x16_t b01  = vorrq_u8(b0, b1);
+    const uint8x8_t  fold = vorr_u8(vget_low_u8(b01), vget_high_u8(b01));
+    const uint8x8_t  out  = vpadd_u8(fold, fold);
+    plane[0 * 32 + row] = vget_lane_u8(out, 0);
+    plane[1 * 32 + row] = vget_lane_u8(out, 1);
+    plane[2 * 32 + row] = vget_lane_u8(out, 2);
+    plane[3 * 32 + row] = vget_lane_u8(out, 3);
+}
+
+static inline uint16x8_t repack_neon_load_f16x8(const uint16_t * src) {
+    return vld1q_u16(src);
+}
+
+#if defined(__aarch64__) || defined(__ARM_FP16_FORMAT_IEEE)
+static inline uint16x8_t repack_neon_load_f16x8(const float * src) {
+    const float16x4_t lo = vcvt_f16_f32(vld1q_f32(src));
+    const float16x4_t hi = vcvt_f16_f32(vld1q_f32(src + 4));
+    return vreinterpretq_u16_f16(vcombine_f16(lo, hi));
+}
+
+static inline uint16x8_t repack_neon_load_f16x8(const ggml_bf16_t * src) {
+    const uint16x8_t  bf   = vld1q_u16((const uint16_t *) src);
+    const float32x4_t f_lo = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(bf), 16));
+    const float32x4_t f_hi = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(bf), 16));
+    const float16x4_t lo   = vcvt_f16_f32(f_lo);
+    const float16x4_t hi   = vcvt_f16_f32(f_hi);
+    return vreinterpretq_u16_f16(vcombine_f16(lo, hi));
+}
+#else
+static inline uint16x8_t repack_neon_load_f16x8(const float * src) {
+    uint16_t tmp[8];
+    for (int i = 0; i < 8; ++i) {
+        tmp[i] = GGML_FP32_TO_FP16(src[i]);
+    }
+    return vld1q_u16(tmp);
+}
+
+static inline uint16x8_t repack_neon_load_f16x8(const ggml_bf16_t * src) {
+    uint16_t tmp[8];
+    for (int i = 0; i < 8; ++i) {
+        tmp[i] = GGML_FP32_TO_FP16(GGML_BF16_TO_FP32(src[i]));
+    }
+    return vld1q_u16(tmp);
+}
+#endif
+
+template <typename SrcT>
+static void repack_neon_f16_tile(uint16_t * tile_dst, const SrcT * src_slice, int64_t row_stride, int ct, int kt) {
+    for (int row0 = 0; row0 < 32; row0 += 4) {
+        const SrcT * r0 = src_slice + (int64_t) (ct * 32 + row0 + 0) * row_stride + kt * 32;
+        const SrcT * r1 = src_slice + (int64_t) (ct * 32 + row0 + 1) * row_stride + kt * 32;
+        const SrcT * r2 = src_slice + (int64_t) (ct * 32 + row0 + 2) * row_stride + kt * 32;
+        const SrcT * r3 = src_slice + (int64_t) (ct * 32 + row0 + 3) * row_stride + kt * 32;
+        for (int cp = 0; cp < 16; cp += 4) {
+            const uint32x4_t v0 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r0 + 2 * cp));
+            const uint32x4_t v1 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r1 + 2 * cp));
+            const uint32x4_t v2 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r2 + 2 * cp));
+            const uint32x4_t v3 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r3 + 2 * cp));
+
+            const uint32x2x2_t t0 = vtrn_u32(vget_low_u32(v0),  vget_low_u32(v1));
+            const uint32x2x2_t t1 = vtrn_u32(vget_high_u32(v0), vget_high_u32(v1));
+            const uint32x2x2_t t2 = vtrn_u32(vget_low_u32(v2),  vget_low_u32(v3));
+            const uint32x2x2_t t3 = vtrn_u32(vget_high_u32(v2), vget_high_u32(v3));
+
+            vst1q_u16(tile_dst + (cp + 0) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t0.val[0], t2.val[0])));
+            vst1q_u16(tile_dst + (cp + 1) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t0.val[1], t2.val[1])));
+            vst1q_u16(tile_dst + (cp + 2) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t1.val[0], t3.val[0])));
+            vst1q_u16(tile_dst + (cp + 3) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t1.val[1], t3.val[1])));
+        }
+    }
+}
+#endif
+
 // --- Generic 32x32 tile grid walkers for block-32 and float types ---
 
+template <typename Fn>
+static void repack_parallel_for_ct(size_t slice_bytes, int n_col_tiles, Fn && fn) {
+    int n_threads = 1;
+    if (slice_bytes > 128 * 1024) {
+        for (int nt : { 4, 3, 2 }) {
+            if (n_col_tiles % nt == 0) {
+                n_threads = nt;
+                break;
+            }
+        }
+    }
+
+    if (n_threads <= 1) {
+        fn(0, n_col_tiles);
+        return;
+    }
+
+    const int chunk = n_col_tiles / n_threads;
+    std::thread workers[3];
+    for (int i = 0; i < n_threads - 1; ++i) {
+        workers[i] = std::thread([&fn, i, chunk]() {
+            fn(i * chunk, (i + 1) * chunk);
+        });
+    }
+    fn((n_threads - 1) * chunk, n_col_tiles);
+    for (int i = 0; i < n_threads - 1; ++i) {
+        workers[i].join();
+    }
+}
+
 template <typename SrcElemT, int ELEMS_PER_KTILE, typename PackTileFn>
 static void repack_to_tiled_grid_32(ggml_tensor * t, const void * data, size_t offset, size_t size,
                                     size_t tile_size, PackTileFn && pack_tile) {
@@ -70,12 +376,14 @@ static void repack_to_tiled_grid_32(ggml_tensor * t, const void * data, size_t o
         const SrcElemT * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * row_stride);
         uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size;
 
-        for (int ct = 0; ct < n_col_tiles; ct++) {
-            for (int kt = 0; kt < n_k_tiles; kt++) {
-                uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
-                pack_tile(tile_dst, src_slice, ne0, ne1, row_stride, ct, kt);
+        repack_parallel_for_ct(matrix_size, n_col_tiles, [&](int ct0, int ct1) {
+            for (int ct = ct0; ct < ct1; ct++) {
+                for (int kt = 0; kt < n_k_tiles; kt++) {
+                    uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
+                    pack_tile(tile_dst, src_slice, ne0, ne1, row_stride, ct, kt);
+                }
             }
-        }
+        });
     }
 }
 
@@ -130,8 +438,24 @@ static void unpack_from_tiled_grid_32(void * data, const ggml_tensor * t, size_t
 static void repack_q4_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
     repack_to_tiled_grid_32<block_q4_0, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_0,
         [](uint8_t * tile_dst, const block_q4_0 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
-            uint8_t tile_quants[32][32];
             ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_4bit_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    scale_dst[row] = src_slice[(ct * 32 + row) * row_stride + kt].d;
+                }
+                return;
+            }
+#endif
+            uint8_t tile_quants[32][32];
             for (int row = 0; row < 32; row++) {
                 const int64_t r = ct * 32 + row;
                 if (r < ne1 && kt < ne0 / 32) {
@@ -167,8 +491,26 @@ static void repack_tiled_q4_0(void * data, const ggml_tensor * t, size_t offset,
 static void repack_q4_1_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
     repack_to_tiled_grid_32<block_q4_1, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_1,
         [](uint8_t * tile_dst, const block_q4_1 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
-            uint8_t tile_quants[32][32];
             ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_4bit_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    const block_q4_1 & src_block = src_slice[(ct * 32 + row) * row_stride + kt];
+                    scale_dst[2 * row + 0] = src_block.d;
+                    scale_dst[2 * row + 1] = src_block.m;
+                }
+                return;
+            }
+#endif
+            uint8_t tile_quants[32][32];
             for (int row = 0; row < 32; row++) {
                 const int64_t r = ct * 32 + row;
                 if (r < ne1 && kt < ne0 / 32) {
@@ -207,8 +549,24 @@ static void repack_tiled_q4_1(void * data, const ggml_tensor * t, size_t offset,
 static void repack_mxfp4_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
     repack_to_tiled_grid_32<block_mxfp4, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_MXFP4,
         [](uint8_t * tile_dst, const block_mxfp4 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
-            uint8_t tile_quants[32][32];
             uint8_t * scale_dst = tile_dst + 512;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_4bit_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    scale_dst[row] = src_slice[(ct * 32 + row) * row_stride + kt].e;
+                }
+                return;
+            }
+#endif
+            uint8_t tile_quants[32][32];
             for (int row = 0; row < 32; row++) {
                 const int64_t r = ct * 32 + row;
                 if (r < ne1 && kt < ne0 / 32) {
@@ -244,6 +602,23 @@ static void repack_tiled_mxfp4(void * data, const ggml_tensor * t, size_t offset
 static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
     repack_to_tiled_grid_32<block_q8_0, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q8_0,
         [](uint8_t * tile_dst, const block_q8_0 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
+            ggml_half * scale_dst = (ggml_half *) (tile_dst + 1024);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = (const uint8_t *) src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_q8_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    scale_dst[row] = src_slice[(ct * 32 + row) * row_stride + kt].d;
+                }
+                return;
+            }
+#endif
             for (int cp = 0; cp < 16; cp++) {
                 const int col0 = cp * 2;
                 const int col1 = col0 + 1;
@@ -254,7 +629,6 @@ static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t offset,
                     tile_dst[cp * 64 + 2 * row + 1] = b ? b->qs[col1] : 0;
                 }
             }
-            ggml_half * scale_dst = (ggml_half *) (tile_dst + 1024);
             for (int row = 0; row < 32; row++) {
                 const int64_t r = ct * 32 + row;
                 scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_slice[r * row_stride + kt].d : 0;
@@ -303,24 +677,33 @@ static void repack_q2_0_tiled(ggml_tensor * t, const void * data, size_t offset,
     for (int64_t slice = 0; slice < t->ne[2] * t->ne[3]; ++slice) {
         const block_q2_0 * src_slice = src + slice * t->ne[1] * blocks_per_row;
         uint8_t * dst_slice = (uint8_t *) t->data + slice * matrix_size;
-        memset(dst_slice, 0, matrix_size);
+        const size_t ct_bytes = (size_t) n_weight_tiles * HTP_TENSOR_TILE_SIZE_Q2_0;
 
-        for (int64_t row_index = 0; row_index < t->ne[1]; ++row_index) {
-            const int64_t ct = row_index / 32;
-            const int row = (int) (row_index % 32);
-            for (int64_t kt = 0; kt < n_weight_tiles; ++kt) {
-                const block_q2_0 * block = src_slice + row_index * blocks_per_row + kt;
-                uint8_t * tile = dst_slice + (ct * n_weight_tiles + kt) * HTP_TENSOR_TILE_SIZE_Q2_0;
-                for (int group = 0; group < 16; ++group) {
-                    for (int byte = 0; byte < 4; ++byte) {
-                        const int k = 4 * group + byte;
-                        const uint8_t code = (block->qs[k / 4] >> (2 * (k % 4))) & 3;
-                        tile[(group >> 2) * 128 + 4 * row + byte] |= code << (2 * (group & 3));
+        repack_parallel_for_ct(matrix_size, (int) n_col_tiles, [&](int ct0, int ct1) {
+            memset(dst_slice + (size_t) ct0 * ct_bytes, 0, (size_t) (ct1 - ct0) * ct_bytes);
+            const int64_t r0 = (int64_t) ct0 * 32;
+            const int64_t r1 = (std::min)((int64_t) ct1 * 32, t->ne[1]);
+            for (int64_t row_index = r0; row_index < r1; ++row_index) {
+                const int64_t ct = row_index / 32;
+                const int row = (int) (row_index % 32);
+                for (int64_t kt = 0; kt < n_weight_tiles; ++kt) {
+                    const block_q2_0 * block = src_slice + row_index * blocks_per_row + kt;
+                    uint8_t * tile = dst_slice + (ct * n_weight_tiles + kt) * HTP_TENSOR_TILE_SIZE_Q2_0;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+                    repack_neon_q2_0_block(block->qs, tile, row);
+#else
+                    for (int group = 0; group < 16; ++group) {
+                        for (int byte = 0; byte < 4; ++byte) {
+                            const int k = 4 * group + byte;
+                            const uint8_t code = (block->qs[k / 4] >> (2 * (k % 4))) & 3;
+                            tile[(group >> 2) * 128 + 4 * row + byte] |= code << (2 * (group & 3));
+                        }
                     }
+#endif
+                    memcpy(tile + 512 + 2 * row, &block->d, sizeof(block->d));
                 }
-                memcpy(tile + 512 + 2 * row, &block->d, sizeof(block->d));
             }
-        }
+        });
     }
 }
 
@@ -368,6 +751,12 @@ static void repack_float_to_f16_tiled(ggml_tensor * t, const void * data, size_t
     repack_to_tiled_grid_32<SrcT, 32>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_F16,
         [&](uint8_t * tile_bytes, const SrcT * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
             uint16_t * tile_dst = (uint16_t *) tile_bytes;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt * 32 + 32 <= ne0) {
+                repack_neon_f16_tile(tile_dst, src_slice, row_stride, ct, kt);
+                return;
+            }
+#endif
             for (int cp = 0; cp < 16; cp++) {
                 const int64_t c0 = kt * 32 + cp * 2;
                 const int64_t c1 = c0 + 1;
@@ -424,20 +813,25 @@ static void repack_kquant_to_tiled_grid(ggml_tensor * t, const void * data, size
     for (int64_t slice_idx = 0; slice_idx < ne2 * ne3; slice_idx++) {
         const BlockT * src_slice = src_matrix + slice_idx * (ne1 * sb_per_row);
         uint8_t * matrix_dst     = (uint8_t *) t->data + slice_idx * matrix_size;
+        const size_t ct_bytes    = (size_t) n_k_tiles * tile_size;
 
-        // Padding rows and OR-ed bit planes require zeroed tiles
-        memset(matrix_dst, 0, matrix_size);
+        repack_parallel_for_ct(matrix_size, n_col_tiles, [&](int ct0, int ct1) {
+            // Padding rows and OR-ed bit planes require zeroed tiles
+            memset(matrix_dst + (size_t) ct0 * ct_bytes, 0, (size_t) (ct1 - ct0) * ct_bytes);
 
-        for (int64_t r = 0; r < ne1; r++) {
-            const int ct  = (int) (r / 32);
-            const int row = (int) (r % 32);
-            const BlockT * src_row = src_slice + r * sb_per_row;
+            const int64_t r0 = (int64_t) ct0 * 32;
+            const int64_t r1 = (std::min)((int64_t) ct1 * 32, ne1);
+            for (int64_t r = r0; r < r1; r++) {
+                const int ct  = (int) (r / 32);
+                const int row = (int) (r % 32);
+                const BlockT * src_row = src_slice + r * sb_per_row;
 
-            for (int kt = 0; kt < n_k_tiles; kt++) {
-                uint8_t * tile = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
-                pack_kt(tile, &src_row[kt / 8], row, kt % 8);
+                for (int kt = 0; kt < n_k_tiles; kt++) {
+                    uint8_t * tile = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
+                    pack_kt(tile, &src_row[kt / 8], row, kt % 8);
+                }
             }
-        }
+        });
     }
 }
 
@@ -493,6 +887,7 @@ static void unpack_kquant_from_tiled_grid(void * data, const ggml_tensor * t, si
 
 // --- K-Quant Helpers & Per-Type Implementations (Q2_K, Q3_K, Q4_K, Q5_K, Q6_K) ---
 
+#if !defined(__ARM_NEON) && !defined(__aarch64__)
 // unsigned 6-bit value (0..63) of element e of a Q6_K block, same bit layout as dequantize_row_q6_K
 static inline uint8_t q6_K_get_quant(const block_q6_K * b, int e) {
     const int c = e / 128;
@@ -510,6 +905,7 @@ static inline uint8_t q6_K_get_quant(const block_q6_K * b, int e) {
     }
     return (uint8_t) (lo | (hi << 4));
 }
+#endif
 
 static void repack_q6_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
     repack_kquant_to_tiled_grid<block_q6_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q6_K,
@@ -519,6 +915,9 @@ static void repack_q6_K_tiled(ggml_tensor * t, const void * data, size_t offset,
             uint8_t * hi_pl   = tile + 512;
             ggml_half * sc_pl = (ggml_half *) (tile + 768);
 
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            repack_neon_q6_k(b, kt_local, lo_pl, hi_pl, row);
+#else
             for (int lk = 0; lk < 32; lk++) {
                 const uint8_t q6 = q6_K_get_quant(b, kt_local * 32 + lk);
                 const int g   = lk >> 2;
@@ -526,6 +925,7 @@ static void repack_q6_K_tiled(ggml_tensor * t, const void * data, size_t offset,
                 lo_pl[(g >> 1) * 128 + pos] |= (uint8_t) ((q6 & 0xF) << ((g & 1) * 4));
                 hi_pl[(g >> 2) * 128 + pos] |= (uint8_t) ((q6 >> 4) << ((g & 3) * 2));
             }
+#endif
             for (int sub = 0; sub < 2; sub++) {
                 sc_pl[sub * 32 + row] = GGML_FP32_TO_FP16(d * (float) b->scales[kt_local * 2 + sub]);
             }
@@ -585,6 +985,7 @@ static void repack_tiled_q6_K(void * data, const ggml_tensor * t, size_t offset,
         });
 }
 
+#if !defined(__ARM_NEON) && !defined(__aarch64__)
 // low 2 bits (0..3) of element e of a Q2_K or Q3_K block, same bit layout as dequantize_row_q2_K / q3_K
 static inline uint8_t q2_3_K_get_low2(const uint8_t * qs, int e) {
     const int c = e / 128;
@@ -600,6 +1001,7 @@ static inline bool q3_K_get_hbit(const block_q3_K * b, int e) {
     const int l = e % 32;
     return (b->hmask[l] >> (c * 4 + j)) & 1;
 }
+#endif
 
 // signed 6-bit scale j (-32..31) of a Q3_K block, same packing as quantize_row_q3_K_ref
 static inline int q3_K_get_scale(const uint8_t * scales, int j) {
@@ -670,6 +1072,10 @@ static void repack_q3_K_tiled(ggml_tensor * t, const void * data, size_t offset,
             uint8_t * neg_pl  = tile + 256;
             ggml_half * sc_pl = (ggml_half *) (tile + 384);
 
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            repack_neon_q2_3_k_low2(b->qs, kt_local, lo_pl, row);
+            repack_neon_q3_k_hmask(b->hmask, kt_local, neg_pl, row);
+#else
             for (int lk = 0; lk < 32; lk++) {
                 const int e   = kt_local * 32 + lk;
                 const int g   = lk >> 2;
@@ -679,6 +1085,7 @@ static void repack_q3_K_tiled(ggml_tensor * t, const void * data, size_t offset,
                     neg_pl[pos] |= (uint8_t) (1 << g);
                 }
             }
+#endif
             for (int sub = 0; sub < 2; sub++) {
                 sc_pl[sub * 32 + row] = GGML_FP32_TO_FP16(d * (float) q3_K_get_scale(b->scales, kt_local * 2 + sub));
             }
@@ -750,11 +1157,15 @@ static void repack_q2_K_tiled(ggml_tensor * t, const void * data, size_t offset,
             uint8_t * sc_pl   = tile + 256;
             ggml_half * dd_pl = (ggml_half *) (tile + 320);
 
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            repack_neon_q2_3_k_low2(b->qs, kt_local, lo_pl, row);
+#else
             for (int lk = 0; lk < 32; lk++) {
                 const int g   = lk >> 2;
                 const int pos = row * 4 + (lk & 3);
                 lo_pl[(g >> 2) * 128 + pos] |= (uint8_t) (q2_3_K_get_low2(b->qs, kt_local * 32 + lk) << ((g & 3) * 2));
             }
+#endif
             for (int sub = 0; sub < 2; sub++) {
                 sc_pl[sub * 32 + row] = b->scales[kt_local * 2 + sub];
             }
@@ -846,23 +1257,46 @@ static inline void pack_k4_scales_mins(uint8_t * scales, ggml_half * d, ggml_hal
     }
 }
 
+template <typename BlockT>
+static inline void repack_k4_quants(uint8_t * tile_dst, const BlockT * b, int row, int kt_local,
+                                    int64_t sb_per_row, bool use_neon) {
+#if defined(__ARM_NEON) || defined(__aarch64__)
+    if (use_neon) {
+        if ((row & 7) == 0) {
+            const uint8_t * qs[8];
+            for (int rr = 0; rr < 8; ++rr) {
+                qs[rr] = b[rr * sb_per_row].qs + (kt_local / 2) * 32;
+            }
+            repack_neon_k4_rows(qs, tile_dst, row, (kt_local & 1) != 0);
+        }
+        return;
+    }
+#else
+    GGML_UNUSED(sb_per_row);
+    GGML_UNUSED(use_neon);
+#endif
+    const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32;
+    const int shift = (kt_local & 1) ? 4 : 0;
+
+    for (int cp = 0; cp < 16; cp++) {
+        const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F;
+        const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F;
+        tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0);
+    }
+}
+
 static void repack_q4_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
+    const int64_t sb_per_row = t->ne[0] / QK_K;
+    const bool    use_neon   = (t->ne[1] % 8 == 0);
     repack_kquant_to_tiled_grid<block_q4_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_K,
-        [](uint8_t * tile_dst, const block_q4_K * b, int row, int kt_local) {
+        [sb_per_row, use_neon](uint8_t * tile_dst, const block_q4_K * b, int row, int kt_local) {
             const float d    = GGML_FP16_TO_FP32(b->d);
             const float dmin = GGML_FP16_TO_FP32(b->dmin);
 
             uint8_t sc, m;
             get_scale_min_k4(kt_local, b->scales, &sc, &m);
 
-            const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32;
-            const int shift = (kt_local & 1) ? 4 : 0;
-
-            for (int cp = 0; cp < 16; cp++) {
-                const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F;
-                const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F;
-                tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0);
-            }
+            repack_k4_quants(tile_dst, b, row, kt_local, sb_per_row, use_neon);
 
             ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
             scale_dst[2 * row + 0] = GGML_FP32_TO_FP16(d * (float) sc);
@@ -900,8 +1334,10 @@ static void repack_tiled_q4_K(void * data, const ggml_tensor * t, size_t offset,
 }
 
 static void repack_q5_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
+    const int64_t sb_per_row = t->ne[0] / QK_K;
+    const bool    use_neon   = (t->ne[1] % 8 == 0);
     repack_kquant_to_tiled_grid<block_q5_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q5_K,
-        [](uint8_t * tile_dst, const block_q5_K * b, int row, int kt_local) {
+        [sb_per_row, use_neon](uint8_t * tile_dst, const block_q5_K * b, int row, int kt_local) {
             const float d     = GGML_FP16_TO_FP32(b->d);
             const float dmin  = GGML_FP16_TO_FP32(b->dmin);
             uint8_t * plane   = tile_dst + 640;
@@ -909,15 +1345,13 @@ static void repack_q5_K_tiled(ggml_tensor * t, const void * data, size_t offset,
             uint8_t sc, m;
             get_scale_min_k4(kt_local, b->scales, &sc, &m);
 
-            const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32;
-            const int shift    = (kt_local & 1) ? 4 : 0;
+            repack_k4_quants(tile_dst, b, row, kt_local, sb_per_row, use_neon);
+
             const uint8_t hbit = (uint8_t) (1 << kt_local);
-
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            repack_neon_q5_k_qh(b->qh, hbit, plane, row);
+#else
             for (int cp = 0; cp < 16; cp++) {
-                const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F;
-                const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F;
-                tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0);
-
                 const int i    = cp / 4;
                 const int lane = (cp % 4) * 32 + row;
                 if (b->qh[2 * cp + 0] & hbit) {
@@ -927,6 +1361,7 @@ static void repack_q5_K_tiled(ggml_tensor * t, const void * data, size_t offset,
                     plane[lane] |= (uint8_t) (1 << (2 * i + 1));
                 }
             }
+#endif
 
             ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
             scale_dst[2 * row + 0] = GGML_FP32_TO_FP16(d * (float) sc);
@@ -1002,31 +1437,37 @@ static void repack_q1_0_tiled(ggml_tensor * t, const void * data, size_t offset,
         const block_q1_0 * src_slice = src_matrix + (slice_idx - start_slice) * ne1 * n_k_tiles;
         uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size;
 
-        for (int ct = 0; ct < n_col_tiles; ct++) {
-            for (int kt = 0; kt < n_k_tiles; kt++) {
-                uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
-                ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
-                memset(tile_dst, 0, tile_size);
+        repack_parallel_for_ct(matrix_size, n_col_tiles, [&](int ct0, int ct1) {
+            for (int ct = ct0; ct < ct1; ct++) {
+                for (int kt = 0; kt < n_k_tiles; kt++) {
+                    uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
+                    ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
+                    memset(tile_dst, 0, tile_size);
 
-                for (int row = 0; row < 32; row++) {
-                    const int64_t r = ct * 32 + row;
-                    if (r >= ne1) {
-                        continue;
-                    }
-
-                    const block_q1_0 & block = src_slice[r * n_k_tiles + kt];
-                    scale_dst[row] = block.d;
-                    for (int k = 0; k < QK1_0; k++) {
-                        int byte_offset;
-                        int bit;
-                        q1_0_tiled_bit_location(k, row, byte_offset, bit);
-                        if (block.qs[k / 8] & (1u << (k % 8))) {
-                            tile_dst[byte_offset] |= (uint8_t) (1u << bit);
+                    for (int row = 0; row < 32; row++) {
+                        const int64_t r = ct * 32 + row;
+                        if (r >= ne1) {
+                            continue;
                         }
+
+                        const block_q1_0 & block = src_slice[r * n_k_tiles + kt];
+                        scale_dst[row] = block.d;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+                        repack_neon_q1_0_block(block.qs, tile_dst, row);
+#else
+                        for (int k = 0; k < QK1_0; k++) {
+                            int byte_offset;
+                            int bit;
+                            q1_0_tiled_bit_location(k, row, byte_offset, bit);
+                            if (block.qs[k / 8] & (1u << (k % 8))) {
+                                tile_dst[byte_offset] |= (uint8_t) (1u << bit);
+                            }
+                        }
+#endif
                     }
                 }
             }
-        }
+        });
     }
 }
 
@@ -1127,11 +1568,11 @@ void htp_repack_tensor_tiled(ggml_tensor * tensor, const void * data, size_t siz
             break;
         case GGML_TYPE_F32:
             repack_float_to_f16_tiled<float>(tensor, data, 0, size,
-                [](float v) { return (uint16_t) ggml_fp32_to_fp16(v); });
+                [](float v) { return (uint16_t) GGML_FP32_TO_FP16(v); });
             break;
         case GGML_TYPE_BF16:
             repack_float_to_f16_tiled<ggml_bf16_t>(tensor, data, 0, size,
-                [](ggml_bf16_t v) { return (uint16_t) ggml_fp32_to_fp16(ggml_bf16_to_fp32(v)); });
+                [](ggml_bf16_t v) { return (uint16_t) GGML_FP32_TO_FP16(GGML_BF16_TO_FP32(v)); });
             break;
         default:
             break;

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

@@ -28,7 +28,7 @@ add_library(${HTP_LIB} SHARED
     unary-ops.c
     sum-rows-ops.c
     softmax-ops.c
-    act-ops.c
+    glu-ops.c
     rope-ops.c
     set-rows-ops.c
     get-rows-ops.c

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

@@ -32,6 +32,7 @@ static inline size_t htp_allreduce_vtcm_buffer_count(
 }
 
 struct htp_allreduce_kernel_params {
+    uint8_t kernel_type;
     int32_t rank;
     int32_t n_ranks;
     int32_t n_threads;
@@ -39,7 +40,6 @@ struct htp_allreduce_kernel_params {
     int32_t elems_per_thread;     // 1D: nelem_per_thread, 2D: nrows_per_thread
     int32_t vtcm_size_per_thread;
     int32_t vtcm_size;
-    int32_t kernel_type;
     int32_t ne0;
     int32_t ne1;
     int32_t row_size_aligned;
@@ -50,6 +50,12 @@ struct htp_allreduce_kernel_params {
     int32_t mode;                 // enum htp_allreduce_mode: FULL or SHARDED_FANOUT
 };
 
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_allreduce_kernel_params) <= 128, "htp_allreduce_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_allreduce_kernel_params) <= 128, "htp_allreduce_kernel_params is too large for kernel_params blob");
+#endif
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -453,7 +453,10 @@ static void merge_runs(
 struct htp_sort_chunk_ctx {
     struct htp_ops_context *              octx;
     const struct htp_sort_kernel_params * kparams;
-    uint8_t *                             vtcm_base;
+    uint8_t *                             phase1_base;
+    size_t                                phase1_slot_size;
+    float *                               merge_values;
+    int32_t *                             merge_indices;
     uint32_t                              real_count[HTP_MAX_NTHREADS];
 };
 
@@ -462,26 +465,27 @@ static void htp_sort_chunk_job(unsigned int n, unsigned int i, void * data) {
     struct htp_ops_context * octx = cctx->octx;
     const struct htp_sort_kernel_params * kparams = cctx->kparams;
     const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * dst  = octx->dst;
 
     uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
-    uint32_t ne00        = (uint32_t) kparams->ne00;
+    uint32_t ne00        = src0->ne[0];
     uint32_t chunk_base  = i * chunk_elems;
 
     uint32_t real_count = (chunk_base < ne00) ? MIN(chunk_elems, ne00 - chunk_base) : 0;
     cctx->real_count[i] = real_count;
 
-    uint8_t * spad = cctx->vtcm_base + (size_t) kparams->phase1_slot_size * i;
+    uint8_t * spad = cctx->phase1_base + cctx->phase1_slot_size * i;
     float *   values_buf  = (float *) spad;
     int32_t * indices_buf = (int32_t *) (spad + chunk_elems * sizeof(float));
 
-    float *   merge_values  = (float *)   (cctx->vtcm_base + kparams->merge_values_off);
-    int32_t * merge_indices = (int32_t *) (cctx->vtcm_base + kparams->merge_indices_off);
+    float *   merge_values  = cctx->merge_values;
+    int32_t * merge_indices = cctx->merge_indices;
 
     struct htp_thread_trace * tr = &octx->ctx->trace[i];
     const dma_addr_t src_addr = src0->data + (size_t) chunk_base * sizeof(float);
 
     if (kparams->is_top_k) {
-        uint32_t k = (uint32_t) kparams->k;
+        uint32_t k = dst->ne[0];
         uint32_t local_k = MIN(k, chunk_elems);
 
         if (real_count == 0) {
@@ -559,17 +563,17 @@ static void htp_argsort_merge_l1_job(unsigned int n, unsigned int i, void * data
     uint32_t count0 = cctx->real_count[c0];
     uint32_t count1 = cctx->real_count[c1];
 
-    size_t slot_size = (size_t) kparams->phase1_slot_size;
+    size_t slot_size = cctx->phase1_slot_size;
     uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
 
-    const float *   val0 = (const float *)   (cctx->vtcm_base + slot_size * c0);
-    const int32_t * idx0 = (const int32_t *) (cctx->vtcm_base + slot_size * c0 + chunk_elems * sizeof(float));
-    const float *   val1 = (const float *)   (cctx->vtcm_base + slot_size * c1);
-    const int32_t * idx1 = (const int32_t *) (cctx->vtcm_base + slot_size * c1 + chunk_elems * sizeof(float));
+    const float *   val0 = (const float *)   (cctx->phase1_base + slot_size * c0);
+    const int32_t * idx0 = (const int32_t *) (cctx->phase1_base + slot_size * c0 + chunk_elems * sizeof(float));
+    const float *   val1 = (const float *)   (cctx->phase1_base + slot_size * c1);
+    const int32_t * idx1 = (const int32_t *) (cctx->phase1_base + slot_size * c1 + chunk_elems * sizeof(float));
 
     uint32_t out_offset = (2 * i) * chunk_elems;
-    float *   out_val = (float *)   (cctx->vtcm_base + kparams->merge_values_off)  + out_offset;
-    int32_t * out_idx = (int32_t *) (cctx->vtcm_base + kparams->merge_indices_off) + out_offset;
+    float *   out_val = cctx->merge_values  + out_offset;
+    int32_t * out_idx = cctx->merge_indices + out_offset;
 
     struct htp_thread_trace * tr = &cctx->octx->ctx->trace[i];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
@@ -587,21 +591,21 @@ static void htp_argsort_merge_l2_job(unsigned int n, unsigned int i, void * data
     struct htp_sort_chunk_ctx * cctx = mctx->cctx;
     const struct htp_sort_kernel_params * kparams = cctx->kparams;
 
-    size_t slot_size = (size_t) kparams->phase1_slot_size;
+    size_t slot_size = cctx->phase1_slot_size;
     uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
 
     uint32_t off0 = (4 * i) * chunk_elems;
     uint32_t off1 = (4 * i + 2) * chunk_elems;
 
-    const float *   val0 = (const float *)   (cctx->vtcm_base + kparams->merge_values_off)  + off0;
-    const int32_t * idx0 = (const int32_t *) (cctx->vtcm_base + kparams->merge_indices_off) + off0;
-    const float *   val1 = (const float *)   (cctx->vtcm_base + kparams->merge_values_off)  + off1;
-    const int32_t * idx1 = (const int32_t *) (cctx->vtcm_base + kparams->merge_indices_off) + off1;
+    const float *   val0 = cctx->merge_values  + off0;
+    const int32_t * idx0 = cctx->merge_indices + off0;
+    const float *   val1 = cctx->merge_values  + off1;
+    const int32_t * idx1 = cctx->merge_indices + off1;
 
     uint32_t count0 = cctx->real_count[4*i + 0] + cctx->real_count[4*i + 1];
     uint32_t count1 = cctx->real_count[4*i + 2] + cctx->real_count[4*i + 3];
 
-    uint8_t * slot_base = cctx->vtcm_base + (4 * i) * slot_size;
+    uint8_t * slot_base = cctx->phase1_base + (4 * i) * slot_size;
     float *   out_val   = (float *) slot_base;
     int32_t * out_idx   = (int32_t *) (slot_base + 4 * chunk_elems * sizeof(float));
 
@@ -614,7 +618,10 @@ static void htp_argsort_merge_l2_job(unsigned int n, unsigned int i, void * data
 struct htp_sort_multi_row_ctx {
     struct htp_ops_context *              octx;
     const struct htp_sort_kernel_params * kparams;
-    uint8_t *                             vtcm_base;
+    uint8_t *                             phase1_base;
+    size_t                                phase1_slot_size;
+    uint32_t                              row_start;
+    uint32_t                              row_end;
     uint32_t                              nrows_per_thread;
 };
 
@@ -624,12 +631,18 @@ struct htp_tensor_row_coord {
     uint32_t i03;
 };
 
-static inline struct htp_tensor_row_coord init_row_coord(uint32_t r, uint32_t ne01, uint32_t ne02) {
+static inline struct htp_tensor_row_coord init_row_coord(
+    uint32_t r,
+    uint32_t ne01,
+    uint32_t ne02,
+    const struct fastdiv_values * div_ne01,
+    const struct fastdiv_values * div_ne01_ne02
+) {
     uint32_t ne01_ne02 = ne01 * ne02;
-    uint32_t i03 = (ne01_ne02 > 0) ? (r / ne01_ne02) : 0;
-    uint32_t rem = (ne01_ne02 > 0) ? (r % ne01_ne02) : 0;
-    uint32_t i02 = (ne01 > 0) ? (rem / ne01) : 0;
-    uint32_t i01 = (ne01 > 0) ? (rem % ne01) : 0;
+    uint32_t i03 = (ne01_ne02 > 0) ? fastdiv(r, div_ne01_ne02) : 0;
+    uint32_t rem = (ne01_ne02 > 0) ? (r - i03 * ne01_ne02) : 0;
+    uint32_t i02 = (ne01 > 0) ? fastdiv(rem, div_ne01) : 0;
+    uint32_t i01 = (ne01 > 0) ? (rem - i02 * ne01) : 0;
     return (struct htp_tensor_row_coord){ i01, i02, i03 };
 }
 
@@ -701,7 +714,7 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
 
-    uint32_t ne00 = (uint32_t) kparams->ne00;
+    uint32_t ne00 = src0->ne[0];
     uint32_t ne01 = src0->ne[1];
     uint32_t ne02 = src0->ne[2];
 
@@ -713,16 +726,16 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
     uint32_t nb2 = dst->nb[2];
     uint32_t nb3 = dst->nb[3];
 
-    uint32_t row_start = (uint32_t) kparams->row_start;
-    uint32_t row_end   = (uint32_t) kparams->row_end;
+    uint32_t row_start = mctx->row_start;
+    uint32_t row_end   = mctx->row_end;
     uint32_t r_start   = row_start + ith * mctx->nrows_per_thread;
     uint32_t r_end     = MIN(r_start + mctx->nrows_per_thread, row_end);
 
     if (r_start >= r_end) return;
 
-    size_t spad_offset = ith * (size_t) (kparams->phase1_slot_size * kparams->n_slots);
-    uint8_t * thread_spad = mctx->vtcm_base + spad_offset;
-    size_t slot_size = (size_t) kparams->phase1_slot_size;
+    size_t slot_size = mctx->phase1_slot_size;
+    size_t spad_offset = ith * slot_size * (size_t) kparams->n_slots;
+    uint8_t * thread_spad = mctx->phase1_base + spad_offset;
 
     uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
     float *   values_buf  = (float *) thread_spad;
@@ -732,14 +745,14 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
     bool is_top_k = kparams->is_top_k != 0;
-    uint32_t k = (uint32_t) kparams->k;
+    uint32_t k = dst->ne[0];
     enum ggml_sort_order order = (enum ggml_sort_order) kparams->order;
 
     const uint32_t src_bytes = ne00 * sizeof(float);
     const uint32_t dst_bytes = (is_top_k ? k : ne00) * sizeof(int32_t);
 
     if (kparams->n_slots == 2) {
-        struct htp_tensor_row_coord prime_coord = init_row_coord(r_start, ne01, ne02);
+        struct htp_tensor_row_coord prime_coord = init_row_coord(r_start, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);
         for (uint32_t r = r_start, spad_idx = 0; r < r_end && spad_idx < 2; r++, spad_idx++) {
             uint8_t * cur_spad = thread_spad + spad_idx * slot_size;
             float *   cur_values  = (float *) cur_spad;
@@ -759,8 +772,8 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
             step_row_coord(&prime_coord, ne01, ne02);
         }
 
-        struct htp_tensor_row_coord dst_coord = init_row_coord(r_start, ne01, ne02);
-        struct htp_tensor_row_coord next_src_coord = init_row_coord(r_start + 2, ne01, ne02);
+        struct htp_tensor_row_coord dst_coord = init_row_coord(r_start, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);
+        struct htp_tensor_row_coord next_src_coord = init_row_coord(r_start + 2, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);
 
         for (uint32_t r = r_start; r < r_end; r++) {
             uint32_t cur_slot = (r - r_start) & 1;
@@ -796,7 +809,7 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
 
         dma_queue_flush(dma_q);
     } else {
-        struct htp_tensor_row_coord coord = init_row_coord(r_start, ne01, ne02);
+        struct htp_tensor_row_coord coord = init_row_coord(r_start, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);
         for (uint32_t r = r_start; r < r_end; r++) {
             dma_addr_t src_addr = src0->data + coord_offset(coord, nb01, nb02, nb03);
             dma_addr_t dst_addr = dst->data + coord_offset(coord, nb1, nb2, nb3);
@@ -847,52 +860,39 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
     uint32_t ne00 = src0->ne[0];
     uint32_t k    = dst->ne[0];
 
-    struct htp_sort_kernel_params kparams_local;
     const struct htp_sort_kernel_params * kparams = (const struct htp_sort_kernel_params *) octx->kernel_params;
 
-    if (kparams->n_threads == 0) {
-        struct htp_sort_vtcm_layout layout;
-        if (!htp_sort_solve_layout(&layout, ne00, total_rows, k, octx->n_threads, octx->ctx->vtcm_size, is_top_k)) {
-            return HTP_STATUS_VTCM_TOO_SMALL;
-        }
-        memset(&kparams_local, 0, sizeof(kparams_local));
-        kparams_local.n_threads         = (int32_t) layout.n_threads;
-        kparams_local.total_rows        = (int32_t) total_rows;
-        kparams_local.row_start         = (int32_t) row_start;
-        kparams_local.row_end           = (int32_t) row_end;
-        kparams_local.ne00              = (int32_t) ne00;
-        kparams_local.k                 = (int32_t) k;
-        kparams_local.order             = is_top_k ? GGML_SORT_ORDER_DESC : (int32_t) octx->op_params[0];
-        kparams_local.is_top_k          = is_top_k ? 1 : 0;
-        kparams_local.use_dma           = 1;
-        kparams_local.chunk_elems       = (int32_t) layout.chunk_elems;
-        kparams_local.n_chunks          = (int32_t) layout.n_chunks;
-        kparams_local.vtcm_size         = (int32_t) layout.total_bytes;
-        kparams_local.phase1_slot_size  = (int32_t) layout.phase1_slot_size;
-        kparams_local.merge_values_off  = (int32_t) layout.merge_values_off;
-        kparams_local.merge_indices_off = (int32_t) layout.merge_indices_off;
-        kparams_local.merge_elems       = (int32_t) layout.merge_elems;
-        kparams_local.n_slots           = (int32_t) layout.n_slots;
-        kparams = &kparams_local;
-    } else if (octx->ctx->mdev.count > 1) {
-        memcpy(&kparams_local, kparams, sizeof(kparams_local));
-        kparams_local.row_start = (int32_t) row_start;
-        kparams_local.row_end   = (int32_t) row_end;
-        kparams = &kparams_local;
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
     }
 
-    if (octx->ctx->vtcm_size < (size_t) kparams->vtcm_size) {
-        FARF(ERROR, "sort: VTCM size too small. Needed %d, have %zu", kparams->vtcm_size, octx->ctx->vtcm_size);
+    struct htp_sort_vtcm_layout layout;
+    htp_sort_vtcm_layout_build(&layout,
+                               total_rows,
+                               (uint32_t) kparams->chunk_elems,
+                               (uint32_t) kparams->n_chunks,
+                               (uint32_t) kparams->merge_elems,
+                               (uint32_t) kparams->n_threads,
+                               (uint32_t) kparams->n_slots);
+
+    if (octx->ctx->vtcm_size < layout.total_bytes) {
+        FARF(ERROR, "sort: VTCM size too small. Needed %zu, have %zu", layout.total_bytes, octx->ctx->vtcm_size);
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
-    uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
+    uint8_t * vtcm_base     = (uint8_t *) octx->ctx->vtcm_base;
+    uint8_t * phase1_base   = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_phase1);
+    float *   merge_values  = VTCM_LAYOUT_PTR(float,   vtcm_base, layout.off_merge_values);
+    int32_t * merge_indices = VTCM_LAYOUT_PTR(int32_t, vtcm_base, layout.off_merge_indices);
 
     if (total_rows == 1 && kparams->n_chunks > 1) {
         struct htp_sort_chunk_ctx cctx;
-        cctx.octx      = octx;
-        cctx.kparams   = kparams;
-        cctx.vtcm_base = vtcm_base;
+        cctx.octx             = octx;
+        cctx.kparams          = kparams;
+        cctx.phase1_base      = phase1_base;
+        cctx.phase1_slot_size = layout.phase1_slot_size;
+        cctx.merge_values     = merge_values;
+        cctx.merge_indices    = merge_indices;
         memset(cctx.real_count, 0, sizeof(cctx.real_count));
 
         work_queue_run(octx->ctx->work_queue, htp_sort_chunk_job, &cctx, (uint32_t) kparams->n_chunks);
@@ -902,8 +902,8 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
             uint32_t local_k     = MIN(k, chunk_elems);
 
             if (k == 1 && ne00 >= 128*1024) {
-                float *   cand_vals = (float *)   (vtcm_base + kparams->merge_values_off);
-                int32_t * cand_idxs = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+                float *   cand_vals = merge_values;
+                int32_t * cand_idxs = merge_indices;
                 float best_val = cand_vals[0];
                 int32_t best_idx = cand_idxs[0];
                 for (int32_t c = 1; c < kparams->n_chunks; c++) {
@@ -918,8 +918,8 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
             }
 
             uint32_t total_candidates = (uint32_t) kparams->n_chunks * local_k;
-            float *   cand_vals = (float *)   (vtcm_base + kparams->merge_values_off);
-            int32_t * cand_idxs = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+            float *   cand_vals = merge_values;
+            int32_t * cand_idxs = merge_indices;
             uint32_t merge_elems = (uint32_t) kparams->merge_elems;
 
             if (merge_elems > total_candidates) {
@@ -943,18 +943,18 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
                 struct htp_argsort_merge_l2_ctx l2ctx = { &cctx, asc };
                 work_queue_run(octx->ctx->work_queue, htp_argsort_merge_l2_job, &l2ctx, 2);
 
-                size_t slot_size = (size_t) kparams->phase1_slot_size;
+                size_t slot_size = layout.phase1_slot_size;
                 uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
 
-                const float *   val0 = (const float *) vtcm_base;
-                const int32_t * idx0 = (const int32_t *) (vtcm_base + 4 * chunk_elems * sizeof(float));
-                const float *   val1 = (const float *) (vtcm_base + 4 * slot_size);
-                const int32_t * idx1 = (const int32_t *) (vtcm_base + 4 * slot_size + 4 * chunk_elems * sizeof(float));
+                const float *   val0 = (const float *) phase1_base;
+                const int32_t * idx0 = (const int32_t *) (phase1_base + 4 * chunk_elems * sizeof(float));
+                const float *   val1 = (const float *) (phase1_base + 4 * slot_size);
+                const int32_t * idx1 = (const int32_t *) (phase1_base + 4 * slot_size + 4 * chunk_elems * sizeof(float));
 
                 uint32_t n01 = cctx.real_count[0] + cctx.real_count[1] + cctx.real_count[2] + cctx.real_count[3];
                 uint32_t n23 = cctx.real_count[4] + cctx.real_count[5] + cctx.real_count[6] + cctx.real_count[7];
 
-                int32_t * final_idx = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+                int32_t * final_idx = merge_indices;
                 struct htp_thread_trace * tr = &octx->ctx->trace[0];
                 htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                 merge_runs(val0, idx0, n01, val1, idx1, n23, NULL, final_idx, asc);
@@ -968,27 +968,27 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
                 uint32_t n01 = cctx.real_count[0] + cctx.real_count[1];
                 uint32_t n23 = cctx.real_count[2] + cctx.real_count[3];
 
-                const float *   val01 = (const float *)   (vtcm_base + kparams->merge_values_off);
-                const int32_t * idx01 = (const int32_t *) (vtcm_base + kparams->merge_indices_off);
-                const float *   val23 = (const float *)   (vtcm_base + kparams->merge_values_off)  + 2 * chunk_elems;
-                const int32_t * idx23 = (const int32_t *) (vtcm_base + kparams->merge_indices_off) + 2 * chunk_elems;
+                const float *   val01 = merge_values;
+                const int32_t * idx01 = merge_indices;
+                const float *   val23 = merge_values  + 2 * chunk_elems;
+                const int32_t * idx23 = merge_indices + 2 * chunk_elems;
 
-                int32_t * final_idx = (int32_t *) vtcm_base;
+                int32_t * final_idx = (int32_t *) phase1_base;
                 struct htp_thread_trace * tr = &octx->ctx->trace[0];
                 htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                 merge_runs(val01, idx01, n01, val23, idx23, n23, NULL, final_idx, asc);
                 htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                 dma_sync_write(octx->ctx->dma[0], dst->data, final_idx, ne00 * sizeof(int32_t));
             } else if (kparams->n_chunks == 2) {
-                size_t slot_size = (size_t) kparams->phase1_slot_size;
+                size_t slot_size = layout.phase1_slot_size;
                 uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
 
-                const float *   val0 = (const float *)   (vtcm_base);
-                const int32_t * idx0 = (const int32_t *) (vtcm_base + chunk_elems * sizeof(float));
-                const float *   val1 = (const float *)   (vtcm_base + slot_size);
-                const int32_t * idx1 = (const int32_t *) (vtcm_base + slot_size + chunk_elems * sizeof(float));
+                const float *   val0 = (const float *)   (phase1_base);
+                const int32_t * idx0 = (const int32_t *) (phase1_base + chunk_elems * sizeof(float));
+                const float *   val1 = (const float *)   (phase1_base + slot_size);
+                const int32_t * idx1 = (const int32_t *) (phase1_base + slot_size + chunk_elems * sizeof(float));
 
-                int32_t * final_idx = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+                int32_t * final_idx = merge_indices;
                 struct htp_thread_trace * tr = &octx->ctx->trace[0];
                 htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                 merge_runs(val0, idx0, cctx.real_count[0], val1, idx1, cctx.real_count[1], NULL, final_idx, asc);
@@ -996,7 +996,7 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
                 dma_sync_write(octx->ctx->dma[0], dst->data, final_idx, ne00 * sizeof(int32_t));
             } else {
                 uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
-                int32_t * final_idx = (int32_t *) (vtcm_base + chunk_elems * sizeof(float));
+                int32_t * final_idx = (int32_t *) (phase1_base + chunk_elems * sizeof(float));
                 dma_sync_write(octx->ctx->dma[0], dst->data, final_idx, ne00 * sizeof(int32_t));
             }
         }
@@ -1006,7 +1006,10 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
     struct htp_sort_multi_row_ctx mctx;
     mctx.octx             = octx;
     mctx.kparams          = kparams;
-    mctx.vtcm_base        = vtcm_base;
+    mctx.phase1_base      = phase1_base;
+    mctx.phase1_slot_size = layout.phase1_slot_size;
+    mctx.row_start        = row_start;
+    mctx.row_end          = row_end;
 
     uint32_t dst_row_bytes = (uint32_t) dst_row_size;
     uint32_t rows_per_chunk = 0;
@@ -1014,13 +1017,13 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
     if (rows_per_chunk == 0) {
         rows_per_chunk = 1;
     }
-    uint32_t dr = (nrows + (uint32_t) kparams->n_threads - 1) / (uint32_t) kparams->n_threads;
+    uint32_t dr = fastdiv(nrows + octx->n_threads - 1, &octx->n_threads_div);
     if (rows_per_chunk > 1) {
         dr = hex_round_up(dr, rows_per_chunk);
     }
     mctx.nrows_per_thread = dr;
 
-    work_queue_run(octx->ctx->work_queue, htp_sort_multi_row_job, &mctx, (uint32_t) kparams->n_threads);
+    work_queue_run(octx->ctx->work_queue, htp_sort_multi_row_job, &mctx, octx->n_threads);
 
     return HTP_STATUS_OK;
 }

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

@@ -7,33 +7,27 @@
 #include <string.h>
 
 #include "hex-fastdiv.h"
+#include "htp-vtcm.h"
 
 struct htp_sort_kernel_params {
-    int32_t  n_threads;
-    int32_t  total_rows;
-    int32_t  row_start;
-    int32_t  row_end;
-    int32_t  ne00;
-    int32_t  k;
-    int32_t  order;             // GGML_SORT_ORDER_ASC (0) or GGML_SORT_ORDER_DESC (1)
-    int32_t  is_top_k;          // 1 if TOP_K, 0 if ARGSORT
-    int32_t  use_dma;           // 1 if DMA enabled
-    int32_t  chunk_elems;
-    int32_t  n_chunks;
-    int32_t  vtcm_size;
-    int32_t  phase1_slot_size;
-    int32_t  merge_values_off;
-    int32_t  merge_indices_off;
-    int32_t  merge_elems;
-    int32_t  n_slots;
-    int32_t  pad[15];
+    uint32_t n_threads;
+    uint32_t order;             // GGML_SORT_ORDER_ASC (0) or GGML_SORT_ORDER_DESC (1)
+    uint32_t is_top_k;          // 1 if TOP_K, 0 if ARGSORT
+    uint32_t chunk_elems;
+    uint32_t n_chunks;
+    uint32_t vtcm_size;
+    uint32_t merge_elems;
+    uint32_t n_slots;
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne01_ne02;
 };
 
 struct htp_sort_vtcm_layout {
     size_t   total_bytes;
     size_t   phase1_slot_size;
-    size_t   merge_values_off;
-    size_t   merge_indices_off;
+    size_t   off_phase1;
+    size_t   off_merge_values;
+    size_t   off_merge_indices;
     uint32_t chunk_elems;
     uint32_t n_chunks;
     uint32_t merge_elems;
@@ -41,6 +35,39 @@ struct htp_sort_vtcm_layout {
     uint32_t n_slots;
 };
 
+static inline void htp_sort_vtcm_layout_build(
+    struct htp_sort_vtcm_layout * L,
+    uint32_t total_rows,
+    uint32_t chunk_elems,
+    uint32_t n_chunks,
+    uint32_t merge_elems,
+    uint32_t n_threads,
+    uint32_t n_slots) {
+
+    memset(L, 0, sizeof(*L));
+    L->chunk_elems = chunk_elems;
+    L->n_chunks    = n_chunks;
+    L->merge_elems = merge_elems;
+    L->n_threads   = n_threads;
+    L->n_slots     = n_slots;
+
+    const size_t values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
+    const size_t indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
+    L->phase1_slot_size       = ((values_size + indices_size) + 255) & ~255;
+
+    size_t off = 0;
+    if (total_rows > 1) {
+        VTCM_LAYOUT_ALLOC(off, off_phase1, L->phase1_slot_size * n_slots * n_threads);
+    } else {
+        VTCM_LAYOUT_ALLOC(off, off_phase1, L->phase1_slot_size * n_chunks * n_slots);
+        const size_t merge_values_size  = ((merge_elems * sizeof(float)) + 127) & ~127;
+        const size_t merge_indices_size = ((merge_elems * sizeof(int32_t)) + 127) & ~127;
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_merge_values,  merge_values_size,  n_chunks > 1);
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_merge_indices, merge_indices_size, n_chunks > 1);
+    }
+    L->total_bytes = off;
+}
+
 static inline bool htp_sort_solve_layout(
     struct htp_sort_vtcm_layout * layout,
     uint32_t ne00,
@@ -75,18 +102,8 @@ static inline bool htp_sort_solve_layout(
             n_threads--;
         }
 
-        size_t total_bytes = spad_per_thread * n_threads;
-        if (total_bytes > vtcm_budget) {
-            return false;
-        }
-
-        layout->total_bytes      = total_bytes;
-        layout->phase1_slot_size = spad_per_slot;
-        layout->chunk_elems      = ne00_padded;
-        layout->n_chunks         = 1;
-        layout->n_threads        = n_threads;
-        layout->n_slots          = n_slots;
-        return true;
+        htp_sort_vtcm_layout_build(layout, total_rows, ne00_padded, 1, 0, n_threads, n_slots);
+        return layout->total_bytes <= vtcm_budget;
     }
 
     uint32_t n_vec = (ne00 + 31) / 32;
@@ -102,17 +119,7 @@ static inline bool htp_sort_solve_layout(
 
     uint32_t chunk_n_vec = n_vec_pow2 / n_chunks;
     uint32_t chunk_elems = chunk_n_vec * 32;
-
-    size_t phase1_values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
-    size_t phase1_indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
-    size_t phase1_slot_size    = ((phase1_values_size + phase1_indices_size) + 255) & ~255;
-    size_t phase1_total_size   = phase1_slot_size * n_chunks;
-
-    size_t merge_values_size  = 0;
-    size_t merge_indices_size = 0;
-    size_t merge_values_off   = phase1_total_size;
-    size_t merge_indices_off  = 0;
-    uint32_t merge_elems      = 0;
+    uint32_t merge_elems = 0;
 
     if (n_chunks > 1) {
         if (is_top_k) {
@@ -125,47 +132,25 @@ static inline bool htp_sort_solve_layout(
         } else {
             merge_elems = n_vec_pow2 * 32;
         }
-        merge_values_size  = ((merge_elems * sizeof(float)) + 127) & ~127;
-        merge_indices_size = ((merge_elems * sizeof(int32_t)) + 127) & ~127;
-        merge_indices_off  = merge_values_off + merge_values_size;
     }
 
-    size_t total_bytes = phase1_total_size + merge_values_size + merge_indices_size;
+    htp_sort_vtcm_layout_build(layout, 1, chunk_elems, n_chunks, merge_elems, n_chunks, 1);
 
-    if (total_bytes > vtcm_budget && n_chunks > 1) {
-        n_chunks = 1;
+    if (layout->total_bytes > vtcm_budget && n_chunks > 1) {
+        n_chunks    = 1;
         chunk_elems = n_vec_pow2 * 32;
-        phase1_values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
-        phase1_indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
-        phase1_slot_size    = ((phase1_values_size + phase1_indices_size) + 255) & ~255;
-        phase1_total_size   = phase1_slot_size;
-        merge_values_size   = 0;
-        merge_indices_size  = 0;
-        merge_values_off    = phase1_total_size;
-        merge_indices_off   = 0;
-        merge_elems         = 0;
-        total_bytes         = phase1_total_size;
+        merge_elems = 0;
+        htp_sort_vtcm_layout_build(layout, 1, chunk_elems, n_chunks, merge_elems, 1, 1);
     }
 
-    if (total_bytes > vtcm_budget) {
+    if (layout->total_bytes > vtcm_budget) {
         return false;
     }
 
-    uint32_t n_slots = 1;
-    if (n_chunks == 1 && phase1_slot_size * 2 <= vtcm_budget) {
-        n_slots = 2;
-        total_bytes = phase1_slot_size * 2;
+    if (n_chunks == 1 && layout->phase1_slot_size * 2 <= vtcm_budget) {
+        htp_sort_vtcm_layout_build(layout, 1, chunk_elems, 1, 0, 1, 2);
     }
 
-    layout->total_bytes       = total_bytes;
-    layout->phase1_slot_size  = phase1_slot_size;
-    layout->merge_values_off  = merge_values_off;
-    layout->merge_indices_off = merge_indices_off;
-    layout->chunk_elems       = chunk_elems;
-    layout->n_chunks          = n_chunks;
-    layout->merge_elems       = merge_elems;
-    layout->n_threads         = n_chunks;
-    layout->n_slots           = n_slots;
     return true;
 }
 

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

@@ -42,6 +42,7 @@ struct htp_binary_context {
     uint32_t nrows_per_thread;
     uint32_t total_rows;
     uint32_t row_start;
+    uint32_t n_srcs;
     size_t   src0_row_size_aligned;
     size_t   src1_row_size_aligned;
     size_t   dst_row_size_aligned;
@@ -380,6 +381,7 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
     size_t dst_spad_half  = layout->dst_spad_half_size;
 
     dma_queue * dma_q = octx->ctx->dma[ith];
+    const uint32_t n_srcs = bctx->n_srcs;
     uint32_t ir_prefetch = start_row;
     int spad_idx = 0;
 
@@ -391,12 +393,7 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
         i02 = fastdiv(rem, &bctx->src0_dim1_div);
         i01 = rem - i02 * ne01;
 
-        uint32_t i13 = (ne13 == 1) ? 0 : i03;
-        uint32_t i12 = (ne12 == 1) ? 0 : i02;
-        uint32_t i11 = (ne11 == 1) ? 0 : i01;
-
         dma_addr_t src0_curr = src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01;
-        dma_addr_t src1_curr = src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11;
         dma_addr_t dst_curr  = dst->data  + i03 * nb3  + i02 * nb2  + i01 * nb1;
 
         uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half;
@@ -405,7 +402,15 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
 
         dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0);
         dma_queue_push(dma_q, dma_make_data(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size);
-        dma_queue_push(dma_q, dma_make_data(s1_spad, src1_curr), bctx->src1_row_size_aligned, nb11, row_size_bytes, current_block_size);
+        for (uint32_t s = 1; s < n_srcs; s++) {
+            const struct htp_tensor * src_s = octx->src[s];
+            uint32_t is3 = (src_s->ne[3] == 1) ? 0 : i03;
+            uint32_t is2 = (src_s->ne[2] == 1) ? 0 : i02;
+            uint32_t is1 = (src_s->ne[1] == 1) ? 0 : i01;
+            dma_addr_t s_curr = src_s->data + is3 * src_s->nb[3] + is2 * src_s->nb[2] + is1 * src_s->nb[1];
+            uint8_t * s_spad = s1_spad + (s - 1) * 2 * src1_spad_half;
+            dma_queue_push(dma_q, dma_make_data(s_spad, s_curr), bctx->src1_row_size_aligned, src_s->nb[1], row_size_bytes, current_block_size);
+        }
         ir_prefetch += current_block_size;
         spad_idx ^= 1;
     }
@@ -422,6 +427,11 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
         compute(d_spad, s0_spad, s1_spad, current_block_size,
                 bctx->dst_row_size_aligned, bctx->src0_row_size_aligned, bctx->src1_row_size_aligned, ne00);
+        for (uint32_t s = 2; s < n_srcs; s++) {
+            uint8_t * ss_spad = (uint8_t *) dma_queue_pop(dma_q).dst;
+            compute(d_spad, d_spad, ss_spad, current_block_size,
+                    bctx->dst_row_size_aligned, bctx->dst_row_size_aligned, bctx->src1_row_size_aligned, ne00);
+        }
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
 
         uint32_t i03, i02, i01, rem;
@@ -440,15 +450,18 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
             p02 = fastdiv(prem, &bctx->src0_dim1_div);
             p01 = prem - p02 * ne01;
 
-            uint32_t p13 = (ne13 == 1) ? 0 : p03;
-            uint32_t p12 = (ne12 == 1) ? 0 : p02;
-            uint32_t p11 = (ne11 == 1) ? 0 : p01;
-
             dma_addr_t s0_next = src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01;
-            dma_addr_t s1_next = src1->data + p13 * nb13 + p12 * nb12 + p11 * nb11;
-
             dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size);
-            dma_queue_push(dma_q, dma_make_data(s1_spad, s1_next), bctx->src1_row_size_aligned, nb11, row_size_bytes, next_block_size);
+
+            for (uint32_t s = 1; s < n_srcs; s++) {
+                const struct htp_tensor * src_s = octx->src[s];
+                uint32_t ps3 = (src_s->ne[3] == 1) ? 0 : p03;
+                uint32_t ps2 = (src_s->ne[2] == 1) ? 0 : p02;
+                uint32_t ps1 = (src_s->ne[1] == 1) ? 0 : p01;
+                dma_addr_t s_next = src_s->data + ps3 * src_s->nb[3] + ps2 * src_s->nb[2] + ps1 * src_s->nb[1];
+                uint8_t * s_spad = s1_spad + (s - 1) * 2 * src1_spad_half;
+                dma_queue_push(dma_q, dma_make_data(s_spad, s_next), bctx->src1_row_size_aligned, src_s->nb[1], row_size_bytes, next_block_size);
+            }
 
             ir_prefetch += next_block_size;
         }
@@ -859,6 +872,7 @@ struct binary_chunked_context {
     uint32_t                 chunk_size;
     uint32_t                 chunks_per_thread;
     uint32_t                 total_chunks;
+    uint32_t                 n_srcs;
     bool                     is_scalar;
     float                    scalar_f32;
     _Float16                 scalar_f16;
@@ -872,7 +886,6 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
     struct binary_chunked_context * ctx = (struct binary_chunked_context *) data;
     struct htp_ops_context * octx = ctx->octx;
     const struct htp_tensor * src0 = octx->src[0];
-    const struct htp_tensor * src1 = octx->src[1];
     const struct htp_tensor * dst  = octx->dst;
 
     const uint32_t start_chunk = ctx->chunks_per_thread * ith;
@@ -895,6 +908,7 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
     uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);
 
     dma_queue * dma_q = octx->ctx->dma[ith];
+    const uint32_t n_srcs = ctx->n_srcs;
     uint32_t prefetch_chunk = start_chunk;
     int spad_idx = 0;
 
@@ -913,9 +927,12 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
         dma_queue_push(dma_q, dma_make_data(d_curr, d_spad), chunk_bytes, chunk_bytes, cur_bytes, 0);
         dma_queue_push(dma_q, dma_make_data(s0_spad, s0_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
         if (!ctx->is_scalar) {
-            dma_addr_t s1_curr = src1->data + (size_t) c_start * elem_size;
             uint8_t * s1_spad = src1_spad_base + spad_idx * chunk_bytes;
-            dma_queue_push(dma_q, dma_make_data(s1_spad, s1_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
+            for (uint32_t s = 1; s < n_srcs; s++) {
+                dma_addr_t s_curr = octx->src[s]->data + (size_t) c_start * elem_size;
+                uint8_t * s_spad  = s1_spad + (s - 1) * 2 * chunk_bytes;
+                dma_queue_push(dma_q, dma_make_data(s_spad, s_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
+            }
         }
 
         prefetch_chunk++;
@@ -943,6 +960,10 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
             }
         } else {
             ctx->compute(d_spad, s0_spad, s1_spad, cur_elems);
+            for (uint32_t s = 2; s < n_srcs; s++) {
+                uint8_t * s_spad = (uint8_t *) dma_queue_pop(dma_q).dst;
+                ctx->compute(d_spad, d_spad, s_spad, cur_elems);
+            }
         }
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) c);
 
@@ -958,8 +979,11 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
             dma_addr_t s0_next = src0->data + (size_t) pc_start * elem_size;
             dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), chunk_bytes, chunk_bytes, p_bytes, 1);
             if (!ctx->is_scalar) {
-                dma_addr_t s1_next = src1->data + (size_t) pc_start * elem_size;
-                dma_queue_push(dma_q, dma_make_data(s1_spad, s1_next), chunk_bytes, chunk_bytes, p_bytes, 1);
+                for (uint32_t s = 1; s < n_srcs; s++) {
+                    dma_addr_t s_next = octx->src[s]->data + (size_t) pc_start * elem_size;
+                    uint8_t * s_spad  = s1_spad + (s - 1) * 2 * chunk_bytes;
+                    dma_queue_push(dma_q, dma_make_data(s_spad, s_next), chunk_bytes, chunk_bytes, p_bytes, 1);
+                }
             }
 
             prefetch_chunk++;
@@ -974,6 +998,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
     const struct htp_tensor * src1 = octx->src[1];
     const struct htp_tensor * dst  = octx->dst;
     const struct htp_binary_kernel_params * kparams = (const struct htp_binary_kernel_params *) octx->kernel_params;
+    const uint32_t n_srcs = kparams->n_srcs > 2 ? kparams->n_srcs : 2;
 
     const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
 
@@ -985,7 +1010,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
     const size_t dst_row_size  = dst->ne[0]  * elem_size;
 
     if (kparams->kernel_type == HTP_BINARY_KERNEL_CHUNKED) {
-        const bool is_scalar = kparams->is_scalar || (src1->ne[0] == 1 && src1->ne[1] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1);
+        const bool is_scalar = kparams->is_scalar;
 
         compute_binary_chunked_t compute = NULL;
         compute_binary_scalar_chunked_f32_t compute_scalar_f32 = NULL;
@@ -1045,12 +1070,16 @@ static int execute_op_binary(struct htp_ops_context * octx) {
         const uint32_t elems_per_line = (src0_type == HTP_TYPE_F32) ? 32 : 64;
 
         if (octx->ctx->mdev.count > 1) {
-            const bool can_split = htp_tensor_mdev_data_aligned(dst) &&
-                                   htp_tensor_mdev_data_aligned(src0) &&
-                                   (is_scalar || htp_tensor_mdev_data_aligned(src1)) &&
-                                   htp_tensor_is_contiguous(dst, elem_size) &&
-                                   htp_tensor_is_contiguous(src0, elem_size) &&
-                                   (is_scalar || htp_tensor_is_contiguous(src1, elem_size));
+            bool can_split = htp_tensor_mdev_data_aligned(dst) &&
+                             htp_tensor_mdev_data_aligned(src0) &&
+                             (is_scalar || htp_tensor_mdev_data_aligned(src1)) &&
+                             htp_tensor_is_contiguous(dst, elem_size) &&
+                             htp_tensor_is_contiguous(src0, elem_size) &&
+                             (is_scalar || htp_tensor_is_contiguous(src1, elem_size));
+            for (uint32_t s = 2; can_split && s < n_srcs; s++) {
+                can_split = htp_tensor_mdev_data_aligned(octx->src[s]) &&
+                            htp_tensor_is_contiguous(octx->src[s], elem_size);
+            }
             const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
                 total_elems, can_split ? elems_per_line : 0,
                 octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
@@ -1072,7 +1101,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
             return HTP_STATUS_VTCM_TOO_SMALL;
         }
 
-        const uint32_t chunk_size = kparams->chunk_size > 0 ? kparams->chunk_size : (32768 / elem_size);
+        const uint32_t chunk_size = kparams->chunk_size;
         const uint32_t total_chunks = (nelem + chunk_size - 1) / chunk_size;
         const uint32_t n_threads = (total_chunks >= 2) ? MIN(octx->n_threads, total_chunks) : 1;
         const uint32_t chunks_per_thread = (total_chunks + n_threads - 1) / n_threads;
@@ -1101,6 +1130,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
             .chunk_size          = chunk_size,
             .chunks_per_thread   = chunks_per_thread,
             .total_chunks        = total_chunks,
+            .n_srcs              = n_srcs,
             .is_scalar           = is_scalar,
             .scalar_f32          = scalar_f32,
             .scalar_f16          = scalar_f16,
@@ -1146,65 +1176,6 @@ static int execute_op_binary(struct htp_ops_context * octx) {
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
-    dma_queue * dma_q = octx->ctx->dma[0];
-    uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src1);
-    if (kparams->kernel_type == HTP_BINARY_KERNEL_ROW_BCAST) {
-        dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1->ne[0] * elem_size, 1);
-    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR ||
-               kparams->kernel_type == HTP_BINARY_KERNEL_REPEAT ||
-               kparams->kernel_type == HTP_BINARY_KERNEL_COMPLEX) {
-        if (!kparams->src1_strided) {
-            const size_t src1_bytes = (size_t) src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3] * elem_size;
-            dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1_bytes, 1);
-        } else {
-            const uint32_t ne11 = src1->ne[1];
-            const uint32_t ne12 = src1->ne[2];
-            const uint32_t ne13 = src1->ne[3];
-            const size_t row_bytes = src1->ne[0] * elem_size;
-            for (uint32_t i13 = 0; i13 < ne13; i13++) {
-                for (uint32_t i12 = 0; i12 < ne12; i12++) {
-                    dma_addr_t src_addr = src1->data + i13 * src1->nb[3] + i12 * src1->nb[2];
-                    uint8_t * dst_ptr = vtcm_src1 + (i13 * ne12 * ne11 + i12 * ne11) * row_bytes;
-                    while (!dma_queue_push(dma_q, dma_make_data(dst_ptr, src_addr),
-                                           row_bytes, src1->nb[1], row_bytes, ne11)) {
-                        dma_queue_pop(dma_q);
-                    }
-                }
-            }
-        }
-    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) {
-        const struct htp_tensor * src2 = octx->src[2];
-        uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
-        dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data),
-                       kparams->src1_row_size_aligned, src1->nb[1],
-                       src1->ne[0] * elem_size, src1->ne[1]);
-        const size_t src2_row_bytes = src2->ne[0] * sizeof(int32_t);
-        dma_queue_push(dma_q, dma_make_data(vtcm_src2, src2->data),
-                       src2_row_bytes, src2->nb[1],
-                       src2_row_bytes, src2->ne[1]);
-    }
-
-    bctx.octx                  = octx;
-    bctx.nrows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
-    bctx.total_rows            = nrows;
-    bctx.row_start             = row_start;
-    bctx.block_max             = bctx.vtcm_layout.rows_per_buffer;
-    bctx.src0_row_size_aligned = src0_row_size_aligned;
-    bctx.src1_row_size_aligned = src1_row_size_aligned;
-    bctx.dst_row_size_aligned  = dst_row_size_aligned;
-    bctx.row_size_bytes        = src0_row_size;
-
-    bctx.src0_dim1_div  = init_fastdiv_values(src0->ne[1]);
-    bctx.src0_dim2_div  = init_fastdiv_values(src0->ne[2]);
-    bctx.src0_dim12_div = init_fastdiv_values(src0->ne[1] * src0->ne[2]);
-
-    bctx.src1_dim1_div  = init_fastdiv_values(src1->ne[1]);
-    bctx.src1_dim2_div  = init_fastdiv_values(src1->ne[2]);
-    bctx.src1_dim3_div  = init_fastdiv_values(src1->ne[3]);
-
-    bctx.split_at_ne01  = kparams->split_at_ne01;
-    bctx.split_at_ne02  = kparams->split_at_ne02;
-
     worker_callback_t worker_func = NULL;
     void * compute_func = NULL;
 
@@ -1322,7 +1293,70 @@ static int execute_op_binary(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    bctx.compute = compute_func;
+    dma_queue * dma_q = octx->ctx->dma[0];
+    uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src1);
+    if (kparams->kernel_type == HTP_BINARY_KERNEL_ROW_BCAST) {
+        dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1->ne[0] * elem_size, 1);
+    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR ||
+               kparams->kernel_type == HTP_BINARY_KERNEL_REPEAT ||
+               kparams->kernel_type == HTP_BINARY_KERNEL_COMPLEX) {
+        if (!kparams->src1_strided) {
+            const size_t src1_bytes = (size_t) src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3] * elem_size;
+            dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1_bytes, 1);
+            if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR && octx->src[2] != NULL) {
+                uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
+                dma_queue_push(dma_q, dma_make_data(vtcm_src2, octx->src[2]->data), bctx.vtcm_layout.src2_size, 0, src1_bytes, 1);
+            }
+        } else {
+            const uint32_t ne11 = src1->ne[1];
+            const uint32_t ne12 = src1->ne[2];
+            const uint32_t ne13 = src1->ne[3];
+            const size_t row_bytes = src1->ne[0] * elem_size;
+            for (uint32_t i13 = 0; i13 < ne13; i13++) {
+                for (uint32_t i12 = 0; i12 < ne12; i12++) {
+                    dma_addr_t src_addr = src1->data + i13 * src1->nb[3] + i12 * src1->nb[2];
+                    uint8_t * dst_ptr = vtcm_src1 + (i13 * ne12 * ne11 + i12 * ne11) * row_bytes;
+                    while (!dma_queue_push(dma_q, dma_make_data(dst_ptr, src_addr),
+                                           row_bytes, src1->nb[1], row_bytes, ne11)) {
+                        dma_queue_pop(dma_q);
+                    }
+                }
+            }
+        }
+    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) {
+        const struct htp_tensor * src2 = octx->src[2];
+        uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
+        dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data),
+                       kparams->src1_row_size_aligned, src1->nb[1],
+                       src1->ne[0] * elem_size, src1->ne[1]);
+        const size_t src2_row_bytes = src2->ne[0] * sizeof(int32_t);
+        dma_queue_push(dma_q, dma_make_data(vtcm_src2, src2->data),
+                       src2_row_bytes, src2->nb[1],
+                       src2_row_bytes, src2->ne[1]);
+    }
+
+    bctx.octx                  = octx;
+    bctx.nrows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    bctx.total_rows            = nrows;
+    bctx.row_start             = row_start;
+    bctx.n_srcs                = n_srcs;
+    bctx.block_max             = bctx.vtcm_layout.rows_per_buffer;
+    bctx.src0_row_size_aligned = src0_row_size_aligned;
+    bctx.src1_row_size_aligned = src1_row_size_aligned;
+    bctx.dst_row_size_aligned  = dst_row_size_aligned;
+    bctx.row_size_bytes        = src0_row_size;
+
+    bctx.src0_dim1_div  = kparams->div_ne01;
+    bctx.src0_dim2_div  = kparams->div_ne02;
+    bctx.src0_dim12_div = kparams->div_ne012;
+
+    bctx.src1_dim1_div  = kparams->div_ne11;
+    bctx.src1_dim2_div  = kparams->div_ne12;
+    bctx.src1_dim3_div  = kparams->div_ne13;
+
+    bctx.split_at_ne01  = kparams->split_at_ne01;
+    bctx.split_at_ne02  = kparams->split_at_ne02;
+    bctx.compute        = compute_func;
 
     if (kparams->kernel_type == HTP_BINARY_KERNEL_ROW_BCAST ||
         kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR ||
@@ -1330,6 +1364,15 @@ static int execute_op_binary(struct htp_ops_context * octx) {
         kparams->kernel_type == HTP_BINARY_KERNEL_COMPLEX ||
         kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) {
         dma_queue_flush(dma_q);
+        if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR && octx->src[2] != NULL && octx->op == HTP_OP_MUL) {
+            uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
+            const uint32_t n_scalars = src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3];
+            if (src0_type == HTP_TYPE_F32) {
+                hvx_mul_f32_aaa(vtcm_src1, vtcm_src1, vtcm_src2, n_scalars);
+            } else if (src0_type == HTP_TYPE_F16) {
+                hvx_mul_f16_aaa(vtcm_src1, vtcm_src1, vtcm_src2, n_scalars);
+            }
+        }
     }
 
     work_queue_run(octx->ctx->work_queue, worker_func, &bctx, n_threads);

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

@@ -6,6 +6,7 @@
 #include <string.h>
 
 #include "hex-common.h"
+#include "hex-fastdiv.h"
 #include "htp-ops.h"
 #include "htp-vtcm.h"
 
@@ -20,7 +21,15 @@ enum htp_binary_kernel_type {
 };
 
 struct htp_binary_kernel_params {
-    uint32_t kernel_type;
+    uint8_t  kernel_type;
+    uint8_t  split_at_ne01;
+    uint8_t  split_at_ne02;
+    uint8_t  src1_strided;
+
+    uint8_t  n_srcs;
+    uint8_t  is_scalar;
+    uint8_t  pad[2];
+
     uint32_t n_threads;
     uint32_t rows_per_buffer;
 
@@ -34,12 +43,14 @@ struct htp_binary_kernel_params {
 
     uint32_t chunk_size;
     uint32_t chunk_bytes;
-    uint32_t is_scalar;
 
-    uint8_t  split_at_ne01;
-    uint8_t  split_at_ne02;
-    uint8_t  src1_strided;
-    uint8_t  pad;
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne02;
+    struct fastdiv_values div_ne012;
+
+    struct fastdiv_values div_ne11;
+    struct fastdiv_values div_ne12;
+    struct fastdiv_values div_ne13;
 };
 
 #if defined(__cplusplus)
@@ -100,6 +111,8 @@ static inline void htp_binary_vtcm_layout_build(
         return;
     }
 
+    const uint32_t n_srcs = kparams->n_srcs > 2 ? kparams->n_srcs : 2;
+
     if (kparams->kernel_type == HTP_BINARY_KERNEL_CHUNKED) {
         const size_t chunk_bytes = kparams->chunk_bytes;
         if (chunk_bytes == 0) {
@@ -107,7 +120,7 @@ static inline void htp_binary_vtcm_layout_build(
         }
 
         L->src0_bytes_per_thread = 2 * chunk_bytes;
-        L->src1_bytes_per_thread = kparams->is_scalar ? 0 : (2 * chunk_bytes);
+        L->src1_bytes_per_thread = kparams->is_scalar ? 0 : (2 * (n_srcs - 1) * chunk_bytes);
         L->dst_bytes_per_thread  = 2 * chunk_bytes;
 
         L->src0_spad_half_size = chunk_bytes;
@@ -135,7 +148,7 @@ static inline void htp_binary_vtcm_layout_build(
     }
 
     const size_t spad_row_total = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE)
-        ? 2 * (kparams->src0_row_size_aligned + kparams->src1_row_size_aligned + kparams->dst_row_size_aligned)
+        ? 2 * (kparams->src0_row_size_aligned + (n_srcs - 1) * kparams->src1_row_size_aligned + kparams->dst_row_size_aligned)
         : 2 * (kparams->src0_row_size_aligned + kparams->dst_row_size_aligned);
 
     const size_t static_vtcm_size = kparams->src1_size + kparams->src2_size;
@@ -143,7 +156,9 @@ static inline void htp_binary_vtcm_layout_build(
         return;
     }
 
-    const size_t rows_per_buffer = (vtcm_size - static_vtcm_size) / (n_threads * spad_row_total);
+    const size_t rows_per_buffer = (kparams->rows_per_buffer > 0)
+        ? kparams->rows_per_buffer
+        : (vtcm_size - static_vtcm_size) / (n_threads * spad_row_total);
     if (rows_per_buffer == 0) {
         return;
     }
@@ -155,11 +170,13 @@ static inline void htp_binary_vtcm_layout_build(
     L->src0_bytes_per_thread = rows_per_buffer * 2 * kparams->src0_row_size_aligned;
     L->dst_bytes_per_thread  = rows_per_buffer * 2 * kparams->dst_row_size_aligned;
     L->src1_bytes_per_thread = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE)
-        ? rows_per_buffer * 2 * kparams->src1_row_size_aligned
+        ? rows_per_buffer * 2 * (n_srcs - 1) * kparams->src1_row_size_aligned
         : 0;
 
     L->src0_spad_half_size = L->src0_bytes_per_thread / 2;
-    L->src1_spad_half_size = L->src1_bytes_per_thread / 2;
+    L->src1_spad_half_size = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE)
+        ? rows_per_buffer * kparams->src1_row_size_aligned
+        : 0;
     L->dst_spad_half_size  = L->dst_bytes_per_thread / 2;
 
     const size_t src0_total = n_threads * L->src0_bytes_per_thread;
@@ -178,4 +195,53 @@ static inline void htp_binary_vtcm_layout_build(
     L->total_bytes = off;
 }
 
+static inline bool htp_binary_solve_chunked(
+    struct htp_binary_vtcm_layout * L,
+    struct htp_binary_kernel_params * kparams,
+    uint32_t total_elems,
+    uint32_t elem_size,
+    uint32_t max_threads,
+    uint32_t n_srcs,
+    bool is_scalar,
+    size_t vtcm_size
+) {
+    const size_t n_bufs = is_scalar ? 2 : (size_t) n_srcs + 1;
+    const size_t scalar_bytes = is_scalar ? 128 : 0;
+    const size_t budget_per_buf = (max_threads > 0 && vtcm_size > scalar_bytes)
+        ? hex_align_down((vtcm_size - scalar_bytes) / (2 * max_threads * n_bufs), 128)
+        : 0;
+    const uint32_t max_chunk_bytes = (budget_per_buf > 0 && budget_per_buf < 32768) ? (uint32_t) budget_per_buf : 32768u;
+    const uint32_t max_chunk_elems = max_chunk_bytes / elem_size;
+    if (max_chunk_elems == 0) {
+        return false;
+    }
+
+    const uint32_t min_chunk_elems = 256 < max_chunk_elems ? 256 : max_chunk_elems;
+    const uint32_t target_chunk_elems = hex_round_up((total_elems + (2 * max_threads) - 1) / (2 * max_threads), 32);
+    const uint32_t chunk_size  = (target_chunk_elems < min_chunk_elems) ? min_chunk_elems
+                               : (target_chunk_elems > max_chunk_elems) ? max_chunk_elems
+                               : target_chunk_elems;
+    const uint32_t chunk_bytes  = hex_round_up(chunk_size * elem_size, 128);
+    const uint32_t total_chunks = (total_elems + chunk_size - 1) / chunk_size;
+    const uint32_t n_threads    = (total_chunks >= 2) ? (max_threads < total_chunks ? max_threads : total_chunks) : 1;
+
+    kparams->kernel_type           = HTP_BINARY_KERNEL_CHUNKED;
+    kparams->n_threads             = n_threads;
+    kparams->n_srcs                = (uint8_t) n_srcs;
+    kparams->rows_per_buffer       = 1;
+    kparams->src1_row_size_aligned = is_scalar ? 0 : kparams->src0_row_size_aligned;
+    kparams->src1_size             = 0;
+    kparams->chunk_size            = chunk_size;
+    kparams->chunk_bytes           = chunk_bytes;
+    kparams->is_scalar             = is_scalar ? 1 : 0;
+
+    htp_binary_vtcm_layout_build(L, kparams, vtcm_size);
+    if (L->total_bytes == 0 || L->total_bytes > vtcm_size) {
+        return false;
+    }
+
+    kparams->vtcm_size = (uint32_t) L->total_bytes;
+    return true;
+}
+
 #endif

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

@@ -294,22 +294,21 @@ static int concat_transposed(struct htp_ops_context * octx, const struct htp_con
 
     const uint32_t n_threads = octx->n_threads;
 
-    // layout precomputed on host; kept for reference:
-    // struct htp_concat_transposed_vtcm_layout layout;
-    // htp_concat_transposed_vtcm_layout_build(&layout, octx->src[0]->ne[0], octx->src[1]->ne[0], type_size, n_threads);
+    struct htp_concat_transposed_vtcm_layout layout;
+    htp_concat_transposed_vtcm_layout_build(&layout, octx->src[0]->ne[0], octx->src[1]->ne[0], type_size, n_threads);
 
     uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
 
     struct htp_concat_context cctx;
     cctx.octx                  = octx;
-    cctx.spad0_base            = vtcm_base;
-    cctx.spad1_base            = vtcm_base + n_threads * kparams->u.transposed.spad0_size_per_thread;
-    cctx.spad0_size_per_thread = kparams->u.transposed.spad0_size_per_thread;
-    cctx.spad1_size_per_thread = kparams->u.transposed.spad1_size_per_thread;
+    cctx.spad0_base            = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_spad0);
+    cctx.spad1_base            = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_spad1);
+    cctx.spad0_size_per_thread = layout.src0_spad_size_per_thread;
+    cctx.spad1_size_per_thread = layout.src1_spad_size_per_thread;
     cctx.row_start             = row_start;
     cctx.nrows                 = nrows;
     cctx.nplanes               = dst->ne[2] * dst->ne[3];
-    cctx.div_ne2               = init_fastdiv_values(dst->ne[2]);
+    cctx.div_ne2               = kparams->u.transposed.div_ne2;
     cctx.nrows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
 
     work_queue_func_t worker_func = (type_size == 4) ? concat_2d_f32_transposed : concat_2d_f16_transposed;

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

@@ -2,6 +2,8 @@
 #define HTP_CONCAT_OPS_H
 
 #include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
 #include "gather-rows.h"
 #include <stdint.h>
 
@@ -13,8 +15,7 @@ enum htp_concat_kernel_type {
 };
 
 struct htp_concat_transposed_params {
-    uint32_t spad0_size_per_thread;
-    uint32_t spad1_size_per_thread;
+    struct fastdiv_values div_ne2;
 };
 
 struct htp_concat_kernel_params {
@@ -40,11 +41,13 @@ _Static_assert(sizeof(struct htp_concat_kernel_params) <= 128, "htp_concat_kerne
 struct htp_concat_transposed_vtcm_layout {
     uint32_t src0_spad_size_per_thread;
     uint32_t src1_spad_size_per_thread;
+    size_t   off_spad0;
+    size_t   off_spad1;
     uint32_t total_bytes;
 };
 
 static inline void htp_concat_transposed_vtcm_layout_build(
-    struct htp_concat_transposed_vtcm_layout * layout,
+    struct htp_concat_transposed_vtcm_layout * L,
     uint32_t src0_ne0,
     uint32_t src1_ne0,
     uint32_t type_size,
@@ -55,9 +58,13 @@ static inline void htp_concat_transposed_vtcm_layout_build(
     uint32_t src1_ne0_padded = hex_round_up(src1_ne0, block_i);
     uint32_t spad0_row_bytes = hex_round_up(src0_ne0 * type_size, 128) + src1_ne0_padded * type_size;
 
-    layout->src0_spad_size_per_thread = block_i * spad0_row_bytes;
-    layout->src1_spad_size_per_thread = src1_ne0_padded * spad1_stride;
-    layout->total_bytes = n_threads * (layout->src0_spad_size_per_thread + layout->src1_spad_size_per_thread);
+    L->src0_spad_size_per_thread = block_i * spad0_row_bytes;
+    L->src1_spad_size_per_thread = src1_ne0_padded * spad1_stride;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_spad0, (size_t) n_threads * L->src0_spad_size_per_thread);
+    VTCM_LAYOUT_ALLOC(off, off_spad1, (size_t) n_threads * L->src1_spad_size_per_thread);
+    L->total_bytes = (uint32_t) off;
 }
 
 #endif // HTP_CONCAT_OPS_H

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

@@ -32,6 +32,7 @@ struct htp_copy_context {
     uint32_t nelem;
     uint32_t elem_per_thread;
 
+    struct htp_copy_convert_vtcm_layout layout;
     uint8_t * vtcm_src0;
     uint8_t * vtcm_dst;
 };
@@ -176,13 +177,14 @@ static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, vo
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                   \
                                                                                              \
     const struct htp_copy_convert_params * cvt = &ct->kparams->u.convert;                    \
-    const uint32_t src0_buf_size = cvt->src0_buf_size;                                       \
-    const uint32_t dst_buf_size  = cvt->dst_buf_size;                                        \
-    const uint32_t src0_row_stride = cvt->src0_row_stride;                                   \
-    const uint32_t dst_row_stride  = cvt->dst_row_stride;                                    \
-    const uint32_t blk_rows        = cvt->blk_rows;                                          \
-    uint8_t * vtcm_src0_base = ct->vtcm_src0 + ith * cvt->spad0_size_per_thread;             \
-    uint8_t * vtcm_dst_base  = ct->vtcm_dst  + ith * cvt->spad1_size_per_thread;             \
+    const struct htp_copy_convert_vtcm_layout * L = &ct->layout;                             \
+    const uint32_t src0_buf_size   = L->src0_buf_size;                                       \
+    const uint32_t dst_buf_size    = L->dst_buf_size;                                        \
+    const uint32_t src0_row_stride = L->src0_row_stride;                                     \
+    const uint32_t dst_row_stride  = L->dst_row_stride;                                      \
+    const uint32_t blk_rows        = L->blk_rows;                                            \
+    uint8_t * vtcm_src0_base = ct->vtcm_src0 + ith * L->spad0_size_per_thread;               \
+    uint8_t * vtcm_dst_base  = ct->vtcm_dst  + ith * L->spad1_size_per_thread;               \
     const uint32_t src0_row_size = ne00 * ct->kparams->src0_type_size;                       \
     const uint32_t dst_row_size  = ne00 * ct->kparams->dst_type_size;                        \
                                                                                              \
@@ -254,17 +256,16 @@ static int cpy_scalar(struct htp_ops_context * octx, const struct htp_copy_kerne
     }
 
     dma_queue * dma_q = octx->ctx->dma[0];
-    dma_addr_t s_vtcm = (dma_addr_t)(uintptr_t) octx->ctx->vtcm_base;
-    dma_addr_t d_vtcm = s_vtcm + VLEN;
+    uint8_t * s_ptr = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, 0);
+    uint8_t * d_ptr = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, VLEN);
+    dma_addr_t s_vtcm = (dma_addr_t)(uintptr_t) s_ptr;
+    dma_addr_t d_vtcm = (dma_addr_t)(uintptr_t) d_ptr;
     const uint32_t s_size = kparams->src0_type_size;
     const uint32_t d_size = kparams->dst_type_size;
 
     dma_queue_push(dma_q, dma_make_data(s_vtcm, src0->data), s_size, s_size, s_size, 1);
     dma_queue_pop(dma_q);
 
-    uint8_t * s_ptr = (uint8_t *) octx->ctx->vtcm_base;
-    uint8_t * d_ptr = s_ptr + VLEN;
-
     const HVX_Vector v_src = hvx_vmem(s_ptr);
     HVX_Vector v_dst;
 
@@ -366,6 +367,8 @@ static int cpy_sameshape_convert(struct htp_ops_context * octx, const struct htp
     const struct htp_copy_convert_params * cvt = &kparams->u.convert;
 
     struct htp_copy_context ct;
+    htp_copy_convert_vtcm_layout_build(&ct.layout, src0->ne[0], kparams->src0_type_size, kparams->dst_type_size, cvt->blk_rows, n_threads);
+
     ct.octx                  = octx;
     ct.kparams               = kparams;
     ct.row_start             = row_start;
@@ -373,8 +376,8 @@ static int cpy_sameshape_convert(struct htp_ops_context * octx, const struct htp
     ct.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
 
     uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
-    ct.vtcm_src0 = vtcm_base;
-    ct.vtcm_dst  = vtcm_base + (size_t) n_threads * cvt->spad0_size_per_thread;
+    ct.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, ct.layout.off_src0);
+    ct.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, ct.layout.off_dst);
 
     work_queue_func_t copy_fun = NULL;
     if (dst->type == HTP_TYPE_F16 && src0->type == HTP_TYPE_F32) {

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

@@ -3,6 +3,7 @@
 
 #include "hex-common.h"
 #include "hex-fastdiv.h"
+#include "htp-vtcm.h"
 #include "gather-rows.h"
 #include <stdint.h>
 
@@ -17,12 +18,6 @@ enum htp_copy_kernel_type {
 };
 
 struct htp_copy_convert_params {
-    uint32_t              src0_buf_size;
-    uint32_t              dst_buf_size;
-    uint32_t              spad0_size_per_thread;
-    uint32_t              spad1_size_per_thread;
-    uint32_t              src0_row_stride;  // row pitch in the VTCM buffers
-    uint32_t              dst_row_stride;
     uint32_t              blk_rows;         // rows moved per DMA descriptor
     struct fastdiv_values div_ne01;
     struct fastdiv_values div_ne02_ne01;
@@ -60,6 +55,9 @@ struct htp_copy_kernel_params {
 #define HTP_COPY_VTCM_BLK_BYTES (32 * 1024)
 
 struct htp_copy_convert_vtcm_layout {
+    size_t   total_bytes;
+    size_t   off_src0;
+    size_t   off_dst;
     uint32_t src0_row_stride;
     uint32_t dst_row_stride;
     uint32_t blk_rows;
@@ -67,11 +65,35 @@ struct htp_copy_convert_vtcm_layout {
     uint32_t dst_buf_size;
     uint32_t spad0_size_per_thread;
     uint32_t spad1_size_per_thread;
-    uint32_t total_bytes;
 };
 
-// ne01 bounds a block: rows of one block share the dim-1 stride.
 static inline void htp_copy_convert_vtcm_layout_build(
+    struct htp_copy_convert_vtcm_layout * L,
+    uint32_t ne00,
+    uint32_t src_type_size,
+    uint32_t dst_type_size,
+    uint32_t blk_rows,
+    uint32_t n_threads) {
+
+    const uint32_t src0_row_stride = hex_round_up(ne00 * src_type_size, 256);
+    const uint32_t dst_row_stride  = hex_round_up(ne00 * dst_type_size, 256);
+
+    L->src0_row_stride       = src0_row_stride;
+    L->dst_row_stride        = dst_row_stride;
+    L->blk_rows              = blk_rows;
+    L->src0_buf_size         = blk_rows * src0_row_stride;
+    L->dst_buf_size          = blk_rows * dst_row_stride;
+    L->spad0_size_per_thread = 2 * L->src0_buf_size;
+    L->spad1_size_per_thread = 2 * L->dst_buf_size;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, (size_t) n_threads * L->spad0_size_per_thread);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  (size_t) n_threads * L->spad1_size_per_thread);
+    L->total_bytes = off;
+}
+
+// ne01 bounds a block: rows of one block share the dim-1 stride.
+static inline void htp_copy_convert_solve_layout(
     struct htp_copy_convert_vtcm_layout * layout,
     uint32_t ne00,
     uint32_t ne01,
@@ -79,10 +101,10 @@ static inline void htp_copy_convert_vtcm_layout_build(
     uint32_t dst_type_size,
     uint32_t n_threads) {
 
-    layout->src0_row_stride = hex_round_up(ne00 * src_type_size, 256);
-    layout->dst_row_stride  = hex_round_up(ne00 * dst_type_size, 256);
+    const uint32_t src0_row_stride = hex_round_up(ne00 * src_type_size, 256);
+    const uint32_t dst_row_stride  = hex_round_up(ne00 * dst_type_size, 256);
 
-    const uint32_t max_row = layout->src0_row_stride > layout->dst_row_stride ? layout->src0_row_stride : layout->dst_row_stride;
+    const uint32_t max_row = src0_row_stride > dst_row_stride ? src0_row_stride : dst_row_stride;
     uint32_t blk_rows = HTP_COPY_VTCM_BLK_BYTES / max_row;
     if (blk_rows > ne01) {
         blk_rows = ne01;
@@ -90,13 +112,8 @@ static inline void htp_copy_convert_vtcm_layout_build(
     if (blk_rows == 0) {
         blk_rows = 1;
     }
-    layout->blk_rows = blk_rows;
 
-    layout->src0_buf_size = blk_rows * layout->src0_row_stride;
-    layout->dst_buf_size  = blk_rows * layout->dst_row_stride;
-    layout->spad0_size_per_thread = 2 * layout->src0_buf_size;
-    layout->spad1_size_per_thread = 2 * layout->dst_buf_size;
-    layout->total_bytes = n_threads * (layout->spad0_size_per_thread + layout->spad1_size_per_thread);
+    htp_copy_convert_vtcm_layout_build(layout, ne00, src_type_size, dst_type_size, blk_rows, n_threads);
 }
 
 #if defined(__cplusplus)

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

@@ -12,9 +12,11 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "htp-vtcm.h"
 #include "hvx-types.h"
 #include "hvx-utils.h"
 #include "dma-queue.h"
+#include "cumsum-ops.h"
 
 #define htp_cumsum_tensors_preamble                         \
     const struct htp_tensor * restrict src0 = octx->src[0]; \
@@ -89,8 +91,8 @@ static inline HVX_Vector hvx_splat_last_f32(HVX_Vector v) {
 }
 
 static inline void hvx_cumsum_row_f32(const float * restrict src, float * restrict dst, uint32_t n) {
-    const uint32_t nvec = n / VLEN_FP32;
-    const uint32_t nloe = n % VLEN_FP32;
+    const uint32_t nvec = n >> 5;
+    const uint32_t nloe = n & (VLEN_FP32 - 1);
 
     HVX_Vector carry = Q6_V_vsplat_R(0);
 
@@ -210,16 +212,19 @@ static void cumsum_thread_f32(unsigned int nth, unsigned int ith, void * data) {
 int op_cumsum_f32(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
+    const struct htp_cumsum_kernel_params * kparams =
+        (const struct htp_cumsum_kernel_params *) octx->kernel_params;
 
-    const uint32_t total_rows      = src0->ne[1] * src0->ne[2] * src0->ne[3];
-    const size_t dst_data_row_size = dst->ne[0] * sizeof(float);
+    const uint32_t total_rows   = kparams->total_rows;
+    const size_t   src_row_size = src0->nb[1];
+    const size_t   dst_row_size = dst->nb[1];
 
     uint32_t row_start = 0;
     uint32_t nrows     = total_rows;
 
     if (octx->ctx->mdev.count > 1) {
         uint32_t rows_per_chunk = 0;
-        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_data_row_size, &rows_per_chunk);
+        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_row_size, &rows_per_chunk);
         const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         row_start = range.start;
         nrows     = range.count;
@@ -229,43 +234,45 @@ int op_cumsum_f32(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t n_threads = octx->n_threads;
 
-    const size_t src_row_size         = src0->nb[1];
-    const size_t dst_row_size         = dst->nb[1];
-    const size_t src_row_size_aligned = hex_round_up(src_row_size, VLEN);
-    const size_t dst_row_size_aligned = hex_round_up(dst_row_size, VLEN);
+    if (!kparams->use_dma && (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
 
-    // 2 ping-pong buffers per thread for src and dst
-    const size_t spad_per_thread = 2 * (src_row_size_aligned + dst_row_size_aligned);
+    struct htp_cumsum_vtcm_layout layout = { 0 };
+    if (kparams->use_dma) {
+        htp_cumsum_vtcm_layout_build(&layout, src_row_size, dst_row_size, n_threads);
 
-    octx->src0_spad.size_per_thread = src_row_size_aligned * 2;
-    octx->dst_spad.size_per_thread  = dst_row_size_aligned * 2;
-
-    octx->src0_spad.size  = n_threads * octx->src0_spad.size_per_thread;
-    octx->dst_spad.size   = n_threads * octx->dst_spad.size_per_thread;
-
-    octx->src0_spad.data  = octx->ctx->vtcm_base;                        octx->src0_spad.src = NULL;
-    octx->dst_spad.data   = octx->src0_spad.data + octx->src0_spad.size; octx->dst_spad.src  = NULL;
+        octx->src0_spad.size_per_thread = layout.src0_bytes_per_thread;
+        octx->dst_spad.size_per_thread  = layout.dst_bytes_per_thread;
+        octx->src0_spad.size = n_threads * layout.src0_bytes_per_thread;
+        octx->dst_spad.size  = n_threads * layout.dst_bytes_per_thread;
+        octx->src0_spad.data = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_src0);
+        octx->dst_spad.data  = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_dst);
+        octx->src0_spad.src  = NULL;
+        octx->dst_spad.src   = NULL;
+    }
 
     struct htp_cumsum_context cctx = {
         .octx                 = octx,
         .src_row_size         = src_row_size,
         .dst_row_size         = dst_row_size,
-        .src_row_size_aligned = src_row_size_aligned,
-        .dst_row_size_aligned = dst_row_size_aligned,
+        .src_row_size_aligned = layout.src_row_size_aligned,
+        .dst_row_size_aligned = layout.dst_row_size_aligned,
         .rows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
         .total_rows           = nrows,
         .row_start            = row_start,
     };
 
-    if (octx->ctx->vtcm_size < spad_per_thread * n_threads) {
-        if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
-            return HTP_STATUS_NO_SUPPORT;
-        }
-        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32, &cctx, n_threads);
-    } else {
+    if (kparams->use_dma) {
         work_queue_run(octx->ctx->work_queue, cumsum_thread_f32_dma, &cctx, n_threads);
+    } else {
+        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32, &cctx, n_threads);
     }
 
     return HTP_STATUS_OK;

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

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

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

@@ -10,10 +10,12 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "htp-vtcm.h"
 #include "hvx-types.h"
 #include "hex-utils.h"
 #include "hvx-copy.h"
 #include "dma-queue.h"
+#include "diag-ops.h"
 
 #define htp_diag_tensors_preamble                           \
     const struct htp_tensor * restrict src0 = octx->src[0]; \
@@ -33,6 +35,7 @@
 
 struct htp_diag_context {
     struct htp_ops_context * octx;
+    const struct htp_diag_kernel_params * kparams;
     size_t          src_batch_size;
     size_t          dst_row_size;
     size_t          src_batch_size_aligned;
@@ -83,8 +86,8 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
     for (uint32_t ib = ib0; ib < ib1; ib++) {
-        const uint32_t i3 = ib / ne02;
-        const uint32_t i2 = ib % ne02;
+        const uint32_t i3 = fastdiv(ib, &dctx->kparams->div_ne02);
+        const uint32_t i2 = ib - i3 * ne02;
 
         const dma_addr_t src_batch = src_data + i3 * nb03 + i2 * nb02;
 
@@ -134,8 +137,8 @@ static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) {
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ib0);
 
     for (uint32_t ib = ib0; ib < ib1; ib++) {
-        const uint32_t i3 = ib / ne02;
-        const uint32_t i2 = ib % ne02;
+        const uint32_t i3 = fastdiv(ib, &dctx->kparams->div_ne02);
+        const uint32_t i2 = ib - i3 * ne02;
 
         const float * restrict src_batch = (const float *)(src_data + i3 * nb03 + i2 * nb02);
 
@@ -155,29 +158,17 @@ static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) {
 int op_diag_f32(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
+    const struct htp_diag_kernel_params * kparams =
+        (const struct htp_diag_kernel_params *) octx->kernel_params;
 
-    const uint32_t total_batches = src0->ne[2] * src0->ne[3];
-    const size_t dst_batch_size  = dst->ne[1] * dst->nb[1];
+    const uint32_t total_batches = kparams->total_batches;
 
     uint32_t batch_start = 0;
     uint32_t nbatches    = total_batches;
 
     if (octx->ctx->mdev.count > 1) {
-        bool can_split = htp_tensor_mdev_data_aligned(dst) && (dst->ne[0] == 1 || dst->nb[0] == sizeof(float)) && !htp_tensor_is_permuted(dst);
-        uint32_t batches_per_chunk = 1;
-        if (can_split) {
-            if (dst->ne[2] > 1 && (dst->nb[2] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0 &&
-                (dst->ne[3] <= 1 || (dst->nb[3] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0)) {
-                batches_per_chunk = 1;
-            } else if (dst->nb[2] == dst_batch_size &&
-                       (dst->ne[3] <= 1 || dst->nb[3] == dst->nb[2] * dst->ne[2])) {
-                batches_per_chunk = (dst_batch_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(dst_batch_size, HEX_L2_LINE_SIZE)) : 1;
-            } else {
-                can_split = false;
-            }
-        }
-
-        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_batches, can_split ? batches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        const bool can_split = (kparams->can_split != 0) && htp_tensor_mdev_data_aligned(dst);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_batches, can_split ? kparams->batches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         batch_start = range.start;
         nbatches    = range.count;
     }
@@ -186,43 +177,49 @@ int op_diag_f32(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    const uint32_t n_threads     = octx->n_threads;
+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nbatches))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
 
-    const size_t src_batch_size         = src0->ne[0] * sizeof(float);
-    const size_t dst_row_size           = dst->ne[0] * sizeof(float);
-    const size_t src_batch_size_aligned = hex_round_up(src_batch_size, VLEN);
-    const size_t dst_row_size_aligned   = hex_round_up(dst_row_size, VLEN);
+    const uint32_t n_threads = octx->n_threads;
 
-    // 1 src buffer + 1 dst row buffer per thread
-    const size_t spad_per_thread = src_batch_size_aligned + dst_row_size_aligned;
+    if (!kparams->use_dma && (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
 
-    octx->src0_spad.size_per_thread = src_batch_size_aligned;
-    octx->dst_spad.size_per_thread  = dst_row_size_aligned;
+    const size_t src_batch_size = src0->ne[0] * sizeof(float);
+    const size_t dst_row_size   = dst->ne[0] * sizeof(float);
 
-    octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread;
-    octx->dst_spad.size  = n_threads * octx->dst_spad.size_per_thread;
+    struct htp_diag_vtcm_layout layout = { 0 };
+    if (kparams->use_dma) {
+        htp_diag_vtcm_layout_build(&layout, src_batch_size, dst_row_size, n_threads);
 
-    octx->src0_spad.data = octx->ctx->vtcm_base;                        octx->src0_spad.src = NULL;
-    octx->dst_spad.data  = octx->src0_spad.data + octx->src0_spad.size; octx->dst_spad.src  = NULL;
+        octx->src0_spad.size_per_thread = layout.src0_bytes_per_thread;
+        octx->dst_spad.size_per_thread  = layout.dst_bytes_per_thread;
+        octx->src0_spad.size = n_threads * layout.src0_bytes_per_thread;
+        octx->dst_spad.size  = n_threads * layout.dst_bytes_per_thread;
+        octx->src0_spad.data = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_src0);
+        octx->dst_spad.data  = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_dst);
+        octx->src0_spad.src  = NULL;
+        octx->dst_spad.src   = NULL;
+    }
 
     struct htp_diag_context dctx = {
         .octx                   = octx,
+        .kparams                = kparams,
         .src_batch_size         = src_batch_size,
         .dst_row_size           = dst_row_size,
-        .src_batch_size_aligned = src_batch_size_aligned,
-        .dst_row_size_aligned   = dst_row_size_aligned,
+        .src_batch_size_aligned = layout.src_batch_size_aligned,
+        .dst_row_size_aligned   = layout.dst_row_size_aligned,
         .batches_per_thread     = fastdiv(nbatches + n_threads - 1, &octx->n_threads_div),
         .total_batches          = nbatches,
         .batch_start            = batch_start,
     };
 
-    if (octx->ctx->vtcm_size < spad_per_thread * n_threads) {
-        if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
-            return HTP_STATUS_NO_SUPPORT;
-        }
-        work_queue_run(octx->ctx->work_queue, diag_thread_f32, &dctx, n_threads);
-    } else {
+    if (kparams->use_dma) {
         work_queue_run(octx->ctx->work_queue, diag_thread_f32_dma, &dctx, n_threads);
+    } else {
+        work_queue_run(octx->ctx->work_queue, diag_thread_f32, &dctx, n_threads);
     }
 
     return HTP_STATUS_OK;

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

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

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

@@ -436,10 +436,10 @@ static inline void dma_sync_write(dma_queue * dma_q, dma_addr_t dst, const void
     }
 }
 
-#define DMA_CACHE_MAX_SIZE 256U
+#define DMA_CACHE_MAX_SIZE 128U
 
 // Fully assoc LRU cache
-typedef struct {
+typedef struct __attribute__((aligned(128))) {
     uint8_t *base;
     uint32_t line_size;
     uint32_t capacity;
@@ -485,7 +485,7 @@ static inline bool dma_cache_fa_push(dma_queue *q, dma_cache_fa *c, dma_addr_t s
 }
 
 // Direct mapped cache
-typedef struct {
+typedef struct __attribute__((aligned(128))) {
     uint8_t *base;
     uint32_t line_size;
     uint32_t capacity;

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

@@ -16,6 +16,7 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "fill-ops.h"
 
 // ggml op_params layout for FILL:
 //   op_params[0] (as float) - the scalar fill value
@@ -30,21 +31,19 @@
                                                \
     const uint32_t nb1 = dst->nb[1];           \
     const uint32_t nb2 = dst->nb[2];           \
-    const uint32_t nb3 = dst->nb[3];           \
-                                               \
-    const uint32_t nr = ne1 * ne2 * ne3;
+    const uint32_t nb3 = dst->nb[3];
 
 struct htp_fill_context {
     struct htp_ops_context * octx;
+    const struct htp_fill_kernel_params * kparams;
     uint32_t nrows_per_thread;
     uint32_t total_rows;  // ne1 * ne2 * ne3
     uint32_t row_start;
-    bool     opt_path;
     HVX_Vector splat_vec;
     uint32_t   elem_size;
 };
 
-static void fill_thread(unsigned int nth, unsigned int ith, void * data) {
+static void fill_thread_contiguous(unsigned int nth, unsigned int ith, void * data) {
     const struct htp_fill_context * fctx = (const struct htp_fill_context *) data;
     struct htp_ops_context        * octx = fctx->octx;
     fill_preamble;
@@ -60,45 +59,59 @@ static void fill_thread(unsigned int nth, unsigned int ith, void * data) {
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
 
-    if (fctx->opt_path) {
-        // Opt path: tensor is fully contiguous, treat as flat array
-        const uint32_t elem_start = ir0 * ne0;
-        const uint32_t elem_end = ir1 * ne0;
-        uint8_t * dst_ptr = (uint8_t *) dst->data + elem_start * fctx->elem_size;
-        hvx_splat_u(dst_ptr, fctx->splat_vec, elem_end - elem_start, fctx->elem_size);
-    } else {
-        // Non-contiguous path: must respect strides
-        for (uint32_t ir = ir0; ir < ir1; ++ir) {
-            const uint32_t i1 = ir % ne1;
-            const uint32_t i2 = (ir / ne1) % ne2;
-            const uint32_t i3 = ir / (ne1 * ne2);
-            uint8_t * dst_ptr = (uint8_t *) dst->data + i1*nb1 + i2*nb2 + i3*nb3;
-            hvx_splat_u(dst_ptr, fctx->splat_vec, ne0, fctx->elem_size);
-        }
+    const uint32_t elem_start = ir0 * ne0;
+    const uint32_t elem_end   = ir1 * ne0;
+    uint8_t * dst_ptr = (uint8_t *) dst->data + elem_start * fctx->elem_size;
+    hvx_splat_u(dst_ptr, fctx->splat_vec, elem_end - elem_start, fctx->elem_size);
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir1);
+    FARF(HIGH, "fill-contig %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
+}
+
+static void fill_thread_strided(unsigned int nth, unsigned int ith, void * data) {
+    const struct htp_fill_context * fctx = (const struct htp_fill_context *) data;
+    struct htp_ops_context        * octx = fctx->octx;
+    const struct htp_fill_kernel_params * kparams = fctx->kparams;
+    fill_preamble;
+
+    // Parallelise over the flat row index spanning ne1*ne2*ne3
+    const uint32_t ir0 = fctx->row_start + fctx->nrows_per_thread * ith;
+    const uint32_t ir1 = MIN(ir0 + fctx->nrows_per_thread, fctx->row_start + fctx->total_rows);
+
+    if (ir0 >= ir1) {
+        return;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
+
+    for (uint32_t ir = ir0; ir < ir1; ++ir) {
+        const uint32_t i3  = fastdiv(ir, &kparams->div_ne1_ne2);
+        const uint32_t rem = ir - i3 * (ne1 * ne2);
+        const uint32_t i2  = fastdiv(rem, &kparams->div_ne1);
+        const uint32_t i1  = rem - i2 * ne1;
+        uint8_t * dst_ptr = (uint8_t *) dst->data + i1*nb1 + i2*nb2 + i3*nb3;
+        hvx_splat_u(dst_ptr, fctx->splat_vec, ne0, fctx->elem_size);
     }
 
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir1);
-    FARF(HIGH, "fill %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
+    FARF(HIGH, "fill-strided %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
 }
 
 int op_fill(struct htp_ops_context * octx) {
-    fill_preamble;
-
-    if (dst->type != HTP_TYPE_F32 && dst->type != HTP_TYPE_F16) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
+    const struct htp_tensor * dst = octx->dst;
+    const struct htp_fill_kernel_params * kparams =
+        (const struct htp_fill_kernel_params *) octx->kernel_params;
 
     if (htp_tensor_is_extended(dst)) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
     uint32_t row_start = 0;
-    uint32_t nrows     = nr;
+    uint32_t nrows     = kparams->total_rows;
 
     if (octx->ctx->mdev.count > 1) {
-        const uint32_t row_size = nb1;
-        const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
-        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(nr, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(kparams->total_rows, htp_tensor_mdev_data_aligned(dst) ? kparams->rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         row_start = range.start;
         nrows     = range.count;
     }
@@ -107,40 +120,43 @@ int op_fill(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    // nr = ne1*ne2*ne3 (flat row count across all outer dims); parallelise over it.
+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t n_threads = octx->n_threads;
 
-    // Optimize if fully contiguous: skip stride arithmetic, treat as flat array
-    const bool opt_path = (nb2 == nb1 * ne1) && (nb3 == nb2 * ne2);
-
-    FARF(HIGH, "fill: (%ux%ux%ux%u) type=%u opt=%d\n",
-         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], dst->type, (int) opt_path);
+    FARF(HIGH, "fill: (%ux%ux%ux%u) type=%u ktype=%u\n",
+         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], dst->type, kparams->kernel_type);
 
     float val_f32 = 0.f;
     memcpy(&val_f32, &octx->op_params[0], sizeof(float));
 
     struct htp_fill_context fctx = {
         .octx             = octx,
+        .kparams          = kparams,
         .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
         .total_rows       = nrows,
         .row_start        = row_start,
-        .opt_path         = opt_path,
+        .elem_size        = kparams->elem_size,
     };
 
     switch (dst->type) {
     case HTP_TYPE_F32:
         fctx.splat_vec = hvx_vec_splat_f32(val_f32);
-        fctx.elem_size = sizeof(float);
         break;
     case HTP_TYPE_F16:
         fctx.splat_vec = hvx_vec_splat_f16((_Float16) val_f32);
-        fctx.elem_size = sizeof(_Float16);
         break;
     default:
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    work_queue_run(octx->ctx->work_queue, fill_thread, &fctx, n_threads);
+    work_queue_func_t worker = (kparams->kernel_type == HTP_FILL_KERNEL_CONTIGUOUS)
+        ? fill_thread_contiguous
+        : fill_thread_strided;
+
+    work_queue_run(octx->ctx->work_queue, worker, &fctx, n_threads);
 
     return HTP_STATUS_OK;
 }

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

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

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

@@ -39,6 +39,26 @@
 // Must be multiple of 32
 #define FLASH_ATTN_BLOCK_SIZE (32 * 2)
 
+// Large scratch lives in ctx->ddr_spad to keep stack usage low.
+// Ops can run on the dspqueue callback thread, which has a small stack.
+struct hvx_fa_ddr_spad {
+    __fp16   slopes[512];
+    uint16_t active_blocks[FA_MAX_KV_BLOCKS];
+    uint16_t thread_active_blocks[HTP_MAX_NTHREADS][FA_MAX_KV_BLOCKS] __attribute__((aligned(128)));
+};
+
+struct hmx_fa_ddr_spad {
+    uint16_t active_blocks[FA_MAX_KV_BLOCKS];
+    __fp16   slopes[] __attribute__((aligned(128))); // [g_br]
+};
+
+static inline void * fa_get_ddr_spad(const struct htp_ops_context * octx, size_t size) {
+    if (!octx->ctx->ddr_spad_base || octx->ctx->ddr_spad_size < size) {
+        return NULL;
+    }
+    return octx->ctx->ddr_spad_base;
+}
+
 struct htp_fa_context {
     const struct htp_ops_context * octx;
 
@@ -61,7 +81,6 @@ struct htp_fa_context {
     uint32_t n_head_log2;
     float m0;
     float m1;
-    __fp16 slopes[512];
 
     uint32_t n_blocks;
     uint32_t n_active_blocks;
@@ -90,7 +109,9 @@ struct htp_fa_context {
     uint8_t *  spad_m;
     uint8_t *  spad_a;
     float *    spad_sinks;
-    uint16_t   active_blocks[FA_MAX_KV_BLOCKS];
+
+    struct hvx_fa_ddr_spad * ddr;
+    dma_cache_dm *           m_cache; // [n_threads]
 
     const struct htp_tensor * k;
     const struct htp_tensor * v;
@@ -156,6 +177,7 @@ struct hmx_fa_context {
     __fp16 *     vtcm_mask_buf;        // VTCM mask buffer [Br * m_line], DMA'd per KV block
     __fp16 *     vtcm_slopes;          // ALiBi slopes [g_br]
     float *      vtcm_sinks;           // Attention sinks
+    __fp16 *     ddr_slopes;           // ALiBi slopes staging buffer in DDR [g_br]
     size_t       row_buf_stride;       // HVX vectors per row buffer (Bc/64)
     size_t       mask_buf_row_stride;  // elements (__fp16) per row in mask buffer
     size_t       q_tile_bytes;
@@ -164,7 +186,7 @@ struct hmx_fa_context {
     size_t       col_vec_bytes;
     size_t       d_tile_bytes;
     bool         mask_broadcast;       // true when mask->ne[2] == 1 (head-independent, single 2D DMA)
-    dma_cache_fa m_cache;
+    dma_cache_fa * m_cache;
 };
 
 // Scan mask rows [q_start, q_start + n_rows_q) and build compact list of KV blocks with any entry != -inf.
@@ -308,11 +330,11 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
     uint8_t * spad_m = factx->spad_m + (mask ? factx->size_m_block * HVX_FA_DMA_CACHE_SIZE : 0) * ith;
     uint8_t * spad_a = factx->spad_a + factx->size_vkq_acc * ith;
 
-    dma_cache_dm m_cache;
-    dma_cache_dm_init(&m_cache, spad_m, factx->size_m_block, HVX_FA_DMA_CACHE_SIZE);
+    dma_cache_dm * m_cache = &factx->m_cache[ith];
+    dma_cache_dm_init(m_cache, spad_m, factx->size_m_block, HVX_FA_DMA_CACHE_SIZE);
 
-    uint16_t thread_active_blocks[FA_MAX_KV_BLOCKS];
-    const uint16_t * active_blocks = factx->shared_active_blocks ? factx->active_blocks : thread_active_blocks;
+    uint16_t * thread_active_blocks = factx->ddr->thread_active_blocks[ith];
+    const uint16_t * active_blocks = factx->shared_active_blocks ? factx->ddr->active_blocks : thread_active_blocks;
     uint32_t n_active_blocks       = factx->n_active_blocks;
     dma_addr_t prev_mp_base        = 0;
 
@@ -385,7 +407,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
             const uint32_t r_iq1 = (r - r_iq3*neq2*neq1 - r_iq2 * neq1);
 
             heads[g] = r_iq2;
-            const __fp16 slope = factx->slopes[r_iq2];
+            const __fp16 slope = factx->ddr->slopes[r_iq2];
             slope_vecs[g] = hvx_vec_splat_f16(slope);
 
             S_vec[g] = hvx_vec_splat_f32(0.0f);
@@ -422,7 +444,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
             // Mask
             if (mask) {
                 const dma_addr_t m_src = mp_base + ic_start * sizeof(__fp16);
-                dma_cache_dm_push(dma_q, &m_cache, ib, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
+                dma_cache_dm_push(dma_q, m_cache, ib, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
             }
         }
 
@@ -682,7 +704,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                     // Mask
                     if (mask) {
                         const dma_addr_t m_src = mp_base + next_ic_start * sizeof(__fp16);
-                        dma_cache_dm_push(dma_q, &m_cache, next_ib, m_src, next_block_size * 2, next_block_size * 2, next_block_size * 2, 1);
+                        dma_cache_dm_push(dma_q, m_cache, next_ib, m_src, next_block_size * 2, next_block_size * 2, next_block_size * 2, 1);
                     }
                 }
             } // end for g
@@ -1887,7 +1909,7 @@ static __attribute__((noinline)) void fa_compute_slopes(
     const float    m0          = factx->m0;
     const float    m1          = factx->m1;
 
-    __fp16 temp_slopes[512] __attribute__((aligned(128)));
+    __fp16 * temp_slopes = factx->ddr_slopes;
     if (G <= 32) {
         // Fast path: Compute G unique slope values in vector registers
         HVX_Vector v_val = hvx_alibi_slopes(kv_head, G, n_head_log2, m0, m1);
@@ -1905,14 +1927,13 @@ static __attribute__((noinline)) void fa_compute_slopes(
         }
     }
 
-    // Allocate stack buffer to avoid scalar writes to VTCM (which generates L2 misses)
-    __fp16 local_slopes[n_rows_g] __attribute__((aligned(128)));
-    for (size_t r = 0; r < n_rows_g; ++r) {
-        local_slopes[r] = temp_slopes[fastmodulo(r, G, &factx->div_G)];
+    // Replicate the G values in DDR to avoid scalar writes to VTCM (which generates L2 misses)
+    for (size_t r = G; r < n_rows_g; ++r) {
+        temp_slopes[r] = temp_slopes[r - G];
     }
 
     // Copy to VTCM slopes using HVX block copy (both are aligned to 128 bytes)
-    hvx_copy_f16_aa((uint8_t *)slopes, (const uint8_t *)local_slopes, n_rows_g);
+    hvx_copy_f16_aa((uint8_t *)slopes, (const uint8_t *)temp_slopes, n_rows_g);
 }
 
 static void fa_push_mask_dma_gqa(
@@ -1958,7 +1979,7 @@ static inline void fa_prefetch_block(dma_queue * dma_q, const struct htp_tensor
     if (mask) {
         if (__builtin_expect(factx->mask_broadcast, true)) {
             const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + prefetch_start * sizeof(__fp16);
-            dma_cache_fa_push(dma_q, &factx->m_cache, ms_src, m_line_bytes, mask->nb[1], prefetch_rows * sizeof(__fp16), n_rows_q);
+            dma_cache_fa_push(dma_q, factx->m_cache, ms_src, m_line_bytes, mask->nb[1], prefetch_rows * sizeof(__fp16), n_rows_q);
         } else {
             fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, prefetch_start, kv_head, G, m_line_bytes, prefetch_rows, n_rows_q, factx);
         }
@@ -2151,7 +2172,14 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
 
     const size_t m_line_bytes = L.m_line_bytes;  // used by the mask DMAs in the KV loop
 
-    dma_cache_fa_init(&factx.m_cache, (uint8_t *) factx.vtcm_mask_buf, L.m_buf_slot_bytes, HMX_FA_DMA_CACHE_SIZE);
+    struct hmx_fa_ddr_spad * ddr = (struct hmx_fa_ddr_spad *) fa_get_ddr_spad(octx, sizeof(struct hmx_fa_ddr_spad) + factx.g_br * sizeof(__fp16));
+    if (!ddr) {
+        return HTP_STATUS_INTERNAL_ERR;
+    }
+
+    factx.ddr_slopes = ddr->slopes;
+    factx.m_cache    = &octx->dma_cache_fa[0];
+    dma_cache_fa_init(factx.m_cache, (uint8_t *) factx.vtcm_mask_buf, L.m_buf_slot_bytes, HMX_FA_DMA_CACHE_SIZE);
 
     // Head-dim padding: the K/V DMA staging buffers and the flat-Q buffer are laid out
     // with padded row strides (size_{k,v,q}_row_padded, covering D_pad columns) but the
@@ -2210,7 +2238,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
             const size_t   g_br_actual = hex_align_up(n_rows_g, HMX_FP16_TILE_N_ROWS);
             const size_t   n_row_tiles = g_br_actual / HMX_FP16_TILE_N_ROWS;
 
-            uint16_t active_blocks[FA_MAX_KV_BLOCKS];
+            uint16_t * active_blocks = ddr->active_blocks;
             uint32_t n_active_blocks = factx.n_kv_blocks;
             if (mask && factx.mask_broadcast && factx.n_kv_blocks > 1) {
                 const size_t scratch_bytes = factx.k_tile_bytes * (factx.pipeline ? 2 : 1);
@@ -2252,7 +2280,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                         if (factx.pipeline && mask) {
                             if (__builtin_expect(factx.mask_broadcast, true)) {
                                 const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + b0_start * sizeof(__fp16);
-                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
+                                dma_cache_fa_push(dma_q, factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
                             } else {
                                 fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, b0_start, kv_head, G, m_line_bytes, b0_rows, n_rows_q, &factx);
                             }
@@ -2453,7 +2481,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                         if (mask) {
                             if (__builtin_expect(factx.mask_broadcast, true)) {
                                 const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
-                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
+                                dma_cache_fa_push(dma_q, factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
                             } else {
                                 fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx);
                             }
@@ -2577,7 +2605,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                         if (factx.pipeline && mask) {
                             if (__builtin_expect(factx.mask_broadcast, true)) {
                                 const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + b0_start * sizeof(__fp16);
-                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
+                                dma_cache_fa_push(dma_q, factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
                             } else {
                                 fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, b0_start, next_kv_head, G, m_line_bytes, b0_rows, n_rows_q, &factx);
                             }
@@ -2680,8 +2708,15 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
     if (n_head > 512) {
         return HTP_STATUS_NO_SUPPORT;
     }
+
+    factx.ddr = (struct hvx_fa_ddr_spad *) fa_get_ddr_spad(octx, sizeof(struct hvx_fa_ddr_spad));
+    if (!factx.ddr) {
+        return HTP_STATUS_INTERNAL_ERR;
+    }
+    factx.m_cache = octx->dma_cache_dm;
+
     for (uint32_t h = 0; h < n_head; ++h) {
-        factx.slopes[h] = (__fp16) ((kparams->max_bias > 0.0f) ? alibi_slope(h, factx.n_head_log2, factx.m0, factx.m1) : 1.0f);
+        factx.ddr->slopes[h] = (__fp16) ((kparams->max_bias > 0.0f) ? alibi_slope(h, factx.n_head_log2, factx.m0, factx.m1) : 1.0f);
     }
 
     // total rows in q
@@ -2754,7 +2789,7 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
     if (has_mask && neq1 == 1 && neq3 == 1 && mask->ne[2] == 1 && factx.n_blocks > 1 && factx.n_blocks <= FA_MAX_KV_BLOCKS) {
         factx.n_active_blocks = fa_scan_active_blocks(
             octx->ctx->dma[0], mask, 0, 1, 0, 0, k->ne[1], FLASH_ATTN_BLOCK_SIZE,
-            factx.n_blocks, factx.spad_k, factx.size_k_block * 2, factx.active_blocks);
+            factx.n_blocks, factx.spad_k, factx.size_k_block * 2, factx.ddr->active_blocks);
         factx.shared_active_blocks = true;
     }
 

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

@@ -35,7 +35,6 @@ enum htp_fa_kernel_type {
 struct htp_fa_kernel_params {
     uint8_t  kernel_type;        // enum htp_fa_kernel_type
     uint8_t  head_split;         // 1 = partition by KV heads in multicore, 0 = token partition
-    uint8_t  flags;              // reserved
     uint8_t  n_threads;          // Number of threads to run
 
     // Common parameters
@@ -49,9 +48,6 @@ struct htp_fa_kernel_params {
     float    logit_softcap;
     uint32_t vtcm_size;
 
-    uint32_t qrows;
-    uint32_t qrows_per_thread;
-    uint32_t qrow_start;
     float    m0;
     float    m1;
     uint32_t n_head_log2;
@@ -67,8 +63,6 @@ struct htp_fa_kernel_params {
     union {
         struct {
             uint32_t g_br;
-            uint32_t row_buf_stride;
-            uint32_t mask_buf_row_stride;
             int32_t  mask_broadcast;
             int32_t  pipeline;
             struct fastdiv_values div_G;
@@ -85,6 +79,8 @@ struct htp_fa_kernel_params {
 
 #if defined(__cplusplus)
 static_assert(sizeof(struct htp_fa_kernel_params) <= 128, "htp_fa_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_fa_kernel_params) <= 128, "htp_fa_kernel_params is too large for kernel_params blob");
 #endif
 
 // VTCM region layout for the HMX flash-attention kernel.

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

@@ -2270,46 +2270,9 @@ static int gated_delta_net_f32_hmx_chunked(
 }
 
 int op_gated_delta_net(struct htp_ops_context * octx) {
-    const struct htp_tensor * q     = octx->src[0];
-    const struct htp_tensor * k     = octx->src[1];
-    const struct htp_tensor * v     = octx->src[2];
-    const struct htp_tensor * g     = octx->src[3];
-    const struct htp_tensor * beta  = octx->src[4];
-    const struct htp_tensor * state = octx->src[5];
-    const struct htp_tensor * dst   = octx->dst;
+    const struct htp_tensor * dst = octx->dst;
 
-    if (q->type != HTP_TYPE_F32 || k->type != HTP_TYPE_F32 || v->type != HTP_TYPE_F32 ||
-        g->type != HTP_TYPE_F32 || beta->type != HTP_TYPE_F32 || state->type != HTP_TYPE_F32 ||
-        dst->type != HTP_TYPE_F32) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    const uint32_t S_v      = v->ne[0];
-    const uint32_t H        = v->ne[1];
-    const uint32_t n_tokens = v->ne[2];
-    const uint32_t n_seqs   = v->ne[3];
-    const uint32_t K        = octx->op_params[0];
-
-    if (S_v == 0 || S_v > HTP_GDN_MAX_SV || H == 0 || n_tokens == 0 || n_seqs == 0) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    if ((g->ne[0] != 1 && g->ne[0] != S_v) || beta->ne[0] != 1) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    if (q->ne[0] != S_v || k->ne[0] != S_v || q->ne[1] == 0 || k->ne[1] == 0 ||
-        q->ne[2] != n_tokens || k->ne[2] != n_tokens || q->ne[3] == 0 || k->ne[3] == 0 ||
-        (n_seqs % q->ne[3]) != 0 || (n_seqs % k->ne[3]) != 0) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    // state holds s0 only: [S_v, S_v, H, n_seqs]
-    if (state->ne[0] != S_v || state->ne[1] != S_v || state->ne[2] != H || state->ne[3] != n_seqs) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    if (dst->ne[0] != S_v * H || dst->ne[1] != n_tokens * n_seqs + S_v * n_seqs * K) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    for (int i = 0; i < 5; i++) {
+    for (int i = 0; i < 6; i++) {
         if (htp_tensor_is_extended(octx->src[i])) {
             return HTP_STATUS_NO_SUPPORT;
         }
@@ -2325,68 +2288,12 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
     }
 
     const struct htp_gdn_kernel_params * kparams = (const struct htp_gdn_kernel_params *) octx->kernel_params;
-    struct htp_gdn_kernel_params kparams_local;
-    if (!kparams || kparams->S_v == 0) {
-        const uint32_t rq3 = n_seqs / q->ne[3];
-        const uint32_t rk3 = n_seqs / k->ne[3];
-        const uint32_t total_rows = H * n_seqs;
-        uint32_t n_threads = (total_rows < octx->n_threads) ? total_rows : octx->n_threads;
-        if (n_threads == 0) {
-            n_threads = 1;
-        }
-
-        memset(&kparams_local, 0, sizeof(kparams_local));
-        kparams_local.n_threads           = n_threads;
-        kparams_local.S_v                 = S_v;
-        kparams_local.H                   = H;
-        kparams_local.n_tokens            = n_tokens;
-        kparams_local.n_seqs              = n_seqs;
-        kparams_local.K                   = K;
-        kparams_local.total_rows          = total_rows;
-        kparams_local.rows_per_thread     = (total_rows + n_threads - 1) / n_threads;
-        const bool can_use_hmx = (octx->ctx->hmx_enabled) &&
-                                 (S_v % 64 == 0) &&
-                                 (n_tokens >= HTP_GDN_MIN_TOKENS) &&
-                                 (g->ne[0] == 1) &&
-                                 (K == 1);
-
-        struct htp_gdn_hmx_vtcm_layout hmx_layout_local;
-        struct htp_gdn_vtcm_layout hvx_layout_local;
-        uint32_t n_heads_batch = 1;
-
-        if (can_use_hmx && htp_gdn_hmx_solve_layout(&hmx_layout_local, S_v, HTP_GDN_CHUNK_SIZE, total_rows, octx->ctx->vtcm_size, n_threads, true, &n_heads_batch)) {
-            kparams_local.kernel_type     = HTP_GDN_KERNEL_HMX_CHUNKED;
-            kparams_local.pipeline        = hmx_layout_local.pipeline ? 1 : 0;
-            kparams_local.chunk_size      = HTP_GDN_CHUNK_SIZE;
-            kparams_local.n_chunks        = (n_tokens + HTP_GDN_CHUNK_SIZE - 1) / HTP_GDN_CHUNK_SIZE;
-            kparams_local.n_heads_batch   = (uint16_t) n_heads_batch;
-            kparams_local.vtcm_size       = (uint32_t) hmx_layout_local.total_bytes;
-            kparams_local.state_aligned   = (uint32_t) hmx_layout_local.state_f32_bytes;
-            kparams_local.vtcm_per_thread = (uint32_t) (hmx_layout_local.total_bytes / (n_threads > 0 ? n_threads : 1));
-        } else {
-            htp_gdn_vtcm_layout_build(&hvx_layout_local, S_v, n_threads);
-            kparams_local.kernel_type     = HTP_GDN_KERNEL_HVX_RECURRENT;
-            kparams_local.pipeline        = 0;
-            kparams_local.n_heads_batch   = 1;
-            kparams_local.state_aligned   = (uint32_t) hvx_layout_local.state_aligned;
-            kparams_local.vtcm_per_thread = (uint32_t) hvx_layout_local.bytes_per_thread;
-            kparams_local.vtcm_size       = (uint32_t) hvx_layout_local.total_bytes;
-        }
-        kparams_local.kda                 = (g->ne[0] == S_v) ? 1 : 0;
-        kparams_local.scale               = 1.0f / sqrtf((float) S_v);
-        kparams_local.state_seq_stride    = (uint32_t) (state->nb[3] / sizeof(float));
-        kparams_local.state_size_per_snap = S_v * S_v * H * n_seqs;
-
-        kparams_local.div_H         = init_fastdiv_values(H);
-        kparams_local.div_q1        = init_fastdiv_values(q->ne[1]);
-        kparams_local.div_k1        = init_fastdiv_values(k->ne[1]);
-        kparams_local.div_rq3       = init_fastdiv_values(rq3);
-        kparams_local.div_rk3       = init_fastdiv_values(rk3);
-        kparams_local.div_n_threads = init_fastdiv_values(n_threads);
-
-        kparams = &kparams_local;
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
     }
 
+    const uint32_t S_v        = kparams->S_v;
+    const uint32_t n_tokens   = kparams->n_tokens;
     const uint32_t total_rows = kparams->total_rows;
     uint32_t row_start = 0;
     uint32_t nrows     = total_rows;
@@ -2416,16 +2323,16 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
         return gated_delta_net_f32_hmx_chunked(octx, kparams, row_start, nrows);
     }
 
-    const uint32_t n_threads = (nrows < kparams->n_threads) ? nrows : kparams->n_threads;
+    const uint32_t n_threads = (nrows < octx->n_threads) ? nrows : octx->n_threads;
 
     struct htp_gdn_context gctx;
     gctx.octx      = octx;
     gctx.kparams   = kparams;
     gctx.row_start = row_start;
     gctx.nrows     = nrows;
-    gctx.vtcm_base = octx->ctx->vtcm_base;
 
     htp_gdn_vtcm_layout_build(&gctx.layout, S_v, n_threads);
+    gctx.vtcm_base = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, gctx.layout.off_state);
 
     if (gctx.layout.total_bytes > octx->ctx->vtcm_size) {
         return HTP_STATUS_VTCM_TOO_SMALL;
@@ -2433,10 +2340,10 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
 
     FARF(HIGH, "gated-delta-net-f32: q(%ux%ux%ux%u) k(%ux%ux%ux%u) v(%ux%ux%ux%u) state(%ux%ux%ux%u) -> (%ux%ux%ux%u) : "
          "vtcm-size %zu n_threads %u\n",
-         q->ne[0], q->ne[1], q->ne[2], q->ne[3],
-         k->ne[0], k->ne[1], k->ne[2], k->ne[3],
-         v->ne[0], v->ne[1], v->ne[2], v->ne[3],
-         state->ne[0], state->ne[1], state->ne[2], state->ne[3],
+         octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3],
+         octx->src[1]->ne[0], octx->src[1]->ne[1], octx->src[1]->ne[2], octx->src[1]->ne[3],
+         octx->src[2]->ne[0], octx->src[2]->ne[1], octx->src[2]->ne[2], octx->src[2]->ne[3],
+         octx->src[5]->ne[0], octx->src[5]->ne[1], octx->src[5]->ne[2], octx->src[5]->ne[3],
          dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
          gctx.layout.total_bytes, n_threads);
 

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

@@ -37,13 +37,7 @@ struct htp_gdn_kernel_params {
     uint32_t K;
 
     uint32_t total_rows;
-    uint32_t row_start;
-    uint32_t nrows;
-    uint32_t rows_per_thread;
-
     uint32_t kda;
-    uint32_t state_aligned;
-    uint32_t vtcm_per_thread;
     uint32_t vtcm_size;
     uint32_t state_seq_stride;
     uint32_t state_size_per_snap;
@@ -55,7 +49,6 @@ struct htp_gdn_kernel_params {
     struct fastdiv_values div_k1;
     struct fastdiv_values div_rq3;
     struct fastdiv_values div_rk3;
-    struct fastdiv_values div_n_threads;
 };
 
 #if defined(__cplusplus)
@@ -67,18 +60,21 @@ _Static_assert(sizeof(struct htp_gdn_kernel_params) <= 128, "htp_gdn_kernel_para
 struct htp_gdn_vtcm_layout {
     size_t state_aligned;
     size_t bytes_per_thread;
+    size_t off_state;
     size_t total_bytes;
 };
 
 static inline void htp_gdn_vtcm_layout_build(
-    struct htp_gdn_vtcm_layout * layout,
+    struct htp_gdn_vtcm_layout * L,
     uint32_t S_v,
     uint32_t n_threads
 ) {
     size_t state_bytes = (size_t) S_v * S_v * sizeof(float);
-    layout->state_aligned = hex_round_up(state_bytes, 128);
-    layout->bytes_per_thread = 2 * layout->state_aligned;
-    layout->total_bytes = layout->bytes_per_thread * n_threads;
+    L->state_aligned = hex_round_up(state_bytes, 128);
+    L->bytes_per_thread = 2 * L->state_aligned;
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_state, L->bytes_per_thread * n_threads);
+    L->total_bytes = off;
 }
 
 struct htp_gdn_hmx_vtcm_layout {

ggml/src/ggml-hexagon/htp/gather-rows.c

@@ -206,14 +206,22 @@ static void gather_rows_thread_f32(unsigned int nth, unsigned int ith, void * da
 
     struct htp_thread_trace * tr = &task->octx->ctx->trace[ith];
 
-    uint8_t * base   = task->octx->ctx->vtcm_base + ith * kparams->slice;
-    HVX_Vector * tab = (HVX_Vector *) base;
-    HVX_Vector * inc = tab + kparams->vg;
-    HVX_Vector * cur = inc + kparams->vg;
-    uint8_t * region_base = (uint8_t *) (cur + kparams->vg);
+    struct htp_gather_rows_vtcm_layout layout;
+    htp_gather_rows_vtcm_layout_build(&layout, kparams->vg, kparams->region_size, kparams->nr_max, kparams->ne, 4);
 
-    uint8_t * region[2] = { region_base, region_base + kparams->buf_stride };
-    uint8_t * out[2]    = { region[0] + kparams->out_off, region[1] + kparams->out_off };
+    uint8_t * base   = VTCM_LAYOUT_PTR(uint8_t, task->octx->ctx->vtcm_base, (size_t) ith * kparams->slice);
+    HVX_Vector * tab = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_tab);
+    HVX_Vector * inc = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_inc);
+    HVX_Vector * cur = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_cur);
+
+    uint8_t * region[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region1),
+    };
+    uint8_t * out[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out1),
+    };
 
     htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, (uint16_t) r_beg);
     gather_rows_init_tables_f32(tab, inc, kparams, a, b);
@@ -279,14 +287,22 @@ static void gather_rows_thread_f16(unsigned int nth, unsigned int ith, void * da
 
     struct htp_thread_trace * tr = &task->octx->ctx->trace[ith];
 
-    uint8_t * base   = task->octx->ctx->vtcm_base + ith * kparams->slice;
-    HVX_Vector * tab = (HVX_Vector *) base;
-    HVX_Vector * inc = tab + kparams->vg;
-    HVX_Vector * cur = inc + kparams->vg;
-    uint8_t * region_base = (uint8_t *) (cur + kparams->vg);
+    struct htp_gather_rows_vtcm_layout layout;
+    htp_gather_rows_vtcm_layout_build(&layout, kparams->vg, kparams->region_size, kparams->nr_max, kparams->ne, 2);
 
-    uint8_t * region[2] = { region_base, region_base + kparams->buf_stride };
-    uint8_t * out[2]    = { region[0] + kparams->out_off, region[1] + kparams->out_off };
+    uint8_t * base   = VTCM_LAYOUT_PTR(uint8_t, task->octx->ctx->vtcm_base, (size_t) ith * kparams->slice);
+    HVX_Vector * tab = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_tab);
+    HVX_Vector * inc = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_inc);
+    HVX_Vector * cur = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_cur);
+
+    uint8_t * region[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region1),
+    };
+    uint8_t * out[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out1),
+    };
 
     htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, (uint16_t) r_beg);
     gather_rows_init_tables_f16(tab, inc, kparams, a, b);

ggml/src/ggml-hexagon/htp/gather-rows.h

@@ -2,6 +2,7 @@
 #define HTP_GATHER_ROWS_H
 
 #include "hex-common.h"
+#include "htp-vtcm.h"
 #include <stdint.h>
 #include <stdbool.h>
 
@@ -19,8 +20,6 @@ struct htp_gather_rows_params {
     uint32_t v1;
     uint32_t span1;
     uint32_t region_size;
-    uint32_t out_off;
-    uint32_t buf_stride;
     uint32_t slice;
     uint32_t total_rows;
     uint32_t rows_per_thread;
@@ -28,6 +27,40 @@ struct htp_gather_rows_params {
     uint8_t  pad[3];
 };
 
+struct htp_gather_rows_vtcm_layout {
+    size_t off_tab;
+    size_t off_inc;
+    size_t off_cur;
+    size_t off_region0;
+    size_t off_out0;
+    size_t off_region1;
+    size_t off_out1;
+    size_t total_per_thread;
+};
+
+static inline void htp_gather_rows_vtcm_layout_build(
+    struct htp_gather_rows_vtcm_layout * L,
+    uint32_t vg,
+    uint32_t region_size,
+    uint32_t nr_max,
+    uint32_t ne,
+    uint32_t elem_size) {
+
+    const size_t vec_bytes      = (size_t) vg * 128;
+    const size_t region_aligned = hex_round_up(region_size, 128);
+    const size_t out_aligned    = hex_round_up(region_aligned + (size_t) nr_max * ne * elem_size, 128) - region_aligned;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_tab,     vec_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_inc,     vec_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_cur,     vec_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_region0, region_aligned);
+    VTCM_LAYOUT_ALLOC(off, off_out0,    out_aligned);
+    VTCM_LAYOUT_ALLOC(off, off_region1, region_aligned);
+    VTCM_LAYOUT_ALLOC(off, off_out1,    out_aligned);
+    L->total_per_thread = off;
+}
+
 static inline bool htp_gather_rows_solve_layout(
     struct htp_gather_rows_params * kparams,
     uint32_t na, uint32_t a_nb1,
@@ -80,8 +113,6 @@ static inline bool htp_gather_rows_solve_layout(
     uint32_t v1          = 0;
     uint32_t span1       = 0;
     uint32_t region_size = 0;
-    uint32_t out_off     = 0;
-    uint32_t buf_stride  = 0;
 
     while (nr_max > 0) {
         v1 = hex_round_up((nr_max - 1) * a_nb1 + na * elem_size, 128);
@@ -98,10 +129,10 @@ static inline bool htp_gather_rows_solve_layout(
             region_size = v1;
         }
 
-        out_off    = hex_round_up(region_size, 128);
-        buf_stride = hex_round_up(out_off + nr_max * ne * elem_size, 128);
+        struct htp_gather_rows_vtcm_layout layout;
+        htp_gather_rows_vtcm_layout_build(&layout, vg, region_size, nr_max, ne, elem_size);
 
-        if ((elem_size != 2 || region_size <= 32768) && fixed + 2 * buf_stride <= slice) {
+        if ((elem_size != 2 || region_size <= 32768) && layout.total_per_thread <= slice) {
             break;
         }
 
@@ -121,8 +152,6 @@ static inline bool htp_gather_rows_solve_layout(
     kparams->v1              = v1;
     kparams->span1           = span1;
     kparams->region_size     = region_size;
-    kparams->out_off         = out_off;
-    kparams->buf_stride      = buf_stride;
     kparams->slice           = slice;
     kparams->total_rows      = total_rows;
     kparams->rows_per_thread = rows_per_thread;

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

@@ -404,7 +404,7 @@ int op_get_rows(struct htp_ops_context * octx) {
 
     if (octx->ctx->mdev.count > 1) {
         uint32_t tasks_per_chunk = 1;
-        htp_tensor_mdev_rows_per_chunk(dst, dst_row_size / dst->ne[0], dst_row_size, &tasks_per_chunk);
+        htp_tensor_mdev_rows_per_chunk(dst, (uint32_t) dst->nb[0], dst_row_size, &tasks_per_chunk);
         const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_tasks, tasks_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         task_start = range.start;
         tasks      = range.count;
@@ -414,7 +414,7 @@ int op_get_rows(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    if (!htp_ops_context_set_n_threads(octx, (uint32_t) kparams->n_threads)) {
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
@@ -428,7 +428,7 @@ int op_get_rows(struct htp_ops_context * octx) {
     grctx.compute_tiled = NULL;
     grctx.task_start = task_start;
     grctx.tasks = tasks;
-    grctx.tasks_per_thread = octx->ctx->mdev.count == 1 ? kparams->tasks_per_thread : fastdiv(tasks + n_threads - 1, &octx->n_threads_div);
+    grctx.tasks_per_thread = fastdiv(tasks + n_threads - 1, &octx->n_threads_div);
     grctx.tile_size   = 0;
     grctx.tile_stride = 0;
     grctx.src0_row_size = src0_row_size;

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

@@ -15,11 +15,10 @@ enum htp_get_rows_kernel_type {
 };
 
 struct htp_get_rows_kernel_params {
-    int32_t  n_threads;
-    int32_t  kernel_type;
-    int32_t  total_tasks;
-    int32_t  tasks_per_thread;
-    int32_t  vtcm_size;
+    uint8_t  kernel_type;
+    uint32_t n_threads;
+    uint32_t total_tasks;
+    uint32_t vtcm_size;
 
     // Fastdiv helpers
     struct fastdiv_values div_ne10;

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

@@ -17,11 +17,12 @@
 #include "hex-common.h"
 #include "htp-tensor.h"
 #include "htp-vtcm.h"
+#include "glu-ops.h"
 
-#define htp_act_preamble                                 \
-    const struct htp_tensor * src0 = actx->octx->src[0]; \
-    const struct htp_tensor * src1 = actx->octx->src[1]; \
-    const struct htp_tensor * dst  = actx->octx->dst;    \
+#define htp_glu_preamble                                 \
+    const struct htp_tensor * src0 = gctx->octx->src[0]; \
+    const struct htp_tensor * src1 = gctx->octx->src[1]; \
+    const struct htp_tensor * dst  = gctx->octx->dst;    \
                                                          \
     const uint32_t ne00 = src0->ne[0];                   \
     const uint32_t ne01 = src0->ne[1];                   \
@@ -53,21 +54,21 @@
     const uint32_t nb2 = dst->nb[2];                     \
     const uint32_t nb3 = dst->nb[3];
 
-struct htp_act_context;
+struct htp_glu_context;
 
 typedef void (*glu_compute_f32_fn_t)(const float * restrict src0,
                                      const float * restrict src1,
                                      float * restrict dst,
                                      const uint32_t num_rows,
-                                     const struct htp_act_context * actx);
+                                     const struct htp_glu_context * gctx);
 
 typedef void (*glu_compute_f16_fn_t)(const _Float16 * restrict src0,
                                      const _Float16 * restrict src1,
                                      _Float16 * restrict dst,
                                      const uint32_t num_rows,
-                                     const struct htp_act_context * actx);
+                                     const struct htp_glu_context * gctx);
 
-struct htp_act_context {
+struct htp_glu_context {
     struct htp_ops_context * octx;
 
     void *                   compute;
@@ -108,53 +109,18 @@ struct htp_act_context {
     size_t                   vtcm_dst_size_per_thread;
 };
 
-struct htp_act_vtcm_layout {
-    size_t total_bytes;
-    size_t off_src0;
-    size_t off_src1;
-    size_t off_dst;
-
-    size_t src0_bytes_per_thread;
-    size_t src1_bytes_per_thread;
-    size_t dst_bytes_per_thread;
-
-    uint32_t vtcm_row_per_thread;
-};
-
-static inline void htp_act_vtcm_layout_build(struct htp_act_vtcm_layout * L,
-                                             size_t                       src0_row_size_aligned,
-                                             size_t                       src1_row_size_aligned,
-                                             size_t                       dst_row_size_aligned,
-                                             uint32_t                     n_threads,
-                                             size_t                       vtcm_size) {
-    const size_t   spad_size_per_row   = src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned;
-    const uint32_t vtcm_row_per_thread = (uint32_t) (vtcm_size / (n_threads * spad_size_per_row));
-
-    L->vtcm_row_per_thread = vtcm_row_per_thread;
-
-    L->src0_bytes_per_thread = src0_row_size_aligned * vtcm_row_per_thread;
-    L->src1_bytes_per_thread = src1_row_size_aligned * vtcm_row_per_thread;
-    L->dst_bytes_per_thread  = dst_row_size_aligned * vtcm_row_per_thread;
-
-    L->off_src0 = 0;
-    L->off_src1 = L->off_src0 + L->src0_bytes_per_thread * n_threads;
-    L->off_dst  = L->off_src1 + L->src1_bytes_per_thread * n_threads;
-
-    L->total_bytes = L->off_dst + L->dst_bytes_per_thread * n_threads;
-}
-
 #define htp_glu_op_preamble                                            \
-    const size_t src0_row_size_aligned = actx->src0_row_size_aligned;  \
-    const size_t src1_row_size_aligned = actx->src1_row_size_aligned;  \
-    const size_t dst_row_size_aligned  = actx->dst_row_size_aligned;   \
-    const int    nc                    = actx->nc;
+    const size_t src0_row_size_aligned = gctx->src0_row_size_aligned;  \
+    const size_t src1_row_size_aligned = gctx->src1_row_size_aligned;  \
+    const size_t dst_row_size_aligned  = gctx->dst_row_size_aligned;   \
+    const int    nc                    = gctx->nc;
 
 // swiglu(x) = x1 * sigmoid(x0)
 static void swiglu_f32(const float * restrict src0,
                                                  const float * restrict src1,
                                                  float * restrict dst,
                                                  const uint32_t num_rows,
-                                                 const struct htp_act_context * actx) {
+                                                 const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -172,10 +138,10 @@ static void swiglu_oai_f32(const float * restrict src0,
                                                      const float * restrict src1,
                                                      float * restrict dst,
                                                      const uint32_t num_rows,
-                                                     const struct htp_act_context * actx) {
+                                                     const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
-    const float alpha = ((const float *) (actx->octx->op_params))[2];
-    const float limit = ((const float *) (actx->octx->op_params))[3];
+    const float alpha = ((const float *) (gctx->octx->op_params))[2];
+    const float limit = ((const float *) (gctx->octx->op_params))[3];
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
         const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
@@ -201,9 +167,9 @@ static void swiglu_clamp_f32(const float * restrict src0,
                                                        const float * restrict src1,
                                                        float * restrict dst,
                                                        const uint32_t                 num_rows,
-                                                       const struct htp_act_context * actx) {
+                                                       const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
-    const float limit = ((const float *) (actx->octx->op_params))[3];
+    const float limit = ((const float *) (gctx->octx->op_params))[3];
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
         const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
@@ -455,7 +421,7 @@ static void geglu_f32(const float * restrict src0,
                                                 const float * restrict src1,
                                                 float * restrict dst,
                                                 const uint32_t num_rows,
-                                                const struct htp_act_context * actx) {
+                                                const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -472,7 +438,7 @@ static void geglu_quick_f32(const float * restrict src0,
                             const float * restrict src1,
                             float * restrict dst,
                             const uint32_t num_rows,
-                            const struct htp_act_context * actx) {
+                            const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -489,7 +455,7 @@ static void geglu_erf_f32(const float * restrict src0,
                           const float * restrict src1,
                           float * restrict dst,
                           const uint32_t num_rows,
-                          const struct htp_act_context * actx) {
+                          const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -501,12 +467,124 @@ static void geglu_erf_f32(const float * restrict src0,
     }
 }
 
+static inline void hvx_reglu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
+    assert((unsigned long) dst  % 128 == 0);
+    assert((unsigned long) src0 % 128 == 0);
+    assert((unsigned long) src1 % 128 == 0);
+
+    HVX_Vector * restrict vdst        = (HVX_Vector *) dst;
+    const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
+    const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
+
+    const uint32_t epv  = 128 / sizeof(float);
+    const uint32_t nvec = n / epv;
+    const uint32_t nloe = n % epv;
+
+    const HVX_Vector v_zero = Q6_V_vzero();
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        HVX_Vector x = Q6_Vsf_vmax_VsfVsf(vsrc0[i], v_zero);
+        vdst[i]      = hvx_vec_mul_f32_f32(x, vsrc1[i]);
+    }
+
+    if (nloe) {
+        HVX_Vector x   = Q6_Vsf_vmax_VsfVsf(vsrc0[i], v_zero);
+        HVX_Vector res = hvx_vec_mul_f32_f32(x, vsrc1[i]);
+        hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(float), res);
+    }
+}
+
+// reglu(x, g) = relu(x) * g
+static void reglu_f32(const float * restrict src0,
+                      const float * restrict src1,
+                      float * restrict dst,
+                      const uint32_t num_rows,
+                      const struct htp_glu_context * gctx) {
+    htp_glu_op_preamble;
+
+    for (uint32_t ib = 0; ib < num_rows; ib++) {
+        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
+        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
+        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
+
+        hvx_reglu_f32_aa(dst_ptr, src0_ptr, src1_ptr, nc);
+    }
+}
+
+static __attribute__((noinline)) void hvx_sigglu_f32_vec(HVX_Vector * restrict vdst,
+                                                         const HVX_Vector * restrict vsrc0,
+                                                         const HVX_Vector * restrict vsrc1,
+                                                         uint32_t nvec) {
+    const HVX_Vector v_one     = hvx_vec_splat_f32(1.0f);
+    const HVX_Vector v_max_exp = hvx_vec_splat_f32(87.0f);
+    const HVX_Vector v_min_exp = hvx_vec_splat_f32(-87.0f);
+
+    for (uint32_t i = 0; i < nvec; i++) {
+        HVX_Vector sig = hvx_vec_fast_sigmoid_f32_guard(vsrc0[i], v_one, v_max_exp, v_min_exp);
+        vdst[i]        = hvx_vec_mul_f32_f32(sig, vsrc1[i]);
+    }
+}
+
+static __attribute__((noinline)) void hvx_sigglu_f32_tail(HVX_Vector * restrict vdst,
+                                                          const HVX_Vector * restrict vsrc0,
+                                                          const HVX_Vector * restrict vsrc1,
+                                                          uint32_t nloe) {
+    const HVX_Vector v_one     = hvx_vec_splat_f32(1.0f);
+    const HVX_Vector v_max_exp = hvx_vec_splat_f32(87.0f);
+    const HVX_Vector v_min_exp = hvx_vec_splat_f32(-87.0f);
+
+    HVX_Vector sig = hvx_vec_fast_sigmoid_f32_guard(vsrc0[0], v_one, v_max_exp, v_min_exp);
+    HVX_Vector res = hvx_vec_mul_f32_f32(sig, vsrc1[0]);
+    hvx_vec_store_a((void *) vdst, nloe * sizeof(float), res);
+}
+
+static inline void hvx_sigglu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
+    assert((unsigned long) dst  % 128 == 0);
+    assert((unsigned long) src0 % 128 == 0);
+    assert((unsigned long) src1 % 128 == 0);
+
+    HVX_Vector * restrict vdst        = (HVX_Vector *) dst;
+    const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
+    const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
+
+    const uint32_t epv  = 128 / sizeof(float);
+    const uint32_t nvec = n / epv;
+    const uint32_t nloe = n % epv;
+
+    if (nvec) {
+        hvx_sigglu_f32_vec(vdst, vsrc0, vsrc1, nvec);
+    }
+    if (nloe) {
+        hvx_sigglu_f32_tail(vdst + nvec, vsrc0 + nvec, vsrc1 + nvec, nloe);
+    }
+}
+
+// sigglu(x, g) = sigmoid(x) * g
+static void sigglu_f32(const float * restrict src0,
+                       const float * restrict src1,
+                       float * restrict dst,
+                       const uint32_t num_rows,
+                       const struct htp_glu_context * gctx) {
+    htp_glu_op_preamble;
+
+    for (uint32_t ib = 0; ib < num_rows; ib++) {
+        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
+        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
+        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
+
+        hvx_sigglu_f32_aa(dst_ptr, src0_ptr, src1_ptr, nc);
+    }
+}
+
 // swiglu(x) = x1 * sigmoid(x0)
 static void swiglu_f16(const _Float16 * restrict src0,
                        const _Float16 * restrict src1,
                        _Float16 * restrict dst,
                        const uint32_t num_rows,
-                       const struct htp_act_context * actx) {
+                       const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -524,10 +602,10 @@ static void swiglu_oai_f16(const _Float16 * restrict src0,
                            const _Float16 * restrict src1,
                            _Float16 * restrict dst,
                            const uint32_t num_rows,
-                           const struct htp_act_context * actx) {
+                           const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
-    const _Float16 alpha = (_Float16) ((const float *) (actx->octx->op_params))[2];
-    const _Float16 limit = (_Float16) ((const float *) (actx->octx->op_params))[3];
+    const _Float16 alpha = (_Float16) ((const float *) (gctx->octx->op_params))[2];
+    const _Float16 limit = (_Float16) ((const float *) (gctx->octx->op_params))[3];
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
         const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
@@ -553,9 +631,9 @@ static void swiglu_clamp_f16(const _Float16 * restrict src0,
                              const _Float16 * restrict src1,
                              _Float16 * restrict dst,
                              const uint32_t                 num_rows,
-                             const struct htp_act_context * actx) {
+                             const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
-    const _Float16 limit = (_Float16) ((const float *) (actx->octx->op_params))[3];
+    const _Float16 limit = (_Float16) ((const float *) (gctx->octx->op_params))[3];
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
         const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
@@ -574,7 +652,7 @@ static void geglu_f16(const _Float16 * restrict src0,
                       const _Float16 * restrict src1,
                       _Float16 * restrict dst,
                       const uint32_t num_rows,
-                      const struct htp_act_context * actx) {
+                      const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -586,56 +664,155 @@ static void geglu_f16(const _Float16 * restrict src0,
     }
 }
 
+static inline void hvx_reglu_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
+    assert((unsigned long) dst  % 128 == 0);
+    assert((unsigned long) src0 % 128 == 0);
+    assert((unsigned long) src1 % 128 == 0);
+
+    HVX_Vector * restrict vdst        = (HVX_Vector *) dst;
+    const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
+    const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
+
+    const uint32_t epv  = 128 / sizeof(_Float16);
+    const uint32_t nvec = n / epv;
+    const uint32_t nloe = n % epv;
+
+    const HVX_Vector v_zero = Q6_V_vzero();
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        HVX_Vector x = Q6_Vhf_vmax_VhfVhf(vsrc0[i], v_zero);
+        vdst[i]      = hvx_vec_mul_f16_f16(x, vsrc1[i]);
+    }
+
+    if (nloe) {
+        HVX_Vector x   = Q6_Vhf_vmax_VhfVhf(vsrc0[i], v_zero);
+        HVX_Vector res = hvx_vec_mul_f16_f16(x, vsrc1[i]);
+        hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(_Float16), res);
+    }
+}
+
+// reglu(x, g) = relu(x) * g
+static void reglu_f16(const _Float16 * restrict src0,
+                      const _Float16 * restrict src1,
+                      _Float16 * restrict dst,
+                      const uint32_t num_rows,
+                      const struct htp_glu_context * gctx) {
+    htp_glu_op_preamble;
+
+    for (uint32_t ib = 0; ib < num_rows; ib++) {
+        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
+        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
+        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
+
+        hvx_reglu_f16_aa(dst_ptr, src0_ptr, src1_ptr, nc);
+    }
+}
+
+static __attribute__((noinline)) void hvx_sigglu_f16_vec(HVX_Vector * restrict vdst,
+                                                         const HVX_Vector * restrict vsrc0,
+                                                         const HVX_Vector * restrict vsrc1,
+                                                         uint32_t nvec) {
+    for (uint32_t i = 0; i < nvec; i++) {
+        HVX_Vector sig = hvx_vec_fast_sigmoid_f16(vsrc0[i]);
+        vdst[i]        = hvx_vec_mul_f16_f16(sig, vsrc1[i]);
+    }
+}
+
+static __attribute__((noinline)) void hvx_sigglu_f16_tail(HVX_Vector * restrict vdst,
+                                                          const HVX_Vector * restrict vsrc0,
+                                                          const HVX_Vector * restrict vsrc1,
+                                                          uint32_t nloe) {
+    HVX_Vector sig = hvx_vec_fast_sigmoid_f16(vsrc0[0]);
+    HVX_Vector res = hvx_vec_mul_f16_f16(sig, vsrc1[0]);
+    hvx_vec_store_a((void *) vdst, nloe * sizeof(_Float16), res);
+}
+
+static inline void hvx_sigglu_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
+    assert((unsigned long) dst  % 128 == 0);
+    assert((unsigned long) src0 % 128 == 0);
+    assert((unsigned long) src1 % 128 == 0);
+
+    HVX_Vector * restrict vdst        = (HVX_Vector *) dst;
+    const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
+    const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
+
+    const uint32_t epv  = 128 / sizeof(_Float16);
+    const uint32_t nvec = n / epv;
+    const uint32_t nloe = n % epv;
+
+    if (nvec) {
+        hvx_sigglu_f16_vec(vdst, vsrc0, vsrc1, nvec);
+    }
+    if (nloe) {
+        hvx_sigglu_f16_tail(vdst + nvec, vsrc0 + nvec, vsrc1 + nvec, nloe);
+    }
+}
+
+// sigglu(x, g) = sigmoid(x) * g
+static void sigglu_f16(const _Float16 * restrict src0,
+                       const _Float16 * restrict src1,
+                       _Float16 * restrict dst,
+                       const uint32_t num_rows,
+                       const struct htp_glu_context * gctx) {
+    htp_glu_op_preamble;
+
+    for (uint32_t ib = 0; ib < num_rows; ib++) {
+        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
+        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
+        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
+
+        hvx_sigglu_f16_aa(dst_ptr, src0_ptr, src1_ptr, nc);
+    }
+}
+
 #define DEFINE_GLU_PER_THREAD(TYPE, SUFFIX)                                                                            \
     static void glu_##SUFFIX##_per_thread(unsigned int nth, unsigned int ith, void * data) {                           \
-        struct htp_act_context * actx = (struct htp_act_context *) data;                                               \
-        htp_act_preamble;                                                                                              \
+        struct htp_glu_context * gctx = (struct htp_glu_context *) data;                                               \
+        htp_glu_preamble;                                                                                              \
                                                                                                                        \
-        struct htp_thread_trace * tr = actx->octx->ctx ? &actx->octx->ctx->trace[ith] : NULL;                          \
+        struct htp_thread_trace * tr = gctx->octx->ctx ? &gctx->octx->ctx->trace[ith] : NULL;                          \
                                                                                                                        \
-        size_t src0_row_size = actx->src0_row_size;                                                                    \
-        size_t src1_row_size = actx->src1_row_size;                                                                    \
-        size_t dst_row_size  = actx->dst_row_size;                                                                     \
+        size_t src0_row_size = gctx->src0_row_size;                                                                    \
+        size_t src1_row_size = gctx->src1_row_size;                                                                    \
+        size_t dst_row_size  = gctx->dst_row_size;                                                                     \
                                                                                                                        \
-        size_t src0_row_stride = actx->src0_row_stride;                                                                \
-        size_t src1_row_stride = actx->src1_row_stride;                                                                \
+        size_t src0_row_stride = gctx->src0_row_stride;                                                                \
+        size_t src1_row_stride = gctx->src1_row_stride;                                                                \
                                                                                                                        \
-        const uint32_t src0_nrows            = actx->src0_nrows;                                                       \
-        const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread;                                            \
+        const uint32_t src0_nrows            = gctx->src0_nrows;                                                       \
+        const uint32_t src0_nrows_per_thread = gctx->src0_nrows_per_thread;                                            \
                                                                                                                        \
-        const uint32_t src0_start_row = actx->row_start + src0_nrows_per_thread * ith;                                 \
-        const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, actx->row_start + src0_nrows);     \
+        const uint32_t src0_start_row = gctx->row_start + src0_nrows_per_thread * ith;                                 \
+        const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, gctx->row_start + src0_nrows);     \
                                                                                                                        \
         /* no work for this thread */                                                                                  \
         if (src0_start_row >= src0_end_row) {                                                                          \
             return;                                                                                                    \
         }                                                                                                              \
                                                                                                                        \
-        const dma_addr_t data_src0 = actx->data_src0;                                                                  \
-        const dma_addr_t data_src1 = actx->data_src1;                                                                  \
-        const dma_addr_t data_dst  = actx->data_dst;                                                                   \
+        const dma_addr_t data_src0 = gctx->data_src0;                                                                  \
+        const dma_addr_t data_src1 = gctx->data_src1;                                                                  \
+        const dma_addr_t data_dst  = gctx->data_dst;                                                                   \
                                                                                                                        \
-        const size_t src0_row_size_aligned = actx->src0_row_size_aligned;                                              \
-        const size_t src1_row_size_aligned = actx->src1_row_size_aligned;                                              \
-        const size_t dst_row_size_aligned  = actx->dst_row_size_aligned;                                               \
+        const size_t src0_row_size_aligned = gctx->src0_row_size_aligned;                                              \
+        const size_t src1_row_size_aligned = gctx->src1_row_size_aligned;                                              \
+        const size_t dst_row_size_aligned  = gctx->dst_row_size_aligned;                                               \
                                                                                                                        \
-        uint8_t * restrict src0_spad_data = actx->vtcm_src0 + (ith * actx->vtcm_src0_size_per_thread);                 \
-        uint8_t * restrict src1_spad_data = actx->vtcm_src1 + (ith * actx->vtcm_src1_size_per_thread);                 \
-        uint8_t * restrict dst_spad_data  = actx->vtcm_dst  + (ith * actx->vtcm_dst_size_per_thread);                  \
+        uint8_t * restrict src0_spad_data = gctx->vtcm_src0 + (ith * gctx->vtcm_src0_size_per_thread);                 \
+        uint8_t * restrict src1_spad_data = gctx->vtcm_src1 + (ith * gctx->vtcm_src1_size_per_thread);                 \
+        uint8_t * restrict dst_spad_data  = gctx->vtcm_dst  + (ith * gctx->vtcm_dst_size_per_thread);                  \
                                                                                                                        \
-        size_t src0_spad_half_size = actx->src0_spad_half_size;                                                        \
-        size_t src1_spad_half_size = actx->src1_spad_half_size;                                                        \
-        size_t dst_spad_half_size  = actx->dst_spad_half_size;                                                         \
+        size_t src0_spad_half_size = gctx->src0_spad_half_size;                                                        \
+        size_t src1_spad_half_size = gctx->src1_spad_half_size;                                                        \
+        size_t dst_spad_half_size  = gctx->dst_spad_half_size;                                                         \
                                                                                                                        \
-        const int BLOCK = actx->block;                                                                                 \
-        if (BLOCK == 0) {                                                                                              \
-            FARF(ERROR, "%s : VTCM reservation %zu is too small, needed %zu\n",                                        \
-                 actx->op_str, actx->vtcm_src0_size_per_thread, src0_row_size_aligned);                                \
-            return;                                                                                                    \
-        }                                                                                                              \
+        const int BLOCK = gctx->block;                                                                                 \
                                                                                                                        \
-        dma_queue * dma_q = actx->octx->ctx->dma[ith];                                                                 \
-        glu_compute_##SUFFIX##_fn_t compute = (glu_compute_##SUFFIX##_fn_t) actx->compute;                             \
+        dma_queue * dma_q = gctx->octx->ctx->dma[ith];                                                                 \
+        glu_compute_##SUFFIX##_fn_t compute = (glu_compute_##SUFFIX##_fn_t) gctx->compute;                             \
                                                                                                                        \
         for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) { \
             const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);                                                 \
@@ -662,7 +839,7 @@ static void geglu_f16(const _Float16 * restrict src0,
             TYPE * src1_spad = (TYPE *) dma_queue_pop(dma_q).dst;                                                      \
                                                                                                                        \
             htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir);                                                     \
-            compute(src0_spad, src1_spad, dst_spad, block_size, actx);                                                 \
+            compute(src0_spad, src1_spad, dst_spad, block_size, gctx);                                                 \
             htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir);                                                      \
                                                                                                                        \
             dma_queue_push(dma_q, dma_make_data(data_dst + (ir * dst_row_size), dst_spad),                             \
@@ -685,16 +862,11 @@ static void geglu_f16(const _Float16 * restrict src0,
 DEFINE_GLU_PER_THREAD(float, f32)
 DEFINE_GLU_PER_THREAD(_Float16, f16)
 
-static int execute_op_activations_f32(struct htp_ops_context * octx) {
+static int execute_op_glu_f32(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * src1 = octx->src[1];
     const struct htp_tensor * dst  = octx->dst;
 
-    if ((dst->ne[0] * SIZEOF_FP32) != dst->nb[1]) {
-        FARF(ERROR, "Non-contiguous dst is not supported at this time \n");
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
     glu_compute_f32_fn_t compute_fn = NULL;
     const char *         op_type    = NULL;
 
@@ -729,11 +901,27 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
             op_type    = "geglu-erf-f32";
             break;
 
+        case HTP_OP_GLU_REGLU:
+            compute_fn = reglu_f32;
+            op_type    = "reglu-f32";
+            break;
+
+        case HTP_OP_GLU_SIGMOID:
+            compute_fn = sigglu_f32;
+            op_type    = "sigglu-f32";
+            break;
+
         default:
-            FARF(ERROR, "Unsupported activations Op %u\n", octx->op);
+            FARF(ERROR, "Unsupported GLU Op %u\n", octx->op);
             return HTP_STATUS_NO_SUPPORT;
     }
 
+    const struct htp_glu_kernel_params * kparams = (const struct htp_glu_kernel_params *) octx->kernel_params;
+
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
     const size_t dst_row_size = dst->ne[0] * SIZEOF_FP32;
 
@@ -762,17 +950,16 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
     const size_t src0_row_stride = src0->nb[1];
     const size_t src1_row_stride = src1 ? src1->nb[1] : src0->nb[1];
 
-    const size_t src0_row_size_aligned = hex_round_up(src0_row_size, VLEN);
-    const size_t src1_row_size_aligned = hex_round_up(src1_row_size, VLEN);
-    const size_t dst_row_size_aligned  = hex_round_up(dst_row_size, VLEN);
+    struct htp_glu_vtcm_layout L;
+    htp_glu_vtcm_layout_build(&L, dst_row_size, n_threads, kparams->vtcm_row_per_thread);
 
-    struct htp_act_vtcm_layout L;
-    htp_act_vtcm_layout_build(&L, src0_row_size_aligned, src1_row_size_aligned, dst_row_size_aligned, n_threads,
-                              octx->ctx->vtcm_size);
+    const size_t src0_row_size_aligned = L.row_size_aligned;
+    const size_t src1_row_size_aligned = L.row_size_aligned;
+    const size_t dst_row_size_aligned  = L.row_size_aligned;
 
     // Make sure the reserved vtcm size is sufficient
-    if (L.vtcm_row_per_thread == 0) {
-        FARF(ERROR, "act-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
+    if (kparams->vtcm_row_per_thread == 0 || octx->ctx->vtcm_size < kparams->vtcm_size) {
+        FARF(ERROR, "glu-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
              (src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned) * n_threads);
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
@@ -789,58 +976,53 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
     }
 
     // Prepare context
-    struct htp_act_context actx;
-    actx.octx    = octx;
-    actx.compute = (void *) compute_fn;
-    actx.op_str  = op_type;
+    struct htp_glu_context gctx;
+    gctx.octx    = octx;
+    gctx.compute = (void *) compute_fn;
+    gctx.op_str  = op_type;
 
-    actx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    gctx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
 
-    actx.src0_row_size = src0_row_size;
-    actx.src1_row_size = src1_row_size;
-    actx.dst_row_size  = dst_row_size;
+    gctx.src0_row_size = src0_row_size;
+    gctx.src1_row_size = src1_row_size;
+    gctx.dst_row_size  = dst_row_size;
 
-    actx.src0_row_size_aligned = src0_row_size_aligned;
-    actx.src1_row_size_aligned = src1_row_size_aligned;
-    actx.dst_row_size_aligned  = dst_row_size_aligned;
+    gctx.src0_row_size_aligned = src0_row_size_aligned;
+    gctx.src1_row_size_aligned = src1_row_size_aligned;
+    gctx.dst_row_size_aligned  = dst_row_size_aligned;
 
-    actx.src0_row_stride = src0_row_stride;
-    actx.src1_row_stride = src1_row_stride;
+    gctx.src0_row_stride = src0_row_stride;
+    gctx.src1_row_stride = src1_row_stride;
 
     uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
-    actx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
-    actx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
-    actx.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
+    gctx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+    gctx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+    gctx.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
 
-    actx.vtcm_src0_size_per_thread = L.src0_bytes_per_thread;
-    actx.vtcm_src1_size_per_thread = L.src1_bytes_per_thread;
-    actx.vtcm_dst_size_per_thread  = L.dst_bytes_per_thread;
+    gctx.vtcm_src0_size_per_thread = (uint32_t) L.src0_bytes_per_thread;
+    gctx.vtcm_src1_size_per_thread = (uint32_t) L.src1_bytes_per_thread;
+    gctx.vtcm_dst_size_per_thread  = (uint32_t) L.dst_bytes_per_thread;
 
-    actx.src0_spad_half_size = L.src0_bytes_per_thread / 2;
-    actx.src1_spad_half_size = L.src1_bytes_per_thread / 2;
-    actx.dst_spad_half_size  = L.dst_bytes_per_thread / 2;
+    gctx.src0_spad_half_size = gctx.vtcm_src0_size_per_thread >> 1;
+    gctx.src1_spad_half_size = gctx.vtcm_src1_size_per_thread >> 1;
+    gctx.dst_spad_half_size  = gctx.vtcm_dst_size_per_thread  >> 1;
 
-    actx.block = actx.src0_spad_half_size / actx.src0_row_size_aligned;
-    actx.src0_nrows = nrows;
-    actx.row_start  = row_start;
+    gctx.block      = L.block;
+    gctx.src0_nrows = nrows;
+    gctx.row_start  = row_start;
 
-    actx.nc = dst->ne[0];
+    gctx.nc = dst->ne[0];
 
     // Addresses and GLU logic
     dma_addr_t data_src0 = src0->data;
     dma_addr_t data_src1 = src1 ? src1->data : 0;
 
-    if (!src1 && (octx->op == HTP_OP_GLU_SWIGLU ||
-                  octx->op == HTP_OP_GLU_SWIGLU_OAI ||
-                  octx->op == HTP_OP_GLU_SWIGLU_CLAMP ||
-                  octx->op == HTP_OP_GLU_GEGLU ||
-                  octx->op == HTP_OP_GLU_GEGLU_QUICK ||
-                  octx->op == HTP_OP_GLU_GEGLU_ERF)) {
+    if (!src1) {
          const int32_t swapped = octx->op_params[1];
          data_src1 = data_src0;
-         actx.src1_row_size = actx.src0_row_size;
+         gctx.src1_row_size = gctx.src0_row_size;
 
-         size_t nc_in_bytes = actx.nc * SIZEOF_FP32;
+         size_t nc_in_bytes = gctx.nc * SIZEOF_FP32;
          if (swapped) {
              data_src0 += nc_in_bytes;
          } else {
@@ -848,24 +1030,19 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
          }
     }
 
-    actx.data_src0 = data_src0;
-    actx.data_src1 = data_src1;
-    actx.data_dst  = dst->data;
+    gctx.data_src0 = data_src0;
+    gctx.data_src1 = data_src1;
+    gctx.data_dst  = dst->data;
 
-    work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f32_per_thread, &actx, n_threads);
+    work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f32_per_thread, &gctx, n_threads);
     return HTP_STATUS_OK;
 }
 
-static int execute_op_activations_f16(struct htp_ops_context * octx) {
+static int execute_op_glu_f16(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * src1 = octx->src[1];
     const struct htp_tensor * dst  = octx->dst;
 
-    if ((dst->ne[0] * SIZEOF_FP16) != dst->nb[1]) {
-        FARF(ERROR, "Non-contiguous dst is not supported at this time \n");
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
     glu_compute_f16_fn_t compute_fn = NULL;
     const char *         op_type    = NULL;
 
@@ -890,11 +1067,27 @@ static int execute_op_activations_f16(struct htp_ops_context * octx) {
             op_type    = "geglu-f16";
             break;
 
+        case HTP_OP_GLU_REGLU:
+            compute_fn = reglu_f16;
+            op_type    = "reglu-f16";
+            break;
+
+        case HTP_OP_GLU_SIGMOID:
+            compute_fn = sigglu_f16;
+            op_type    = "sigglu-f16";
+            break;
+
         default:
-            FARF(ERROR, "Unsupported activations Op %u\n", octx->op);
+            FARF(ERROR, "Unsupported GLU Op %u\n", octx->op);
             return HTP_STATUS_NO_SUPPORT;
     }
 
+    const struct htp_glu_kernel_params * kparams = (const struct htp_glu_kernel_params *) octx->kernel_params;
+
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
     const size_t dst_row_size = dst->ne[0] * SIZEOF_FP16;
 
@@ -923,17 +1116,16 @@ static int execute_op_activations_f16(struct htp_ops_context * octx) {
     const size_t src0_row_stride = src0->nb[1];
     const size_t src1_row_stride = src1 ? src1->nb[1] : src0->nb[1];
 
-    const size_t src0_row_size_aligned = hex_round_up(src0_row_size, VLEN);
-    const size_t src1_row_size_aligned = hex_round_up(src1_row_size, VLEN);
-    const size_t dst_row_size_aligned  = hex_round_up(dst_row_size, VLEN);
+    struct htp_glu_vtcm_layout L;
+    htp_glu_vtcm_layout_build(&L, dst_row_size, n_threads, kparams->vtcm_row_per_thread);
 
-    struct htp_act_vtcm_layout L;
-    htp_act_vtcm_layout_build(&L, src0_row_size_aligned, src1_row_size_aligned, dst_row_size_aligned, n_threads,
-                              octx->ctx->vtcm_size);
+    const size_t src0_row_size_aligned = L.row_size_aligned;
+    const size_t src1_row_size_aligned = L.row_size_aligned;
+    const size_t dst_row_size_aligned  = L.row_size_aligned;
 
     // Make sure the reserved vtcm size is sufficient
-    if (L.vtcm_row_per_thread == 0) {
-        FARF(ERROR, "act-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
+    if (kparams->vtcm_row_per_thread == 0 || octx->ctx->vtcm_size < kparams->vtcm_size) {
+        FARF(ERROR, "glu-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
              (src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned) * n_threads);
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
@@ -950,56 +1142,53 @@ static int execute_op_activations_f16(struct htp_ops_context * octx) {
     }
 
     // Prepare context
-    struct htp_act_context actx;
-    actx.octx    = octx;
-    actx.compute = (void *) compute_fn;
-    actx.op_str  = op_type;
+    struct htp_glu_context gctx;
+    gctx.octx    = octx;
+    gctx.compute = (void *) compute_fn;
+    gctx.op_str  = op_type;
 
-    actx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    gctx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
 
-    actx.src0_row_size = src0_row_size;
-    actx.src1_row_size = src1_row_size;
-    actx.dst_row_size  = dst_row_size;
+    gctx.src0_row_size = src0_row_size;
+    gctx.src1_row_size = src1_row_size;
+    gctx.dst_row_size  = dst_row_size;
 
-    actx.src0_row_size_aligned = src0_row_size_aligned;
-    actx.src1_row_size_aligned = src1_row_size_aligned;
-    actx.dst_row_size_aligned  = dst_row_size_aligned;
+    gctx.src0_row_size_aligned = src0_row_size_aligned;
+    gctx.src1_row_size_aligned = src1_row_size_aligned;
+    gctx.dst_row_size_aligned  = dst_row_size_aligned;
 
-    actx.src0_row_stride = src0_row_stride;
-    actx.src1_row_stride = src1_row_stride;
+    gctx.src0_row_stride = src0_row_stride;
+    gctx.src1_row_stride = src1_row_stride;
 
     uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
-    actx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
-    actx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
-    actx.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
+    gctx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+    gctx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+    gctx.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
 
-    actx.vtcm_src0_size_per_thread = L.src0_bytes_per_thread;
-    actx.vtcm_src1_size_per_thread = L.src1_bytes_per_thread;
-    actx.vtcm_dst_size_per_thread  = L.dst_bytes_per_thread;
+    gctx.vtcm_src0_size_per_thread = (uint32_t) L.src0_bytes_per_thread;
+    gctx.vtcm_src1_size_per_thread = (uint32_t) L.src1_bytes_per_thread;
+    gctx.vtcm_dst_size_per_thread  = (uint32_t) L.dst_bytes_per_thread;
 
-    actx.src0_spad_half_size = L.src0_bytes_per_thread / 2;
-    actx.src1_spad_half_size = L.src1_bytes_per_thread / 2;
-    actx.dst_spad_half_size  = L.dst_bytes_per_thread / 2;
+    gctx.src0_spad_half_size = gctx.vtcm_src0_size_per_thread >> 1;
+    gctx.src1_spad_half_size = gctx.vtcm_src1_size_per_thread >> 1;
+    gctx.dst_spad_half_size  = gctx.vtcm_dst_size_per_thread  >> 1;
 
-    actx.block = actx.src0_spad_half_size / actx.src0_row_size_aligned;
-    actx.src0_nrows = nrows;
-    actx.row_start  = row_start;
+    gctx.block      = L.block;
+    gctx.src0_nrows = nrows;
+    gctx.row_start  = row_start;
 
-    actx.nc = dst->ne[0];
+    gctx.nc = dst->ne[0];
 
     // Addresses and GLU logic
     dma_addr_t data_src0 = src0->data;
     dma_addr_t data_src1 = src1 ? src1->data : 0;
 
-    if (!src1 && (octx->op == HTP_OP_GLU_SWIGLU ||
-                  octx->op == HTP_OP_GLU_SWIGLU_OAI ||
-                  octx->op == HTP_OP_GLU_SWIGLU_CLAMP ||
-                  octx->op == HTP_OP_GLU_GEGLU)) {
+    if (!src1) {
          const int32_t swapped = octx->op_params[1];
          data_src1 = data_src0;
-         actx.src1_row_size = actx.src0_row_size;
+         gctx.src1_row_size = gctx.src0_row_size;
 
-         size_t nc_in_bytes = actx.nc * SIZEOF_FP16;
+         size_t nc_in_bytes = gctx.nc * SIZEOF_FP16;
          if (swapped) {
              data_src0 += nc_in_bytes;
          } else {
@@ -1007,22 +1196,22 @@ static int execute_op_activations_f16(struct htp_ops_context * octx) {
          }
     }
 
-    actx.data_src0 = data_src0;
-    actx.data_src1 = data_src1;
-    actx.data_dst  = dst->data;
+    gctx.data_src0 = data_src0;
+    gctx.data_src1 = data_src1;
+    gctx.data_dst  = dst->data;
 
-    work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f16_per_thread, &actx, n_threads);
+    work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f16_per_thread, &gctx, n_threads);
     return HTP_STATUS_OK;
 }
 
 
-int op_activations(struct htp_ops_context * octx) {
+int op_glu(struct htp_ops_context * octx) {
     switch (octx->src[0]->type) {
         case HTP_TYPE_F32:
-            return execute_op_activations_f32(octx);
+            return execute_op_glu_f32(octx);
 
         case HTP_TYPE_F16:
-            return execute_op_activations_f16(octx);
+            return execute_op_glu_f16(octx);
 
         default:
             return HTP_STATUS_NO_SUPPORT;

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

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

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

@@ -17,6 +17,10 @@
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
 #endif
 
+#ifndef HEX_L2_LINE_SIZE
+#define HEX_L2_LINE_SIZE 128
+#endif
+
 static inline uint32_t hex_ceil_pow2(uint32_t x) {
     if (x <= 1) { return 1; }
     int p = 2;

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

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

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

@@ -12,8 +12,10 @@
 #include <assert.h>
 #include <dspqueue.h>
 #include <stdatomic.h>
+#include <stddef.h>
 #include <stdint.h>
 #include <stdbool.h>
+#include <string.h>
 
 #ifndef HTP_MAX_NTHREADS
 #define HTP_MAX_NTHREADS 10
@@ -66,14 +68,18 @@ struct htp_ops_context {
     // TODO convert these to an array
     struct htp_spad src0_spad;
     struct htp_spad src1_spad;
-    struct htp_spad src2_spad;
-    struct htp_spad src3_spad;
     struct htp_spad dst_spad;
 
     uint32_t              flags;
     uint32_t              n_threads;
     struct fastdiv_values n_threads_div;
     int                   status;
+
+    // Must stay last: not cleared by htp_ops_context_reset(), ops init the caches they use
+    union {
+        dma_cache_dm dma_cache_dm[HTP_MAX_NTHREADS];
+        dma_cache_fa dma_cache_fa[HTP_MAX_NTHREADS];
+    } __attribute__((aligned(128)));
 };
 
 // Main context for htp DSP backend
@@ -118,6 +124,13 @@ struct htp_context {
     size_t                 footprint;
 };
 
+static inline void htp_ops_context_reset(struct htp_ops_context * octx, struct htp_context * ctx) {
+    memset(octx, 0, offsetof(struct htp_ops_context, dma_cache_dm));
+    octx->ctx           = ctx;
+    octx->n_threads     = ctx->n_threads;
+    octx->n_threads_div = ctx->n_threads_div;
+}
+
 static inline bool htp_ops_context_set_n_threads(struct htp_ops_context * octx, uint32_t n_threads) {
     if (n_threads == 0 || n_threads > octx->ctx->n_threads) {
         return false;
@@ -148,7 +161,7 @@ int op_unary(struct htp_ops_context * octx);
 int op_sum(struct htp_ops_context * octx);
 int op_sum_rows(struct htp_ops_context * octx);
 int op_argmax(struct htp_ops_context * octx);
-int op_activations(struct htp_ops_context * octx);
+int op_glu(struct htp_ops_context * octx);
 int op_softmax(struct htp_ops_context * octx);
 int op_add_id(struct htp_ops_context * octx);
 int op_rope(struct htp_ops_context * octx);

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

@@ -38,11 +38,6 @@ enum htp_data_type {
     HTP_TYPE_INVALID
 };
 
-// Constants for internal types
-#define QK_Q8_0_TILED  128  // 32x32 Q8_0 tiled layout
-
-
-
 // Do not reorder first 4 (used as an index)
 enum htp_op_code {
     HTP_OP_MUL = 0,
@@ -115,6 +110,8 @@ enum htp_op_code {
     HTP_OP_UNARY_GELU_QUICK,
     HTP_OP_IM2COL_3D,
     HTP_OP_HADAMARD,
+    HTP_OP_GLU_REGLU,
+    HTP_OP_GLU_SIGMOID,
 
     HTP_OP_INVALID
 };
@@ -125,8 +122,8 @@ enum htp_op_code {
 #define HTP_OP_MAX_PARAMS  16   // aka GGML_MAX_OP_PARAMS
 #define HTP_OP_MAX_KERN_PARAMS 32
 
-#define HTP_OP_MAX_BUFS    16
-#define HTP_OP_MAX_TENSORS 8192 // must stay under 64K (uint16)
+#define HTP_OP_MAX_BUFS    64
+#define HTP_OP_MAX_TENSORS 65534 // must stay under 64K (uint16)
 
 #define HTP_FENCE_TIMEOUT  (1000000000ULL)
 #define HTP_FENCE_SLOT_SIZE 128
@@ -136,9 +133,10 @@ enum htp_op_code {
 #define HTP_MMAP_MAX_VMEM  (2147483648u)
 
 enum htp_tensor_flags {
-    HTP_TENSOR_WEIGHT  = (1U << 0), // Tensor buffer model weight data (not compute)
-    HTP_TENSOR_TILED   = (1U << 1), // Tensor is in tiled format
-    HTP_TENSOR_FENCE   = (1U << 2)  // Tensor is synchronization fence (explicitly managed)
+    HTP_TENSOR_WEIGHT   = (1U << 0), // Tensor buffer model weight data (not compute)
+    HTP_TENSOR_TILED    = (1U << 1), // Tensor is in tiled format
+    HTP_TENSOR_FENCE    = (1U << 2), // Tensor is synchronization fence (explicitly managed)
+    HTP_TENSOR_FUSEABLE = (1U << 3), // Tensor can be fused (host-only)
 };
 
 enum htp_buf_flags {

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

@@ -233,7 +233,7 @@ static inline void quantize_block_f32_q8_0_tiled(float * restrict x, uint8_t * r
 
 static void quantize_row_f32_q8_0_tiled(float * restrict x, uint8_t * restrict y, uint32_t k) {
     assert(k % 32 == 0);
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (k + qk - 1) / qk;
 
     for (uint32_t i = 0; i < nb; i++) {
@@ -244,7 +244,7 @@ static void quantize_row_f32_q8_0_tiled(float * restrict x, uint8_t * restrict y
 
 static void quantize_row_f32_q8_1_tiled(float * restrict x, uint8_t * restrict y, uint32_t k) {
     assert(k % 32 == 0);
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (k + qk - 1) / qk;
 
     for (uint32_t i = 0; i < nb; i++) {
@@ -255,7 +255,7 @@ static void quantize_row_f32_q8_1_tiled(float * restrict x, uint8_t * restrict y
 
 static void quantize_row_f32_q8_1_s16_tiled(float * restrict x, uint8_t * restrict y, uint32_t k) {
     assert(k % 32 == 0);
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (k + qk - 1) / qk;
 
     for (uint32_t i = 0; i < nb; i++) {
@@ -310,7 +310,7 @@ static inline void quantize_block_f16_f16_tiled(const __fp16 * restrict x, uint8
 
 static void quantize_row_f32_f16_tiled(const float * restrict x, uint8_t * restrict y, uint32_t k) {
     assert(k % 32 == 0);
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (k + qk - 1) / qk;
 
     for (uint32_t i = 0; i < nb; i++) {
@@ -321,7 +321,7 @@ static void quantize_row_f32_f16_tiled(const float * restrict x, uint8_t * restr
 
 static void quantize_row_f16_f16_tiled(const __fp16 * restrict x, uint8_t * restrict y, uint32_t k) {
     assert(k % 32 == 0);
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (k + qk - 1) / qk;
 
     for (uint32_t i = 0; i < nb; i++) {
@@ -1963,7 +1963,7 @@ static inline void quantize_f32_q8_0_tiled_block_kernel(
     uint32_t c
 ) {
     (void) tmp_data;
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (ne0 + qk - 1) / qk;
 
     for (uint32_t ib = ib_first; ib < ib_last; ++ib) {
@@ -1993,7 +1993,7 @@ static inline void quantize_f32_q8_1_tiled_block_kernel(
     uint32_t c
 ) {
     (void) tmp_data;
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (ne0 + qk - 1) / qk;
 
     for (uint32_t ib = ib_first; ib < ib_last; ++ib) {
@@ -2023,7 +2023,7 @@ static inline void quantize_f32_q8_1_s16_tiled_block_kernel(
     uint32_t c
 ) {
     (void) tmp_data;
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (ne0 + qk - 1) / qk;
 
     for (uint32_t ib = ib_first; ib < ib_last; ++ib) {
@@ -2053,7 +2053,7 @@ static inline void quantize_f32_f16_tiled_block_kernel(
     uint32_t c
 ) {
     (void) tmp_data;
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (ne0 + qk - 1) / qk;
 
     for (uint32_t ib = ib_first; ib < ib_last; ++ib) {
@@ -2083,7 +2083,7 @@ static inline void quantize_f16_f16_tiled_block_kernel(
     uint32_t c
 ) {
     (void) tmp_data;
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (ne0 + qk - 1) / qk;
 
     for (uint32_t ib = ib_first; ib < ib_last; ++ib) {

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

@@ -11,7 +11,8 @@
 static inline void hvx_fast_rms_norm_f32(const uint8_t * restrict src,
                                          uint8_t * restrict dst,
                                          const int num_elems,
-                                         float     epsilon) {
+                                         float     epsilon,
+                                         float     scale) {
 
     const HVX_Vector * restrict v_src = (HVX_Vector *) src;
     HVX_Vector * restrict v_dst       = (HVX_Vector *) dst;
@@ -48,6 +49,9 @@ static inline void hvx_fast_rms_norm_f32(const uint8_t * restrict src,
 
     // Scale full vectors
     HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v));
+    if (scale != 0.0f) {
+        scale_v = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(scale_v, hvx_vec_splat_f32(scale)));
+    }
 
     #pragma unroll(4)
     for (int i = 0; i < nvec; i++) {
@@ -72,7 +76,8 @@ static inline void hvx_fast_rms_norm_mul_f32(const uint8_t * restrict src,
                                              const uint8_t * restrict weight,
                                              uint8_t * restrict dst,
                                              const int num_elems,
-                                             float     epsilon) {
+                                             float     epsilon,
+                                             float     scale) {
     const HVX_Vector * restrict v_src    = (const HVX_Vector *) src;
     const HVX_Vector * restrict v_weight = (const HVX_Vector *) weight;
     HVX_Vector * restrict v_dst          = (HVX_Vector *) dst;
@@ -109,6 +114,9 @@ static inline void hvx_fast_rms_norm_mul_f32(const uint8_t * restrict src,
 
     // Scale and multiply
     HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v));
+    if (scale != 0.0f) {
+        scale_v = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(scale_v, hvx_vec_splat_f32(scale)));
+    }
 
     #pragma unroll(4)
     for (int i = 0; i < nvec; i++) {

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

@@ -68,6 +68,14 @@ static inline HVX_Vector hvx_vec_tanh_f32(HVX_Vector x) {
     return Q6_Vsf_equals_Vqf32(res);
 }
 
+static inline HVX_Vector hvx_vec_silu_f32(HVX_Vector x) {
+    HVX_Vector one     = hvx_vec_splat_f32(1.0f);
+    HVX_Vector max_exp = hvx_vec_splat_f32(87.f);
+    HVX_Vector min_exp = hvx_vec_splat_f32(-87.f);
+    HVX_Vector sig     = hvx_vec_fast_sigmoid_f32_guard(x, one, max_exp, min_exp);
+    return Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(x, sig));
+}
+
 #define hvx_sigmoid_loop_body(dst_type, src_type, vec_store)                                  \
     do {                                                                                      \
         dst_type * restrict vdst = (dst_type *) dst;                                          \

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

@@ -768,19 +768,13 @@ int op_im2col(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-
     const struct htp_tensor * dst = octx->dst;
 
     uint32_t row_base = 0;
     uint32_t nrows    = kparams->total_rows;
     if (octx->ctx->mdev.count > 1) {
-        const uint32_t row_size = dst->nb[2];
-        const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
         const struct htp_tensor_mdev_range range =
-            htp_tensor_mdev_partition(kparams->total_rows, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0,
+            htp_tensor_mdev_partition(kparams->total_rows, htp_tensor_mdev_data_aligned(dst) ? kparams->rows_per_chunk : 0,
                                       octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         row_base = range.start;
         nrows    = range.count;
@@ -790,20 +784,27 @@ int op_im2col(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
+    if (!htp_ops_context_set_n_threads(octx, MIN((uint32_t) kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    struct htp_im2col_vtcm_layout layout;
+    htp_im2col_vtcm_layout_build(&layout, kparams->src_row_bytes, kparams->dst_row_bytes, kparams->n_threads);
+
     struct htp_im2col_context ictx = { 0 };
     uint8_t * const base = octx->ctx->vtcm_base;
 
-    const uint32_t pth = MIN(octx->n_threads, nrows);
+    const uint32_t pth = octx->n_threads;
     ictx.octx                   = octx;
     ictx.pe_row_base            = row_base;
     ictx.pe_nrows               = nrows;
-    ictx.pe_rows_per_thread     = (nrows + pth - 1) / pth;
+    ictx.pe_rows_per_thread     = fastdiv(nrows + pth - 1, &octx->n_threads_div);
     ictx.pe_src_row_bytes       = kparams->src_row_bytes;
     ictx.pe_dst_row_bytes       = kparams->dst_row_bytes;
-    ictx.pe_vtcm_src            = VTCM_LAYOUT_PTR(uint8_t, base, kparams->off_src);
-    ictx.pe_vtcm_dst            = VTCM_LAYOUT_PTR(uint8_t, base, kparams->off_dst);
-    ictx.pe_src_size_per_thread = kparams->src_bytes_per_thread;
-    ictx.pe_dst_size_per_thread = kparams->dst_bytes_per_thread;
+    ictx.pe_vtcm_src            = VTCM_LAYOUT_PTR(uint8_t, base, layout.off_src);
+    ictx.pe_vtcm_dst            = VTCM_LAYOUT_PTR(uint8_t, base, layout.off_dst);
+    ictx.pe_src_size_per_thread = layout.src_bytes_per_thread;
+    ictx.pe_dst_size_per_thread = layout.dst_bytes_per_thread;
     ictx.pe_owb                 = kparams->owb;
     ictx.pe_wb                  = kparams->wb;
     ictx.pe_nbpr                = kparams->nbpr;
@@ -821,9 +822,10 @@ int op_im2col(struct htp_ops_context * octx) {
     ictx.patch_stride           = kparams->patch_stride;
 
     for (uint32_t ith = 0; ith < pth; ith++) {
-        uint32_t r = row_base + ictx.pe_rows_per_thread * ith;
-        ictx.thread_ioh[ith] = r % kparams->OH;
-        ictx.thread_in[ith]  = r / kparams->OH;
+        const uint32_t r  = row_base + ictx.pe_rows_per_thread * ith;
+        const uint32_t in = fastdiv(r, &kparams->div_OH);
+        ictx.thread_ioh[ith] = r - in * kparams->OH;
+        ictx.thread_in[ith]  = in;
     }
 
     ictx.dst_elem = dst->type == HTP_TYPE_F16 ? sizeof(__fp16) : sizeof(float);
@@ -850,19 +852,13 @@ int op_im2col_3d(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-
     const struct htp_tensor * dst = octx->dst;
 
     uint32_t row_base = 0;
     uint32_t nrows    = kparams->total_rows;
     if (octx->ctx->mdev.count > 1) {
-        const uint32_t row_size = dst->nb[2];
-        const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
         const struct htp_tensor_mdev_range range =
-            htp_tensor_mdev_partition(kparams->total_rows, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0,
+            htp_tensor_mdev_partition(kparams->total_rows, htp_tensor_mdev_data_aligned(dst) ? kparams->rows_per_chunk : 0,
                                       octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         row_base = range.start;
         nrows    = range.count;
@@ -872,20 +868,27 @@ int op_im2col_3d(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
+    if (!htp_ops_context_set_n_threads(octx, MIN((uint32_t) kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    struct htp_im2col_vtcm_layout layout;
+    htp_im2col_vtcm_layout_build(&layout, kparams->src_row_bytes, kparams->dst_row_bytes, kparams->n_threads);
+
     struct htp_im2col_context ictx = { 0 };
     uint8_t * const base = octx->ctx->vtcm_base;
 
-    const uint32_t pth = MIN(octx->n_threads, nrows);
+    const uint32_t pth = octx->n_threads;
     ictx.octx                   = octx;
     ictx.pe_row_base            = row_base;
     ictx.pe_nrows               = nrows;
-    ictx.pe_rows_per_thread     = (nrows + pth - 1) / pth;
+    ictx.pe_rows_per_thread     = fastdiv(nrows + pth - 1, &octx->n_threads_div);
     ictx.pe_src_row_bytes       = kparams->src_row_bytes;
     ictx.pe_dst_row_bytes       = kparams->dst_row_bytes;
-    ictx.pe_vtcm_src            = VTCM_LAYOUT_PTR(uint8_t, base, kparams->off_src);
-    ictx.pe_vtcm_dst            = VTCM_LAYOUT_PTR(uint8_t, base, kparams->off_dst);
-    ictx.pe_src_size_per_thread = kparams->src_bytes_per_thread;
-    ictx.pe_dst_size_per_thread = kparams->dst_bytes_per_thread;
+    ictx.pe_vtcm_src            = VTCM_LAYOUT_PTR(uint8_t, base, layout.off_src);
+    ictx.pe_vtcm_dst            = VTCM_LAYOUT_PTR(uint8_t, base, layout.off_dst);
+    ictx.pe_src_size_per_thread = layout.src_bytes_per_thread;
+    ictx.pe_dst_size_per_thread = layout.dst_bytes_per_thread;
     ictx.pe_owb                 = kparams->owb;
     ictx.pe_wb                  = kparams->wb;
     ictx.pe_nbpr                = kparams->nbpr;
@@ -903,11 +906,12 @@ int op_im2col_3d(struct htp_ops_context * octx) {
     ictx.patch_stride           = kparams->patch_stride;
 
     for (uint32_t ith = 0; ith < pth; ith++) {
-        uint32_t rem = row_base + ictx.pe_rows_per_thread * ith;
-        ictx.thread_ioh[ith] = rem % kparams->OH;
-        rem /= kparams->OH;
-        ictx.thread_iod[ith] = rem % kparams->OD;
-        ictx.thread_in[ith]  = rem / kparams->OD;
+        const uint32_t rem  = row_base + ictx.pe_rows_per_thread * ith;
+        const uint32_t q_oh = fastdiv(rem, &kparams->div_OH);
+        ictx.thread_ioh[ith] = rem - q_oh * kparams->OH;
+        const uint32_t in   = fastdiv(q_oh, &kparams->div_OD);
+        ictx.thread_iod[ith] = q_oh - in * kparams->OD;
+        ictx.thread_in[ith]  = in;
     }
 
     ictx.dst_elem = dst->type == HTP_TYPE_F16 ? sizeof(__fp16) : sizeof(float);

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

@@ -2,6 +2,8 @@
 #define HTP_IM2COL_OPS_H
 
 #include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
 #include <stdint.h>
 
 enum htp_im2col_kernel_type {
@@ -11,38 +13,33 @@ enum htp_im2col_kernel_type {
 };
 
 struct htp_im2col_vtcm_layout {
-    uint32_t off_src;
-    uint32_t off_dst;
+    size_t   total_bytes;
+    size_t   off_src;
+    size_t   off_dst;
     uint32_t src_bytes_per_thread;
     uint32_t dst_bytes_per_thread;
-    uint32_t total_bytes;
 };
 
 static inline void htp_im2col_vtcm_layout_build(
-    struct htp_im2col_vtcm_layout * layout,
+    struct htp_im2col_vtcm_layout * L,
     uint32_t src_row_bytes,
     uint32_t dst_row_bytes,
     uint32_t n_threads) {
 
-    layout->src_bytes_per_thread = 2 * src_row_bytes + 256;
-    layout->dst_bytes_per_thread = 2 * dst_row_bytes;
+    L->src_bytes_per_thread = 2 * src_row_bytes + 256;
+    L->dst_bytes_per_thread = 2 * dst_row_bytes;
 
-    layout->off_src     = 0;
-    layout->off_dst     = layout->off_src + layout->src_bytes_per_thread * n_threads;
-    layout->total_bytes = layout->off_dst + layout->dst_bytes_per_thread * n_threads;
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src, (size_t) L->src_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst, (size_t) L->dst_bytes_per_thread * n_threads);
+    L->total_bytes = off;
 }
 
 struct htp_im2col_kernel_params {
     uint8_t  kernel_type;
-    uint8_t  is_2d;
-    uint8_t  is_3d;
     uint8_t  n_threads;
 
     uint32_t vtcm_size;
-    uint32_t off_src;
-    uint32_t off_dst;
-    uint32_t src_bytes_per_thread;
-    uint32_t dst_bytes_per_thread;
     uint32_t src_row_bytes;
     uint32_t dst_row_bytes;
 
@@ -57,6 +54,10 @@ struct htp_im2col_kernel_params {
     uint32_t OD, OH, OW;
     uint32_t patch_stride;
     uint32_t total_rows;
+    uint32_t rows_per_chunk;
+
+    struct fastdiv_values div_OH;
+    struct fastdiv_values div_OD;
 };
 
 #if defined(__cplusplus)

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

@@ -955,7 +955,9 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_GLU_GEGLU:
         case HTP_OP_GLU_GEGLU_QUICK:
         case HTP_OP_GLU_GEGLU_ERF:
-            return op_activations(octx);
+        case HTP_OP_GLU_REGLU:
+        case HTP_OP_GLU_SIGMOID:
+            return op_glu(octx);
 
         case HTP_OP_SOFTMAX:
             return op_softmax(octx);
@@ -1244,8 +1246,6 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_buf_desc * bufs
 
     octx->src0_spad.src = NULL;
     octx->src1_spad.src = NULL;
-    octx->src2_spad.src = NULL;
-    octx->src3_spad.src = NULL;
     octx->dst_spad.src  = NULL;
 
     return octx->status;
@@ -1306,10 +1306,7 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
     prep_tensors(ctx, bufs, tens, n_tens);
 
     struct htp_ops_context *octx = &ctx->octx;
-    memset(octx, 0, sizeof(*octx));
-    octx->n_threads     = ctx->n_threads;
-    octx->n_threads_div = ctx->n_threads_div;
-    octx->ctx           = ctx;
+    htp_ops_context_reset(octx, ctx);
 
     mdev_group_init(ctx, req);
 

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

@@ -1816,7 +1816,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
                 return HTP_STATUS_NO_SUPPORT;
             }
 
-            const uint32_t qk = QK_Q8_0_TILED;
+            const uint32_t qk = HTP_MM_ACT_BLOCK_K;
             const uint32_t nb = (ne10 + qk - 1) / qk;
             const uint32_t total_nb = act_nrows * nb;
 
@@ -1886,7 +1886,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
     mmctx->vtcm_src0_stride = src0_row_size_padded;
     mmctx->vtcm_act_stride  = act_row_size;
     if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW) {
-        mmctx->vtcm_act_raw_stride = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+        mmctx->vtcm_act_raw_stride = htp_mm_act_raw_row_size(ne10);
     } else if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM) {
         mmctx->vtcm_act_raw_stride = hex_round_up(ne10 * sizeof(float), 128);
     } else {
@@ -1904,7 +1904,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
             if (need_quant) {
                 hvx_mm_transfer_act_dma(octx, kparams, act, mmctx->vtcm_act_raw, mmctx->vtcm_act_raw_stride, m_start, cur_m_rows);
 
-                const uint32_t qk = QK_Q8_0_TILED;
+                const uint32_t qk = HTP_MM_ACT_BLOCK_K;
                 const uint32_t nb = (ne10 + qk - 1) / qk;
                 const uint32_t total_nb = cur_m_rows * nb;
                 uint32_t quant_tasks;
@@ -3914,7 +3914,7 @@ static int hvx_mm_matmul_id(
     const struct htp_tensor * restrict ids = octx->src[2];
     const size_t src0_row_size = nb01;
 
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (ne10 + qk - 1) / qk;
     const uint32_t total_nb = act_nrows * nb;
 
@@ -3966,7 +3966,7 @@ static int hvx_mm_matmul_id(
 
     mmctx->vtcm_src0_stride    = src0_row_size_padded;
     mmctx->vtcm_act_stride     = act_row_size;
-    mmctx->vtcm_act_raw_stride = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+    mmctx->vtcm_act_raw_stride = htp_mm_act_raw_row_size(ne10);
 
     mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div);
     mmctx->vtcm_act_size_per_thread  = L.act_bytes;
@@ -4081,7 +4081,7 @@ static int hvx_mm_matmul_id_nx(
     const struct htp_tensor * restrict ids  = octx->src[n_weights + 1];
     const size_t src0_row_size = src0->nb[1];
 
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (act->ne[0] + qk - 1) / qk;
     const uint32_t total_nb = act_nrows * nb;
 
@@ -4124,7 +4124,7 @@ static int hvx_mm_matmul_id_nx(
 
     mmctx->vtcm_src0_stride    = 0;
     mmctx->vtcm_act_stride     = act_row_size;
-    mmctx->vtcm_act_raw_stride = hex_round_up(act->ne[0] * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+    mmctx->vtcm_act_raw_stride = htp_mm_act_raw_row_size(act->ne[0]);
 
     mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div);
     mmctx->vtcm_act_size_per_thread  = L.act_bytes;
@@ -4474,7 +4474,7 @@ static int op_matmul_nx_impl(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (act->ne[0] + qk - 1) / qk;
     const uint32_t total_nb = act_nrows * nb;
 
@@ -4518,7 +4518,7 @@ static int op_matmul_nx_impl(struct htp_ops_context * octx) {
 
     mmctx->vtcm_src0_stride    = 0;
     mmctx->vtcm_act_stride     = act_row_size;
-    mmctx->vtcm_act_raw_stride = hex_round_up(act->ne[0] * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+    mmctx->vtcm_act_raw_stride = htp_mm_act_raw_row_size(act->ne[0]);
 
     mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div);
     mmctx->vtcm_act_size_per_thread  = L.act_bytes;

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

@@ -25,6 +25,7 @@ extern "C" {
 #define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q2_0  640
 
 // --- Activation Tiled Block Sizes (including padding) ---
+#define HTP_MM_ACT_BLOCK_K             128
 #define HTP_MM_ACT_TILE_SIZE_Q8_0      1152
 #define HTP_MM_ACT_TILE_SIZE_Q8_1      1280
 #define HTP_MM_ACT_TILE_SIZE_F16       2048
@@ -74,10 +75,7 @@ struct htp_mm_kernel_params {
     int32_t  aligned_tile_size;  // Aligned weight tile size (padded to 128)
     int32_t  act_row_size;       // Row size for activation scratchpad
     int32_t  vtcm_size;          // Total required scratchpad size in VTCM
-    int32_t  vtcm_src0_size;     // src0 scratchpad size in VTCM
-    int32_t  vtcm_act_size;      // activation scratchpad size in VTCM
     int32_t  vtcm_bias_size;     // bias scratchpad size in VTCM (fused only)
-    int32_t  vtcm_dst_size;      // dst scratchpad size in VTCM
 
     // Precomputed division values
     struct fastdiv_values div_ne12_ne1;
@@ -192,19 +190,19 @@ static inline bool htp_mm_weight_has_offset(int weight_type) {
 
 // --- Activation/Row Size Helpers ---
 static inline size_t htp_mm_q8_0_tiled_row_size(uint32_t ne) {
-    const uint32_t ne_padded = ((ne + 127) / 128) * 128;
+    const uint32_t ne_padded = hex_round_up(ne, HTP_MM_ACT_BLOCK_K);
     const uint32_t nb_32 = ne_padded / 32;
     return nb_32 * HTP_MM_ACT_TILE_SIZE_Q8_0;
 }
 
 static inline size_t htp_mm_q8_1_tiled_row_size(uint32_t ne) {
-    const uint32_t ne_padded = ((ne + 127) / 128) * 128;
+    const uint32_t ne_padded = hex_round_up(ne, HTP_MM_ACT_BLOCK_K);
     const uint32_t nb_32 = ne_padded / 32;
     return nb_32 * HTP_MM_ACT_TILE_SIZE_Q8_1;
 }
 
 static inline size_t htp_mm_f16_tiled_row_size(uint32_t ne) {
-    const uint32_t ne_padded = ((ne + 127) / 128) * 128;
+    const uint32_t ne_padded = hex_round_up(ne, HTP_MM_ACT_BLOCK_K);
     const uint32_t nb_32 = ne_padded / 32;
     return nb_32 * HTP_MM_ACT_TILE_SIZE_F16;
 }
@@ -216,6 +214,10 @@ static inline size_t htp_mm_tiled_act_row_size(int weight_type, uint32_t ne) {
     return htp_mm_weight_has_offset(weight_type) ? htp_mm_q8_1_tiled_row_size(ne) : htp_mm_q8_0_tiled_row_size(ne);
 }
 
+static inline size_t htp_mm_act_raw_row_size(uint32_t ne) {
+    return hex_round_up(ne, HTP_MM_ACT_BLOCK_K) * sizeof(float);
+}
+
 static inline size_t htp_mm_round_up(size_t n, size_t m) {
     return ((n + m - 1) / m) * m;
 }
@@ -445,7 +447,7 @@ static inline void htp_mm_hvx_vtcm_layout_build(
 
         size_t tiled_act_row_size = htp_mm_tiled_act_row_size(wtype, ne10);
         size_t q_act_sz = hex_round_up(tiled_act_row_size * act_nrows, 128);
-        size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+        size_t raw_row_size = htp_mm_act_raw_row_size(ne10);
 
         src0_sz    = weight_sz_per_thread * n_threads; // shared single-weight prefetch buffer
         act_sz     = q_act_sz;                         // quantized activation buffer
@@ -461,7 +463,7 @@ static inline void htp_mm_hvx_vtcm_layout_build(
 
         act_sz = htp_mm_round_up(act_row_size_tiled * act_nrows, 256);
 
-        size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+        size_t raw_row_size = htp_mm_act_raw_row_size(ne10);
 
         src0_sz    = src0_sz_per_thread * n_threads;
         dst_sz     = 0;
@@ -502,7 +504,7 @@ static inline void htp_mm_hvx_vtcm_layout_build(
 
                 size_t dst_slice_per_thread = (dst_nrows > 0 && act_nrows == 1) ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0;
                 dst_sz = dst_slice_per_thread * n_threads;
-                size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+                size_t raw_row_size = htp_mm_act_raw_row_size(ne10);
                 act_raw_sz = hex_round_up(raw_row_size * act_nrows, 128);
                 break;
             }
@@ -582,7 +584,7 @@ static inline bool htp_mm_hvx_solve_vtcm_params(
 
     size_t eff_row_size = row_size;
     if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
-        eff_row_size += hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+        eff_row_size += htp_mm_act_raw_row_size(ne10);
     }
 
     uint32_t m_chunk = (uint32_t) (avail_act / eff_row_size);

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

@@ -17,18 +17,28 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "htp-vtcm.h"
+#include "pad-ops.h"
 
 /* Circular wrap: maps any integer x into [0, n) */
-static inline uint32_t wrap_around(int32_t x, uint32_t n) {
-    return (uint32_t)(((x % (int32_t)n) + (int32_t)n) % (int32_t)n);
+static inline uint32_t wrap_around(int32_t x, uint32_t n, const struct fastdiv_values * div_n) {
+    if (x >= 0) {
+        const uint32_t ux = (uint32_t) x;
+        return ux < n ? ux : fastmodulo(ux, n, div_n);
+    }
+    const uint32_t ux  = (uint32_t) (-x);
+    const uint32_t rem = ux < n ? ux : fastmodulo(ux, n, div_n);
+    return rem ? (n - rem) : 0;
 }
 
 /* Decompose a flat dst row index into (i1, i2, i3) */
 static inline void pad_decompose_row(uint32_t ir, uint32_t ne1, uint32_t ne2,
+                                     const struct htp_pad_kernel_params * kparams,
                                      uint32_t *i1, uint32_t *i2, uint32_t *i3) {
-    *i1 = ir % ne1;
-    *i2 = (ir / ne1) % ne2;
-    *i3 = ir / (ne1 * ne2);
+    *i3 = fastdiv(ir, &kparams->div_ne1_ne2);
+    const uint32_t rem = ir - (*i3) * (ne1 * ne2);
+    *i2 = fastdiv(rem, &kparams->div_ne1);
+    *i1 = rem - (*i2) * ne1;
 }
 
 /* Return non-zero if row (i1,i2,i3) falls in the non-padded interior */
@@ -62,25 +72,28 @@ static inline dma_addr_t pad_src_row_data(const struct htp_tensor * src,
 
 /* Compute the DDR src row pointer for a circular row (wrap-around indexing) */
 static inline const uint8_t * pad_circ_src_row_ptr(const struct htp_tensor * src,
+                                                    const struct htp_pad_kernel_params * kparams,
                                                     uint32_t i1, uint32_t i2, uint32_t i3,
                                                     int32_t lp1, int32_t lp2, int32_t lp3) {
     return (const uint8_t *) src->data
-        + wrap_around((int32_t)i1 - lp1, src->ne[1]) * src->nb[1]
-        + wrap_around((int32_t)i2 - lp2, src->ne[2]) * src->nb[2]
-        + wrap_around((int32_t)i3 - lp3, src->ne[3]) * src->nb[3];
+        + wrap_around((int32_t)i1 - lp1, src->ne[1], &kparams->div_ne01) * src->nb[1]
+        + wrap_around((int32_t)i2 - lp2, src->ne[2], &kparams->div_ne02) * src->nb[2]
+        + wrap_around((int32_t)i3 - lp3, src->ne[3], &kparams->div_ne03) * src->nb[3];
 }
 
 static inline dma_addr_t pad_circ_src_row_data(const struct htp_tensor * src,
+                                               const struct htp_pad_kernel_params * kparams,
                                                uint32_t i1, uint32_t i2, uint32_t i3,
                                                int32_t lp1, int32_t lp2, int32_t lp3) {
     return src->data
-        + wrap_around((int32_t)i1 - lp1, src->ne[1]) * src->nb[1]
-        + wrap_around((int32_t)i2 - lp2, src->ne[2]) * src->nb[2]
-        + wrap_around((int32_t)i3 - lp3, src->ne[3]) * src->nb[3];
+        + wrap_around((int32_t)i1 - lp1, src->ne[1], &kparams->div_ne01) * src->nb[1]
+        + wrap_around((int32_t)i2 - lp2, src->ne[2], &kparams->div_ne02) * src->nb[2]
+        + wrap_around((int32_t)i3 - lp3, src->ne[3], &kparams->div_ne03) * src->nb[3];
 }
 
 struct htp_pad_context {
     struct htp_ops_context * octx;
+    const struct htp_pad_kernel_params * kparams;
 
     int32_t  lp0, rp0;
     int32_t  lp1, rp1;
@@ -152,7 +165,7 @@ static void pad_job_per_thread_hvx(unsigned int nth, unsigned int ith, void * da
 
     for (uint32_t dst_row = row_start; dst_row < row_end; dst_row++) {
         uint32_t i1, i2, i3;
-        pad_decompose_row(dst_row, ne1, ne2, &i1, &i2, &i3);
+        pad_decompose_row(dst_row, ne1, ne2, pctx->kparams, &i1, &i2, &i3);
 
         uint8_t * dst_ptr = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
 
@@ -219,7 +232,7 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void
             dst_row_size, dst_row_size_aligned, dst_row_size, 0);
 
         uint32_t i1, i2, i3;
-        pad_decompose_row(ir, ne1, ne2, &i1, &i2, &i3);
+        pad_decompose_row(ir, ne1, ne2, pctx->kparams, &i1, &i2, &i3);
         const int interior = pad_is_interior(i1, i2, i3,
                                              lp1, rp1, ne1,
                                              lp2, rp2, ne2,
@@ -246,7 +259,7 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void
         uint8_t * src_spad_cur = (uint8_t *) dma_queue_pop(dma_q).dst;
 
         uint32_t i1, i2, i3;
-        pad_decompose_row(ir, ne1, ne2, &i1, &i2, &i3);
+        pad_decompose_row(ir, ne1, ne2, pctx->kparams, &i1, &i2, &i3);
 
         const dma_addr_t dst_data = dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
 
@@ -278,7 +291,7 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void
         const uint32_t next_row = ir + 2;
         if (next_row < row_end) {
             uint32_t ni1, ni2, ni3;
-            pad_decompose_row(next_row, ne1, ne2, &ni1, &ni2, &ni3);
+            pad_decompose_row(next_row, ne1, ne2, pctx->kparams, &ni1, &ni2, &ni3);
             const int next_interior = pad_is_interior(ni1, ni2, ni3,
                                                       lp1, rp1, ne1,
                                                       lp2, rp2, ne2,
@@ -315,10 +328,10 @@ static void pad_job_per_thread_hvx_circular(unsigned int nth, unsigned int ith,
 
     for (uint32_t dst_row = row_start; dst_row < row_end; dst_row++) {
         uint32_t i1, i2, i3;
-        pad_decompose_row(dst_row, ne1, ne2, &i1, &i2, &i3);
+        pad_decompose_row(dst_row, ne1, ne2, pctx->kparams, &i1, &i2, &i3);
 
         uint8_t       * dst_ptr = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
-        const uint8_t * src_row = pad_circ_src_row_ptr(src, i1, i2, i3, lp1, lp2, lp3);
+        const uint8_t * src_row = pad_circ_src_row_ptr(src, pctx->kparams, i1, i2, i3, lp1, lp2, lp3);
 
         if (nb00 == type_size) {
 
@@ -393,8 +406,8 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i
             dst_row_size, dst_row_size_aligned, dst_row_size, 0);
 
         uint32_t pi1, pi2, pi3;
-        pad_decompose_row(ir, ne1, ne2, &pi1, &pi2, &pi3);
-        const dma_addr_t src_data = pad_circ_src_row_data(src, pi1, pi2, pi3, lp1, lp2, lp3);
+        pad_decompose_row(ir, ne1, ne2, pctx->kparams, &pi1, &pi2, &pi3);
+        const dma_addr_t src_data = pad_circ_src_row_data(src, pctx->kparams, pi1, pi2, pi3, lp1, lp2, lp3);
         dma_queue_push(dma_q,
             dma_make_data(src_spad_cur, src_data),
             src_row_size_aligned, src_row_size, src_row_size, 1);
@@ -411,7 +424,7 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i
         uint8_t * src_spad_cur = (uint8_t *) dma_queue_pop(dma_q).dst;
 
         uint32_t i1, i2, i3;
-        pad_decompose_row(ir, ne1, ne2, &i1, &i2, &i3);
+        pad_decompose_row(ir, ne1, ne2, pctx->kparams, &i1, &i2, &i3);
         const dma_addr_t dst_data = dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
 
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
@@ -455,8 +468,8 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i
         const uint32_t next_row = ir + 2;
         if (next_row < row_end) {
             uint32_t nri1, nri2, nri3;
-            pad_decompose_row(next_row, ne1, ne2, &nri1, &nri2, &nri3);
-            const dma_addr_t next_src_data = pad_circ_src_row_data(src, nri1, nri2, nri3, lp1, lp2, lp3);
+            pad_decompose_row(next_row, ne1, ne2, pctx->kparams, &nri1, &nri2, &nri3);
+            const dma_addr_t next_src_data = pad_circ_src_row_data(src, pctx->kparams, nri1, nri2, nri3, lp1, lp2, lp3);
             dma_queue_push(dma_q,
                 dma_make_data(src_spad_cur, next_src_data),
                 src_row_size_aligned, src_row_size, src_row_size, 1);
@@ -475,15 +488,8 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i
 int op_pad(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
-
-    // Only F32 supported
-    size_t type_size;
-    switch (src0->type) {
-        case HTP_TYPE_F32: type_size = 4; break;
-        default:
-            FARF(ERROR, "pad-hvx: unsupported type %u\n", src0->type);
-            return HTP_STATUS_NO_SUPPORT;
-    }
+    const struct htp_pad_kernel_params * kparams =
+        (const struct htp_pad_kernel_params *) octx->kernel_params;
 
     const int32_t lp0 = octx->op_params[0];
     const int32_t rp0 = octx->op_params[1];
@@ -495,11 +501,10 @@ int op_pad(struct htp_ops_context * octx) {
     const int32_t rp3 = octx->op_params[7];
     const int32_t circular = octx->op_params[8];
 
-    const uint32_t ne0  = dst->ne[0];
-    const uint32_t ne00 = src0->ne[0];
-
-    const uint32_t total_dst_rows = dst->ne[1] * dst->ne[2] * dst->ne[3];
-    const size_t dst_row_size     = (size_t)ne0 * type_size;
+    const uint32_t total_dst_rows = kparams->total_rows;
+    const size_t   type_size      = kparams->elem_size;
+    const size_t   src_row_size   = src0->ne[0] * type_size;
+    const size_t   dst_row_size   = dst->ne[0] * type_size;
 
     uint32_t row_start = 0;
     uint32_t nrows     = total_dst_rows;
@@ -516,60 +521,68 @@ int op_pad(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t n_threads = octx->n_threads;
 
-    const size_t src_row_size         = (size_t)ne00 * type_size;
-    const size_t src_row_size_aligned = hex_round_up(src_row_size, VLEN);
-    const size_t dst_row_size_aligned = hex_round_up(dst_row_size, VLEN);
-
-    // Total VTCM needed: 2 buffers (ping+pong) for src and dst, per thread
-    const size_t vtcm_needed = (size_t)n_threads * 2 * (src_row_size_aligned + dst_row_size_aligned);
-
-    const int use_dma = (src0->nb[0] == (uint32_t)type_size) && (ne00 >= 512) &&
-                        (octx->ctx->vtcm_size >= vtcm_needed);
-
-    if (!use_dma && (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) {
+    if (!kparams->use_dma && (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    if (use_dma) {
-        octx->src0_spad.size_per_thread = 2 * src_row_size_aligned;
-        octx->dst_spad.size_per_thread  = 2 * dst_row_size_aligned;
-        octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread;
-        octx->dst_spad.size  = n_threads * octx->dst_spad.size_per_thread;
-        octx->src0_spad.data = octx->ctx->vtcm_base;
-        octx->dst_spad.data  = octx->src0_spad.data + octx->src0_spad.size;
+    struct htp_pad_vtcm_layout layout = { 0 };
+    if (kparams->use_dma) {
+        htp_pad_vtcm_layout_build(&layout, src_row_size, dst_row_size, n_threads);
+
+        octx->src0_spad.size_per_thread = layout.src0_bytes_per_thread;
+        octx->dst_spad.size_per_thread  = layout.dst_bytes_per_thread;
+        octx->src0_spad.size = n_threads * layout.src0_bytes_per_thread;
+        octx->dst_spad.size  = n_threads * layout.dst_bytes_per_thread;
+        octx->src0_spad.data = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_src0);
+        octx->dst_spad.data  = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_dst);
         octx->src0_spad.src  = NULL;
         octx->dst_spad.src   = NULL;
     }
 
     struct htp_pad_context pctx = {
-        .octx             = octx,
+        .octx                 = octx,
+        .kparams              = kparams,
         .lp0 = lp0, .rp0 = rp0,
         .lp1 = lp1, .rp1 = rp1,
         .lp2 = lp2, .rp2 = rp2,
         .lp3 = lp3, .rp3 = rp3,
-        .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
-        .total_dst_rows   = nrows,
-        .row_start        = row_start,
-        .type_size        = type_size,
+        .nrows_per_thread     = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
+        .total_dst_rows       = nrows,
+        .row_start            = row_start,
+        .type_size            = type_size,
         .src_row_size         = src_row_size,
-        .src_row_size_aligned = src_row_size_aligned,
+        .src_row_size_aligned = layout.src_row_size_aligned,
         .dst_row_size         = dst_row_size,
-        .dst_row_size_aligned = dst_row_size_aligned,
+        .dst_row_size_aligned = layout.dst_row_size_aligned,
     };
 
     FARF(HIGH, "pad-hvx%s%s: (%ux%ux%ux%u) -> (%ux%ux%ux%u) pads=(%d,%d,%d,%d,%d,%d,%d,%d)\n",
-         circular ? "-circ" : "",
-         use_dma   ? "-dma"  : "",
+         circular         ? "-circ" : "",
+         kparams->use_dma ? "-dma"  : "",
          src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3],
          dst->ne[0],  dst->ne[1],  dst->ne[2],  dst->ne[3],
          lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
 
-    if      (circular && use_dma) { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_circular_dma, &pctx, n_threads); }
-    else if (circular)            { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_circular,     &pctx, n_threads); }
-    else if (use_dma)             { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_dma,          &pctx, n_threads); }
-    else                          { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx,              &pctx, n_threads); }
+    switch (kparams->kernel_type) {
+        case HTP_PAD_KERNEL_HVX_CIRCULAR_DMA:
+            work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_circular_dma, &pctx, n_threads);
+            break;
+        case HTP_PAD_KERNEL_HVX_CIRCULAR:
+            work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_circular,     &pctx, n_threads);
+            break;
+        case HTP_PAD_KERNEL_HVX_DMA:
+            work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_dma,          &pctx, n_threads);
+            break;
+        default:
+            work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx,              &pctx, n_threads);
+            break;
+    }
 
     return HTP_STATUS_OK;
 }

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

@@ -0,0 +1,65 @@
+#ifndef HTP_PAD_OPS_H
+#define HTP_PAD_OPS_H
+
+#include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
+#include <stdint.h>
+
+enum htp_pad_kernel_type {
+    HTP_PAD_KERNEL_HVX              = 0,
+    HTP_PAD_KERNEL_HVX_DMA          = 1,
+    HTP_PAD_KERNEL_HVX_CIRCULAR     = 2,
+    HTP_PAD_KERNEL_HVX_CIRCULAR_DMA = 3,
+};
+
+struct htp_pad_vtcm_layout {
+    size_t total_bytes;
+    size_t off_src0;
+    size_t off_dst;
+    size_t src_row_size_aligned;
+    size_t dst_row_size_aligned;
+    size_t src0_bytes_per_thread;
+    size_t dst_bytes_per_thread;
+};
+
+static inline void htp_pad_vtcm_layout_build(
+    struct htp_pad_vtcm_layout * L,
+    size_t                       src_row_size,
+    size_t                       dst_row_size,
+    uint32_t                     n_threads
+) {
+    L->src_row_size_aligned  = hex_round_up(src_row_size, 128);
+    L->dst_row_size_aligned  = hex_round_up(dst_row_size, 128);
+    L->src0_bytes_per_thread = 2 * L->src_row_size_aligned;
+    L->dst_bytes_per_thread  = 2 * L->dst_row_size_aligned;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread * n_threads);
+
+    L->total_bytes = off;
+}
+
+struct htp_pad_kernel_params {
+    uint8_t  kernel_type;
+    uint32_t n_threads;
+    uint32_t total_rows;
+    uint32_t use_dma;
+    uint32_t elem_size;
+    uint32_t vtcm_size;
+
+    struct fastdiv_values div_ne1;
+    struct fastdiv_values div_ne1_ne2;
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne02;
+    struct fastdiv_values div_ne03;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_pad_kernel_params) <= 128, "htp_pad_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_pad_kernel_params) <= 128, "htp_pad_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_PAD_OPS_H

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

@@ -247,6 +247,7 @@ typedef void (*pool_plane_fn_t)(const float * src, float * dst, const struct htp
 struct pool_2d_context {
     struct htp_ops_context * octx;
     const struct htp_pool_2d_kernel_params * kparams;
+    struct htp_pool_vtcm_layout layout;
     pool_plane_fn_t pool_plane;
     uint32_t n_threads;
     uint32_t plane_start;
@@ -257,6 +258,7 @@ struct pool_2d_context {
 static void pool_2d_thread(unsigned int nth, unsigned int ith, void * data) {
     struct pool_2d_context * ctx = (struct pool_2d_context *) data;
     const struct htp_pool_2d_kernel_params * p = ctx->kparams;
+    const struct htp_pool_vtcm_layout * L = &ctx->layout;
     const struct htp_tensor * src0 = ctx->octx->src[0];
     const struct htp_tensor * dst = ctx->octx->dst;
     pool_plane_fn_t pool_plane = ctx->pool_plane;
@@ -271,15 +273,14 @@ static void pool_2d_thread(unsigned int nth, unsigned int ith, void * data) {
     struct htp_thread_trace * tr = &ctx->octx->ctx->trace[ith];
     dma_queue * dma_queue = ctx->octx->ctx->dma[ith];
 
-    const uint32_t src_spad_half = p->src_plane_bytes_aligned;
-    const uint32_t dst_spad_half = p->dst_plane_bytes_aligned;
-    const uint32_t src_bytes_per_thread = 2 * src_spad_half;
-    const uint32_t dst_bytes_per_thread = 2 * dst_spad_half;
-    const size_t off_dst = (size_t) ctx->n_threads * src_bytes_per_thread;
+    const uint32_t src_plane_bytes = L->src_plane_bytes;
+    const uint32_t dst_plane_bytes = L->dst_plane_bytes;
+    const uint32_t src_spad_half   = (uint32_t) L->src_spad_half_size;
+    const uint32_t dst_spad_half   = (uint32_t) L->dst_spad_half_size;
 
     uint8_t * vtcm_base = (uint8_t *) ctx->octx->ctx->vtcm_base;
-    uint8_t * src_spad  = vtcm_base + ith * src_bytes_per_thread;
-    uint8_t * dst_spad  = vtcm_base + off_dst + ith * dst_bytes_per_thread;
+    uint8_t * src_spad  = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, L->off_src) + ith * L->src_bytes_per_thread;
+    uint8_t * dst_spad  = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, L->off_dst) + ith * L->dst_bytes_per_thread;
 
     float * srcb2[2] = { (float *) src_spad, (float *) (src_spad + src_spad_half) };
     float * dstb2[2] = { (float *) dst_spad, (float *) (dst_spad + dst_spad_half) };
@@ -290,14 +291,14 @@ static void pool_2d_thread(unsigned int nth, unsigned int ith, void * data) {
     for (uint32_t i = 0; i < total && i < 2; ++i) {
         dma_queue_push(dma_queue,
                        dma_make_data(dst->data, dstb2[i]),
-                       p->dst_plane_bytes, dst_spad_half,
-                       p->dst_plane_bytes, 0);
+                       dst_plane_bytes, dst_spad_half,
+                       dst_plane_bytes, 0);
 
-        const dma_addr_t src_addr = src0->data + (first + i) * p->src_plane_bytes;
+        const dma_addr_t src_addr = src0->data + (first + i) * src_plane_bytes;
         dma_queue_push(dma_queue,
                        dma_make_data(srcb2[i], src_addr),
-                       src_spad_half, p->src_plane_bytes,
-                       p->src_plane_bytes, 1);
+                       src_spad_half, src_plane_bytes,
+                       src_plane_bytes, 1);
     }
 
     for (uint32_t i = 0; i < total; ++i) {
@@ -313,18 +314,18 @@ static void pool_2d_thread(unsigned int nth, unsigned int ith, void * data) {
         pool_plane(srcb, dstb, p);
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) plane);
 
-        const dma_addr_t dst_addr = dst->data + plane * p->dst_plane_bytes;
+        const dma_addr_t dst_addr = dst->data + plane * dst_plane_bytes;
         dma_queue_push(dma_queue,
                        dma_make_data(dst_addr, dstb),
-                       p->dst_plane_bytes, dst_spad_half,
-                       p->dst_plane_bytes, 1);
+                       dst_plane_bytes, dst_spad_half,
+                       dst_plane_bytes, 1);
 
         if (i + 2 < total) {
-            const dma_addr_t next_src_addr = src0->data + (plane + 2) * p->src_plane_bytes;
+            const dma_addr_t next_src_addr = src0->data + (plane + 2) * src_plane_bytes;
             dma_queue_push(dma_queue,
                            dma_make_data(srcb, next_src_addr),
-                           src_spad_half, p->src_plane_bytes,
-                           p->src_plane_bytes, 1);
+                           src_spad_half, src_plane_bytes,
+                           src_plane_bytes, 1);
         }
     }
 
@@ -337,22 +338,15 @@ static void pool_2d_thread(unsigned int nth, unsigned int ith, void * data) {
 }
 
 int op_pool_2d(struct htp_ops_context * octx) {
-    const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst = octx->dst;
     const struct htp_pool_2d_kernel_params * p =
         (const struct htp_pool_2d_kernel_params *) octx->kernel_params;
 
-    if (src0->type != HTP_TYPE_F32 || dst->type != HTP_TYPE_F32 ||
-        (p->pool_op != HTP_POOL_AVG && p->pool_op != HTP_POOL_MAX)) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
     uint32_t plane_start = 0;
     uint32_t plane_count = p->planes;
     if (octx->ctx->mdev.count > 1) {
-        const uint32_t planes_per_chunk = (p->dst_plane_bytes > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(p->dst_plane_bytes, HEX_L2_LINE_SIZE)) : 1;
         const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
-            plane_count, htp_tensor_mdev_data_aligned(dst) ? planes_per_chunk : 0,
+            plane_count, htp_tensor_mdev_data_aligned(dst) ? p->planes_per_chunk : 0,
             octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         plane_start = range.start;
         plane_count = range.count;
@@ -369,16 +363,12 @@ int op_pool_2d(struct htp_ops_context * octx) {
     const uint32_t planes_per_thread = fastdiv(plane_count + n_threads - 1, &octx->n_threads_div);
 
     pool_plane_fn_t pool_plane;
-    if (p->global_path) {
-        pool_plane = pool_plane_global;
-    } else if (p->block_path) {
-        pool_plane = pool_plane_block;
-    } else if (p->narrow_path) {
-        pool_plane = pool_plane_hvx_narrow;
-    } else if (p->fast_path) {
-        pool_plane = pool_plane_hvx;
-    } else {
-        pool_plane = pool_plane_general;
+    switch (p->kernel_type) {
+        case HTP_POOL_KERNEL_GLOBAL:     pool_plane = pool_plane_global;     break;
+        case HTP_POOL_KERNEL_BLOCK:      pool_plane = pool_plane_block;      break;
+        case HTP_POOL_KERNEL_HVX_NARROW: pool_plane = pool_plane_hvx_narrow; break;
+        case HTP_POOL_KERNEL_HVX:        pool_plane = pool_plane_hvx;        break;
+        default:                         pool_plane = pool_plane_general;    break;
     }
 
     struct pool_2d_context ctx = {
@@ -390,6 +380,7 @@ int op_pool_2d(struct htp_ops_context * octx) {
         .plane_count = plane_count,
         .planes_per_thread = planes_per_thread,
     };
+    htp_pool_vtcm_layout_build(&ctx.layout, p->src_x, p->src_y, p->dst_x, p->dst_y, n_threads);
     work_queue_run(octx->ctx->work_queue, pool_2d_thread, &ctx, n_threads);
     return HTP_STATUS_OK;
 }

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

@@ -7,8 +7,18 @@
 #include <string.h>
 
 #include "hex-common.h"
+#include "htp-vtcm.h"
+
+enum htp_pool_kernel_type {
+    HTP_POOL_KERNEL_GENERAL = 0,
+    HTP_POOL_KERNEL_HVX,
+    HTP_POOL_KERNEL_HVX_NARROW,
+    HTP_POOL_KERNEL_GLOBAL,
+    HTP_POOL_KERNEL_BLOCK,
+};
 
 struct htp_pool_2d_kernel_params {
+    uint8_t  kernel_type;
     uint32_t src_x;
     uint32_t src_y;
     uint32_t dst_x;
@@ -19,17 +29,10 @@ struct htp_pool_2d_kernel_params {
     uint32_t stride_y;
     int32_t  pad_x;
     int32_t  pad_y;
-    uint32_t src_plane_bytes;
-    uint32_t dst_plane_bytes;
-    uint32_t src_plane_bytes_aligned;
-    uint32_t dst_plane_bytes_aligned;
     uint32_t n_threads;
     uint32_t planes;
+    uint32_t planes_per_chunk;
     uint32_t pool_op;
-    uint32_t fast_path;
-    uint32_t narrow_path;
-    uint32_t global_path;
-    uint32_t block_path;
     uint32_t avg_divide_count;
     uint32_t ox_lo;
     uint32_t ox_hi;
@@ -46,12 +49,36 @@ struct htp_pool_vtcm_layout {
     size_t   total_bytes;
     size_t   off_src;
     size_t   off_dst;
+    uint32_t src_plane_bytes;
+    uint32_t dst_plane_bytes;
     size_t   src_bytes_per_thread;
     size_t   dst_bytes_per_thread;
     size_t   src_spad_half_size;
     size_t   dst_spad_half_size;
 };
 
+static inline void htp_pool_vtcm_layout_build(
+    struct htp_pool_vtcm_layout * L,
+    uint32_t src_x,
+    uint32_t src_y,
+    uint32_t dst_x,
+    uint32_t dst_y,
+    uint32_t n_threads
+) {
+    // Full-plane double buffering (256 bytes guard space for vector loads)
+    L->src_plane_bytes      = src_x * src_y * sizeof(float);
+    L->dst_plane_bytes      = dst_x * dst_y * sizeof(float);
+    L->src_spad_half_size   = hex_round_up(L->src_plane_bytes + 256, 128);
+    L->dst_spad_half_size   = hex_round_up(L->dst_plane_bytes, 128);
+    L->src_bytes_per_thread = 2 * L->src_spad_half_size;
+    L->dst_bytes_per_thread = 2 * L->dst_spad_half_size;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src, L->src_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst, L->dst_bytes_per_thread * n_threads);
+    L->total_bytes = off;
+}
+
 static inline bool htp_pool_solve_layout(
     struct htp_pool_vtcm_layout * layout,
     uint32_t src_x,
@@ -61,25 +88,8 @@ static inline bool htp_pool_solve_layout(
     uint32_t n_threads,
     size_t   vtcm_budget
 ) {
-    // Full-plane double buffering (256 bytes guard space for vector loads)
-    const size_t src_plane_bytes   = (size_t) src_x * src_y * sizeof(float);
-    const size_t dst_plane_bytes   = (size_t) dst_x * dst_y * sizeof(float);
-    const size_t src_plane_aligned = hex_round_up((uint32_t) src_plane_bytes + 256, 128);
-    const size_t dst_plane_aligned = hex_round_up((uint32_t) dst_plane_bytes, 128);
-
-    const size_t spad_per_thread = 2 * (src_plane_aligned + dst_plane_aligned);
-    if (spad_per_thread * n_threads > vtcm_budget) {
-        return false;
-    }
-
-    layout->src_spad_half_size   = src_plane_aligned;
-    layout->dst_spad_half_size   = dst_plane_aligned;
-    layout->src_bytes_per_thread = 2 * src_plane_aligned;
-    layout->dst_bytes_per_thread = 2 * dst_plane_aligned;
-    layout->off_src              = 0;
-    layout->off_dst              = layout->src_bytes_per_thread * n_threads;
-    layout->total_bytes          = layout->off_dst + layout->dst_bytes_per_thread * n_threads;
-    return true;
+    htp_pool_vtcm_layout_build(layout, src_x, src_y, dst_x, dst_y, n_threads);
+    return layout->total_bytes <= vtcm_budget;
 }
 
 #endif // HTP_POOL_OPS_H

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

@@ -16,14 +16,11 @@
 #include "hex-profile.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "repeat-ops.h"
 
 struct htp_repeat_context {
-    struct htp_ops_context * octx;
-
-    uint32_t nr0;
-    uint32_t nr1;
-    uint32_t nr2;
-    uint32_t nr3;
+    struct htp_ops_context *               octx;
+    const struct htp_repeat_kernel_params * kparams;
 
     uint32_t nrows_per_thread;
     uint32_t total_dst_rows;  // ne1 * ne2 * ne3
@@ -35,6 +32,7 @@ struct htp_repeat_context {
 static void repeat_job_per_thread(unsigned int nth, unsigned int ith, void * data) {
     const struct htp_repeat_context * rctx = (const struct htp_repeat_context *) data;
     struct htp_ops_context * octx = rctx->octx;
+    const struct htp_repeat_kernel_params * kparams = rctx->kparams;
     const struct htp_tensor * src = octx->src[0];
     const struct htp_tensor * dst = octx->dst;
 
@@ -58,29 +56,30 @@ static void repeat_job_per_thread(unsigned int nth, unsigned int ith, void * dat
     const uint32_t nb2 = dst->nb[2];
     const uint32_t nb3 = dst->nb[3];
 
-    const uint32_t nr0 = rctx->nr0;
-    const uint32_t nr1 = rctx->nr1;
-    const uint32_t nr2 = rctx->nr2;
-    const uint32_t nr3 = rctx->nr3;
+    const uint32_t nr0 = kparams->nr0;
 
     const size_t row_bytes = ne00 * rctx->type_size;
 
     const uint32_t row_start = rctx->row_start + rctx->nrows_per_thread * ith;
     const uint32_t row_end   = MIN(row_start + rctx->nrows_per_thread, rctx->row_start + rctx->total_dst_rows);
+    if (row_start >= row_end) {
+        return;
+    }
 
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, row_start);
 
     for (uint32_t dst_row = row_start; dst_row < row_end; dst_row++) {
         // Decompose flat dst row index into (i1, i2, i3)
-        const uint32_t i1 = dst_row % ne1;
-        const uint32_t i2 = (dst_row / ne1) % ne2;
-        const uint32_t i3 = dst_row / (ne1 * ne2);
+        const uint32_t dst_row_div_ne1 = fastdiv(dst_row, &kparams->div_ne1);
+        const uint32_t i1              = dst_row - dst_row_div_ne1 * ne1;
+        const uint32_t i3              = fastdiv(dst_row, &kparams->div_ne12);
+        const uint32_t i2              = dst_row_div_ne1 - i3 * ne2;
 
         // Map to source indices (tiling)
-        const uint32_t k1 = i1 % ne01;
-        const uint32_t k2 = i2 % ne02;
-        const uint32_t k3 = i3 % ne03;
+        const uint32_t k1 = fastmodulo(i1, ne01, &kparams->div_ne01);
+        const uint32_t k2 = fastmodulo(i2, ne02, &kparams->div_ne02);
+        const uint32_t k3 = fastmodulo(i3, ne03, &kparams->div_ne03);
 
         const uint8_t * src_row = (const uint8_t *) src->data + k1 * nb01 + k2 * nb02 + k3 * nb03;
         uint8_t * dst_base      = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
@@ -104,15 +103,6 @@ int op_repeat(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
 
-    // Validate that dst dims are multiples of src dims
-    if (dst->ne[0] % src0->ne[0] != 0 ||
-        dst->ne[1] % src0->ne[1] != 0 ||
-        dst->ne[2] % src0->ne[2] != 0 ||
-        dst->ne[3] % src0->ne[3] != 0) {
-        FARF(ERROR, "repeat: dst dims must be multiples of src dims\n");
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-
     size_t type_size;
     switch (src0->type) {
         case HTP_TYPE_F32: type_size = 4; break;
@@ -126,6 +116,13 @@ int op_repeat(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
+    const struct htp_repeat_kernel_params * kparams =
+        (const struct htp_repeat_kernel_params *) octx->kernel_params;
+
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t total_dst_rows  = dst->ne[1] * dst->ne[2] * dst->ne[3];
     const size_t dst_row_size = dst->ne[0] * type_size;
 
@@ -148,10 +145,7 @@ int op_repeat(struct htp_ops_context * octx) {
 
     struct htp_repeat_context rctx = {
         .octx             = octx,
-        .nr0              = dst->ne[0] / src0->ne[0],
-        .nr1              = dst->ne[1] / src0->ne[1],
-        .nr2              = dst->ne[2] / src0->ne[2],
-        .nr3              = dst->ne[3] / src0->ne[3],
+        .kparams          = kparams,
         .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
         .total_dst_rows   = nrows,
         .row_start        = row_start,
@@ -161,7 +155,7 @@ int op_repeat(struct htp_ops_context * octx) {
     FARF(HIGH, "repeat: (%ux%ux%ux%u) -> (%ux%ux%ux%u) nr=(%u,%u,%u,%u)\n",
          src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3],
          dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         rctx.nr0, rctx.nr1, rctx.nr2, rctx.nr3);
+         kparams->nr0, kparams->nr1, kparams->nr2, kparams->nr3);
 
     work_queue_run(octx->ctx->work_queue, repeat_job_per_thread, &rctx, n_threads);
 

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

@@ -0,0 +1,27 @@
+#ifndef HTP_REPEAT_OPS_H
+#define HTP_REPEAT_OPS_H
+
+#include "hex-fastdiv.h"
+#include <stdint.h>
+
+struct htp_repeat_kernel_params {
+    uint32_t n_threads;
+    uint32_t nr0;
+    uint32_t nr1;
+    uint32_t nr2;
+    uint32_t nr3;
+
+    struct fastdiv_values div_ne1;
+    struct fastdiv_values div_ne12;
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne02;
+    struct fastdiv_values div_ne03;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_repeat_kernel_params) <= 128, "htp_repeat_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_repeat_kernel_params) <= 128, "htp_repeat_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_REPEAT_OPS_H

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

@@ -17,16 +17,15 @@
 #include "hex-profile.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "roll-ops.h"
 
 struct htp_roll_context {
     struct htp_ops_context * octx;
+    const struct htp_roll_kernel_params * kparams;
 
     uint32_t row_start;
     uint32_t nrows;
     uint32_t nrows_per_thread;
-
-    struct fastdiv_values div_ne1;
-    struct fastdiv_values div_ne2_ne1;
 };
 
 static inline uint32_t htp_roll_wrap(int32_t i, uint32_t ne) {
@@ -39,30 +38,32 @@ static inline uint32_t htp_roll_wrap(int32_t i, uint32_t ne) {
     return (uint32_t) i;
 }
 
-#define htp_roll_preamble                             \
-    const struct htp_tensor * src0 = octx->src[0];    \
-    const struct htp_tensor * dst  = octx->dst;       \
-                                                      \
-    const uint32_t ne0 = dst->ne[0];                  \
-    const uint32_t ne1 = dst->ne[1];                  \
-    const uint32_t ne2 = dst->ne[2];                  \
-    const uint32_t ne3 = dst->ne[3];                  \
-                                                      \
-    const uint32_t nb01 = src0->nb[1];                \
-    const uint32_t nb02 = src0->nb[2];                \
-    const uint32_t nb03 = src0->nb[3];                \
-                                                      \
-    const uint32_t nb1 = dst->nb[1];                  \
-    const uint32_t nb2 = dst->nb[2];                  \
-    const uint32_t nb3 = dst->nb[3];                  \
-                                                      \
-    const int32_t s0 = octx->op_params[0];            \
-    const int32_t s1 = octx->op_params[1];            \
-    const int32_t s2 = octx->op_params[2];            \
-    const int32_t s3 = octx->op_params[3];            \
-                                                      \
-    const uint32_t i0_src0 = htp_roll_wrap(-s0, ne0); \
-    const uint32_t n0      = ne0 - i0_src0;
+#define htp_roll_preamble                                            \
+    const struct htp_tensor * src0 = octx->src[0];                   \
+    const struct htp_tensor * dst  = octx->dst;                      \
+    const struct htp_roll_kernel_params * kparams =                  \
+        (const struct htp_roll_kernel_params *) octx->kernel_params; \
+                                                                     \
+    const uint32_t ne0 = dst->ne[0];                                 \
+    const uint32_t ne1 = dst->ne[1];                                 \
+    const uint32_t ne2 = dst->ne[2];                                 \
+    const uint32_t ne3 = dst->ne[3];                                 \
+                                                                     \
+    const uint32_t nb01 = src0->nb[1];                               \
+    const uint32_t nb02 = src0->nb[2];                               \
+    const uint32_t nb03 = src0->nb[3];                               \
+                                                                     \
+    const uint32_t nb1 = dst->nb[1];                                 \
+    const uint32_t nb2 = dst->nb[2];                                 \
+    const uint32_t nb3 = dst->nb[3];                                 \
+                                                                     \
+    const int32_t s0 = octx->op_params[0];                           \
+    const int32_t s1 = octx->op_params[1];                           \
+    const int32_t s2 = octx->op_params[2];                           \
+    const int32_t s3 = octx->op_params[3];                           \
+                                                                     \
+    const uint32_t i0_src0 = kparams->i0_src0;                       \
+    const uint32_t n0      = kparams->n0;
 
 #define htp_roll_dma_preamble dma_queue * q = octx->ctx->dma[0];
 
@@ -91,10 +92,10 @@ static inline void roll_dma_push_rows(dma_queue *               q,
                                       uint32_t                  src_row,
                                       uint32_t                  nrows,
                                       uint32_t                  row_size,
-                                      uint32_t                  i0_src0) {
+                                      uint32_t                  i0_src0,
+                                      uint32_t                  n0) {
     const dma_addr_t dst_base = dst->data + (size_t) dst_row * row_size;
     const dma_addr_t src_base = src0->data + (size_t) src_row * row_size;
-    const uint32_t   n0       = src0->ne[0] - i0_src0;
 
     roll_dma_push(q, dst_base, src_base + (size_t) i0_src0 * sizeof(float),
                   row_size, row_size, n0 * sizeof(float), nrows);
@@ -125,7 +126,7 @@ static int roll_dma_f32_contiguous(struct htp_ops_context * octx) {
     const uint32_t row_size = ne0 * sizeof(float);
 
     if (s1 == 0 && s2 == 0 && s3 == 0) {
-        roll_dma_push_rows(q, dst, src0, 0, 0, ne1 * ne2 * ne3, row_size, i0_src0);
+        roll_dma_push_rows(q, dst, src0, 0, 0, ne1 * ne2 * ne3, row_size, i0_src0, n0);
         dma_queue_flush(q);
         return HTP_STATUS_OK;
     }
@@ -139,9 +140,9 @@ static int roll_dma_f32_contiguous(struct htp_ops_context * octx) {
             const uint32_t n2_first = ne2 - i2_src0;
 
             roll_dma_push_rows(q, dst, src0, dst_row0, src_row0, n2_first * ne1,
-                               row_size, i0_src0);
+                               row_size, i0_src0, n0);
             roll_dma_push_rows(q, dst, src0, dst_row0 + n2_first * ne1, i03 * ne2 * ne1,
-                               i2_src0 * ne1, row_size, i0_src0);
+                               i2_src0 * ne1, row_size, i0_src0, n0);
         }
 
         dma_queue_flush(q);
@@ -159,9 +160,9 @@ static int roll_dma_f32_contiguous(struct htp_ops_context * octx) {
             const uint32_t src_row0 = (i03 * ne2 + i02) * ne1;
 
             roll_dma_push_rows(q, dst, src0, dst_row0, src_row0 + i1_src0,
-                               n1_first, row_size, i0_src0);
+                               n1_first, row_size, i0_src0, n0);
             roll_dma_push_rows(q, dst, src0, dst_row0 + n1_first, src_row0,
-                               i1_src0, row_size, i0_src0);
+                               i1_src0, row_size, i0_src0, n0);
         }
     }
 
@@ -214,9 +215,9 @@ static void roll_thread_f32(unsigned int nth, unsigned int ith, void * data) {
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, row_start);
 
     for (uint32_t row = row_start; row < row_end; row++) {
-        const uint32_t i3  = fastdiv(row, &rctx->div_ne2_ne1);
+        const uint32_t i3  = fastdiv(row, &kparams->div_ne2_ne1);
         const uint32_t rem = row - i3 * ne2 * ne1;
-        const uint32_t i2  = fastdiv(rem, &rctx->div_ne1);
+        const uint32_t i2  = fastdiv(rem, &kparams->div_ne1);
         const uint32_t i1  = rem - i2 * ne1;
 
         const uint32_t i01 = htp_roll_wrap((int32_t) i1 - s1, ne1);
@@ -245,23 +246,14 @@ static void roll_thread_f32(unsigned int nth, unsigned int ith, void * data) {
 int execute_op_roll_f32(struct htp_ops_context * octx) {
     htp_roll_preamble;
 
-    if (src0->type != HTP_TYPE_F32 || dst->type != HTP_TYPE_F32) {
-        FARF(ERROR, "roll: unsupported type %u -> %u\n", src0->type, dst->type);
-        return HTP_STATUS_NO_SUPPORT;
+    if (kparams->kernel_type == HTP_ROLL_KERNEL_DMA_CONTIGUOUS) {
+        return roll_dma_f32_contiguous(octx);
+    }
+    if (kparams->kernel_type == HTP_ROLL_KERNEL_DMA_STRIDED) {
+        return roll_dma_f32_strided(octx);
     }
 
-    if (src0->nb[0] != sizeof(float) || dst->nb[0] != sizeof(float)) {
-        FARF(ERROR, "roll: unsupported nb0 %u -> %u\n", src0->nb[0], dst->nb[0]);
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    if (src0->ne[0] != ne0 || src0->ne[1] != ne1 ||
-        src0->ne[2] != ne2 || src0->ne[3] != ne3) {
-        FARF(ERROR, "roll: shape mismatch\n");
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-
-    const uint32_t total_rows = ne1 * ne2 * ne3;
+    const uint32_t total_rows = kparams->total_rows;
     const size_t dst_row_size = ne0 * sizeof(float);
 
     uint32_t row_start = 0;
@@ -279,25 +271,21 @@ int execute_op_roll_f32(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    if (octx->ctx->mdev.count <= 1) {
-        if (htp_tensor_is_contiguous(src0, sizeof(float)) && htp_tensor_is_contiguous(dst, sizeof(float))) {
-            return roll_dma_f32_contiguous(octx);
-        }
-        return roll_dma_f32_strided(octx);
-    }
-
     if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t n_threads = octx->n_threads;
     struct htp_roll_context rctx = {
         .octx             = octx,
+        .kparams          = kparams,
         .row_start        = row_start,
         .nrows            = nrows,
         .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
-        .div_ne1          = init_fastdiv_values(dst->ne[1]),
-        .div_ne2_ne1      = init_fastdiv_values(dst->ne[2] * dst->ne[1]),
     };
 
     work_queue_run(octx->ctx->work_queue, roll_thread_f32, &rctx, n_threads);

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

@@ -0,0 +1,30 @@
+#ifndef HTP_ROLL_OPS_H
+#define HTP_ROLL_OPS_H
+
+#include "hex-fastdiv.h"
+#include <stdint.h>
+
+enum htp_roll_kernel_type {
+    HTP_ROLL_KERNEL_DMA_CONTIGUOUS = 1,
+    HTP_ROLL_KERNEL_DMA_STRIDED    = 2,
+    HTP_ROLL_KERNEL_HVX            = 3,
+};
+
+struct htp_roll_kernel_params {
+    uint8_t  kernel_type;
+    uint32_t n_threads;
+    uint32_t total_rows;
+    uint32_t i0_src0;
+    uint32_t n0;
+
+    struct fastdiv_values div_ne1;
+    struct fastdiv_values div_ne2_ne1;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_roll_kernel_params) <= 128, "htp_roll_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_roll_kernel_params) <= 128, "htp_roll_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_ROLL_OPS_H

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

@@ -700,15 +700,6 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) {
     const struct htp_tensor * src2 = octx->src[2];
     const struct htp_tensor * dst  = octx->dst;
 
-    switch (octx->op) {
-        case HTP_OP_ROPE:
-            break;
-
-        default:
-            FARF(ERROR, "Unsupported Op %u\n", octx->op);
-            return HTP_STATUS_NO_SUPPORT;
-    }
-
     const struct htp_rope_kernel_params * kparams = (const struct htp_rope_kernel_params *) octx->kernel_params;
     if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
         return HTP_STATUS_INVAL_PARAMS;
@@ -746,20 +737,23 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) {
     const size_t dst_row_size    = dst->ne[0] * sizeof(float);
     const size_t dst_row_stride  = dst->nb[1];
 
+    struct htp_rope_vtcm_layout layout;
+    htp_rope_vtcm_layout_build(&layout, src0->ne[0], kparams->n_threads, src2 ? src2->ne[0] : 0);
+
     struct htp_rope_context rctx;
     memset(&rctx, 0, sizeof(struct htp_rope_context));
 
     rctx.octx                  = octx;
-    rctx.vtcm_base             = (uint8_t *) octx->ctx->vtcm_base;
-    rctx.spad_per_thread       = kparams->spad_per_thread;
-    rctx.theta_cache_offset    = kparams->theta_cache_offset;
+    rctx.vtcm_base             = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_spad);
+    rctx.spad_per_thread       = layout.bytes_per_thread;
+    rctx.theta_cache_offset    = layout.theta_cache_size_aligned;
 
     if (src2) {
         dma_queue * dma_q = octx->ctx->dma[0];
         const size_t ff_size = src2->ne[0] * sizeof(float);
-        float * vtcm_freq_factors = (float *) (rctx.vtcm_base + kparams->freq_factors_offset);
+        float * vtcm_freq_factors = VTCM_LAYOUT_PTR(float, octx->ctx->vtcm_base, layout.off_freq_factors);
         dma_queue_push(dma_q, dma_make_data(vtcm_freq_factors, src2->data),
-                       kparams->freq_factors_size, 0, ff_size, 1);
+                       layout.freq_factors_size_aligned, 0, ff_size, 1);
         dma_queue_pop(dma_q);
         rctx.freq_factors = vtcm_freq_factors;
     }
@@ -791,7 +785,7 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) {
     rctx.src0_row_stride       = src0_row_stride;
     rctx.dst_row_size          = dst_row_size;
     rctx.dst_row_stride        = dst_row_stride;
-    rctx.src0_row_size_aligned = kparams->src0_row_size_aligned;
+    rctx.src0_row_size_aligned = layout.src0_row_size_aligned;
 
     rctx.src0_nrows            = nrows;
     rctx.nrows                 = nrows;

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

@@ -3,6 +3,7 @@
 
 #include "hex-common.h"
 #include "hex-fastdiv.h"
+#include "htp-vtcm.h"
 
 #define HTP_ROPE_SPAD_BLOCK  8
 #define HTP_ROPE_SPAD_NSLOTS 4
@@ -10,14 +11,7 @@
 
 struct htp_rope_kernel_params {
     uint32_t n_threads;
-    uint32_t src0_nrows;
-    uint32_t src0_nrows_per_thread;
     uint32_t vtcm_size;
-    uint32_t spad_per_thread;
-    uint32_t theta_cache_offset;
-    uint32_t src0_row_size_aligned;
-    uint32_t freq_factors_offset;
-    uint32_t freq_factors_size;
 
     struct fastdiv_values div_ne2_ne1;
     struct fastdiv_values div_ne1;
@@ -31,6 +25,8 @@ _Static_assert(sizeof(struct htp_rope_kernel_params) <= 128, "htp_rope_kernel_pa
 
 struct htp_rope_vtcm_layout {
     size_t total_bytes;
+    size_t off_spad;
+    size_t off_freq_factors;
     size_t bytes_per_thread;
     size_t theta_cache_size_aligned;
     size_t src0_row_size_aligned;
@@ -38,7 +34,7 @@ struct htp_rope_vtcm_layout {
 };
 
 static inline void htp_rope_vtcm_layout_build(
-    struct htp_rope_vtcm_layout * layout,
+    struct htp_rope_vtcm_layout * L,
     uint32_t ne00,
     uint32_t n_threads,
     uint32_t n_freq_factors
@@ -48,11 +44,15 @@ static inline void htp_rope_vtcm_layout_build(
     const size_t theta_cache_size_aligned  = hex_round_up((uint32_t) src0_row_size, 256);
     const size_t freq_factors_size_aligned = hex_round_up(n_freq_factors * sizeof(float), 256);
 
-    layout->src0_row_size_aligned     = src0_row_size_aligned;
-    layout->theta_cache_size_aligned  = theta_cache_size_aligned;
-    layout->freq_factors_size_aligned = freq_factors_size_aligned;
-    layout->bytes_per_thread          = theta_cache_size_aligned + HTP_ROPE_SPAD_NROWS * src0_row_size_aligned;
-    layout->total_bytes               = layout->bytes_per_thread * n_threads + freq_factors_size_aligned;
+    L->src0_row_size_aligned     = src0_row_size_aligned;
+    L->theta_cache_size_aligned  = theta_cache_size_aligned;
+    L->freq_factors_size_aligned = freq_factors_size_aligned;
+    L->bytes_per_thread          = theta_cache_size_aligned + HTP_ROPE_SPAD_NROWS * src0_row_size_aligned;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_spad, L->bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_freq_factors, freq_factors_size_aligned, n_freq_factors > 0);
+    L->total_bytes = off;
 }
 
 static inline uint8_t * rope_spad_slot(uint8_t * base, uint32_t slot, size_t row_size_aligned) {

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

@@ -79,8 +79,10 @@ static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsig
     const uint32_t ir1 = MIN(ir0 + dr, srctx->task_start + srctx->tasks);                                        \
     dma_queue * dma_q = octx->ctx->dma[ith];                                                                     \
     const struct htp_set_rows_vtcm_layout * vtcm_layout = &srctx->vtcm_layout;                                   \
-    uint8_t * vtcm_src0 = srctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread;   \
-    uint8_t * vtcm_dst  = srctx->vtcm_base + vtcm_layout->off_dst  + ith * vtcm_layout->dst_bytes_per_thread;    \
+    uint8_t * vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, srctx->vtcm_base, vtcm_layout->off_src0) +                    \
+                          ith * vtcm_layout->src0_bytes_per_thread;                                              \
+    uint8_t * vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, srctx->vtcm_base, vtcm_layout->off_dst) +                     \
+                          ith * vtcm_layout->dst_bytes_per_thread;                                               \
     const uint32_t src0_row_size = ne00 * sizeof(float);                                                         \
     const uint32_t dst_row_size  = htp_tensor_get_row_size(octx->dst->type, ne00);                               \
     const uint32_t nrows_per_thread = ir1 - ir0;                                                                 \
@@ -188,14 +190,6 @@ int op_set_rows(struct htp_ops_context * octx) {
     const struct htp_set_rows_kernel_params * kparams = (const struct htp_set_rows_kernel_params *)octx->kernel_params;
     set_rows_preamble;
 
-    if (octx->src[0]->type != HTP_TYPE_F32) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    if (octx->dst->type != HTP_TYPE_F32 && octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_Q8_0) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
     if (htp_tensor_is_extended(octx->src[1])) {
         return HTP_STATUS_NO_SUPPORT;
     }

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

@@ -2,18 +2,13 @@
 #define HTP_SET_ROWS_OPS_H
 
 #include "hex-fastdiv.h"
+#include "htp-tensor.h"
+#include "htp-vtcm.h"
 
 struct htp_set_rows_kernel_params {
-    int32_t  n_threads;
-    int32_t  total_tasks;
-    int32_t  tasks_per_thread;
-    int32_t  vtcm_size;
-
-    // Fastdiv helpers
-    struct fastdiv_values div_ne11;
-    struct fastdiv_values div_ne12;
-    struct fastdiv_values div_tasks_per_thread;
-    struct fastdiv_values div_ne02;
+    uint32_t n_threads;
+    uint32_t total_tasks;
+    uint32_t vtcm_size;
 };
 
 struct htp_set_rows_vtcm_layout {
@@ -29,40 +24,27 @@ struct htp_set_rows_vtcm_layout {
 };
 
 static inline void htp_set_rows_vtcm_layout_build(
-    struct htp_set_rows_vtcm_layout * vtcm_layout,
+    struct htp_set_rows_vtcm_layout * L,
     int dst_type,
     uint32_t ne00,
     uint32_t n_threads) {
 
-    size_t src0_row_size = ne00 * 4;
-    size_t dst_row_size = 0;
-    switch (dst_type) {
-        case 0: // HTP_TYPE_F32
-            dst_row_size = ne00 * 4;
-            break;
-        case 1: // HTP_TYPE_F16
-            dst_row_size = ne00 * 2;
-            break;
-        case 8: // HTP_TYPE_Q8_0
-            dst_row_size = (ne00 / 32) * 34;
-            break;
-        default:
-            dst_row_size = 0;
-            break;
-    }
+    size_t src0_row_size = ne00 * sizeof(float);
+    size_t dst_row_size  = htp_tensor_get_row_size(dst_type, ne00);
 
     size_t src0_row_size_aligned = (src0_row_size + 255) & ~255;
     size_t dst_row_size_aligned  = (dst_row_size + 255) & ~255;
 
-    vtcm_layout->src0_spad_half_size = src0_row_size_aligned;
-    vtcm_layout->dst_spad_half_size  = dst_row_size_aligned;
+    L->src0_spad_half_size = src0_row_size_aligned;
+    L->dst_spad_half_size  = dst_row_size_aligned;
 
-    vtcm_layout->src0_bytes_per_thread = src0_row_size_aligned * 2;
-    vtcm_layout->dst_bytes_per_thread  = dst_row_size_aligned * 2;
+    L->src0_bytes_per_thread = src0_row_size_aligned * 2;
+    L->dst_bytes_per_thread  = dst_row_size_aligned * 2;
 
-    vtcm_layout->off_src0 = 0;
-    vtcm_layout->off_dst  = vtcm_layout->off_src0 + vtcm_layout->src0_bytes_per_thread * n_threads;
-    vtcm_layout->total_bytes = vtcm_layout->off_dst + vtcm_layout->dst_bytes_per_thread * n_threads;
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread  * n_threads);
+    L->total_bytes = off;
 }
 
 #if defined(__cplusplus)

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

@@ -496,7 +496,7 @@ static int execute_op_softmax_f32(struct htp_ops_context * octx) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
-    if (kparams->kernel_id >= HTP_SOFTMAX_KERNEL_COUNT) {
+    if (kparams->kernel_type >= HTP_SOFTMAX_KERNEL_COUNT) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
@@ -532,34 +532,38 @@ static int execute_op_softmax_f32(struct htp_ops_context * octx) {
     const uint32_t n_threads = octx->n_threads;
     uint8_t * const vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
 
-    const uint32_t off_src0 = 0;
-    const uint32_t off_dst  = off_src0 + kparams->vtcm_src0_size_per_thread * kparams->n_threads;
-    const uint32_t off_src1 = off_dst  + kparams->vtcm_dst_size_per_thread  * kparams->n_threads;
+    struct htp_softmax_vtcm_layout layout;
+    htp_softmax_vtcm_layout_build(&layout,
+                                  src0->ne[0],
+                                  kparams->use_src1 ? octx->src[1]->ne[0] : 0,
+                                  kparams->use_src1 != 0,
+                                  kparams->use_f16 != 0,
+                                  kparams->n_threads);
 
     struct htp_softmax_context smctx = {
         .octx                  = octx,
         .kparams               = kparams,
-        .compute               = (void *) softmax_kernels[kparams->kernel_id],
+        .compute               = (void *) softmax_kernels[kparams->kernel_type],
 
         .data_src0             = src0->data,
         .data_src1             = kparams->use_src1 ? octx->src[1]->data : 0,
         .data_dst              = dst->data,
 
-        .vtcm_src0             = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, off_src0),
-        .vtcm_dst              = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, off_dst),
-        .vtcm_src1             = VTCM_LAYOUT_PTR_OPTIONAL(uint8_t, vtcm_base, off_src1, kparams->use_src1),
+        .vtcm_src0             = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_src0),
+        .vtcm_dst              = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_dst),
+        .vtcm_src1             = VTCM_LAYOUT_PTR_OPTIONAL(uint8_t, vtcm_base, layout.off_src1, kparams->use_src1),
 
-        .vtcm_src0_size_per_thread = kparams->vtcm_src0_size_per_thread,
-        .vtcm_src1_size_per_thread = kparams->vtcm_src1_size_per_thread,
-        .vtcm_dst_size_per_thread  = kparams->vtcm_dst_size_per_thread,
+        .vtcm_src0_size_per_thread = (uint32_t) layout.src0_bytes_per_thread,
+        .vtcm_src1_size_per_thread = (uint32_t) layout.src1_bytes_per_thread,
+        .vtcm_dst_size_per_thread  = (uint32_t) layout.dst_bytes_per_thread,
 
-        .src0_spad_half_size   = kparams->src0_spad_half_size,
-        .src1_spad_half_size   = kparams->src1_spad_half_size,
-        .dst_spad_half_size    = kparams->dst_spad_half_size,
+        .src0_spad_half_size   = (uint32_t) layout.src0_spad_half_size,
+        .src1_spad_half_size   = (uint32_t) layout.src1_spad_half_size,
+        .dst_spad_half_size    = (uint32_t) layout.dst_spad_half_size,
 
-        .src0_row_size_aligned = kparams->src0_row_size_aligned,
-        .src1_row_size_aligned = kparams->src1_row_size_aligned,
-        .dst_row_size_aligned  = kparams->dst_row_size_aligned,
+        .src0_row_size_aligned = (uint32_t) layout.src0_spad_half_size,
+        .src1_row_size_aligned = (uint32_t) layout.src1_spad_half_size,
+        .dst_row_size_aligned  = (uint32_t) layout.dst_spad_half_size,
 
         .use_f16               = kparams->use_f16 != 0,
 

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

@@ -7,8 +7,9 @@
 #include <math.h>
 #include "hex-fastdiv.h"
 #include "hex-common.h"
+#include "htp-vtcm.h"
 
-enum htp_softmax_kernel_id {
+enum htp_softmax_kernel_type {
     HTP_SOFTMAX_KERNEL_NOMASK = 0,
     HTP_SOFTMAX_KERNEL_MASK_F32,
     HTP_SOFTMAX_KERNEL_MASK_F16,
@@ -16,28 +17,14 @@ enum htp_softmax_kernel_id {
 };
 
 struct htp_softmax_kernel_params {
+    uint8_t  kernel_type;
     uint32_t n_threads;
-    uint32_t src0_nrows;
-    uint32_t src0_nrows_per_thread;
     uint32_t vtcm_size;
 
-    uint32_t vtcm_src0_size_per_thread;
-    uint32_t vtcm_src1_size_per_thread;
-    uint32_t vtcm_dst_size_per_thread;
-
-    uint32_t src0_row_size_aligned;
-    uint32_t src1_row_size_aligned;
-    uint32_t dst_row_size_aligned;
-
-    uint32_t src0_spad_half_size;
-    uint32_t src1_spad_half_size;
-    uint32_t dst_spad_half_size;
-
     uint32_t n_head;
     uint32_t n_head_log2;
     uint32_t use_src1;
     uint32_t use_f16;
-    uint32_t kernel_id;
 
     float    scale;
     float    max_bias;
@@ -72,7 +59,7 @@ struct htp_softmax_vtcm_layout {
 };
 
 static inline void htp_softmax_vtcm_layout_build(
-    struct htp_softmax_vtcm_layout * layout,
+    struct htp_softmax_vtcm_layout * L,
     uint32_t ne00,
     uint32_t ne10,
     bool use_src1,
@@ -87,20 +74,21 @@ static inline void htp_softmax_vtcm_layout_build(
     size_t dst_row_size_aligned  = hex_round_up(dst_row_size,  128);
     size_t src1_row_size_aligned = use_src1 ? hex_round_up(src1_row_size, 128) : 0;
 
-    layout->src0_spad_half_size = src0_row_size_aligned;
-    layout->dst_spad_half_size  = dst_row_size_aligned;
-    layout->src1_spad_half_size = src1_row_size_aligned;
+    L->src0_spad_half_size = src0_row_size_aligned;
+    L->dst_spad_half_size  = dst_row_size_aligned;
+    L->src1_spad_half_size = src1_row_size_aligned;
 
     // Double buffering: 2 half-buffers per thread
-    layout->src0_bytes_per_thread = src0_row_size_aligned * 2;
-    layout->dst_bytes_per_thread  = dst_row_size_aligned  * 2;
-    layout->src1_bytes_per_thread = src1_row_size_aligned * 2;
+    L->src0_bytes_per_thread = src0_row_size_aligned * 2;
+    L->dst_bytes_per_thread  = dst_row_size_aligned  * 2;
+    L->src1_bytes_per_thread = src1_row_size_aligned * 2;
 
-    layout->off_src0 = 0;
-    layout->off_dst  = layout->off_src0 + layout->src0_bytes_per_thread * n_threads;
-    layout->off_src1 = layout->off_dst  + layout->dst_bytes_per_thread  * n_threads;
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread  * n_threads);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_src1, L->src1_bytes_per_thread * n_threads, use_src1);
 
-    layout->total_bytes = layout->off_src1 + layout->src1_bytes_per_thread * n_threads;
+    L->total_bytes = off;
 }
 
 #endif // HTP_SOFTMAX_OPS_H

ggml/src/ggml-hexagon/htp/solve-tri-ops.c

@@ -13,14 +13,16 @@
 #include "htp-tensor.h"
 #include "hvx-types.h"
 #include "hvx-utils.h"
+#include "solve-tri-ops.h"
 
 struct htp_solve_tri_context {
-    struct htp_ops_context * octx;
-    uint32_t                 jobs_per_thread;
-    uint32_t                 total_jobs;
-    uint32_t                 job_start;
-    uint32_t                 k_chunks;
-    uint32_t                 col_block;
+    struct htp_ops_context *               octx;
+    const struct htp_solve_tri_kernel_params * kparams;
+    uint32_t                               jobs_per_thread;
+    uint32_t                               total_jobs;
+    uint32_t                               job_start;
+    uint32_t                               k_chunks;
+    uint32_t                               col_block;
 };
 
 static inline void solve_tri_row_scalar(const float * A_row,
@@ -80,6 +82,7 @@ static inline void solve_tri_row_hvx(const float * A_row,
 static void solve_tri_batch_thread_f32(unsigned int nth, unsigned int ith, void * data) {
     struct htp_solve_tri_context * sctx = (struct htp_solve_tri_context *) data;
     struct htp_ops_context *       octx = sctx->octx;
+    const struct htp_solve_tri_kernel_params * kparams = sctx->kparams;
 
     const struct htp_tensor * src0 = octx->src[0];  // A
     const struct htp_tensor * src1 = octx->src[1];  // B
@@ -91,7 +94,7 @@ static void solve_tri_batch_thread_f32(unsigned int nth, unsigned int ith, void
     const uint32_t ne02 = src0->ne[2];
 
     const uint32_t col_block = VLEN_FP32;
-    const uint32_t k_full    = (k / col_block) * col_block;
+    const uint32_t k_full    = k & ~(col_block - 1);
 
     const uint32_t start_batch = sctx->job_start + sctx->jobs_per_thread * ith;
     const uint32_t end_batch   = MIN(start_batch + sctx->jobs_per_thread, sctx->job_start + sctx->total_jobs);
@@ -100,7 +103,7 @@ static void solve_tri_batch_thread_f32(unsigned int nth, unsigned int ith, void
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start_batch);
 
     for (uint32_t batch = start_batch; batch < end_batch; ++batch) {
-        const uint32_t i03 = batch / ne02;
+        const uint32_t i03 = fastdiv(batch, &kparams->div_ne02);
         const uint32_t i02 = batch - i03 * ne02;
 
         const float * A_batch =
@@ -141,6 +144,7 @@ static void solve_tri_batch_thread_f32(unsigned int nth, unsigned int ith, void
 static void solve_tri_chunk_thread_f32(unsigned int nth, unsigned int ith, void * data) {
     struct htp_solve_tri_context * sctx = (struct htp_solve_tri_context *) data;
     struct htp_ops_context *       octx = sctx->octx;
+    const struct htp_solve_tri_kernel_params * kparams = sctx->kparams;
 
     const struct htp_tensor * src0 = octx->src[0];  // A
     const struct htp_tensor * src1 = octx->src[1];  // B
@@ -158,10 +162,10 @@ static void solve_tri_chunk_thread_f32(unsigned int nth, unsigned int ith, void
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start_job);
 
     for (uint32_t job = start_job; job < end_job; ++job) {
-        const uint32_t batch = job / sctx->k_chunks;
+        const uint32_t batch = fastdiv(job, &kparams->div_k_chunks);
         const uint32_t chunk = job - batch * sctx->k_chunks;
 
-        const uint32_t i03 = batch / ne02;
+        const uint32_t i03 = fastdiv(batch, &kparams->div_ne02);
         const uint32_t i02 = batch - i03 * ne02;
 
         const float * A_batch =
@@ -198,36 +202,17 @@ int op_solve_tri(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];  // A
     const struct htp_tensor * src1 = octx->src[1];  // B
     const struct htp_tensor * dst  = octx->dst;     // X
-
-    if (src0->type != HTP_TYPE_F32 || src1->type != HTP_TYPE_F32 || dst->type != HTP_TYPE_F32) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    // left=true, lower=true, uni=false only
-    if (src0->ne[0] != src0->ne[1]) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-    if (src0->ne[1] != src1->ne[1]) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-    if (src0->ne[2] != src1->ne[2] || src0->ne[3] != src1->ne[3]) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-    if (dst->ne[0] != src1->ne[0] || dst->ne[1] != src1->ne[1] || dst->ne[2] != src1->ne[2] ||
-        dst->ne[3] != src1->ne[3]) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
+    const struct htp_solve_tri_kernel_params * kparams =
+        (const struct htp_solve_tri_kernel_params *) octx->kernel_params;
 
     if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(src1) || htp_tensor_is_extended(dst)) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    const uint32_t k = src1->ne[0];
-
-    const uint32_t col_block     = VLEN_FP32;
-    const uint32_t k_chunks      = (k + col_block - 1) / col_block;
-    const uint32_t total_batches = src0->ne[2] * src0->ne[3];
-    const bool     batched       = total_batches >= (uint32_t) octx->n_threads;
+    const uint32_t col_block     = kparams->col_block;
+    const uint32_t k_chunks      = kparams->k_chunks;
+    const uint32_t total_batches = kparams->total_batches;
+    const bool     batched       = (kparams->batched != 0);
 
     FARF(HIGH, "solve-tri: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : batched %d\n",
          src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3],
@@ -239,9 +224,7 @@ int op_solve_tri(struct htp_ops_context * octx) {
         uint32_t njobs     = total_batches;
 
         if (octx->ctx->mdev.count > 1) {
-            const uint32_t batch_size = dst->nb[2];
-            const uint32_t batches_per_chunk = (batch_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(batch_size, HEX_L2_LINE_SIZE)) : 1;
-            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_batches, htp_tensor_mdev_data_aligned(dst) ? batches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_batches, htp_tensor_mdev_data_aligned(dst) ? kparams->batches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
             job_start = range.start;
             njobs     = range.count;
         }
@@ -250,11 +233,15 @@ int op_solve_tri(struct htp_ops_context * octx) {
             return HTP_STATUS_OK;
         }
 
-        // Batch-level parallelism
+        if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, njobs))) {
+            return HTP_STATUS_INVAL_PARAMS;
+        }
+
         const uint32_t n_threads = octx->n_threads;
 
         struct htp_solve_tri_context sctx = {
             .octx            = octx,
+            .kparams         = kparams,
             .jobs_per_thread = fastdiv(njobs + n_threads - 1, &octx->n_threads_div),
             .total_jobs      = njobs,
             .job_start       = job_start,
@@ -264,8 +251,7 @@ int op_solve_tri(struct htp_ops_context * octx) {
 
         work_queue_run(octx->ctx->work_queue, solve_tri_batch_thread_f32, &sctx, n_threads);
     } else {
-        // Chunk-level parallelism
-        const uint32_t total_jobs = total_batches * k_chunks;
+        const uint32_t total_jobs = kparams->total_jobs;
 
         uint32_t job_start = 0;
         uint32_t njobs     = total_jobs;
@@ -281,10 +267,15 @@ int op_solve_tri(struct htp_ops_context * octx) {
             return HTP_STATUS_OK;
         }
 
+        if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, njobs))) {
+            return HTP_STATUS_INVAL_PARAMS;
+        }
+
         const uint32_t n_threads = octx->n_threads;
 
         struct htp_solve_tri_context sctx = {
             .octx            = octx,
+            .kparams         = kparams,
             .jobs_per_thread = fastdiv(njobs + n_threads - 1, &octx->n_threads_div),
             .total_jobs      = njobs,
             .job_start       = job_start,

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

@@ -0,0 +1,26 @@
+#ifndef HTP_SOLVE_TRI_OPS_H
+#define HTP_SOLVE_TRI_OPS_H
+
+#include "hex-fastdiv.h"
+#include <stdint.h>
+
+struct htp_solve_tri_kernel_params {
+    uint32_t n_threads;
+    uint32_t total_batches;
+    uint32_t total_jobs;
+    uint32_t k_chunks;
+    uint32_t col_block;
+    uint32_t batched;
+    uint32_t batches_per_chunk;
+
+    struct fastdiv_values div_ne02;
+    struct fastdiv_values div_k_chunks;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_solve_tri_kernel_params) <= 128, "htp_solve_tri_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_solve_tri_kernel_params) <= 128, "htp_solve_tri_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_SOLVE_TRI_OPS_H

ggml/src/ggml-hexagon/htp/ssm-conv.c

@@ -135,6 +135,13 @@ static inline void hvx_ssm_conv_unpack_to_T(const float * raw, float * T, uint32
     }
 }
 
+static __attribute__((noinline)) void hvx_ssm_conv_silu_vtcm(float * restrict buf, uint32_t n_vecs) {
+    HVX_Vector * restrict vbuf = (HVX_Vector *) buf;
+    for (uint32_t i = 0; i < n_vecs; ++i) {
+        vbuf[i] = hvx_vec_silu_f32(vbuf[i]);
+    }
+}
+
 // Decode dot product specialization for d_conv == 4: multiply in the raw channel-major layout,
 // then deinterleave the products so each vector holds one tap of 32 channels, and sum.
 // Keeps both operands in DMA layout - no transpose, no scratch.
@@ -189,9 +196,10 @@ static void ssm_conv_thread_f32_decode(unsigned int nth, unsigned int ith, void
 
     dma_queue * dma_q = octx->ctx->dma[ith];
 
-    const uint32_t d_conv  = kparams->d_conv;
-    const uint32_t d_inner = kparams->d_inner;
-    const uint32_t n_s     = kparams->n_s;
+    const uint32_t d_conv     = kparams->d_conv;
+    const uint32_t d_inner    = kparams->d_inner;
+    const uint32_t n_s        = kparams->n_s;
+    const bool     apply_silu = kparams->silu != 0;
 
     const uint32_t dr  = scctx->nrows_per_thread;
     const uint32_t ir0 = scctx->row_start + dr * ith;
@@ -232,7 +240,7 @@ static void ssm_conv_thread_f32_decode(unsigned int nth, unsigned int ith, void
     // raw_dot keeps both operands in the DMA layout, so src1 needs no prep pass
     const bool raw_dot = (d_conv == 4) && (d_inner_per_thread % VLEN_FP32 == 0);
 
-    const uint32_t n_tiles   = (d_inner_per_thread + d_inner_tile - 1) / d_inner_tile;
+    const uint32_t n_tiles   = fastdiv(d_inner_per_thread + d_inner_tile - 1, &kparams->div_d_inner_tile);
     const uint32_t n_chunks  = n_s * n_tiles;
     const size_t   row_bytes = d_conv * sizeof(float);
 
@@ -295,21 +303,18 @@ static void ssm_conv_thread_f32_decode(unsigned int nth, unsigned int ith, void
             hvx_ssm_conv_unpack_to_T(src0_tile_raw[c & 1], src0_T, tile_n, tile_stride, d_conv);
 
             for (uint32_t cb = 0; cb < tile_n; cb += VLEN_FP32) {
-                const uint32_t cb_n = MIN(VLEN_FP32, tile_n - cb);
                 HVX_Vector acc = hvx_vec_splat_f32(0.0f);
                 for (uint32_t j = 0; j < d_conv; ++j) {
                     HVX_Vector x = *(const HVX_Vector *)(src0_T + j * tile_stride + cb);
                     HVX_Vector w = *(const HVX_Vector *)(src1_T + j * d_inner_stride + tile_off + cb);
                     acc          = Q6_Vqf32_vadd_Vqf32Vqf32(acc, Q6_Vqf32_vmpy_VsfVsf(x, w));
                 }
-                HVX_Vector y = Q6_Vsf_equals_Vqf32(acc);
-                if (cb_n == VLEN_FP32) {
-                    *(HVX_Vector *)(out + cb) = y;
-                } else {
-                    hvx_vec_store_u(out + cb, cb_n * sizeof(float), y);
-                }
+                *(HVX_Vector *)(out + cb) = Q6_Vsf_equals_Vqf32(acc);
             }
         }
+        if (apply_silu) {
+            hvx_ssm_conv_silu_vtcm(out, hex_round_up(tile_n, VLEN_FP32) / VLEN_FP32);
+        }
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i3);
 
         const dma_addr_t dst_ddr        = dst->data + i3 * dst_stride_seq_bytes + (ir0 + tile_off) * sizeof(float);
@@ -352,11 +357,12 @@ static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void
 
     dma_queue * dma_q = octx->ctx->dma[ith];
 
-    const uint32_t d_conv  = kparams->d_conv;
-    const uint32_t d_inner = kparams->d_inner;
-    const uint32_t n_t     = kparams->n_t;
-    const uint32_t n_s     = kparams->n_s;
-    const uint32_t ncs     = src0->ne[0];
+    const uint32_t d_conv     = kparams->d_conv;
+    const uint32_t d_inner    = kparams->d_inner;
+    const uint32_t n_t        = kparams->n_t;
+    const uint32_t n_s        = kparams->n_s;
+    const uint32_t ncs        = src0->ne[0];
+    const bool     apply_silu = kparams->silu != 0;
 
     const uint32_t dr  = scctx->nrows_per_thread;
     const uint32_t ir0 = scctx->row_start + dr * ith;
@@ -397,7 +403,7 @@ static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void
 
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
-    const uint32_t n_tiles   = (d_inner_per_thread + d_inner_tile - 1) / d_inner_tile;
+    const uint32_t n_tiles   = fastdiv(d_inner_per_thread + d_inner_tile - 1, &kparams->div_d_inner_tile);
     const uint32_t n_chunks  = n_s * n_tiles;
     const size_t   row_bytes = ncs * sizeof(float);
 
@@ -513,10 +519,14 @@ static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void
             htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) tile_off);
         }
 
+        if (apply_silu) {
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) tile_off);
+            hvx_ssm_conv_silu_vtcm(out, n_t * (dst_tile_stride / C_TILE));
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) tile_off);
+        }
+
         // Writeback dst_tile VTCM -> DDR via 2D DMA
-        dma_queue_push(dma_q,
-                       dma_make_data(dst->data + i3 * dst_stride_seq_bytes + (ir0 + tile_off) * sizeof(float),
-                                     (uint8_t *) out),
+        dma_queue_push(dma_q, dma_make_data(dst->data + i3 * dst_stride_seq_bytes + (ir0 + tile_off) * sizeof(float), (uint8_t *) out),
                        dst_stride_token_bytes, dst_tile_stride * sizeof(float), tile_n * sizeof(float), n_t);
 
         if (c + 2 < n_chunks) {
@@ -547,13 +557,8 @@ int op_ssm_conv_f32(struct htp_ops_context * octx) {
     const struct htp_tensor * src1 = octx->src[1];
     const struct htp_tensor * dst  = octx->dst;
 
-    if (src0->type != HTP_TYPE_F32 || src1->type != HTP_TYPE_F32 || dst->type != HTP_TYPE_F32) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
     const struct htp_ssm_conv_kernel_params * kparams = (const struct htp_ssm_conv_kernel_params *) octx->kernel_params;
 
-
     if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
@@ -582,26 +587,30 @@ int op_ssm_conv_f32(struct htp_ops_context * octx) {
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
-    const uint32_t n_threads = octx->n_threads;
+    const uint32_t n_threads          = octx->n_threads;
+    const uint32_t raw_rpt            = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    const uint32_t d_inner_per_thread = hex_round_up(raw_rpt, VLEN_FP32);
 
-    octx->src0_spad.size_per_thread = kparams->vtcm_src0_size_per_thread;
-    octx->src1_spad.size_per_thread = kparams->vtcm_src1_size_per_thread;
-    octx->dst_spad.size_per_thread  = kparams->vtcm_dst_size_per_thread;
+    struct htp_ssm_conv_vtcm_layout L;
+    htp_ssm_conv_vtcm_layout_build(&L, d_inner_per_thread, kparams->d_conv, kparams->n_t, kparams->d_inner_tile, n_threads);
 
-    octx->src0_spad.size = kparams->vtcm_src0_size;
-    octx->src1_spad.size = kparams->vtcm_src1_size;
-    octx->dst_spad.size  = kparams->vtcm_dst_size;
+    uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
 
-    octx->src0_spad.data = octx->ctx->vtcm_base;
-    octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size;
-    octx->dst_spad.data  = octx->src1_spad.data + octx->src1_spad.size;
+    octx->src0_spad.size_per_thread = L.src0_bytes_per_thread;
+    octx->src1_spad.size_per_thread = L.src1_bytes_per_thread;
+    octx->dst_spad.size_per_thread  = L.dst_bytes_per_thread;
+
+    octx->src0_spad.size = L.src0_bytes;
+    octx->src1_spad.size = L.src1_bytes;
+    octx->dst_spad.size  = L.dst_bytes;
+
+    octx->src0_spad.data = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+    octx->src1_spad.data = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+    octx->dst_spad.data  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
     octx->src0_spad.src  = NULL;
     octx->src1_spad.src  = NULL;
     octx->dst_spad.src   = NULL;
 
-    const uint32_t raw_rpt            = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
-    const uint32_t d_inner_per_thread = hex_round_up(raw_rpt, VLEN_FP32);
-
     struct htp_ssm_conv_context scctx = {
         .octx             = octx,
         .kparams          = kparams,

ggml/src/ggml-hexagon/htp/ssm-conv.h

@@ -3,8 +3,10 @@
 
 #include <stdint.h>
 
+#include "hex-common.h"
 #include "hex-fastdiv.h"
 #include "htp-ops.h"
+#include "htp-vtcm.h"
 
 struct htp_ssm_conv_kernel_params {
     uint32_t n_threads;
@@ -14,16 +16,62 @@ struct htp_ssm_conv_kernel_params {
     uint32_t n_s;
     uint32_t d_inner_tile;
 
-    uint32_t vtcm_src0_size_per_thread;
-    uint32_t vtcm_src1_size_per_thread;
-    uint32_t vtcm_dst_size_per_thread;
-
-    uint32_t vtcm_src0_size;
-    uint32_t vtcm_src1_size;
-    uint32_t vtcm_dst_size;
     uint32_t vtcm_size;
+    uint32_t silu;
+
+    struct fastdiv_values div_d_inner_tile;
 };
 
+struct htp_ssm_conv_vtcm_layout {
+    size_t total_bytes;
+    size_t off_src0;
+    size_t off_src1;
+    size_t off_dst;
+    size_t src0_bytes_per_thread;
+    size_t src1_bytes_per_thread;
+    size_t dst_bytes_per_thread;
+    size_t src0_bytes;
+    size_t src1_bytes;
+    size_t dst_bytes;
+};
+
+static inline void htp_ssm_conv_vtcm_layout_build(
+    struct htp_ssm_conv_vtcm_layout * L,
+    uint32_t d_inner_per_thread,
+    uint32_t d_conv,
+    uint32_t n_t,
+    uint32_t d_inner_tile,
+    uint32_t n_threads
+) {
+    const size_t src1_raw_bytes = hex_round_up((size_t) d_inner_per_thread * d_conv * sizeof(float), 128) + 128;
+    const size_t src1_T_bytes   = hex_round_up((size_t) d_conv * d_inner_per_thread * sizeof(float), 128);
+
+    L->src1_bytes_per_thread = src1_raw_bytes + src1_T_bytes;
+
+    if (n_t == 1) {
+        const size_t src0_tile_raw_bytes = hex_round_up((size_t) d_inner_tile * d_conv * sizeof(float), 128);
+        const size_t src0_T_bytes        = hex_round_up((size_t) d_conv * d_inner_tile * sizeof(float), 128);
+        L->src0_bytes_per_thread         = 2 * src0_tile_raw_bytes + src0_T_bytes;
+        L->dst_bytes_per_thread          = 2 * hex_round_up((size_t) d_inner_tile * sizeof(float), 128);
+    } else {
+        const uint32_t ncs          = d_conv - 1 + n_t;
+        const size_t   src0_block_T = hex_round_up((size_t) ncs * 32 * sizeof(float), 128);
+        const size_t   src0_tile_raw = hex_round_up((size_t) d_inner_tile * ncs * sizeof(float), 128);
+        L->src0_bytes_per_thread    = 2 * src0_tile_raw + src0_block_T;
+        L->dst_bytes_per_thread     = 2 * hex_round_up((size_t) d_inner_tile * n_t * sizeof(float), 128);
+    }
+
+    L->src0_bytes = L->src0_bytes_per_thread * n_threads;
+    L->src1_bytes = L->src1_bytes_per_thread * n_threads;
+    L->dst_bytes  = L->dst_bytes_per_thread * n_threads;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_src1, L->src1_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes);
+    L->total_bytes = off;
+}
+
 #if defined(__cplusplus)
 static_assert(sizeof(struct htp_ssm_conv_kernel_params) <= 128, "htp_ssm_conv_kernel_params is too large for kernel_params blob");
 #else

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

@@ -18,6 +18,7 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "sum-rows-ops.h"
 
 #define sum_rows_preamble                         \
     const struct htp_tensor *src0 = octx->src[0]; \
@@ -52,7 +53,6 @@ struct sum_rows_context {
     size_t          dst_stride;
     uint32_t        rows_per_thread;
     uint32_t        total_rows;
-    bool            opt_path;
 };
 
 static void sum_rows_thread_f32(unsigned int nth, unsigned int ith, void *data) {
@@ -71,7 +71,6 @@ static void sum_rows_thread_f32(unsigned int nth, unsigned int ith, void *data)
     const size_t   src_stride = smctx->src_stride;
     const size_t   dst_stride = smctx->dst_stride;
     const uint32_t ne00       = smctx->ne00;
-    const bool     opt_path   = smctx->opt_path;
 
     const float * restrict src_th = (const float *) (smctx->src_data + (start_row * src_stride));
     float       * restrict dst_th = (float *)       (smctx->dst_data + (start_row * dst_stride));
@@ -90,11 +89,7 @@ static void sum_rows_thread_f32(unsigned int nth, unsigned int ith, void *data)
             hex_l2fetch((const uint8_t *) src_local + src_stride, src_stride, src_stride, 1);
         }
 
-        if (opt_path) {
-            *dst_local = hvx_reduce_sum_f32_a((const uint8_t *) src_local, ne00);
-        } else {
-            *dst_local = hvx_reduce_sum_f32((const uint8_t *) src_local, ne00);
-        }
+        *dst_local = hvx_reduce_sum_f32((const uint8_t *) src_local, ne00);
     }
 
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start_row);
@@ -111,6 +106,13 @@ int op_sum_rows(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
+    const struct htp_sum_rows_kernel_params * kparams =
+        (const struct htp_sum_rows_kernel_params *) octx->kernel_params;
+
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t src0_nrows      = ne01 * ne02 * ne03;
     const size_t dst_data_row_size = dst->ne[0] * sizeof(float);
 
@@ -129,14 +131,9 @@ int op_sum_rows(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    const uint32_t n_threads = octx->n_threads;
+    const uint32_t n_threads       = octx->n_threads;
     const uint32_t rows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
 
-    bool opt_path = false;
-    if ((0 == hex_is_aligned((void *) src0->data, VLEN)) && !(nb01 & (VLEN - 1))) {
-        opt_path = true;
-    }
-
     struct sum_rows_context smctx = {
         .octx            = octx,
         .src_data        = (const uint8_t *) src0->data + row_start * nb01,
@@ -146,7 +143,6 @@ int op_sum_rows(struct htp_ops_context * octx) {
         .dst_stride      = nb1,
         .rows_per_thread = rows_per_thread,
         .total_rows      = nrows,
-        .opt_path        = opt_path,
     };
 
     work_queue_run(octx->ctx->work_queue, sum_rows_thread_f32, &smctx, n_threads);
@@ -197,6 +193,13 @@ int op_sum(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
+    const struct htp_sum_rows_kernel_params * kparams =
+        (const struct htp_sum_rows_kernel_params *) octx->kernel_params;
+
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     if (octx->ctx->mdev.count > 1 && octx->ctx->mdev.idx > 0) {
         return HTP_STATUS_OK;
     }
@@ -207,9 +210,8 @@ int op_sum(struct htp_ops_context * octx) {
         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);
+    const uint32_t n_threads        = octx->n_threads;
+    const uint32_t elems_per_thread = kparams->elems_per_thread;
 
     struct sum_context sctx = {
         .octx             = octx,
@@ -324,6 +326,13 @@ int op_argmax(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
+    const struct htp_sum_rows_kernel_params * kparams =
+        (const struct htp_sum_rows_kernel_params *) octx->kernel_params;
+
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t ne00 = src0->ne[0];
     const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
 
@@ -341,9 +350,8 @@ int op_argmax(struct htp_ops_context * octx) {
             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);
+        const uint32_t n_threads        = octx->n_threads;
+        const uint32_t elems_per_thread = kparams->elems_per_thread;
 
         struct argmax_context actx = {
             .octx             = octx,
@@ -389,8 +397,8 @@ int op_argmax(struct htp_ops_context * octx) {
         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;
+    const uint32_t n_threads       = octx->n_threads;
+    const uint32_t rows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
 
     struct argmax_context actx = {
         .octx             = octx,

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

@@ -0,0 +1,17 @@
+#ifndef HTP_SUM_ROWS_OPS_H
+#define HTP_SUM_ROWS_OPS_H
+
+#include <stdint.h>
+
+struct htp_sum_rows_kernel_params {
+    uint32_t n_threads;
+    uint32_t elems_per_thread;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_sum_rows_kernel_params) <= 128, "htp_sum_rows_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_sum_rows_kernel_params) <= 128, "htp_sum_rows_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_SUM_ROWS_OPS_H