llama : optimize vocabulary loading (#30679)

* llama : optimize vocabulary loading

* Update src/unicode.h

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>

* CR : added reserve in gguf_kv_to_str

* CR : trailing space fix

---------

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>

.github/workflows/build-webgpu.yml

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

common/common.cpp

@@ -1175,6 +1175,7 @@ static const std::map<common_decision_type, std::string> COMMON_DECISION_TYPE_NA
     { COMMON_DECISION_TYPE_LFM2_D1,        "lfm2-d1"       },
     { COMMON_DECISION_TYPE_LFM2_D1_OMNI,   "lfm2-d1-omni"  },
     { COMMON_DECISION_TYPE_GRANITE_SWITCH, "granite-switch"},
+    { COMMON_DECISION_TYPE_H2O_LIGHTNING,  "h2o-lightning" },
     // for future models, ALWAYS add org name prefix to avoid name collisions, e.g. "orgname-modelname"
 };
 
@@ -1234,6 +1235,9 @@ common_gguf_info common_get_gguf_info(const std::string & fname) {
         info.n_ctx_train = gguf_get_val_u32(gguf_ctx.get(), ctx_id);
     }
 
+    // same key as llama_model_chat_template(model, nullptr)
+    info.has_chat_tmpl = gguf_find_key(gguf_ctx.get(), "tokenizer.chat_template") >= 0;
+
     return info;
 }
 

common/common.h

@@ -975,6 +975,7 @@ enum common_decision_type {
     COMMON_DECISION_TYPE_LFM2_D1, // same as openjev, the labels depend on the question type
     COMMON_DECISION_TYPE_LFM2_D1_OMNI, // same as laya, other prompt layout
     COMMON_DECISION_TYPE_GRANITE_SWITCH, // same as openjev, run by the decision aLoRA of granite switch
+    COMMON_DECISION_TYPE_H2O_LIGHTNING, // same as openjev, the labels are read after an answer prefill
     COMMON_DECISION_TYPE_UNKNOWN, // a decision model of a type that is not supported
 };
 
@@ -984,6 +985,7 @@ common_decision_type common_get_decision_type(const struct llama_model * model);
 struct common_gguf_info {
     common_decision_type decision_type = COMMON_DECISION_TYPE_UNKNOWN; // UNKNOWN if the file is missing, unreadable, or invalid
     uint32_t             n_ctx_train   = 0;                            // 0 if unknown
+    bool                 has_chat_tmpl = false;                        // has the default tokenizer.chat_template
 };
 
 common_gguf_info common_get_gguf_info(const std::string & fname);

common/speculative.cpp

@@ -16,6 +16,7 @@
 #include <cassert>
 #include <cmath>
 #include <cstring>
+#include <fstream>
 #include <iomanip>
 #include <map>
 #include <cinttypes>
@@ -2152,9 +2153,6 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
 
     uint16_t n_draft;
 
-    bool save_dynamic;
-    bool save_static;
-
     struct seq_info {
         size_t cache_size = 0; // number of tokens in n-gram cache
 
@@ -2168,54 +2166,72 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
     common_speculative_impl_ngram_cache(
             const common_params_speculative & params,
             uint32_t n_seq,
-            uint16_t n_draft,
-            const std::string & path_static,
-            const std::string & path_dynamic,
-            bool save_dynamic,
-            bool save_static)
+            uint16_t n_draft)
         : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, n_seq, n_draft)
         , params(params.ngram_cache)
         , n_draft(n_draft)
-        , save_dynamic(save_dynamic)
-        , save_static(save_static)
     {
+        const auto & ngram = params.ngram_cache;
+
         SPC_TRC("%s", "adding speculative implementation 'ngram-cache'\n");
         SPC_TRC("- n_draft=%d, cache_static=%s, cache_dynamic=%s\n",
                 n_draft,
-                path_static.empty() ? "none" : path_static.c_str(),
-                path_dynamic.empty() ? "none" : path_dynamic.c_str());
+                ngram.lookup_cache_static.empty() ? "none" : ngram.lookup_cache_static.c_str(),
+                ngram.lookup_cache_dynamic.empty() ? "none" : ngram.lookup_cache_dynamic.c_str());
 
         sinfos.resize(n_seq);
 
-        if (!path_static.empty()) {
+        if (!ngram.lookup_cache_static.empty()) {
             try {
-                auto ngram_cache_static = common_ngram_cache_load(path_static);
+                auto ngram_cache_static = common_ngram_cache_load(ngram.lookup_cache_static);
 
                 for (auto & sinfo : sinfos) {
                     sinfo.ngram_cache_static = ngram_cache_static;
                 }
-            } catch (...) {
-                SPC_ERR("failed to open static lookup cache: %s", path_static.c_str());
+            } catch (const std::ifstream::failure &) {
+                SPC_ERR("failed to open static lookup cache: %s\n", ngram.lookup_cache_static.c_str());
                 GGML_ABORT("Couldn't read static lookup cache");
             }
         }
 
-        if (!path_dynamic.empty()) {
+        if (!ngram.lookup_cache_dynamic.empty()) {
             try {
-                auto ngram_cache_dynamic = common_ngram_cache_load(path_dynamic);
+                auto ngram_cache_dynamic = common_ngram_cache_load(ngram.lookup_cache_dynamic);
 
                 for (auto & sinfo : sinfos) {
                     sinfo.ngram_cache_dynamic = ngram_cache_dynamic;
                 }
-            } catch (...) {
-                SPC_ERR("failed to open dynamic lookup cache: %s", path_dynamic.c_str());
-                GGML_ABORT("Couldn't read dynamic lookup cache");
+            } catch (const std::ifstream::failure &) {
+                std::error_code ec;
+                const bool exists = std::filesystem::exists(ngram.lookup_cache_dynamic, ec);
+                if (!exists && !ec) {
+                    // missing file is created on destruction
+                    SPC_WRN("dynamic lookup cache file does not exist: %s, starting with empty cache\n",
+                            ngram.lookup_cache_dynamic.c_str());
+                } else {
+                    SPC_ERR("failed to open dynamic lookup cache: %s\n", ngram.lookup_cache_dynamic.c_str());
+                    GGML_ABORT("Couldn't read dynamic lookup cache");
+                }
             }
         }
     }
 
-    void begin(llama_seq_id /*seq_id*/, const llama_tokens & /*prompt*/) override {
-        // noop
+    ~common_speculative_impl_ngram_cache() override {
+        if (params.lookup_cache_dynamic.empty() || sinfos.empty()) {
+            return;
+        }
+
+        auto & dynamic_cache = sinfos.front().ngram_cache_dynamic;
+        for (auto & sinfo : sinfos) {
+            common_ngram_cache_merge(dynamic_cache, sinfo.ngram_cache_context);
+        }
+
+        SPC_INF("saving dynamic lookup cache at %s\n", params.lookup_cache_dynamic.c_str());
+        common_ngram_cache_save(dynamic_cache, params.lookup_cache_dynamic);
+    }
+
+    void begin(llama_seq_id seq_id, const llama_tokens & /*prompt*/) override {
+        sinfos[seq_id].cache_size = 0;
     }
 
     void draft_one(
@@ -2226,22 +2242,6 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
 
         const auto & prompt = *dparams.prompt;
 
-        if (sinfo.cache_size < prompt.size() + 1) {
-            llama_tokens tokens_new;
-            tokens_new.reserve(prompt.size() + 1 - sinfo.cache_size);
-            for (size_t j = sinfo.cache_size; j < prompt.size(); ++j) {
-                tokens_new.push_back(prompt[j]);
-            }
-            tokens_new.push_back(dparams.id_last); // add the last token
-
-            // Update context ngram cache with new dparams.prompt:
-            common_ngram_cache_update(
-                    sinfo.ngram_cache_context,
-                    LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX,
-                    tokens_new, tokens_new.size(), false);
-            sinfo.cache_size = prompt.size() + 1;
-        }
-
         llama_tokens inp;
         inp.reserve(prompt.size() + 1);
         for (size_t j = 0; j < prompt.size(); ++j) {
@@ -2249,6 +2249,16 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
         }
         inp.push_back(dparams.id_last);
 
+        if (sinfo.cache_size < inp.size()) {
+            // Update context ngram cache with the full history so n-grams
+            // ending at the new tokens have their preceding context
+            common_ngram_cache_update(
+                    sinfo.ngram_cache_context,
+                    LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX,
+                    inp, inp.size() - sinfo.cache_size, false);
+            sinfo.cache_size = inp.size();
+        }
+
         result.push_back(dparams.id_last);
 
         common_ngram_cache_draft(
@@ -2309,20 +2319,12 @@ static common_ngram_map get_common_ngram_map(
     return common_ngram_map(size_key, size_value, key_only, min_hits);
 }
 
-static common_speculative_impl_ngram_cache create_state_ngram_cache(
+static std::unique_ptr<common_speculative_impl_ngram_cache> create_state_ngram_cache(
         const common_speculative_config & config,
-        uint32_t n_seq,
-        const std::string & path_static,
-        const std::string & path_dynamic) {
+        uint32_t n_seq) {
     uint16_t n_draft = 8; // TODO get from config?
 
-    // TODO bool param in common/common.h to set save_static/save_dynamic?
-    bool save_static = false;
-    bool save_dynamic = false;
-
-    common_speculative_impl_ngram_cache state(config.params, n_seq, n_draft, path_static, path_dynamic, save_static, save_dynamic);
-
-    return state;
+    return std::make_unique<common_speculative_impl_ngram_cache>(config.params, n_seq, n_draft);
 }
 
 std::string common_speculative_type_name_str(const std::vector<common_speculative_type> & types) {
@@ -2822,11 +2824,7 @@ common_speculative * common_speculative_init(common_params_speculative & params,
                 break;
             }
             case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: {
-                auto state = create_state_ngram_cache(
-                        config, n_seq,
-                        params.ngram_cache.lookup_cache_static,
-                        params.ngram_cache.lookup_cache_dynamic);
-                impls.push_back(std::make_unique<common_speculative_impl_ngram_cache>(state));
+                impls.push_back(create_state_ngram_cache(config, n_seq));
                 break;
             }
             default:

conversion/__init__.py

@@ -200,7 +200,9 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "MiniMaxM3SparseForCausalLM": "minimax",
     "MiniMaxM3SparseForConditionalGeneration": "minimax",
     "Ministral3ForCausalLM": "mistral3",
+    "Ministral3Model": "mistral3",
     "Mistral3ForConditionalGeneration": "mistral3",
+    "Mistral3Model": "mistral3",
     "VoxtralRealtimeForConditionalGeneration": "mistral3",
     "MistralForCausalLM": "llama",
     "MixtralForCausalLM": "llama",
@@ -226,6 +228,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "MuseSparkForConditionalGeneration": "muse_spark",
     "OpenELMForCausalLM": "openelm",
     "OpenJevModel": "qwen",
+    "H2OLightningModel": "h2o_lightning",
     "OrionForCausalLM": "orion",
     "PLMForCausalLM": "plm",
     "PLaMo2ForCausalLM": "plamo",
@@ -357,6 +360,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
     "MiniCPMV4_6ForConditionalGeneration": "minicpm",
     "MiniCPMV4_7ForConditionalGeneration": "minicpm",
     "Mistral3ForConditionalGeneration": "llava",
+    "Mistral3Model": "llava",
     "NemotronH_Omni_Reasoning_V3": "nemotron",
     "NemotronH_Nano_VL_V2": "nemotron",
     "MuseGlimmerForConditionalGeneration": "muse_glimmer",
@@ -375,6 +379,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
     "Qwen3VLForConditionalGeneration": "qwen3vl",
     "Qwen3VLMoeForConditionalGeneration": "qwen3vl",
     "OpenJevModel": "qwen3vl",
+    "H2OLightningModel": "h2o_lightning",
     "Qwen3_5ForConditionalGeneration": "qwen3vl",
     "Qwen3_5MoeForConditionalGeneration": "qwen3vl",
     "Qwen4ExpForConditionalGeneration": "qwen4exp",

conversion/gemma.py

@@ -999,6 +999,8 @@ class Gemma4VisionAudioModel(MmprojModel):
         assert self.hparams_vision is not None
         self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.GEMMA4V)
         self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams_vision.get("layer_norm_eps", 1e-6))
+        assert self.hparams_vision.get("hidden_activation", "gelu_pytorch_tanh") == "gelu_pytorch_tanh"
+        self.gguf_writer.add_vision_use_gelu(True)
 
         # audio params
         if self.has_audio_encoder:

conversion/h2o_lightning.py

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

conversion/lfm2.py

@@ -374,6 +374,16 @@ class LFM2VLModel(MmprojModel):
             if resize_algo is None:
                 raise ValueError(f"unsupported resample: {resample}")
             self.gguf_writer.add_vision_image_resize_algo(resize_algo)
+        # tiling params from processor_config.json (do_image_splitting=False -> single tile, i.e. no tiling)
+        preproc = self.preprocessor_config
+        if preproc.get("do_image_splitting", True):
+            self.gguf_writer.add_vision_preproc_min_tiles(preproc.get("min_tiles", 2))
+            self.gguf_writer.add_vision_preproc_max_tiles(preproc.get("max_tiles", 10))
+        else:
+            self.gguf_writer.add_vision_preproc_min_tiles(1)
+            self.gguf_writer.add_vision_preproc_max_tiles(1)
+        self.gguf_writer.add_vision_preproc_tile_size(preproc.get("tile_size", 512))
+        self.gguf_writer.add_vision_preproc_max_pixels_tolerance(preproc.get("max_pixels_tolerance", 2.0))
 
     @classmethod
     def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:

conversion/llava.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import json
 
-from typing import Iterable, TYPE_CHECKING
+from typing import Callable, Iterable, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from torch import Tensor
@@ -128,3 +128,24 @@ class LlavaVisionModel(MmprojModel):
             yield from super().modify_tensors(img_break_embd, name, bid)
 
         return # skip other tensors
+
+
+@ModelBase.register("Mistral3Model")
+@ModelBase.example("nvidia/Nemotron-3.5-Embed-1B-BF16")
+class Mistral3EmbeddingVisionModel(LlavaVisionModel):
+    def __init__(self, dir_model, *args, **kwargs):
+        hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False)
+        text_config = hparams.get("text_config", hparams)
+        if text_config.get("is_causal") is not False:
+            raise NotImplementedError("Expected a Nemotron Embed model for Mistral3Model, but is_causal is not False")
+        super().__init__(dir_model, *args, **kwargs)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        if item[0].endswith("embed_tokens.weight"):
+            return ModelBase.filter_tensors(item)
+        return super().filter_tensors(item)
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_vision_preproc_image_size(self.preprocessor_config["size"]["longest_edge"])

conversion/mistral3.py

@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, gguf
+from .base import ModelBase, TextModel, get_model_architecture, gguf
 
 from .deepseek import DeepseekV2Model
 from .llama import LlamaModel
@@ -13,9 +13,11 @@ from .llama import LlamaModel
 
 @ModelBase.register(
     "Mistral3ForConditionalGeneration",
+    "Mistral3Model",
     "Ministral3ForCausalLM",
+    "Ministral3Model",
 )
-@ModelBase.example("mistralai/Mistral-Small-3.1-24B-Instruct-2503", "hf-tiny-v2/tiny-random-Ministral3ForCausalLM")
+@ModelBase.example("mistralai/Mistral-Small-3.1-24B-Instruct-2503", "hf-tiny-v2/tiny-random-Ministral3ForCausalLM", "nvidia/Nemotron-3-Embed-1B-BF16", "nvidia/Nemotron-3.5-Embed-1B-BF16")
 class Mistral3Model(TextModel):
     class Ministral3Model(LlamaModel):
         model_arch = gguf.MODEL_ARCH.MISTRAL3
@@ -29,6 +31,18 @@ class Mistral3Model(TextModel):
                 self.gguf_writer.add_rope_scaling_yarn_log_mul(rope_params["mscale_all_dim"])
                 self.gguf_writer.add_attn_temperature_scale(rope_params["llama_4_scaling_beta"])
 
+    class Mistral3EmbedNemotronModel(Ministral3Model):
+        model_arch = gguf.MODEL_ARCH.MISTRAL3
+
+        def set_vocab(self):
+            super().set_vocab()
+            self.gguf_writer.add_add_bos_token(False)
+            self.gguf_writer.add_add_eos_token(False)
+
+        def set_gguf_parameters(self):
+            super().set_gguf_parameters()
+            self._try_set_pooling_type()
+
     class Mistral4Model(DeepseekV2Model):
         model_arch = gguf.MODEL_ARCH.MISTRAL4
         skip_mtp = False # model contains no MTP layers, so no need to skip
@@ -42,12 +56,20 @@ class Mistral3Model(TextModel):
     model_arch = gguf.MODEL_ARCH.MISTRAL3 # unused
     impl: TextModel
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self, dir_model, *args, **kwargs):
+        hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False)
+        hf_arch = get_model_architecture(hparams, self.model_type)
+        if hf_arch in ("Ministral3Model", "Mistral3Model"):
+            text_config = hparams.get("text_config", hparams)
+            if text_config.get("is_causal") is not False:
+                raise NotImplementedError(f"Expected a Nemotron Embed model for {hf_arch}, but is_causal is not False")
+        super().__init__(dir_model, *args, **kwargs)
         if self.hparams.get("model_type") == "mistral4":
-            self.impl = Mistral3Model.Mistral4Model(*args, **kwargs)
+            self.impl = Mistral3Model.Mistral4Model(dir_model, *args, **kwargs)
+        elif self.hf_arch in ("Ministral3Model", "Mistral3Model"):
+            self.impl = Mistral3Model.Mistral3EmbedNemotronModel(dir_model, *args, **kwargs)
         else:
-            self.impl = Mistral3Model.Ministral3Model(*args, **kwargs)
+            self.impl = Mistral3Model.Ministral3Model(dir_model, *args, **kwargs)
 
     def set_vocab(self):
         self.impl.set_vocab()

docs/backend/zDNN.md

@@ -2,6 +2,7 @@
 
 > [!WARNING]
 > **Note:** zDNN is **not** the same as ZenDNN.
+>
 > - **zDNN** (this page): IBM's Deep Neural Network acceleration library for IBM Z & LinuxONE Mainframes
 > - **ZenDNN**: AMD's deep learning library for AMD EPYC CPUs ([see ZenDNN documentation](ZenDNN.md))
 
@@ -46,6 +47,14 @@ The IBM zDNN backend has the following CMake options that control the behaviour
 | --------------------- | ------------- | -------------------------------------------------------------------------------------- |
 | `GGML_ZDNN_HIGH_PREC` | `0`           | Set to `1` to requantize Q8_0 weights in smaller groups, for lower loss at lower speed |
 
+## Performance Tunables
+
+| Flag Option             | Default Value | Description                                            |
+| ----------------------- | ------------- | ------------------------------------------------------ |
+| `--flash-attn` or `-fa` | `auto`        | Set to `on` to run attention on the NNPA (Recommended) |
+
+Note: Set `-fa on` explicitly. Flash Attention fuses each layer's attention (the KQ matrix multiplication, softmax and KQV matrix multiplication) into one `FLASH_ATTN_EXT` operation, which zDNN runs on the NNPA instead of the CPU. The speedup grows with the context length. With `auto`, llama.cpp currently turns Flash Attention off, because the operation runs on a different device than the layer. zDNN runs it when the KV cache is F16 (the default), the model uses no attention sinks, ALiBi or logit softcapping, and there is one KV cache stream (one sequence, or `--kv-unified`). Otherwise it falls back to the CPU.
+
 ## 1. Install zDNN Library
 
 Note: A minimum of libzdnn 1.1.2 is required for the IBM zDNN backend to link correctly.

ggml/src/ggml-backend-meta.cpp

@@ -774,6 +774,13 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
         return src_ss[0];
     };
 
+    auto handle_pool_1d = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+        if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) {
+            return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
+        }
+        return src_ss[0];
+    };
+
     auto handle_flash_attn_ext = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
         GGML_ASSERT(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
 
@@ -990,12 +997,14 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             case GGML_OP_CONV_3D:
             case GGML_OP_CONV_2D_DW:
             case GGML_OP_CONV_TRANSPOSE_2D:
-            case GGML_OP_POOL_1D:
             case GGML_OP_POOL_2D:
             case GGML_OP_POOL_2D_BACK:
             case GGML_OP_UPSCALE: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ true);
             } break;
+            case GGML_OP_POOL_1D: {
+                split_state = handle_pool_1d(src_ss);
+            } break;
             case GGML_OP_PAD: {
                 split_state = handle_pad(src_ss);
             } break;

ggml/src/ggml-backend.cpp

@@ -2226,6 +2226,7 @@ bool ggml_op_alloc_size_may_expand(enum ggml_op op) {
         case GGML_OP_CUMSUM:
         case GGML_OP_ARGSORT:
         case GGML_OP_TOP_K:
+        case GGML_OP_ARGMAX:
             return true;
         default:
             return false;

ggml/src/ggml-cpu/simd-mappings.h

@@ -42,6 +42,7 @@ extern "C" {
     #define GGML_CPU_COMPUTE_FP32_TO_FP16(x) neon_compute_fp32_to_fp16(x)
 
     #define GGML_CPU_FP16_TO_FP32(x) GGML_CPU_COMPUTE_FP16_TO_FP32(x)
+    #define GGML_CPU_FP32_TO_FP16(x) GGML_CPU_COMPUTE_FP32_TO_FP16(x)
 
     static inline float neon_compute_fp16_to_fp32(ggml_fp16_t h) {
         __fp16 tmp;
@@ -150,9 +151,15 @@ extern float ggml_table_f32_ue4m3[1 << 8];
 #define GGML_CPU_FP32_TO_FP16(x) GGML_CPU_COMPUTE_FP32_TO_FP16(x)
 #endif
 
-// On ARM NEON, it's quicker to directly convert x -> x instead of calling into ggml_lookup_fp16_to_fp32,
-// so we define GGML_CPU_FP16_TO_FP32 and GGML_CPU_FP32_TO_FP16 elsewhere for NEON.
-// This is also true for POWER9.
+// a branch that defines a GGML_CPU_COMPUTE_* macro must also select it, or the fallback takes over silently
+#if defined(GGML_CPU_COMPUTE_FP16_TO_FP32) && !defined(GGML_CPU_FP16_TO_FP32) && !defined(__F16C__) // see: https://github.com/ggml-org/llama.cpp/issues/30342
+#error "a branch defined GGML_CPU_COMPUTE_FP16_TO_FP32 without selecting it as GGML_CPU_FP16_TO_FP32"
+#endif
+#if defined(GGML_CPU_COMPUTE_FP32_TO_FP16) && !defined(GGML_CPU_FP32_TO_FP16)
+#error "a branch defined GGML_CPU_COMPUTE_FP32_TO_FP16 without selecting it as GGML_CPU_FP32_TO_FP16"
+#endif
+
+// fall back to the lookup table / the portable software conversion for whichever direction no branch selected
 #if !defined(GGML_CPU_FP16_TO_FP32)
 inline static float ggml_lookup_fp16_to_fp32(ggml_fp16_t f) {
     uint16_t s;

ggml/src/ggml-cpu/vec.cpp

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

ggml/src/ggml-cpu/vec.h

@@ -44,7 +44,6 @@ inline static ggml_float ggml_sve_sum_f32x2(svfloat32_t sum_lo, svfloat32_t sum_
 #endif
 
 #define GGML_GELU_FP16
-#define GGML_GELU_QUICK_FP16
 
 #define GGML_SOFT_MAX_UNROLL 4
 #define GGML_VEC_DOT_UNROLL  2
@@ -1054,11 +1053,7 @@ inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float *
     }
 }
 #else
-inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) {
-    for (int i = 0; i < n; ++i) {
-        y[i] = ggml_gelu_quick_f32(x[i]);
-    }
-}
+void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x);
 #endif
 
 // Sigmoid Linear Unit (SiLU) function
@@ -1143,6 +1138,13 @@ inline static svfloat32_t ggml_v_silu(svbool_t pg, svfloat32_t x) {
     return svdiv_f32_x(pg, x, one_plus_exp_neg_x);
 }
 
+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static svfloat32_t ggml_v_gelu_quick(svbool_t pg, svfloat32_t x) {
+    const svfloat32_t one = svdup_n_f32_x(pg, 1.0f);
+    const svfloat32_t exp_cx = ggml_v_expf(pg, svmul_n_f32_x(pg, x, GELU_QUICK_COEF));
+    return svdiv_f32_x(pg, x, svadd_f32_x(pg, one, exp_cx));
+}
+
 #elif defined(__ARM_NEON) && defined(__aarch64__)
 
 // adapted from arm limited optimized routine
@@ -1182,6 +1184,13 @@ inline static float32x4_t ggml_v_silu(float32x4_t x) {
     return vdivq_f32(x, one_plus_exp_neg_x);
 }
 
+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static float32x4_t ggml_v_gelu_quick(float32x4_t x) {
+    const float32x4_t one = vdupq_n_f32(1.0f);
+    const float32x4_t exp_cx = ggml_v_expf(vmulq_n_f32(x, GELU_QUICK_COEF));
+    return vdivq_f32(x, vaddq_f32(one, exp_cx));
+}
+
 #elif defined(__AVX512F__) && defined(__AVX512DQ__)
 
 // adapted from arm limited optimized routine
@@ -1225,6 +1234,13 @@ inline static __m512 ggml_v_silu(__m512 x) {
     return _mm512_div_ps(x, one_plus_exp_neg_x);
 }
 
+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static __m512 ggml_v_gelu_quick(__m512 x) {
+    const __m512 one = _mm512_set1_ps(1);
+    const __m512 exp_cx = ggml_v_expf(_mm512_mul_ps(x, _mm512_set1_ps(GELU_QUICK_COEF)));
+    return _mm512_div_ps(x, _mm512_add_ps(one, exp_cx));
+}
+
 #elif defined(__AVX2__) && defined(__FMA__)
 
 // adapted from arm limited optimized routine
@@ -1280,6 +1296,13 @@ inline static __m256 ggml_v_silu(__m256 x) {
     return _mm256_div_ps(x, one_plus_exp_neg_x);
 }
 
+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static __m256 ggml_v_gelu_quick(__m256 x) {
+    const __m256 one = _mm256_set1_ps(1);
+    const __m256 exp_cx = ggml_v_expf(_mm256_mul_ps(x, _mm256_set1_ps(GELU_QUICK_COEF)));
+    return _mm256_div_ps(x, _mm256_add_ps(one, exp_cx));
+}
+
 #elif defined(__SSE2__) // __AVX2__ / __ARM_NEON
 
 #if defined(__FMA__)
@@ -1334,6 +1357,13 @@ inline static __m128 ggml_v_silu(__m128 x) {
     return _mm_div_ps(x, one_plus_exp_neg_x);
 }
 
+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static __m128 ggml_v_gelu_quick(__m128 x) {
+    const __m128 one = _mm_set1_ps(1);
+    const __m128 exp_cx = ggml_v_expf(_mm_mul_ps(x, _mm_set1_ps(GELU_QUICK_COEF)));
+    return _mm_div_ps(x, _mm_add_ps(one, exp_cx));
+}
+
 #elif defined(__VXE__) || defined(__VXE2__)
 
 // adapted from arm limited optimized routine
@@ -1374,6 +1404,13 @@ inline static float32x4_t ggml_v_silu(float32x4_t v_x) {
     return vec_div(v_x, v_one_plus_exp_neg_x);
 }
 
+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static float32x4_t ggml_v_gelu_quick(float32x4_t x) {
+    const float32x4_t one = vec_splats(1.0f);
+    const float32x4_t exp_cx = ggml_v_expf(vec_mul(x, vec_splats(GELU_QUICK_COEF)));
+    return vec_div(x, vec_add(one, exp_cx));
+}
+
 #elif defined(__riscv_v_intrinsic)
 
 // adapted from arm limited optimized routine
@@ -1426,6 +1463,12 @@ inline static vfloat32m2_t ggml_v_silu_m2(vfloat32m2_t x, int vl) {
     return __riscv_vfdiv_vv_f32m2(x, one_plus_exp_neg_x, vl);
 }
 
+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static vfloat32m2_t ggml_v_gelu_quick_m2(vfloat32m2_t x, int vl) {
+    const vfloat32m2_t exp_cx = ggml_v_expf_m2(__riscv_vfmul_vf_f32m2(x, GELU_QUICK_COEF, vl), vl);
+    return __riscv_vfdiv_vv_f32m2(x, __riscv_vfadd_vf_f32m2(exp_cx, 1.0f, vl), vl);
+}
+
 #endif // __ARM_NEON / __AVX2__ / __SSE2__ / __VXE__ / __riscv_v_intrinsic
 
 inline static void ggml_vec_silu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
@@ -1571,11 +1614,7 @@ inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float
     }
 }
 #else
-inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) {
-    for (int i = 0; i < n; ++i) {
-        y[i] = ggml_gelu_quick_f32(x[i]) * g[i];
-    }
-}
+void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g);
 #endif
 
 inline static void ggml_vec_geglu_quick_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) {

ggml/src/ggml-cuda/common.cuh

@@ -1551,11 +1551,11 @@ struct ggml_backend_cuda_context {
     // pool
     std::unique_ptr<ggml_cuda_pool> pools[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS];
 
-    static std::unique_ptr<ggml_cuda_pool> new_pool_for_device(int device, int stream_no);
+    static std::unique_ptr<ggml_cuda_pool> new_pool_for_device(ggml_backend_cuda_context & ctx, int device, int stream_no);
 
     ggml_cuda_pool & pool(int device) {
         if (pools[device][curr_stream_no] == nullptr) {
-            pools[device][curr_stream_no] = new_pool_for_device(device, curr_stream_no);
+            pools[device][curr_stream_no] = new_pool_for_device(*this, device, curr_stream_no);
         }
         return *pools[device][curr_stream_no];
     }

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

@@ -457,6 +457,7 @@ const ggml_cuda_device_info & ggml_cuda_info() {
 struct ggml_cuda_pool_leg : public ggml_cuda_pool {
     static const int MAX_BUFFERS = 256;
 
+    ggml_backend_cuda_context & ctx;
     int device;
     struct ggml_cuda_buffer {
         void * ptr = nullptr;
@@ -466,7 +467,8 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
     ggml_cuda_buffer buffer_pool[MAX_BUFFERS] = {};
     size_t pool_size = 0;
 
-    explicit ggml_cuda_pool_leg(int device) :
+    ggml_cuda_pool_leg(ggml_backend_cuda_context & ctx, int device) :
+        ctx(ctx),
         device(device) {
     }
 
@@ -488,6 +490,20 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
         }
     }
 
+    // captured CUDA graphs can still use freed buffers, so they must be captured again
+    void invalidate_cuda_graphs() {
+#ifdef USE_CUDA_GRAPH
+        int n_captured = 0;
+        for (const auto & [key, graph] : ctx.cuda_graphs) {
+            n_captured += graph->instance != nullptr;
+        }
+        if (n_captured > 0) {
+            GGML_LOG_DEBUG(GGML_CUDA_NAME " pool[%d]: memory freed, %d captured CUDA graphs must be captured again\n", device, n_captured);
+        }
+        ctx.cuda_graphs.clear();
+#endif // USE_CUDA_GRAPH
+    }
+
     void * alloc(size_t size, size_t * actual_size) override {
 #ifdef DEBUG_CUDA_MALLOC
         int nnz = 0;
@@ -538,6 +554,7 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
                            device, look_ahead_size/1024.0/1024.0, cached_bytes/1024.0/1024.0);
             CUDA_CHECK(cudaDeviceSynchronize());
             clear_pool();
+            invalidate_cuda_graphs();
             err = ggml_cuda_device_malloc(&ptr, look_ahead_size, device);
             if (err == cudaSuccess) {
                 GGML_LOG_DEBUG(GGML_CUDA_NAME " pool[%d]: retry succeeded\n", device);
@@ -566,6 +583,7 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
         ggml_cuda_set_device(device);
         CUDA_CHECK(cudaFree(ptr));
         pool_size -= size;
+        invalidate_cuda_graphs();
     }
 };
 
@@ -720,14 +738,15 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
 };
 #endif // defined(GGML_USE_VMM)
 
-std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(int                  device,
-                                                                               [[maybe_unused]] int stream_no) {
+std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(ggml_backend_cuda_context & ctx,
+                                                                               int                         device,
+                                                                               [[maybe_unused]] int        stream_no) {
 #if defined(GGML_USE_VMM)
     if (ggml_cuda_info().devices[device].vmm) {
         return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_vmm(device));
     }
 #endif // defined(GGML_USE_VMM)
-    return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_leg(device));
+    return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_leg(ctx, device));
 }
 
 // destroying a cuBLAS handle while a graph is being captured in a different thread can result in a CUDA error
@@ -2014,6 +2033,8 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c
     const ggml_tensor * src0 = dst->src[0];
     const ggml_tensor * src1 = dst->src[1];
 
+    const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
+
     if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
         return true;
     }
@@ -2032,7 +2053,7 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c
         return false;
     }
 
-    if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
+    if (ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
         return false;
     }
 
@@ -2050,6 +2071,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *
     GGML_TENSOR_BINARY_OP_LOCALS
 
     const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
+    const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
 
     // [TAG_MUL_MAT_ID_CUDA_GRAPHS]
     if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
@@ -2074,7 +2096,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *
             return;
         }
 
-        if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
+        if (ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
             ggml_cuda_mul_mat_f(ctx, src0, src1, ids, dst);
             return;
         }

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

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

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/main.c

@@ -316,9 +316,17 @@ static int vtcm_release_callback(unsigned int rctx, void * state) {
     return 0;
 }
 
-static int vtcm_alloc(struct htp_context * ctx) {
+static unsigned int vtcm_query_size(void) {
     unsigned int vtcm_size = 8 * 1024 * 1024;  // 8MB default
-    HAP_compute_res_query_VTCM(0, &vtcm_size, NULL, NULL, NULL);
+    unsigned int vtcm_avail;
+    compute_res_vtcm_page_t total_layout, avail_layout;
+    // some images (e.g. FuSa cDSP) fail the query when the layout out-params are NULL
+    HAP_compute_res_query_VTCM(0, &vtcm_size, &total_layout, &vtcm_avail, &avail_layout);
+    return vtcm_size;
+}
+
+static int vtcm_alloc(struct htp_context * ctx) {
+    unsigned int vtcm_size = vtcm_query_size();
 
     compute_res_attr_t attr;
     HAP_compute_res_attr_init(&attr);
@@ -770,9 +778,7 @@ AEEResult htp_iface_hwinfo(remote_handle64 handle, uint32_t * n_threads, uint32_
     *n_hvx     = n_hvx_val;
     *n_hmx     = 1;
 
-    uint32_t vtcm_sz = 8 * 1024 * 1024; // 8MB default fallback
-    HAP_compute_res_query_VTCM(0, (unsigned int *)&vtcm_sz, NULL, NULL, NULL);
-    *vtcm_size = vtcm_sz;
+    *vtcm_size = vtcm_query_size();
 
     return AEE_SUCCESS;
 }

ggml/src/ggml-metal/ggml-metal-context.m

@@ -539,8 +539,14 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *
                 descriptor.outputURL = [NSURL fileURLWithPath:path];
 
                 NSError * error = nil;
-                if (![[MTLCaptureManager sharedCaptureManager] startCaptureWithDescriptor:descriptor error:&error]) {
+                const bool ok = [[MTLCaptureManager sharedCaptureManager] startCaptureWithDescriptor:descriptor error:&error];
+
+                [descriptor release];
+
+                if (!ok) {
                     GGML_LOG_ERROR("%s: error: unable to start capture '%s' (did you set METAL_CAPTURE_ENABLED=1 ?)\n", __func__, [[error localizedDescription] UTF8String]);
+                    [ctx->capture_scope release];
+                    ctx->capture_scope = nil;
                 } else {
                     [ctx->capture_scope beginScope];
                     ctx->capture_started = true;
@@ -600,6 +606,8 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *
         // enter here only when capturing in order to wait for all computation to finish
         // otherwise, we leave the graph to compute asynchronously
         if (use_capture && ctx->capture_started) {
+            enum ggml_status res = GGML_STATUS_SUCCESS;
+
             // wait for completion and check status of each command buffer
             // needed to detect if the device ran out-of-memory for example (#1881)
             {
@@ -613,11 +621,11 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *
                         GGML_LOG_INFO("error: %s\n", [[cmd_buf error].localizedDescription UTF8String]);
                     }
 
-                    return GGML_STATUS_FAILED;
+                    res = GGML_STATUS_FAILED;
                 }
             }
 
-            for (int i = 0; i < n_cb; ++i) {
+            for (int i = 0; i < n_cb && res == GGML_STATUS_SUCCESS; ++i) {
                 id<MTLCommandBuffer> cmd_buf = ctx->cmd_bufs[i].obj;
                 [cmd_buf waitUntilCompleted];
 
@@ -627,8 +635,8 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *
                     if (status == MTLCommandBufferStatusError) {
                         GGML_LOG_INFO("error: %s\n", [[cmd_buf error].localizedDescription UTF8String]);
                     }
-
-                    return GGML_STATUS_FAILED;
+                    res = GGML_STATUS_FAILED;
+                    break;
                 }
 
                 id<MTLCommandBuffer> next_buffer = (i + 1 < n_cb ? ctx->cmd_bufs[i + 1].obj : nil);
@@ -643,7 +651,8 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *
 
                 if (ctx->abort_callback && ctx->abort_callback(ctx->abort_callback_data)) {
                     GGML_LOG_INFO("%s: command buffer %d aborted", __func__, i);
-                    return GGML_STATUS_ABORTED;
+                    res = GGML_STATUS_ABORTED;
+                    break;
                 }
 
                 [next_buffer commit];
@@ -652,7 +661,14 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *
             [ctx->capture_scope endScope];
             [[MTLCaptureManager sharedCaptureManager] stopCapture];
 
+            [ctx->capture_scope release];
+            ctx->capture_scope = nil;
+
             ctx->capture_started = false;
+
+            if (res != GGML_STATUS_SUCCESS) {
+                return res;
+            }
         }
     }
 

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -1567,6 +1567,28 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax(ggml_meta
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax_part(ggml_metal_library_t lib) {
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, "kernel_argmax_part_f32");
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, "kernel_argmax_part_f32", "kernel_argmax_part_f32", nullptr);
+    }
+
+    res.smem = 32*(sizeof(float) + sizeof(int32_t));
+
+    return res;
+}
+
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax_final(ggml_metal_library_t lib) {
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, "kernel_argmax_final_f32");
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, "kernel_argmax_final_f32", "kernel_argmax_final_f32", nullptr);
+    }
+
+    res.smem = 32*(sizeof(float) + sizeof(int32_t));
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort(ggml_metal_library_t lib, const ggml_tensor * op) {
     assert(op->op == GGML_OP_ARGSORT);
 

ggml/src/ggml-metal/ggml-metal-device.h

@@ -145,6 +145,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax            (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax_part       (ggml_metal_library_t lib);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax_final      (ggml_metal_library_t lib);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort           (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge     (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht              (ggml_metal_library_t lib, int n, enum ggml_type tsrc);

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -16,6 +16,9 @@
 
 #define N_MM_NPART_AMAX 256
 
+#define N_ARGMAX_NPART_MAX 256
+#define N_ARGMAX_CHUNK 4096
+
 // kernel parameters for mat-vec threadgroups
 //
 // N_R0: number of src0 rows to process per simdgroup
@@ -1379,8 +1382,10 @@ typedef struct {
 } ggml_metal_kargs_pool_1d;
 
 typedef struct {
-     int64_t ne00;
+    int64_t ne00;
     uint64_t nb01;
+    int64_t npart;
+    int64_t ne01;
 } ggml_metal_kargs_argmax;
 
 typedef struct {

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -2319,13 +2319,18 @@ int ggml_metal_op_pool_1d(ggml_metal_op_t ctx, int idx) {
 
     auto pipeline = ggml_metal_library_get_pipeline_pool_1d(lib, op, op_pool);
 
-    const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), (int) np);
-    const int ntg = (np + nth - 1) / nth;
+    const int nth = std::min(256, std::max(1, k0));
+    const int ntg = (int) np;
+
+    const size_t smem = op_pool == GGML_OP_POOL_AVG ?
+        2*32*sizeof(float) :
+        32*sizeof(float);
 
     ggml_metal_encoder_set_pipeline(enc, pipeline);
     ggml_metal_encoder_set_bytes   (enc, &args_pool_1d, sizeof(args_pool_1d),  0);
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         2);
+    ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
 
     ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1);
 
@@ -5421,6 +5426,25 @@ int ggml_metal_op_timestep_embedding(ggml_metal_op_t ctx, int idx) {
     return 1;
 }
 
+size_t ggml_metal_op_argmax_extra(const ggml_tensor * op) {
+    assert(op->op == GGML_OP_ARGMAX);
+
+    const int64_t ne00 = op->src[0]->ne[0];
+    const int64_t ne01 = op->src[0]->ne[1];
+
+    int npart = 1;
+    if (ne00 > N_ARGMAX_CHUNK) {
+        npart = (int) std::min<int64_t>(N_ARGMAX_NPART_MAX, (ne00 + N_ARGMAX_CHUNK - 1)/N_ARGMAX_CHUNK);
+    }
+
+    if (npart == 1) {
+        return 0;
+    }
+
+    // partial max + partial argmax for every row
+    return (sizeof(float) + sizeof(int32_t)) * npart * ne01;
+}
+
 int ggml_metal_op_argmax(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);
 
@@ -5432,22 +5456,61 @@ int ggml_metal_op_argmax(ggml_metal_op_t ctx, int idx) {
     GGML_TENSOR_LOCALS( int32_t, ne,  op,         ne);
     GGML_TENSOR_LOCALS(uint64_t, nb,  op,         nb);
 
+    const int64_t nrows = ggml_nrows(op->src[0]);
+
+    int npart = 1;
+    if (ne00 > N_ARGMAX_CHUNK) {
+        npart = (int) std::min<int64_t>(N_ARGMAX_NPART_MAX, (ne00 + N_ARGMAX_CHUNK - 1)/N_ARGMAX_CHUNK);
+    }
+
     ggml_metal_kargs_argmax args = {
-        /*.ne00 = */ ne00,
-        /*.nb01 = */ nb01,
+        /*.ne00  = */ ne00,
+        /*.nb01  = */ nb01,
+        /*.npart = */ npart,
+        /*.ne01  = */ ne01,
     };
 
-    auto pipeline = ggml_metal_library_get_pipeline_argmax(lib, op);
+    const size_t smem = 32*(sizeof(float) + sizeof(int32_t));
 
-    const int64_t nrows = ggml_nrows(op->src[0]);
+    if (npart > 1) {
+        auto pipeline_part = ggml_metal_library_get_pipeline_argmax_part(lib);
+        auto pipeline_final = ggml_metal_library_get_pipeline_argmax_final(lib);
+
+        ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op);
+        ggml_metal_buffer_id bid_tmp = bid_dst;
+        bid_tmp.offs += ggml_nbytes(op);
+
+        ggml_metal_encoder_set_pipeline(enc, pipeline_part);
+        ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
+        ggml_metal_encoder_set_buffer  (enc, bid_tmp, 2);
+        ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
+        ggml_metal_encoder_dispatch_threadgroups(enc, npart, nrows, 1, 256, 1, 1);
+
+        ggml_metal_op_concurrency_reset(ctx);
+
+        int nth = 32; // SIMD width
+        while (nth < npart && nth < 256) {
+            nth *= 2;
+        }
+
+        ggml_metal_encoder_set_pipeline(enc, pipeline_final);
+        ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
+        ggml_metal_encoder_set_buffer  (enc, bid_tmp, 1);
+        ggml_metal_encoder_set_buffer  (enc, bid_dst, 2);
+        ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
+        ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, nth, 1, 1);
+
+        return 1;
+    }
+
+    auto pipeline = ggml_metal_library_get_pipeline_argmax(lib, op);
 
     int nth = 32; // SIMD width
     while (nth < ne00 && nth*ne01*ne02*ne03 < 256) {
         nth *= 2;
     }
 
-    const size_t smem = pipeline.smem;
-
     ggml_metal_encoder_set_pipeline(enc, pipeline);
     ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), 1);

ggml/src/ggml-metal/ggml-metal-ops.h

@@ -47,6 +47,8 @@ size_t ggml_metal_op_flash_attn_ext_extra_tmp(const struct ggml_tensor * op);
 size_t ggml_metal_op_flash_attn_ext_extra_kv_f16(const struct ggml_tensor * op);
 size_t ggml_metal_op_flash_attn_ext_extra_idx(const struct ggml_tensor * op);
 
+size_t ggml_metal_op_argmax_extra(const struct ggml_tensor * op);
+
 int ggml_metal_op_concat            (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_repeat            (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_acc               (ggml_metal_op_t ctx, int idx);

ggml/src/ggml-metal/ggml-metal.cpp

@@ -245,6 +245,10 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_
             {
                 res = 2*sizeof(int32_t)*ggml_nelements(tensor->src[0]);
             } break;
+        case GGML_OP_ARGMAX:
+            {
+                res += ggml_metal_op_argmax_extra(tensor);
+            } break;
         default:
             break;
     }

ggml/src/ggml-metal/kernels/misc.metal

@@ -60,6 +60,140 @@ kernel void kernel_argmax_f32(
     dst_i32[tgpig] = arg_val;
 }
 
+kernel void kernel_argmax_part_f32(
+        constant ggml_metal_kargs_argmax & args,
+        device   const char * src0,
+        device         char * tmp,
+        threadgroup    char * shmem [[threadgroup(0)]],
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint  sgitg[[simdgroup_index_in_threadgroup]],
+        uint  tiisg[[thread_index_in_simdgroup]],
+        uint3  ntg[[threads_per_threadgroup]]) {
+    const int row  = tgpig.y;
+    const int part = tgpig.x;
+    const int ne00 = (int) args.ne00;
+    const int npart = (int) args.npart;
+    const int chunk = (ne00 + npart - 1) / npart;
+    const int start = part * chunk;
+    const int end   = min(start + chunk, ne00);
+
+    device const float * x_row = (device const float *) ((device const char *) src0 + row * args.nb01);
+
+    float   lmax = -INFINITY;
+    int32_t larg = -1;
+
+    for (int i00 = start + tpitg.x; i00 < end; i00 += ntg.x) {
+        if (x_row[i00] > lmax) {
+            lmax = x_row[i00];
+            larg = i00;
+        }
+    }
+
+    // find the argmax value in the block
+    float max_val = simd_max(lmax);
+    int32_t arg_val = simd_max(select(-1, larg, lmax == max_val));
+
+    threadgroup   float * shared_maxval = (threadgroup   float *) shmem;
+    threadgroup int32_t * shared_argmax = (threadgroup int32_t *) shmem + N_SIMDWIDTH;
+
+    if (ntg.x > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            shared_maxval[tiisg] = -INFINITY;
+            shared_argmax[tiisg] = -1;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            shared_maxval[sgitg] = max_val;
+            shared_argmax[sgitg] = arg_val;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        max_val = shared_maxval[tiisg];
+        arg_val = shared_argmax[tiisg];
+
+        float max_val_reduced   = simd_max(max_val);
+        int32_t arg_val_reduced = simd_max(select(-1, arg_val, max_val == max_val_reduced));
+
+        max_val = max_val_reduced;
+        arg_val = arg_val_reduced;
+    }
+
+    if (tpitg.x == 0) {
+        device float * tmp_max = (device float *) tmp;
+        device int32_t * tmp_arg = (device int32_t *) tmp + npart * (int) args.ne01;
+
+        tmp_max[row*npart + part] = max_val;
+        tmp_arg[row*npart + part] = arg_val;
+    }
+}
+
+kernel void kernel_argmax_final_f32(
+        constant ggml_metal_kargs_argmax & args,
+        device   const char * tmp,
+        device         char * dst,
+        threadgroup    char * shmem [[threadgroup(0)]],
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint  sgitg[[simdgroup_index_in_threadgroup]],
+        uint  tiisg[[thread_index_in_simdgroup]],
+        uint3  ntg[[threads_per_threadgroup]]) {
+    const int row = tgpig.x;
+    const int npart = (int) args.npart;
+
+    device const float * tmp_max = (device const float *) tmp;
+    device const int32_t * tmp_arg = (device const int32_t *) tmp + npart * (int) args.ne01;
+
+    float   lmax = -INFINITY;
+    int32_t larg = -1;
+
+    for (int i = tpitg.x; i < npart; i += ntg.x) {
+        if (tmp_max[row*npart + i] > lmax) {
+            lmax = tmp_max[row*npart + i];
+            larg = tmp_arg[row*npart + i];
+        }
+    }
+
+    // find the argmax value in the block
+    float max_val = simd_max(lmax);
+    int32_t arg_val = simd_max(select(-1, larg, lmax == max_val));
+
+    threadgroup   float * shared_maxval = (threadgroup   float *) shmem;
+    threadgroup int32_t * shared_argmax = (threadgroup int32_t *) shmem + N_SIMDWIDTH;
+
+    if (ntg.x > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            shared_maxval[tiisg] = -INFINITY;
+            shared_argmax[tiisg] = -1;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            shared_maxval[sgitg] = max_val;
+            shared_argmax[sgitg] = arg_val;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        max_val = shared_maxval[tiisg];
+        arg_val = shared_argmax[tiisg];
+
+        float max_val_reduced   = simd_max(max_val);
+        int32_t arg_val_reduced = simd_max(select(-1, arg_val, max_val == max_val_reduced));
+
+        max_val = max_val_reduced;
+        arg_val = arg_val_reduced;
+    }
+
+    if (tpitg.x == 0) {
+        ((device int32_t *) dst)[row] = arg_val;
+    }
+}
+
 kernel void kernel_diag_f32(
         constant ggml_metal_kargs_diag & args,
         device   const char * src0,

ggml/src/ggml-metal/kernels/pool.metal

@@ -84,65 +84,117 @@ kernel void kernel_pool_1d_max_f32(
         constant        ggml_metal_kargs_pool_1d & args,
         device  const   float * src,
         device          float * dst,
-        uint            gid [[thread_position_in_grid]]
-) {
-
-    if (gid >= args.np) {
+        threadgroup    char * shmem [[threadgroup(0)]],
+        uint tgpig[[threadgroup_position_in_grid]],
+        uint tpitg[[thread_position_in_threadgroup]],
+        uint sgitg[[simdgroup_index_in_threadgroup]],
+        uint tiisg[[thread_index_in_simdgroup]],
+        uint ntg[[threads_per_threadgroup]]) {
+    if (tgpig >= args.np) {
         return;
     }
 
-    const int ow  = (int)gid % args.OW;
-    const int row = (int)gid / args.OW;
+    const int out = (int) tgpig;
+    const int ow  = out % (int) args.OW;
+    const int row = out / (int) args.OW;
 
-    const int base = ow * args.s0 - args.p0;
+    const int base     = ow * args.s0 - args.p0;
+    const int src_off  = row * args.IW;
+    const int dst_off  = row * args.OW;
 
-    float acc = -INFINITY;
+    float lmax = -INFINITY;
 
-    const int src_off = row * args.IW;
-    const int dst_off = row * args.OW;
-
-    for (int ki = 0; ki < args.k0; ++ki) {
-        int j = base + ki;
-        if (j < 0 || j >= args.IW){
-            continue;
+    for (int ki = tpitg; ki < args.k0; ki += ntg) {
+        const int j = base + ki;
+        if (j >= 0 && j < args.IW) {
+            lmax = max(lmax, src[src_off + j]);
         }
-        float v = src[src_off + j];
-        acc = max(acc, v);
     }
 
-    dst[dst_off + ow] = acc;
+    float max_val = simd_max(lmax);
+
+    threadgroup float * shared_max = (threadgroup float *) shmem;
+
+    if (ntg > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            shared_max[tiisg] = -INFINITY;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            shared_max[sgitg] = max_val;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        max_val = shared_max[tiisg];
+        max_val = simd_max(max_val);
+    }
+
+    if (tpitg == 0) {
+        dst[dst_off + ow] = max_val;
+    }
 }
 
 kernel void kernel_pool_1d_avg_f32(
         constant        ggml_metal_kargs_pool_1d & args,
         device  const   float * src,
         device          float * dst,
-        uint            gid [[thread_position_in_grid]]
-) {
-
-    if (gid >= args.np) {
+        threadgroup    char * shmem [[threadgroup(0)]],
+        uint tgpig[[threadgroup_position_in_grid]],
+        uint tpitg[[thread_position_in_threadgroup]],
+        uint sgitg[[simdgroup_index_in_threadgroup]],
+        uint tiisg[[thread_index_in_simdgroup]],
+        uint ntg[[threads_per_threadgroup]]) {
+    if (tgpig >= args.np) {
         return;
     }
 
-    const int ow  = (int)gid % args.OW;
-    const int row = (int)gid / args.OW;
+    const int out = (int) tgpig;
+    const int ow  = out % (int) args.OW;
+    const int row = out / (int) args.OW;
 
-    const int base = ow * args.s0 - args.p0;
+    const int base     = ow * args.s0 - args.p0;
+    const int src_off  = row * args.IW;
+    const int dst_off  = row * args.OW;
 
-    float acc = 0.0f;
-    int   cnt = 0;
+    float lsum = 0.0f;
+    int   lcnt = 0;
 
-    const int src_off = row * args.IW;
-    const int dst_off = row * args.OW;
-
-    for (int ki = 0; ki < args.k0; ++ki) {
+    for (int ki = tpitg; ki < args.k0; ki += ntg) {
         const int j = base + ki;
-        if (j < 0 || j >= args.IW) {
-            continue;
+        if (j >= 0 && j < args.IW) {
+            lsum += src[src_off + j];
+            lcnt += 1;
         }
-        acc += src[src_off + j];
-        cnt += 1;
     }
 
-    dst[dst_off + ow] = (cnt > 0) ? (acc / (float)cnt) : 0.0f;
+    float sum = simd_sum(lsum);
+    float cnt = simd_sum((float) lcnt);
+
+    threadgroup float * shared_sum = (threadgroup float *) shmem;
+    threadgroup float * shared_cnt = (threadgroup float *) shmem + N_SIMDWIDTH;
+
+    if (ntg > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            shared_sum[tiisg] = 0.0f;
+            shared_cnt[tiisg] = 0.0f;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            shared_sum[sgitg] = sum;
+            shared_cnt[sgitg] = cnt;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        sum = shared_sum[tiisg];
+        cnt = shared_cnt[tiisg];
+
+        sum = simd_sum(sum);
+        cnt = simd_sum(cnt);
+    }
+
+    if (tpitg == 0) {
+        dst[dst_off + ow] = (cnt > 0.0f) ? (sum / cnt) : 0.0f;
+    }
 }

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

@@ -831,6 +831,7 @@ struct ggml_backend_opencl_context {
     cl_program program_mul_mv_id_mxfp4_f32_flat;
     cl_program program_mul_mm_f32_f32_l4_lm;
     cl_program program_mul_mm_f16_f32_l4_lm;
+    cl_program program_mul_mm_f16_f32_l4_lm_n8;
     cl_program program_mul_mm_q8_0_f32_l4_lm;
 
     cl_kernel kernel_add, kernel_add_row, kernel_add_f16, kernel_add_row_f16;
@@ -1052,6 +1053,7 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_mul_mm_f32_f32_l4_lm;
     cl_kernel kernel_gemv_f32_f32_mc;  // multi-column (small-N) f32 GEMV for spec/MTP verify
     cl_kernel kernel_mul_mm_f16_f32_l4_lm;
+    cl_kernel kernel_mul_mm_f16_f32_l4_lm_n8;
     cl_kernel kernel_mul_mm_q1_0_f32_l4_lm;
     cl_kernel kernel_mul_mm_q4_0_f32_l4_lm;
     cl_kernel kernel_mul_mm_q4_1_f32_l4_lm;
@@ -1241,12 +1243,15 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_gemm_noshuffle_q4_1_f32;
     cl_kernel kernel_gemm_noshuffle_q8_0_f32, kernel_gemm_noshuffle_q8_0_f32_bin;
     cl_kernel kernel_gemm_noshuffle_q8_0_q8_1_dp4a = nullptr;  // dp4a (int8) dense q8_0 prefill GEMM (opt-in)
+    cl_kernel kernel_gemm_noshuffle_q8_0_q8_1_dp4a_act_vload4 = nullptr;
     cl_kernel kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg = nullptr;  // q8_0 dense dp4a, weights via texture (opt-in)
     cl_kernel kernel_gemm_noshuffle_q8_0_q8_1_dp4a_ila_a8_bin = nullptr;
     cl_kernel kernel_gemv_noshuffle_q8_0_f32;
     cl_kernel kernel_gemv_noshuffle_q8_0_f32_splitk;  // split-K across WGs (small-M decode)
     cl_kernel kernel_gemm_noshuffle_q1_0_f32;
     cl_kernel kernel_gemv_noshuffle_q1_0_f32;
+    cl_kernel kernel_gemv_noshuffle_q1_0_f32_nsg1 = nullptr;
+    int min_m_gemv_nsg1_q1_0 = 0;
     cl_kernel kernel_gemv_noshuffle_q4_k_f32;
     cl_kernel kernel_gemv_noshuffle_q4_k_f32_o4;  // 4-output-per-WI, long-vocab lm_head
     cl_kernel kernel_gemv_noshuffle_q4_k_f32_tiled;  // tiled-wide layout (opt-in)
@@ -1260,9 +1265,17 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_ila_a8_bin;
     cl_kernel kernel_gemv_noshuffle_q4_k_f32_32b_trans;
     cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a = nullptr;  // dp4a (int8) dense prefill GEMM
+    cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_act_vload4 = nullptr;
     cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg = nullptr;  // dp4a dense prefill GEMM, weights via texture (X1 opt-in)
+    cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow = nullptr;
+    cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow_wimg = nullptr;
+    int tile_n_dp4a_q4_k = 32;
+    int tile_n_dp4a_narrow_q4_k = 32;
+    int max_n_dp4a_narrow_q4_k = 16;
     cl_kernel kernel_gemm_noshuffle_q5_k_q8_1_dp4a = nullptr;  // dp4a (int8) dense q5_K prefill GEMM
+    cl_kernel kernel_gemm_noshuffle_q5_k_q8_1_dp4a_act_vload4 = nullptr;
     cl_kernel kernel_gemm_noshuffle_q6_k_q8_1_dp4a = nullptr;  // dp4a (int8) dense q6_K prefill GEMM
+    cl_kernel kernel_gemm_noshuffle_q6_k_q8_1_dp4a_act_vload4 = nullptr;
     cl_kernel kernel_quant_a_q8_1;                    // plain activation q8_1 pre-pass
     cl_kernel kernel_gemm_noshuffle_q4_k_f32_r1;
     cl_kernel kernel_gemm_noshuffle_q4_k_f32_kimg;
@@ -1296,6 +1309,7 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_gemm_noshuffle_iq4_nl_f32;
     cl_kernel kernel_gemm_noshuffle_iq4_nl_q8_1_dp4a = nullptr;  // dp4a (int8) dense IQ4_NL prefill GEMM
     cl_kernel kernel_gemm_noshuffle_q4_0_q8_1_dp4a = nullptr;  // dp4a (int8) dense q4_0 prefill GEMM
+    cl_kernel kernel_gemm_noshuffle_q4_0_q8_1_dp4a_act_vload4 = nullptr;
     cl_kernel kernel_gemm_coopk_q4_0_q8_1_dp4a_c4 = nullptr;
     cl_kernel kernel_gemm_coopk_q4_0_q8_1_dp4a_c2 = nullptr;
     int       nsg_coopk_dp4a_q4_0  = 8;
@@ -2781,6 +2795,11 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
             build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
 
         CL_CHECK((backend_ctx->kernel_mul_mm_f16_f32_l4_lm = clCreateKernel(backend_ctx->program_mul_mm_f16_f32_l4_lm, "kernel_mul_mm_f16_f32_l4_lm", &err), err));
+
+        const std::string narrow_opts = compile_opts + " -DBN=8 -DTN=1 -DKERNEL_NAME_LM=kernel_mul_mm_f16_f32_l4_lm_n8";
+        backend_ctx->program_mul_mm_f16_f32_l4_lm_n8 =
+            build_program_from_source(backend_ctx, kernel_src.c_str(), narrow_opts);
+        CL_CHECK((backend_ctx->kernel_mul_mm_f16_f32_l4_lm_n8 = clCreateKernel(backend_ctx->program_mul_mm_f16_f32_l4_lm_n8, "kernel_mul_mm_f16_f32_l4_lm_n8", &err), err));
         GGML_LOG_CONT(".");
     }
 
@@ -3983,6 +4002,16 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
 
         CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q1_0_f32 = clCreateKernel(prog, "kernel_gemv_noshuffle_q1_0_f32", &err), err));
         CL_CHECK(clReleaseProgram(prog));
+
+        const char * nsg1_min_m_env = getenv("GGML_OPENCL_Q1_0_GEMV_NSG1_MIN_M");
+        backend_ctx->min_m_gemv_nsg1_q1_0 = nsg1_min_m_env ? atoi(nsg1_min_m_env) :
+            (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E ? 8192 : 0);
+        if (backend_ctx->min_m_gemv_nsg1_q1_0 > 0) {
+            const std::string nsg1_opts = CL_gemv_compile_opts + " -DN_SIMDGROUP=1";
+            cl_program prog_nsg1 = build_program_from_source(backend_ctx, kernel_src_CL_gemv_general.c_str(), nsg1_opts);
+            CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q1_0_f32_nsg1 = clCreateKernel(prog_nsg1, "kernel_gemv_noshuffle_q1_0_f32", &err), err));
+            CL_CHECK(clReleaseProgram(prog_nsg1));
+        }
         GGML_LOG_CONT(".");
     }
 
@@ -4323,6 +4352,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
 #endif
         cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_0_q8_1_dp4a", &err), err));
+        CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a_act_vload4 = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_0_q8_1_dp4a_act_vload4", &err), err));
         CL_CHECK(clReleaseProgram(prog));
         GGML_LOG_CONT(".");
     }
@@ -4453,16 +4483,33 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
 #else
         const std::string kernel_src = read_file("gemm_noshuffle_q4_k_q8_1_dp4a.cl");
 #endif
-        // Per-device dp4a dense tile. The X2-tuned TILESIZE_N=32 over-occupies LDS on
-        // X1 (1152 B/WG -> few resident WGs); TILESIZE_N=8 (288 B) lifts occupancy on
-        // X1, byte-identical. X2E keeps 32. Env override wins.
-        int q4k_dp4a_ts = (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) ? 8 : 32;
-        if (const char * e = getenv("GGML_OPENCL_Q4K_DP4A_TS")) { q4k_dp4a_ts = atoi(e); }
-        std::string dp4a_opts = compile_opts + " -DTILESIZE_N=" + std::to_string(q4k_dp4a_ts);
+        // X1 uses a smaller N tile to reduce local memory and register use.
+        int tile_n = (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) ? 8 : 32;
+        if (const char * e = getenv("GGML_OPENCL_Q4K_DP4A_TS")) { tile_n = atoi(e); }
+        GGML_ASSERT(tile_n >= 4 && tile_n <= 64 && tile_n % 4 == 0);
+        backend_ctx->tile_n_dp4a_q4_k = tile_n;
+        std::string dp4a_opts = compile_opts + " -DTILESIZE_N=" + std::to_string(tile_n);
         cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), dp4a_opts);
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_k_q8_1_dp4a", &err), err));
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg", &err), err));
+        CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_act_vload4 = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_k_q8_1_dp4a_act_vload4", &err), err));
         CL_CHECK(clReleaseProgram(prog));
+
+        // Compile a smaller N tile to avoid unused columns in small batches.
+        int tile_n_narrow = MIN(tile_n, 16);
+        if (const char * e = getenv("GGML_OPENCL_Q4_K_DP4A_TS_NARROW")) { tile_n_narrow = atoi(e); }
+        GGML_ASSERT(tile_n_narrow >= 4 && tile_n_narrow <= 64 && tile_n_narrow % 4 == 0);
+        if (const char * e = getenv("GGML_OPENCL_Q4_K_DP4A_NARROW_MAX")) {
+            backend_ctx->max_n_dp4a_narrow_q4_k = atoi(e);
+        }
+        backend_ctx->tile_n_dp4a_narrow_q4_k = tile_n_narrow;
+        if (tile_n_narrow != tile_n) {
+            std::string narrow_opts = compile_opts + " -DTILESIZE_N=" + std::to_string(tile_n_narrow);
+            cl_program narrow_prog = build_program_from_source(backend_ctx, kernel_src.c_str(), narrow_opts);
+            CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow = clCreateKernel(narrow_prog, "kernel_gemm_noshuffle_q4_k_q8_1_dp4a", &err), err));
+            CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow_wimg = clCreateKernel(narrow_prog, "kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg", &err), err));
+            CL_CHECK(clReleaseProgram(narrow_prog));
+        }
         GGML_LOG_CONT(".");
     }
 
@@ -4478,6 +4525,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
         cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q8_0_q8_1_dp4a", &err), err));
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg = clCreateKernel(prog, "kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg", &err), err));
+        CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a_act_vload4 = clCreateKernel(prog, "kernel_gemm_noshuffle_q8_0_q8_1_dp4a_act_vload4", &err), err));
         CL_CHECK(clReleaseProgram(prog));
         GGML_LOG_CONT(".");
     }
@@ -4493,6 +4541,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
 #endif
         cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q5_k_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q5_k_q8_1_dp4a", &err), err));
+        CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q5_k_q8_1_dp4a_act_vload4 = clCreateKernel(prog, "kernel_gemm_noshuffle_q5_k_q8_1_dp4a_act_vload4", &err), err));
         CL_CHECK(clReleaseProgram(prog));
         GGML_LOG_CONT(".");
     }
@@ -4508,6 +4557,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
 #endif
         cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q6_k_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q6_k_q8_1_dp4a", &err), err));
+        CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q6_k_q8_1_dp4a_act_vload4 = clCreateKernel(prog, "kernel_gemm_noshuffle_q6_k_q8_1_dp4a_act_vload4", &err), err));
         CL_CHECK(clReleaseProgram(prog));
         GGML_LOG_CONT(".");
     }
@@ -9531,22 +9581,6 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
                     op->src[0]->ne[1] >= 32768) {   // vocab-scale weight; no FFN/attn weight is this tall
                     return false;
                 }
-                // The generic mul_mv (GEMV) kernels are wrong for large-batch prefill on
-                // Adreno. A quant mul_mat only avoids the GEMV when it reaches the Adreno
-                // trans-weight GEMM, which needs both a GEMM kernel for the type and
-                // use_adreno_kernels(). Decline the large-N shapes that would otherwise
-                // fall through to the GEMV.
-                {
-                    const ggml_type t = op->src[0]->type;
-                    const bool type_has_gemm = (t == GGML_TYPE_Q4_0 || t == GGML_TYPE_Q4_1 ||
-                                                t == GGML_TYPE_IQ4_NL || t == GGML_TYPE_Q8_0 ||
-                                                t == GGML_TYPE_Q4_K  || t == GGML_TYPE_Q5_K  ||
-                                                t == GGML_TYPE_Q6_K);
-                    const bool uses_gemm = type_has_gemm && use_adreno_kernels(backend_ctx, op->src[0]);
-                    if (!uses_gemm && op->src[1]->ne[1] >= 512) {
-                        return false;
-                    }
-                }
                 return op->src[1]->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]);
             } else if (op->src[0]->type == GGML_TYPE_Q8_0) {
                 // ggml_cl_mul_mat_q8_0_f32_adreno now honors src1/dst view_offs (the
@@ -19834,7 +19868,9 @@ static void ggml_cl_mul_mat_q1_0_f32_adreno(ggml_backend_t backend, const ggml_t
         img_desc.buffer = b_sub_buf;
         CL_CHECK((b_img = clCreateImage(context, CL_MEM_READ_ONLY, &img_fmt, &img_desc, NULL, &err), err));
 
-        kernel = backend_ctx->kernel_gemv_noshuffle_q1_0_f32;
+        // Tall M has enough row parallelism to skip the K split and its reduction.
+        const bool use_nsg1 = backend_ctx->min_m_gemv_nsg1_q1_0 > 0 && M >= backend_ctx->min_m_gemv_nsg1_q1_0;
+        kernel = use_nsg1 ? backend_ctx->kernel_gemv_noshuffle_q1_0_f32_nsg1 : backend_ctx->kernel_gemv_noshuffle_q1_0_f32;
 
         int r2 = 1;
         int r3 = 1;
@@ -19856,8 +19892,9 @@ static void ggml_cl_mul_mat_q1_0_f32_adreno(ggml_backend_t backend, const ggml_t
         CL_CHECK(clSetKernelArg(kernel, 14, sizeof(int),      &r3));
 
         size_t wavesize = backend_ctx->adreno_wave_size;
-        size_t local_work_size[]  = { wavesize, 4, 1 };
-        size_t global_work_size[] = { CEIL_DIV(M, wavesize)*wavesize, 4, 1 };
+        const size_t nsg = use_nsg1 ? 1 : 4;
+        size_t local_work_size[]  = { wavesize, nsg, 1 };
+        size_t global_work_size[] = { CEIL_DIV(M, wavesize)*wavesize, nsg, 1 };
 
         backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
 
@@ -20321,7 +20358,14 @@ static void ggml_cl_mul_mat_q4_0_f32_adreno(ggml_backend_t backend, const ggml_t
         // per-lane K-walk) for small M. Layout stride is fixed (4 uints/block), so only
         // the K-split count changes; the mc3 kernel reads it via get_local_size(1). The
         // ne1==1 base kernel hardcodes N_SIMDGROUP=4, so it always stays at 4.
-        const int mc3_nsg = (use_q40_mc3 && ne01 < 4096) ? 8 : 4;
+        int mc3_nsg = (use_q40_mc3 && ne01 < 4096) ? 8 : 4;
+        if (use_q40_mc3) {
+            // Register pressure can lower the kernel limit below the device limit.
+            const size_t max_workgroup_size = backend_ctx->get_kernel_workgroup_size(kernel);
+            while (mc3_nsg > 1 && (size_t)(64 * mc3_nsg) > max_workgroup_size) {
+                mc3_nsg /= 2;
+            }
+        }
         size_t local_work_size[3] = {64, (size_t)mc3_nsg, 1};
         size_t global_work_size[3] = {(size_t)CEIL_DIV(ne01/2, 64)*64, (size_t)mc3_nsg, 1};
 
@@ -20396,11 +20440,13 @@ static void ggml_cl_mul_mat_q4_0_f32_adreno(ggml_backend_t backend, const ggml_t
             return;
         }
 
-        // dp4a (int8) dense prefill GEMM, default off
+        static const char * q4_0_act_vload4_env = getenv("GGML_OPENCL_Q4_0_DP4A_ACT_VLOAD4");
+        const bool q4_0_act_vload4_on = q4_0_act_vload4_env == nullptr || atoi(q4_0_act_vload4_env) != 0;
+        // Enable dense prefill by default on X2E when the uint4 activation tile is enabled.
         static const char * q4_0_dense_dp4a_env = getenv("GGML_OPENCL_Q4_0_DENSE_DP4A");
         bool q4_0_dense_dp4a_on = q4_0_dense_dp4a_env
             ? (atoi(q4_0_dense_dp4a_env) != 0)
-            : false;
+            : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E && q4_0_act_vload4_on);
         // dot prod has to be available
         q4_0_dense_dp4a_on = backend_ctx->has_integer_dot && q4_0_dense_dp4a_on;
 
@@ -20427,7 +20473,9 @@ static void ggml_cl_mul_mat_q4_0_f32_adreno(ggml_backend_t backend, const ggml_t
             size_t q_global[1] = { (size_t)(((n_blocks + 63) / 64) * 64) };
             backend_ctx->enqueue_ndrange_kernel(qk, 1, q_global, q_local, dst);
 
-            cl_kernel dk = backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a;
+            cl_kernel dk = q4_0_act_vload4_on
+                         ? backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a_act_vload4
+                         : backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a;
             int ai = 0;
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem),   &extra0_q4_0->q));
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem),   &extra0_q4_0->d));
@@ -21720,8 +21768,13 @@ static void ggml_cl_mul_mat_q8_0_f32_adreno(ggml_backend_t backend, const ggml_t
                 }
             }
 
-            cl_kernel dk = use_wimg ? backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg
-                                    : backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a;
+            static const char * q8_0_act_vload4_env = getenv("GGML_OPENCL_Q8_0_DP4A_ACT_VLOAD4");
+            const bool q8_0_act_vload4_on = q8_0_act_vload4_env == nullptr || atoi(q8_0_act_vload4_env) != 0;
+            cl_kernel dk = use_wimg
+                ? backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg
+                : q8_0_act_vload4_on
+                    ? backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a_act_vload4
+                    : backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a;
             int ai = 0;
             if (use_wimg) {
                 CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &q8_q_img));
@@ -22667,8 +22720,18 @@ static void ggml_cl_mul_mat_q4_k_f32_adreno(ggml_backend_t backend, const ggml_t
                 }
             }
 
-            cl_kernel dk = use_wimg ? backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg
-                                    : backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a;
+            const bool use_narrow = backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow != nullptr
+                                 && N <= backend_ctx->max_n_dp4a_narrow_q4_k;
+            static const char * q4_k_act_vload4_env = getenv("GGML_OPENCL_Q4_K_DP4A_ACT_VLOAD4");
+            const bool q4_k_act_vload4_on = q4_k_act_vload4_env == nullptr || atoi(q4_k_act_vload4_env) != 0;
+            cl_kernel dk = use_narrow
+                ? (use_wimg ? backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow_wimg
+                            : backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow)
+                : use_wimg
+                    ? backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg
+                    : q4_k_act_vload4_on
+                        ? backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_act_vload4
+                        : backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a;
             int ai = 0;
             if (use_wimg) {
                 CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &q4k_q_img));
@@ -22689,12 +22752,9 @@ static void ggml_cl_mul_mat_q4_k_f32_adreno(ggml_backend_t backend, const ggml_t
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_uchar), &mask_d6));
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_uchar), &mask_d4));
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_uchar), &mask_hi2));
-            // Must match the compile-time TILESIZE_N chosen at program build (per-device,
-            // X1E=8 else 32; env override). Same inputs -> same value.
-            int q4k_dp4a_ts = (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) ? 8 : 32;
-            if (const char * e = getenv("GGML_OPENCL_Q4K_DP4A_TS")) q4k_dp4a_ts = atoi(e);
+            const int tile_n = use_narrow ? backend_ctx->tile_n_dp4a_narrow_q4_k : backend_ctx->tile_n_dp4a_q4_k;
             size_t d_local[3]  = { 64, 1, 1 };
-            size_t d_global[3] = { 64, (size_t)(M / 64), (size_t)CEIL_DIV(N, q4k_dp4a_ts) };
+            size_t d_global[3] = { 64, (size_t)(M / 64), (size_t)CEIL_DIV(N, tile_n) };
             backend_ctx->enqueue_ndrange_kernel(dk, 3, d_global, d_local, dst);
 
             if (q4k_q_img != nullptr) {
@@ -23333,7 +23393,11 @@ static void ggml_cl_mul_mat_q6_K_f32_adreno(ggml_backend_t backend, const ggml_t
             size_t q_global[1] = { (size_t)(((n_blocks + 63) / 64) * 64) };
             backend_ctx->enqueue_ndrange_kernel(qk, 1, q_global, q_local, dst);
 
-            cl_kernel dk = backend_ctx->kernel_gemm_noshuffle_q6_k_q8_1_dp4a;
+            static const char * q6_k_act_vload4_env = getenv("GGML_OPENCL_Q6_K_DP4A_ACT_VLOAD4");
+            const bool q6_k_act_vload4_on = q6_k_act_vload4_env == nullptr || atoi(q6_k_act_vload4_env) != 0;
+            cl_kernel dk = q6_k_act_vload4_on
+                         ? backend_ctx->kernel_gemm_noshuffle_q6_k_q8_1_dp4a_act_vload4
+                         : backend_ctx->kernel_gemm_noshuffle_q6_k_q8_1_dp4a;
             int ai = 0;
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem),   &extra0_q6_K->ql));
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem),   &extra0_q6_K->qh));
@@ -23889,7 +23953,11 @@ static void ggml_cl_mul_mat_q5_K_f32_adreno(ggml_backend_t backend, const ggml_t
             size_t q_global[1] = { (size_t)(((n_blocks + 63) / 64) * 64) };
             backend_ctx->enqueue_ndrange_kernel(qk, 1, q_global, q_local, dst);
 
-            cl_kernel dk = backend_ctx->kernel_gemm_noshuffle_q5_k_q8_1_dp4a;
+            static const char * q5_k_act_vload4_env = getenv("GGML_OPENCL_Q5_K_DP4A_ACT_VLOAD4");
+            const bool q5_k_act_vload4_on = q5_k_act_vload4_env == nullptr || atoi(q5_k_act_vload4_env) != 0;
+            cl_kernel dk = q5_k_act_vload4_on
+                         ? backend_ctx->kernel_gemm_noshuffle_q5_k_q8_1_dp4a_act_vload4
+                         : backend_ctx->kernel_gemm_noshuffle_q5_k_q8_1_dp4a;
             int ai = 0;
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem),   &extra0_q5_k->q));
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem),   &extra0_q5_k->qh));
@@ -24815,7 +24883,9 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
                     return;
                 }
 #endif
-                kernel = backend_ctx->kernel_mul_mm_f16_f32_l4_lm;
+                static const char * mm_narrow_n_env = getenv("GGML_OPENCL_MM_NARROW_N");
+                const bool use_narrow_n = (mm_narrow_n_env == nullptr || mm_narrow_n_env[0] != '0') && ne11 >= 2 && ne11 <= 8;
+                kernel = use_narrow_n ? backend_ctx->kernel_mul_mm_f16_f32_l4_lm_n8 : backend_ctx->kernel_mul_mm_f16_f32_l4_lm;
                 nth0 = 128; // calculated as (BM*BN)/(TM*TN)
 
                 int batch_stride_a = ne00*ne01;
@@ -24874,8 +24944,9 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
                 CL_CHECK(clSetKernelArg(kernel, 17, sizeof(int),      &r2));
                 CL_CHECK(clSetKernelArg(kernel, 18, sizeof(int),      &r3));
 
-                // 64 is block tile size BM and BN - change here when BM and BN in the kernel are changed.
-                size_t global_work_size[] = {(size_t)(CEIL_DIV(ne01, 64)*nth0), (size_t)(CEIL_DIV(ne11, 64)), (size_t)ne12*ne13};
+                // BN=8/TN=1 keeps 128 threads for small batches, including attention KQ/KQV.
+                const int bn = use_narrow_n ? 8 : 64;
+                size_t global_work_size[] = {(size_t)(CEIL_DIV(ne01, 64)*nth0), (size_t)(CEIL_DIV(ne11, bn)), (size_t)ne12*ne13};
                 size_t local_work_size[] = {(size_t)nth0, 1, 1};
 
                 backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);

ggml/src/ggml-opencl/kernels/gemm_noshuffle_q4_0_q8_1_dp4a.cl

@@ -125,3 +125,129 @@ kernel void kernel_gemm_noshuffle_q4_0_q8_1_dp4a(
     }
 #undef NGROUPS
 }
+
+// Stage each 32-K activation block as two uint4 values.
+inline int dot8_q8a_v(uint8 qw, uint4 a0, uint4 a1) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a0.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a0.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a0.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a0.w, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a1.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a1.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a1.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a1.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_q4_0_q8_1_dp4a_act_vload4(
+        __global const ushort * src0_q,
+        __global const half   * src0_d,
+        __global const uint   * src1_qa,
+        __global const half   * src1_da,
+        __global const half   * src1_sa,
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+    __local half sh_s[TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        acc[g] = (float4)(0.0f);
+    }
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub = step >> 5;
+
+        const float d_w = (float)src0_d[rrow + sub * (uint)m];
+
+        const uint qsbase = rrow + (step >> 2) * (uint)m;
+        uint8 qw;
+        qw.s0 = EXP4(src0_q[qsbase + 0 * m]);
+        qw.s1 = EXP4(src0_q[qsbase + 1 * m]);
+        qw.s2 = EXP4(src0_q[qsbase + 2 * m]);
+        qw.s3 = EXP4(src0_q[qsbase + 3 * m]);
+        qw.s4 = EXP4(src0_q[qsbase + 4 * m]);
+        qw.s5 = EXP4(src0_q[qsbase + 5 * m]);
+        qw.s6 = EXP4(src0_q[qsbase + 6 * m]);
+        qw.s7 = EXP4(src0_q[qsbase + 7 * m]);
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+            sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+#define DOTV(T) dot8_q8a_v(qw, sh_qa4[T][0], sh_qa4[T][1])
+            rf.s0 = (float)DOTV(b+0);
+            rf.s1 = (float)DOTV(b+1);
+            rf.s2 = (float)DOTV(b+2);
+            rf.s3 = (float)DOTV(b+3);
+#undef DOTV
+
+            acc[g] += d_w * (LD4(sh_d, b) * rf - 8.0f * LD4(sh_s, b));
+        }
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) {
+            dst[(c0 + 0) * (uint)m + row] = a.s0;
+        }
+        if (c0 + 1 < (uint)n_no_padding) {
+            dst[(c0 + 1) * (uint)m + row] = a.s1;
+        }
+        if (c0 + 2 < (uint)n_no_padding) {
+            dst[(c0 + 2) * (uint)m + row] = a.s2;
+        }
+        if (c0 + 3 < (uint)n_no_padding) {
+            dst[(c0 + 3) * (uint)m + row] = a.s3;
+        }
+    }
+#undef NGROUPS
+}

ggml/src/ggml-opencl/kernels/gemm_noshuffle_q4_k_q8_1_dp4a.cl

@@ -279,3 +279,141 @@ kernel void kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg(
     }
 #undef NGROUPS
 }
+
+// Stage each 32-K activation block as two uint4 values.
+inline int dot8_q8a_v(uint8 qw, uint4 a0, uint4 a1) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a0.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a0.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a0.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a0.w, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a1.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a1.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a1.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a1.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_q4_k_q8_1_dp4a_act_vload4(
+        __global const ushort * src0_q,
+        __global const uchar  * src0_s,
+        __global const half   * src0_d,
+        __global const half   * src0_dm,
+        __global const uint   * src1_qa,
+        __global const half   * src1_da,
+        __global const half   * src1_sa,
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k,
+        uchar  mask_d6,
+        uchar  mask_d4,
+        uchar  mask_hi2
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+    __local half sh_s[TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        acc[g] = (float4)(0.0f);
+    }
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub     = step >> 5;
+        const uint sb_idx  = step / QK_K;
+        const uint sub_idx = sub & 7;
+
+        const float dd  = (float)src0_d [rrow + sb_idx * m];
+        const float dmm = (float)src0_dm[rrow + sb_idx * m];
+        global const uchar * sc = src0_s + sb_idx * K_SCALE_SIZE * (uint)m + rrow;
+        uchar sv, mn;
+        get_scale_min_k4(sub_idx, sc, (uint)m, &sv, &mn, mask_d6, mask_d4, mask_hi2);
+        const float scale = dd  * (float)sv;
+        const float minv  = dmm * (float)mn;
+
+        const uint wbase = rrow + (step >> 2) * (uint)m;
+        uint8 qw;
+        qw.s0 = EXP4(src0_q[wbase + 0 * m]);
+        qw.s1 = EXP4(src0_q[wbase + 1 * m]);
+        qw.s2 = EXP4(src0_q[wbase + 2 * m]);
+        qw.s3 = EXP4(src0_q[wbase + 3 * m]);
+        qw.s4 = EXP4(src0_q[wbase + 4 * m]);
+        qw.s5 = EXP4(src0_q[wbase + 5 * m]);
+        qw.s6 = EXP4(src0_q[wbase + 6 * m]);
+        qw.s7 = EXP4(src0_q[wbase + 7 * m]);
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+            sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+#define DOTV(T) dot8_q8a_v(qw, sh_qa4[T][0], sh_qa4[T][1])
+            rf.s0 = (float)DOTV(b+0);
+            rf.s1 = (float)DOTV(b+1);
+            rf.s2 = (float)DOTV(b+2);
+            rf.s3 = (float)DOTV(b+3);
+#undef DOTV
+            acc[g] += scale * LD4(sh_d, b) * rf - minv * LD4(sh_s, b);
+        }
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) {
+            dst[(c0 + 0) * (uint)m + row] = a.s0;
+        }
+        if (c0 + 1 < (uint)n_no_padding) {
+            dst[(c0 + 1) * (uint)m + row] = a.s1;
+        }
+        if (c0 + 2 < (uint)n_no_padding) {
+            dst[(c0 + 2) * (uint)m + row] = a.s2;
+        }
+        if (c0 + 3 < (uint)n_no_padding) {
+            dst[(c0 + 3) * (uint)m + row] = a.s3;
+        }
+    }
+#undef NGROUPS
+}

ggml/src/ggml-opencl/kernels/gemm_noshuffle_q5_k_q8_1_dp4a.cl

@@ -162,3 +162,141 @@ kernel void kernel_gemm_noshuffle_q5_k_q8_1_dp4a(
     }
 #undef NGROUPS
 }
+
+// Stage each 32-K activation block as two uint4 values.
+inline int dot8_q8a_v(uint8 qw, uint4 a0, uint4 a1) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a0.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a0.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a0.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a0.w, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a1.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a1.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a1.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a1.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_q5_k_q8_1_dp4a_act_vload4(
+        __global const ushort * src0_q,
+        __global const uchar  * src0_qh,
+        __global const uchar  * src0_s,
+        __global const half   * src0_d,
+        __global const half   * src0_dm,
+        __global const uint   * src1_qa,
+        __global const half   * src1_da,
+        __global const half   * src1_sa,
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k,
+        uchar  mask_d6,
+        uchar  mask_d4,
+        uchar  mask_hi2
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint num_superblocks = (uint)k / QK_K;
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+    __local half sh_s[TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        acc[g] = (float4)(0.0f);
+    }
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub     = step >> 5;
+        const uint sb_idx  = step / QK_K;
+        const uint sub_idx = sub & 7;
+
+        const float dd  = (float)src0_d [rrow + sb_idx * m];
+        const float dmm = (float)src0_dm[rrow + sb_idx * m];
+        global const uchar * sc = src0_s + rrow * num_superblocks * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
+        uchar sv, mn;
+        get_scale_min_k4(sub_idx, sc, &sv, &mn, mask_d6, mask_d4, mask_hi2);
+        const float scale = dd  * (float)sv;
+        const float minv  = dmm * (float)mn;
+
+        const uint wbase  = rrow + (step >> 2) * (uint)m;
+        const uint qhbase = rrow + (step >> 3) * (uint)m;
+        uint8 qw;
+#define QWU(u) ( EXP4((uint)src0_q[wbase + (uint)(u) * m]) \
+               | EXP1( (uint)((src0_qh[qhbase + (uint)((u) >> 1) * m] >> (((u) & 1) * 4)) & 0x0Fu) ) )
+        qw.s0 = QWU(0); qw.s1 = QWU(1); qw.s2 = QWU(2); qw.s3 = QWU(3);
+        qw.s4 = QWU(4); qw.s5 = QWU(5); qw.s6 = QWU(6); qw.s7 = QWU(7);
+#undef QWU
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+            sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+#define DOTV(T) dot8_q8a_v(qw, sh_qa4[T][0], sh_qa4[T][1])
+            rf.s0 = (float)DOTV(b+0);
+            rf.s1 = (float)DOTV(b+1);
+            rf.s2 = (float)DOTV(b+2);
+            rf.s3 = (float)DOTV(b+3);
+#undef DOTV
+            acc[g] += scale * LD4(sh_d, b) * rf - minv * LD4(sh_s, b);
+        }
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) {
+            dst[(c0 + 0) * (uint)m + row] = a.s0;
+        }
+        if (c0 + 1 < (uint)n_no_padding) {
+            dst[(c0 + 1) * (uint)m + row] = a.s1;
+        }
+        if (c0 + 2 < (uint)n_no_padding) {
+            dst[(c0 + 2) * (uint)m + row] = a.s2;
+        }
+        if (c0 + 3 < (uint)n_no_padding) {
+            dst[(c0 + 3) * (uint)m + row] = a.s3;
+        }
+    }
+#undef NGROUPS
+}

ggml/src/ggml-opencl/kernels/gemm_noshuffle_q6_k_q8_1_dp4a.cl

@@ -142,3 +142,117 @@ kernel void kernel_gemm_noshuffle_q6_k_q8_1_dp4a(
     }
 #undef NGROUPS
 }
+
+// Stage each 32-K activation block as two uint4 values.
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_q6_k_q8_1_dp4a_act_vload4(
+        __global const ushort * src0_ql,
+        __global const uchar  * src0_qh,
+        __global const ushort * src0_s,
+        __global const half   * src0_d,
+        __global const uint   * src1_qa,
+        __global const half   * src1_da,
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        acc[g] = (float4)(0.0f);
+    }
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub    = step >> 5;
+        const uint sb_idx = step / QK_K;
+
+        const float dd = (float)src0_d[rrow + sb_idx * m];
+        const char2 sc = as_char2(src0_s[rrow + sub * m]);
+        const float scale0 = dd * (float)sc.s0;
+        const float scale1 = dd * (float)sc.s1;
+
+        const uint wbase = rrow + (step >> 2) * (uint)m;
+        uint qw[8];
+        #pragma unroll
+        for (int u = 0; u < 8; ++u) {
+            const uint o  = wbase + (uint)u * (uint)m;
+            qw[u] = SIGN6(EXP4((uint)src0_ql[o]) | EXP2((uint)src0_qh[o]));
+        }
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+            #define DOT_TOK(j) { \
+                const uint4 a0 = sh_qa4[b + (j)][0]; \
+                const uint4 a1 = sh_qa4[b + (j)][1]; \
+                const int raw1 = dot4_q8a(qw[0], qw[1], qw[2], qw[3], a0.x, a0.y, a0.z, a0.w); \
+                const int raw2 = dot4_q8a(qw[4], qw[5], qw[6], qw[7], a1.x, a1.y, a1.z, a1.w); \
+                rf.s##j = scale0 * (float)raw1 + scale1 * (float)raw2; \
+            }
+            DOT_TOK(0); DOT_TOK(1); DOT_TOK(2); DOT_TOK(3);
+            #undef DOT_TOK
+            const float4 ad = (float4)((float)sh_d[b+0], (float)sh_d[b+1], (float)sh_d[b+2], (float)sh_d[b+3]);
+            acc[g] += ad * rf;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) {
+            dst[(c0 + 0) * (uint)m + row] = a.s0;
+        }
+        if (c0 + 1 < (uint)n_no_padding) {
+            dst[(c0 + 1) * (uint)m + row] = a.s1;
+        }
+        if (c0 + 2 < (uint)n_no_padding) {
+            dst[(c0 + 2) * (uint)m + row] = a.s2;
+        }
+        if (c0 + 3 < (uint)n_no_padding) {
+            dst[(c0 + 3) * (uint)m + row] = a.s3;
+        }
+    }
+#undef NGROUPS
+}

ggml/src/ggml-opencl/kernels/gemm_noshuffle_q8_0_q8_1_dp4a.cl

@@ -210,3 +210,125 @@ kernel void kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg(
     }
 #undef NGROUPS
 }
+
+// Stage each 32-K activation block as two uint4 values.
+inline int dot8_q8a_v(uint8 qw, uint4 a0, uint4 a1) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a0.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a0.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a0.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a0.w, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a1.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a1.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a1.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a1.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_q8_0_q8_1_dp4a_act_vload4(
+        __global const uint  * src0_q,
+        __global const half  * src0_d,
+        __global const uint  * src1_qa,
+        __global const half  * src1_da,
+        __global       float * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        acc[g] = (float4)(0.0f);
+    }
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub = step >> 5;
+
+        const float d_w = (float)src0_d[rrow + sub * (uint)m];
+
+        const uint wbase = rrow + (step >> 2) * (uint)m;
+        uint8 qw;
+        qw.s0 = src0_q[wbase + 0 * m];
+        qw.s1 = src0_q[wbase + 1 * m];
+        qw.s2 = src0_q[wbase + 2 * m];
+        qw.s3 = src0_q[wbase + 3 * m];
+        qw.s4 = src0_q[wbase + 4 * m];
+        qw.s5 = src0_q[wbase + 5 * m];
+        qw.s6 = src0_q[wbase + 6 * m];
+        qw.s7 = src0_q[wbase + 7 * m];
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+#define DOTV(T) dot8_q8a_v(qw, sh_qa4[T][0], sh_qa4[T][1])
+            rf.s0 = (float)DOTV(b+0);
+            rf.s1 = (float)DOTV(b+1);
+            rf.s2 = (float)DOTV(b+2);
+            rf.s3 = (float)DOTV(b+3);
+#undef DOTV
+            acc[g] += d_w * LD4(sh_d, b) * rf;
+        }
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) {
+            dst[(c0 + 0) * (uint)m + row] = a.s0;
+        }
+        if (c0 + 1 < (uint)n_no_padding) {
+            dst[(c0 + 1) * (uint)m + row] = a.s1;
+        }
+        if (c0 + 2 < (uint)n_no_padding) {
+            dst[(c0 + 2) * (uint)m + row] = a.s2;
+        }
+        if (c0 + 3 < (uint)n_no_padding) {
+            dst[(c0 + 3) * (uint)m + row] = a.s3;
+        }
+    }
+#undef NGROUPS
+}

ggml/src/ggml-opencl/kernels/gemv_noshuffle_q1_0_f32.cl

@@ -8,7 +8,9 @@
 #endif
 
 #define QK1_0 128
+#ifndef N_SIMDGROUP
 #define N_SIMDGROUP 4
+#endif
 
 #define dequantizeBlockAccum_q1(total, bits, scale, regB, lb)                                       \
     total += (2.0f*(float)((bits >>  0) & 1u) - 1.0f) * scale * sub_group_broadcast(regB.s0, lb+0); \
@@ -106,15 +108,19 @@ __kernel void kernel_gemv_noshuffle_q1_0_f32(
         dequantizeBlockAccum_q1(totalSum, regA.s3, scale, regB, 12);
     }
 
-    // reduction in local memory, assumes #wave = N_SIMDGROUP = 4
-    local float reduceLM[SIMDGROUP_WIDTH * 3];
-    if (groupId == 1) reduceLM[SIMDGROUP_WIDTH * 0 + slid] = totalSum;
-    if (groupId == 2) reduceLM[SIMDGROUP_WIDTH * 1 + slid] = totalSum;
-    if (groupId == 3) reduceLM[SIMDGROUP_WIDTH * 2 + slid] = totalSum;
+#if N_SIMDGROUP > 1
+    // Sum the K slices from the other subgroups.
+    local float reduceLM[SIMDGROUP_WIDTH * (N_SIMDGROUP - 1)];
+    if (groupId > 0) {
+        reduceLM[SIMDGROUP_WIDTH * (groupId - 1) + slid] = totalSum;
+    }
     barrier(CLK_LOCAL_MEM_FENCE);
-    if (groupId == 0) totalSum += reduceLM[SIMDGROUP_WIDTH * 0 + slid];
-    if (groupId == 0) totalSum += reduceLM[SIMDGROUP_WIDTH * 1 + slid];
-    if (groupId == 0) totalSum += reduceLM[SIMDGROUP_WIDTH * 2 + slid];
+    if (groupId == 0) {
+        for (uint w = 0; w < N_SIMDGROUP - 1; ++w) {
+            totalSum += reduceLM[SIMDGROUP_WIDTH * w + slid];
+        }
+    }
+#endif
 
     if (groupId == 0) {
         dst = (global float*)((global char*)dst + offsetd);
@@ -122,6 +128,8 @@ __kernel void kernel_gemv_noshuffle_q1_0_f32(
         // so when ne01 is not a multiple of the wave size the tail work-items run past
         // row ne01 and would overrun dst into the adjacent tensor. No-op / byte-identical
         // when ne01 is wave-aligned (no padding).
-        if (gid < M) dst[gid] = totalSum;
+        if (gid < M) {
+            dst[gid] = totalSum;
+        }
     }
 }

ggml/src/ggml-opencl/kernels/mul_mm_f16_f32_l4_lm.cl

@@ -4,12 +4,20 @@
 #define LOAD_VEC_B 4
 
 #define BM 64
+#ifndef BN
 #define BN 64
+#endif
 #define BK 16
 #define TM 4
+#ifndef TN
 #define TN 8
+#endif
 
-kernel void kernel_mul_mm_f16_f32_l4_lm(
+#ifndef KERNEL_NAME_LM
+#define KERNEL_NAME_LM kernel_mul_mm_f16_f32_l4_lm
+#endif
+
+kernel void KERNEL_NAME_LM(
     global half4 * src0,
     ulong offset0,
     global float4 * src1,
@@ -93,7 +101,8 @@ kernel void kernel_mul_mm_f16_f32_l4_lm(
             }
         }
 
-        for (int l = 0; l < BN; l += loadstride_b) {
+        // BN=8 uses only the threads that load a column inside the tile.
+        for (int l = 0; loadc_b + l < BN; l += loadstride_b) {
             if (ic*BN + loadc_b + l < ne11) {
                 const int idx = pos_b + (loadc_b + l) * stride_b / LOAD_VEC_B + loadr_b;
                 buf_b[(loadr_b * LOAD_VEC_B + 0) * BN + loadc_b + l] = src1[idx].s0;

ggml/src/ggml-opencl/kernels/mul_mv_mxfp4_f32_flat.cl

@@ -132,7 +132,7 @@ kernel void kernel_mul_mv_mxfp4_f32_flat(
         global float4 * y4 = (global float4 *)yb;
 
         #pragma unroll
-        for (short row = 0; row < N_R0_MXFP4; row++) {
+        for (short row = 0; row < N_R0_MXFP4 && first_row + row < ne0; row++) {
             uchar xb_e = x_e[row * nb + ib];
 #ifdef SRC0Q_IMG
             ushort4 xb_q = as_ushort4(read_imageui(src0_q, (offset_q + row * nb + ib) * 2 + it).xy);

ggml/src/ggml-opencl/kernels/rope.cl

@@ -520,7 +520,7 @@ kernel void kernel_rope_multi_f16(
     src0 = (global void*)((global char*)src0 + offset0);
     src1 = (global int*)((global char*)src1 + offset1);
     src2 = (global float*)((global char*)src2 + offset2);
-    dst = (global float*)((global char*)dst + offsetd);
+    dst = (global half*)((global char*)dst + offsetd);
 
     int i3 = get_group_id(2);
     int i2 = get_group_id(1);
@@ -714,7 +714,7 @@ kernel void kernel_rope_vision_f16(
     src0 = (global void*)((global char*)src0 + offset0);
     src1 = (global int*)((global char*)src1 + offset1);
     src2 = (global float*)((global char*)src2 + offset2);
-    dst = (global float*)((global char*)dst + offsetd);
+    dst = (global half*)((global char*)dst + offsetd);
 
     int i3 = get_group_id(2);
     int i2 = get_group_id(1);

ggml/src/ggml-sycl/convert.cpp

@@ -440,7 +440,8 @@ static void dequantize_row_iq3_xxs_sycl_reorder(const void * vx, dst_t * y, cons
         sycl::nd_range<3>(sycl::range<3>(1, 1, nb) * sycl::range<3>(1, 1, 32), sycl::range<3>(1, 1, 32)),
         [=](sycl::nd_item<3> item_ct1) {
             dequantize_block_iq3_xxs_reorder(
-                vx, y, item_ct1, nb, iq3xxs_grid, ksigns_iq2xs, kmask_iq2xs);
+                vx, y + item_ct1.get_group(2) * QK_K, item_ct1.get_group(2), item_ct1, nb,
+                iq3xxs_grid, ksigns_iq2xs, kmask_iq2xs);
         });
 }
 
@@ -472,7 +473,8 @@ static void dequantize_row_iq3_s_sycl_reorder(const void * vx, dst_t * y, const
     stream->parallel_for(
         sycl::nd_range<3>(sycl::range<3>(1, 1, nb) * sycl::range<3>(1, 1, 32), sycl::range<3>(1, 1, 32)),
         [=](sycl::nd_item<3> item_ct1) {
-            dequantize_block_iq3_s_reorder(vx, y, item_ct1, nb, kmask_iq2xs, iq3s_grid);
+            dequantize_block_iq3_s_reorder(vx, y + item_ct1.get_group(2) * QK_K, item_ct1.get_group(2),
+                                           item_ct1, nb, kmask_iq2xs, iq3s_grid);
         });
 }
 

ggml/src/ggml-sycl/dequantize.hpp

@@ -1481,12 +1481,11 @@ static void dequantize_block_iq3_xxs(const void * __restrict__ vx, dst_t * __res
 }
 
 template<typename dst_t>
-static void dequantize_block_iq3_xxs_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy,
+static void dequantize_block_iq3_xxs_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy, const int64_t i,
                                              const sycl::nd_item<3> & item_ct1, const int64_t nb,
                                              const uint32_t * iq3xxs_grid,
                                              const uint8_t * ksigns_iq2xs,
                                              const uint8_t * kmask_iq2xs) {
-    const int64_t i   = item_ct1.get_group(2);
     const int64_t tid = item_ct1.get_local_id(2);
 #if QK_K == 256
     const int64_t il = tid / 8;
@@ -1501,7 +1500,7 @@ static void dequantize_block_iq3_xxs_reorder(const void * __restrict__ vx, dst_t
     const uint32_t aux32 = gas[0] | (gas[1] << 16);
     const float dl = (float) d[i] * (0.5f + (aux32 >> 28)) * 0.5f;
     const uint8_t signs = ksigns_iq2xs[(aux32 >> (7 * il)) & 127];
-    dst_t * y = yy + i * QK_K + 32 * ib + 8 * il;
+    dst_t * y = yy + 32 * ib + 8 * il;
 #pragma unroll
     for (int j = 0; j < 4; ++j) {
         y[j+0] = dl * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f);
@@ -1510,6 +1509,7 @@ static void dequantize_block_iq3_xxs_reorder(const void * __restrict__ vx, dst_t
 #else
     GGML_UNUSED(vx);
     GGML_UNUSED(yy);
+    GGML_UNUSED(i);
     GGML_UNUSED(item_ct1);
     GGML_UNUSED(nb);
     GGML_UNUSED(iq3xxs_grid);
@@ -1551,10 +1551,9 @@ dequantize_block_iq3_s(const void *__restrict__ vx, dst_t *__restrict__ yy,
 
 template <typename dst_t>
 __dpct_inline__ static void
-dequantize_block_iq3_s_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy,
+dequantize_block_iq3_s_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy, const int64_t i,
                                const sycl::nd_item<3> & item_ct1, const int64_t nb,
                                const uint8_t * kmask_iq2xs, const uint32_t * iq3s_grid) {
-    const int64_t i   = item_ct1.get_group(2);
     const int64_t tid = item_ct1.get_local_id(2);
 #if QK_K == 256
     const int64_t il = tid / 8;
@@ -1574,7 +1573,7 @@ dequantize_block_iq3_s_reorder(const void * __restrict__ vx, dst_t * __restrict_
         iq3s_grid + (qs_ib[2*il+1] | ((qh[ib] << (7 - 2*il)) & 256)));
     const float dl = (float) d * (1 + 2 * ((scales[ib/2] >> (4 * (ib%2))) & 0xf));
     const uint8_t signs = signs_base[4 * ib + il];
-    dst_t * y = yy + i * QK_K + 32 * ib + 8 * il;
+    dst_t * y = yy + 32 * ib + 8 * il;
 #pragma unroll
     for (int j = 0; j < 4; ++j) {
         y[j+0] = dl * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f);
@@ -1583,6 +1582,7 @@ dequantize_block_iq3_s_reorder(const void * __restrict__ vx, dst_t * __restrict_
 #else
     GGML_UNUSED(vx);
     GGML_UNUSED(yy);
+    GGML_UNUSED(i);
     GGML_UNUSED(item_ct1);
     GGML_UNUSED(nb);
     GGML_UNUSED(kmask_iq2xs);

ggml/src/ggml-sycl/fattn-sparse.cpp

@@ -202,7 +202,8 @@ bool ggml_sycl_flash_attn_ext_sparse(ggml_backend_sycl_context & ctx, ggml_tenso
     uint8_t *    d_V    = v_alloc.get();
     sycl::half * d_mask = m_alloc.get();
 
-    SYCL_CHECK(CHECK_TRY_ERROR(stream->memset(d_cnt, 0, sizeof(int32_t))));
+    // reset with a kernel: on Level Zero a queue memset is not always ordered before the next kernel
+    stream->single_task([=]() { *d_cnt = 0; });
 
     sparse_fa_compact_mask(stream, (const sycl::half *) mask->data,
                            d_idx, d_cnt, n_kv, n_kv_g);

ggml/src/ggml-sycl/fattn.cpp

@@ -281,6 +281,22 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const
 void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     ggml_sycl_set_device(ctx.device);
 
+    ggml_tensor * Q = dst->src[0];
+    ggml_tensor * K = dst->src[1];
+    ggml_tensor * V = dst->src[2];
+    if (K && Q && K->ne[0] != Q->ne[0] && Q->ne[0] > 0 && K->ne[0] % Q->ne[0] == 0) {
+        const int64_t d = Q->ne[0];
+        const int64_t n_heads = (K->ne[0] * K->ne[2]) / d;
+        K->ne[0] = d;
+        K->ne[2] = n_heads;
+    }
+    if (V && Q && V->ne[0] != Q->ne[0] && Q->ne[0] > 0 && V->ne[0] % Q->ne[0] == 0) {
+        const int64_t d = Q->ne[0];
+        const int64_t n_heads = (V->ne[0] * V->ne[2]) / d;
+        V->ne[0] = d;
+        V->ne[2] = n_heads;
+    }
+
     // sparse nodes are gathered down to n_kv_max rows and re-dispatched here
     if (ggml_sycl_flash_attn_ext_sparse(ctx, dst)) {
         return;

ggml/src/ggml-sycl/fused-gemm.cpp

@@ -860,7 +860,7 @@ static void fused_dequant_gemm_tile(
     }
 }
 
-template <typename S, typename block_q_t>
+template <typename S, typename block_q_t, bool reordered = false>
 static void fused_dequant_gemm_launch(const void * src0, const typename S::tsb * packed, float * dst, const int M,
                                       const int N, const int Npad, const int K, const int ldd,
                                       const int64_t groups_n, const int64_t groups_m, dpct::queue_ptr stream) {
@@ -871,8 +871,8 @@ static void fused_dequant_gemm_launch(const void * src0, const typename S::tsb *
             sycl::nd_range<2>(sycl::range<2>(groups_n, groups_m * S::WG_SIZE), sycl::range<2>(1, S::WG_SIZE)),
             [=](sycl::nd_item<2> item) [[sycl::reqd_sub_group_size(S::SG)]] {
                 const int n0 = item.get_group(0) * S::BN;
-                fused_dequant_gemm_tile<S, block_q_t, false>((const block_q_t *) src0, packed, dst, M, Npad, K, ldd,
-                                                      n0, n0, N, tile_a, tile_c, item);
+                fused_dequant_gemm_tile<S, block_q_t, reordered>((const block_q_t *) src0, packed, dst, M, Npad, K,
+                                                      ldd, n0, n0, N, tile_a, tile_c, item);
             });
     });
 }
@@ -974,7 +974,7 @@ template <typename F> static bool fg_visit_type(ggml_type type, bool reordered,
 }
 
 template <typename S>
-static bool fg_fused_run(ggml_type src0_type, const void * src0, const void * src1, ggml_type src1_type,
+static bool fg_fused_run(ggml_type src0_type, bool src0_reordered, const void * src0, const void * src1, ggml_type src1_type,
                          float * dst, int64_t M, int64_t N, int64_t K, int64_t ldd, ggml_sycl_pool & pool,
                          dpct::queue_ptr stream) {
     const int64_t groups_n = (N + S::BN - 1) / S::BN;
@@ -992,21 +992,27 @@ static bool fg_fused_run(ggml_type src0_type, const void * src0, const void * sr
     }
 
     const typename S::tsb * packed = packed_b.get();
-    return fg_visit_type(src0_type, false, [&](auto tag) {
+    return fg_visit_type(src0_type, src0_reordered, [&](auto tag) {
         using block_q_t = typename decltype(tag)::type;
         if constexpr (fg_plain_ok<block_q_t>()) {
-            fused_dequant_gemm_launch<S, block_q_t>(src0, packed, dst, (int) M, (int) N, Npad, (int) K, (int) ldd,
-                                                    groups_n, groups_m, stream);
+            fused_dequant_gemm_launch<S, block_q_t, decltype(tag)::reordered>(src0, packed, dst, (int) M, (int) N,
+                                                                              Npad, (int) K, (int) ldd, groups_n,
+                                                                              groups_m, stream);
         }
     });
 }
 
-bool ggml_sycl_fused_dequant_gemm(ggml_type src0_type, const void * src0, const void * src1, ggml_type src1_type,
+bool ggml_sycl_fused_dequant_gemm(ggml_type src0_type, bool src0_reordered, const void * src0, const void * src1, ggml_type src1_type,
                                   int32_t src1_prec, float * dst, int64_t M, int64_t N, int64_t K, int64_t ldd,
                                   ggml_sycl_pool & pool, dpct::queue_ptr stream) {
-    // every BN columns dequantize A again, so wide N is left to the library GEMM
+    // every BN columns dequantize A again, so wide N is left to the library GEMM.
+    // src0_reordered: the persistent weight reorder may have been installed on this
+    // tensor by an earlier small-batch mul_mat; the A stage then has to decode the
+    // SoA layout. Types without an SoA stage (fg_soa<T>::supported == false)
+    // decline via fg_visit_layout and the caller falls through to the
+    // reorder-aware conversion route.
     bool plain_ok = false;
-    fg_visit_type(src0_type, false, [&](auto tag) { plain_ok = fg_plain_ok<typename decltype(tag)::type>(); });
+    fg_visit_type(src0_type, src0_reordered, [&](auto tag) { plain_ok = fg_plain_ok<typename decltype(tag)::type>(); });
     if (g_ggml_sycl_dynamic_precision == GGML_SYCL_DYNAMIC_PRECISION_F32 ||
         !ggml_sycl_xmx_gather_type_enabled(src0_type) || !plain_ok) {
         return false;
@@ -1023,7 +1029,7 @@ bool ggml_sycl_fused_dequant_gemm(ggml_type src0_type, const void * src0, const
     }
     bool launched = false;
     fg_visit_combo(combo, [&](auto s) {
-        launched = fg_fused_run<decltype(s)>(src0_type, src0, src1, src1_type, dst, M, N, K, ldd, pool, stream);
+        launched = fg_fused_run<decltype(s)>(src0_type, src0_reordered, src0, src1, src1_type, dst, M, N, K, ldd, pool, stream);
     });
     return launched;
 }

ggml/src/ggml-sycl/fused-gemm.hpp

@@ -69,7 +69,7 @@ inline bool ggml_sycl_xmx_gather_type_enabled(ggml_type src0_type) {
 // (op_params[3], [TAG_GGML_PREC]); the accumulator is f32, which meets any request.
 // q8_0 and the k-quants are not handled here, only in the grouped path below.
 // Returns false when the case is not handled (type, device, precision, or shape).
-bool ggml_sycl_fused_dequant_gemm(ggml_type src0_type, const void * src0, const void * src1, ggml_type src1_type,
+bool ggml_sycl_fused_dequant_gemm(ggml_type src0_type, bool src0_reordered, const void * src0, const void * src1, ggml_type src1_type,
                                   int32_t src1_prec, float * dst, int64_t M, int64_t N, int64_t K, int64_t ldd,
                                   ggml_sycl_pool & pool, dpct::queue_ptr stream);
 

ggml/src/ggml-sycl/getrows.cpp

@@ -301,6 +301,52 @@ static void get_rows_sycl_reorder_k(const ggml_tensor *src0, const ggml_tensor *
     GGML_UNUSED(src0);
 }
 
+// IQ3_S / IQ3_XXS in the persistent reorder layout: one work-group of 32 threads
+// dequantizes one QK_K block of one requested row, reusing the reorder-aware
+// block dequantizers that the dequantize-based mul_mat path already uses.
+template <typename F>
+static void get_rows_sycl_reorder_iq3(ggml_backend_sycl_context & ctx, const ggml_tensor * src0,
+                                      const ggml_tensor * src1, ggml_tensor * dst,
+                                      const int32_t * src1_dd, float * dst_dd,
+                                      F dequantize_block) {
+    GGML_TENSOR_BINARY_OP_LOCALS
+
+    // the persistent reorder is only installed for 2D weights
+    GGML_ASSERT(ne02 == 1 && ne03 == 1);
+    GGML_ASSERT(ne00 % QK_K == 0);
+
+    const int64_t nblk    = ne00 / QK_K;
+    const int64_t nb      = nblk * ne01; // blocks in the whole tensor
+    const int64_t ngroups = ne10 * ne11 * ne12 * nblk;
+
+    const size_t s1  = nb1 / ggml_element_size(dst);
+    const size_t s2  = nb2 / ggml_element_size(dst);
+    const size_t s3  = nb3 / ggml_element_size(dst);
+
+    const size_t s10 = nb10 / ggml_element_size(src1);
+    const size_t s11 = nb11 / ggml_element_size(src1);
+    const size_t s12 = nb12 / ggml_element_size(src1);
+
+    queue_ptr stream = ctx.stream();
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, ngroups * 32), sycl::range<3>(1, 1, 32)),
+        [=](sycl::nd_item<3> item_ct1) {
+            const int64_t g   = item_ct1.get_group(2);
+            const int64_t b   = g % nblk;
+            const int64_t r   = g / nblk;
+            const int64_t i10 = r % ne10;
+            const int64_t i11 = (r / ne10) % ne11;
+            const int64_t i12 = r / (ne10 * ne11);
+
+            const int64_t i01 = src1_dd[i10*s10 + i11*s11 + i12*s12];
+
+            float * y = dst_dd + i10*s1 + i11*s2 + i12*s3 + b*QK_K;
+            dequantize_block(y, i01*nblk + b, item_ct1, nb);
+        });
+
+    GGML_UNUSED(dst);
+}
+
 template <int qk, int qr, dequantize_kernel_f32_t dq>
 static void get_rows_sycl_f32(ggml_backend_sycl_context & ctx, const ggml_tensor *src0, const ggml_tensor *src1,
                               ggml_tensor *dst, const void *src0_dd,
@@ -460,7 +506,9 @@ void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     GGML_ASSERT(dst->nb[0] == ggml_type_size(dst->type));
 
     const int32_t * src1_i32 = (const int32_t *) dst->src[1]->data;
-    // a weight shared with MUL_MAT (e.g. tied embeddings) can be in reorder layout
+    // a weight shared with MUL_MAT (e.g. tied embeddings) can be in the
+    // persistent reorder layout (incl. IQ3_S / IQ3_XXS since the SYCL IQ3
+    // reorder work)
     const bool src0_reordered = dst->src[0]->extra &&
         ((const ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder;
     /* TODO: Refactor and remove duplicates */
@@ -506,6 +554,15 @@ void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
             src1_i32, (float *)dst->data, ctx.stream());
             break;
         case GGML_TYPE_IQ3_XXS:
+            if (src0_reordered) {
+                const void * vx = dst->src[0]->data;
+                get_rows_sycl_reorder_iq3(ctx, dst->src[0], dst->src[1], dst, src1_i32, (float *) dst->data,
+                    [=](float * y, int64_t iblk, const sycl::nd_item<3> & item_ct1, int64_t nb) {
+                        dequantize_block_iq3_xxs_reorder(vx, y, iblk, item_ct1, nb,
+                                                         iq3xxs_grid, ksigns_iq2xs, kmask_iq2xs);
+                    });
+                break;
+            }
             get_rows_sycl<QK_K, 1, dequantize_iq3_xxs>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
             src1_i32, (float *)dst->data, ctx.stream());
             break;
@@ -518,6 +575,15 @@ void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
             src1_i32, (float *)dst->data, ctx.stream());
             break;
         case GGML_TYPE_IQ3_S:
+            if (src0_reordered) {
+                const void * vx = dst->src[0]->data;
+                get_rows_sycl_reorder_iq3(ctx, dst->src[0], dst->src[1], dst, src1_i32, (float *) dst->data,
+                    [=](float * y, int64_t iblk, const sycl::nd_item<3> & item_ct1, int64_t nb) {
+                        dequantize_block_iq3_s_reorder(vx, y, iblk, item_ct1, nb,
+                                                       kmask_iq2xs, iq3s_grid);
+                    });
+                break;
+            }
             get_rows_sycl<QK_K, 1, dequantize_iq3_s>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
             src1_i32, (float *)dst->data, ctx.stream());
             break;

ggml/src/ggml-sycl/ggml-sycl.cpp

@@ -3154,9 +3154,16 @@ inline void ggml_sycl_op_mul_mat_sycl(
 #endif
 
     // dequantize inside the GEMM instead of writing the f16 weights out and reading them back; src1
-    // goes in its own type, so there is no separate conversion pass
+    // goes in its own type, so there is no separate conversion pass. The A stage
+    // decodes the weight bytes in place, so an installed persistent reorder must be
+    // visible to it: pass the flag through; types without an SoA A-stage decline
+    // inside and this falls through to the reorder-aware conversion route below.
+    bool src0_reordered = false;
+    if (const auto * src0_extra = static_cast<const ggml_tensor_extra_gpu *>(src0->extra)) {
+        src0_reordered = src0_extra->optimized_feature.reorder;
+    }
     if (ggml_is_quantized(src0->type) && ggml_is_contiguous(src0) && row_diff == src0->ne[1] &&
-        ggml_sycl_fused_dequant_gemm(src0->type, src0_dd_i, src1_ddf_i, src1->type, ggml_sycl_src1_prec(dst), dst_dd_i,
+        ggml_sycl_fused_dequant_gemm(src0->type, src0_reordered, src0_dd_i, src1_ddf_i, src1->type, ggml_sycl_src1_prec(dst), dst_dd_i,
                                      row_diff, src1_ncols, ne10, ldc, ctx.pool(), stream)) {
         return;
     }

ggml/src/ggml-zdnn/common.hpp

@@ -6,6 +6,8 @@
 
 #include "zdnn.h"
 
+#include "scratch.hpp"
+
 #include <algorithm>
 #include <cstdlib>
 #include <future>
@@ -14,9 +16,8 @@
 #include <utility>
 #include <vector>
 
-#define GGML_ZDNN_NAME          "zDNN"
-#define GGML_ZDNN_VERSION       ZDNN_VERNUM
-#define GGML_ZDNN_SCRATCH_SLOTS 3  // int8 inputs, DFloat16 partial results, partial results for grouped matmul
+#define GGML_ZDNN_NAME    "zDNN"
+#define GGML_ZDNN_VERSION ZDNN_VERNUM
 
 #define ZDNN_CHECK(stmt)                \
     do {                                \
@@ -68,44 +69,100 @@ struct ggml_backend_zdnn_context {
     // zero bias for stacked quantized matmuls, keyed by (n_stacks, ne0)
     std::map<std::pair<int64_t, int64_t>, std::unique_ptr<ggml_backend_zdnn_buffer>> zero_bias;
 
-    std::vector<int8_t> scratch_q;
-    std::vector<float>  scratch_f32;
-
-    void * scratch_ztensor[GGML_ZDNN_SCRATCH_SLOTS]      = {};  // 4K aligned
-    size_t scratch_ztensor_size[GGML_ZDNN_SCRATCH_SLOTS] = {};
+    ggml_zdnn_pool pool;  // scratch memory of the ops
 };
 
-template <typename T>
-static inline T * ggml_zdnn_resize_scratch_max(std::vector<T> & scratch, int64_t n) {
-    if ((int64_t)scratch.size() < n) {
-        scratch.resize(n);
-    }
-
-    return scratch.data();
-}
-
 //
 // Scratch buffers
 //
 
-static inline void ggml_zdnn_init_scratch_ztensor(ggml_backend_zdnn_context * ctx,
-                                                  zdnn_tensor_desc          * pre_tfm_desc,
-                                                  zdnn_tensor_desc          * tfm_desc,
-                                                  zdnn_ztensor              * ztensor,
-                                                  int32_t                     slot) {
+// the ztensor uses the memory of scratch, so scratch must outlive it
+static inline void ggml_zdnn_init_scratch_ztensor(ggml_backend_zdnn_context     * ctx,
+                                                  zdnn_tensor_desc              * pre_tfm_desc,
+                                                  zdnn_tensor_desc              * tfm_desc,
+                                                  zdnn_ztensor                  * ztensor,
+                                                  ggml_zdnn_pool_alloc<uint8_t> & scratch) {
 
     const uint64_t size = zdnn_getsize_ztensor(tfm_desc);
 
-    if (ctx->scratch_ztensor_size[slot] < size) {
-        free(ctx->scratch_ztensor[slot]);
-        ctx->scratch_ztensor[slot]      = std::aligned_alloc(4096, GGML_PAD(size, 4096));
-        ctx->scratch_ztensor_size[slot] = GGML_PAD(size, 4096);
-        GGML_ASSERT(ctx->scratch_ztensor[slot] != nullptr);
+    zdnn_init_quantized_ztensor(pre_tfm_desc, tfm_desc, 1.0f, 0.0f, ztensor);
+    ztensor->buffer      = scratch.alloc(ctx->pool, size);
+    ztensor->buffer_size = size;
+}
+
+// same as ggml_zdnn_init_scratch_ztensor, for a DLFLOAT16 ztensor
+static inline void ggml_zdnn_init_scratch_ztensor_dlf16(ggml_backend_zdnn_context     * ctx,
+                                                        zdnn_tensor_desc              * pre_tfm_desc,
+                                                        zdnn_tensor_desc              * tfm_desc,
+                                                        zdnn_ztensor                  * ztensor,
+                                                        ggml_zdnn_pool_alloc<uint8_t> & scratch) {
+
+    const uint64_t size = zdnn_getsize_ztensor(tfm_desc);
+
+    zdnn_init_ztensor(pre_tfm_desc, tfm_desc, ztensor);
+    ztensor->buffer      = scratch.alloc(ctx->pool, size);
+    ztensor->buffer_size = size;
+}
+
+//
+// Stacked ztensors
+//
+
+// view of the stacks [s0, s0 + n) of a stacked ztensor, the stacks are the outermost dim of its buffer
+// the view descs must outlive the view
+static inline zdnn_ztensor ggml_zdnn_stack_view(zdnn_tensor_desc * pre_tfm_desc,
+                                                zdnn_tensor_desc * tfm_desc,
+                                          const zdnn_ztensor     & src,
+                                                int64_t            n_stacks,
+                                                int64_t            s0,
+                                                int64_t            n) {
+
+    GGML_ASSERT(src.buffer_size % n_stacks == 0);
+    const uint64_t stack_bytes = src.buffer_size / n_stacks;
+
+    *pre_tfm_desc = *src.pre_transformed_desc;
+    *tfm_desc     = *src.transformed_desc;
+
+    switch (pre_tfm_desc->layout) {
+        case ZDNN_3DS:
+            {
+                pre_tfm_desc->dim3 = n;
+            } break;
+        case ZDNN_2DS:
+            {
+                pre_tfm_desc->dim2 = n;
+            } break;
+        default:
+            GGML_ABORT("%s: unsupported stacked layout", __func__);
     }
 
-    zdnn_init_quantized_ztensor(pre_tfm_desc, tfm_desc, 1.0f, 0.0f, ztensor);
-    ztensor->buffer      = ctx->scratch_ztensor[slot];
-    ztensor->buffer_size = size;
+    tfm_desc->dim4 = n;
+    GGML_ASSERT(zdnn_getsize_ztensor(tfm_desc) == n*stack_bytes);
+
+    zdnn_ztensor view = src;
+    view.pre_transformed_desc = pre_tfm_desc;
+    view.transformed_desc     = tfm_desc;
+    view.buffer               = (char *) src.buffer + s0*stack_bytes;
+    view.buffer_size          = n * stack_bytes;
+    return view;
+}
+
+// zero bias of n_stacks stacks, created once per shape as it never changes
+static inline const zdnn_ztensor & ggml_zdnn_zero_bias(ggml_backend_zdnn_context * ctx, int64_t n_stacks, int64_t ne0) {
+    std::unique_ptr<ggml_backend_zdnn_buffer> & bias = ctx->zero_bias[{ n_stacks, ne0 }];
+
+    if (!bias) {
+        bias = std::make_unique<ggml_backend_zdnn_buffer>();
+
+        std::vector<float> zeros(n_stacks * ne0, 0.0f);
+
+        zdnn_init_pre_transformed_desc(ZDNN_2DS, FP32, &bias->pre_tfm_desc, n_stacks, ne0);
+        ZDNN_CHECK(zdnn_generate_transformed_desc(&bias->pre_tfm_desc, &bias->tfm_desc));
+        ZDNN_CHECK(zdnn_init_ztensor_with_malloc(&bias->pre_tfm_desc, &bias->tfm_desc, &bias->ztensor));
+        ZDNN_CHECK(zdnn_transform_ztensor(&bias->ztensor, zeros.data()));
+    }
+
+    return bias->ztensor;
 }
 
 //

ggml/src/ggml-zdnn/fattn.cpp

@@ -0,0 +1,255 @@
+#include "ggml.h"
+#include "fattn.hpp"
+
+#include <algorithm>
+#include <cstring>
+
+// DLFLOAT16 has no infinity, so masked scores use the lowest FP16 value, exp of which is still 0
+static inline float ggml_zdnn_fattn_mask_value(ggml_fp16_t m) {
+    return std::max(GGML_FP16_TO_FP32(m), -65504.0f);
+}
+
+bool ggml_zdnn_supports_flash_attn_ext(const ggml_backend_zdnn_device_context * ctx_dev, const ggml_tensor * op) {
+    const ggml_tensor * q    = op->src[0];
+    const ggml_tensor * k    = op->src[1];
+    const ggml_tensor * v    = op->src[2];
+    const ggml_tensor * mask = op->src[3];
+
+    float max_bias      = 0.0f;
+    float logit_softcap = 0.0f;
+    memcpy(&max_bias,      (const float *) op->op_params + 1, sizeof(float));
+    memcpy(&logit_softcap, (const float *) op->op_params + 2, sizeof(float));
+
+    // sinks, alibi, softcap and sparse k/v sets are not implemented
+    if (op->src[4] != nullptr || max_bias != 0.0f || logit_softcap != 0.0f || ggml_get_op_params_i32(op, 4) != 0) {
+        return false;
+    }
+
+    if (op->type != GGML_TYPE_F32 || q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16) {
+        return false;
+    }
+
+    if (mask != nullptr && (mask->type != GGML_TYPE_F16 || mask->ne[2] != 1 || mask->ne[3] != 1)) {
+        return false;
+    }
+
+    if (q->ne[3] != 1 || k->ne[3] != 1 || v->ne[3] != 1 || k->ne[2] != v->ne[2] || q->ne[2] % k->ne[2] != 0) {
+        return false;
+    }
+
+    // rows are gathered with their strides, but each row has to be contiguous
+    if (q->nb[0] != sizeof(float) || k->nb[0] != sizeof(ggml_fp16_t) || v->nb[0] != sizeof(ggml_fp16_t)) {
+        return false;
+    }
+
+    const int64_t max_dim = ctx_dev->max_size;
+    const int64_t group   = q->ne[2] / k->ne[2];
+
+    return q->ne[0] <= max_dim && v->ne[0] <= max_dim && k->ne[1] <= max_dim &&
+           q->ne[2] <= max_dim && group * q->ne[1] <= max_dim;
+}
+
+void ggml_zdnn_flash_attn_ext(ggml_backend_zdnn_context * ctx, ggml_tensor * dst) {
+    const ggml_tensor * q    = dst->src[0];
+    const ggml_tensor * k    = dst->src[1];
+    const ggml_tensor * v    = dst->src[2];
+    const ggml_tensor * mask = dst->src[3];
+
+    const int64_t d_k       = q->ne[0];
+    const int64_t n_q       = q->ne[1];
+    const int64_t n_head    = q->ne[2];
+    const int64_t n_kv      = k->ne[1];
+    const int64_t n_head_kv = k->ne[2];
+    const int64_t d_v       = v->ne[0];
+    const int64_t group     = n_head / n_head_kv;
+
+    float scale = 1.0f;
+    memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float));
+
+    const int32_t n_threads = n_q*n_kv >= 65536 ? ctx->n_threads : 1;
+
+    // one query runs all heads in one call with the mask row as the K*Q bias, more queries run per head and add the mask
+    const bool single = n_q == 1;
+
+    const int64_t n_mask = mask ? (single ? n_head_kv*n_kv : n_q*n_kv) : 0;
+
+    ggml_zdnn_pool_alloc<float>       q_host_alloc(ctx->pool, n_head*n_q*d_k);
+    ggml_zdnn_pool_alloc<float>       o_host_alloc(ctx->pool, n_head*n_q*d_v);
+    ggml_zdnn_pool_alloc<float>       mask_host_alloc(ctx->pool, n_mask);
+    ggml_zdnn_pool_alloc<ggml_fp16_t> k_host_alloc(ctx->pool, n_head_kv*n_kv*d_k);
+    ggml_zdnn_pool_alloc<ggml_fp16_t> v_host_alloc(ctx->pool, n_head_kv*n_kv*d_v);
+
+    float       * q_host    = q_host_alloc.get();
+    float       * o_host    = o_host_alloc.get();
+    float       * mask_host = mask_host_alloc.get();
+    ggml_fp16_t * k_host    = k_host_alloc.get();
+    ggml_fp16_t * v_host    = v_host_alloc.get();
+
+    // q is [n_head][n_q][d_k] with the scale folded in, k and v are [n_head_kv][n_kv][d]
+    ggml_zdnn_parallel_for(n_threads, n_head, [&](int64_t h_start, int64_t h_end) {
+        for (int64_t h = h_start; h < h_end; h++) {
+            for (int64_t iq = 0; iq < n_q; iq++) {
+                const float * src = (const float *)((const char *) q->data + iq*q->nb[1] + h*q->nb[2]);
+                float       * row = q_host + (h*n_q + iq)*d_k;
+                for (int64_t i0 = 0; i0 < d_k; i0++) {
+                    row[i0] = src[i0] * scale;
+                }
+            }
+        }
+    });
+
+    ggml_zdnn_parallel_for(n_threads, n_head_kv*n_kv, [&](int64_t r_start, int64_t r_end) {
+        for (int64_t r = r_start; r < r_end; r++) {
+            const int64_t hk = r / n_kv;
+            const int64_t ik = r % n_kv;
+            memcpy(k_host + r*d_k, (const char *) k->data + ik*k->nb[1] + hk*k->nb[2], d_k*sizeof(ggml_fp16_t));
+            memcpy(v_host + r*d_v, (const char *) v->data + ik*v->nb[1] + hk*v->nb[2], d_v*sizeof(ggml_fp16_t));
+        }
+    });
+
+    zdnn_tensor_desc q_pre_tfm_desc, q_tfm_desc;
+    zdnn_tensor_desc k_pre_tfm_desc, k_tfm_desc;
+    zdnn_tensor_desc v_pre_tfm_desc, v_tfm_desc;
+    zdnn_tensor_desc o_pre_tfm_desc, o_tfm_desc;
+    zdnn_ztensor     q_ztensor, k_ztensor, v_ztensor, o_ztensor;
+
+    ggml_zdnn_pool_alloc<uint8_t> q_ztensor_alloc, k_ztensor_alloc, v_ztensor_alloc, o_ztensor_alloc;
+
+    // a single query keeps the heads of a kv head in one stack, which is the same memory order as one head per stack
+    if (single) {
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &q_pre_tfm_desc, n_head_kv, group, d_k);
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &o_pre_tfm_desc, n_head_kv, group, d_v);
+    } else {
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &q_pre_tfm_desc, n_head, n_q, d_k);
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &o_pre_tfm_desc, n_head, n_q, d_v);
+    }
+    zdnn_init_pre_transformed_desc(ZDNN_3DS, FP16, &k_pre_tfm_desc, n_head_kv, n_kv, d_k);
+    zdnn_init_pre_transformed_desc(ZDNN_3DS, FP16, &v_pre_tfm_desc, n_head_kv, n_kv, d_v);
+
+    ZDNN_CHECK(zdnn_generate_transformed_desc(&q_pre_tfm_desc, &q_tfm_desc));
+    ZDNN_CHECK(zdnn_generate_transformed_desc(&k_pre_tfm_desc, &k_tfm_desc));
+    ZDNN_CHECK(zdnn_generate_transformed_desc(&v_pre_tfm_desc, &v_tfm_desc));
+    ZDNN_CHECK(zdnn_generate_transformed_desc(&o_pre_tfm_desc, &o_tfm_desc));
+
+    ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &q_pre_tfm_desc, &q_tfm_desc, &q_ztensor, q_ztensor_alloc);
+    ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &k_pre_tfm_desc, &k_tfm_desc, &k_ztensor, k_ztensor_alloc);
+    ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &v_pre_tfm_desc, &v_tfm_desc, &v_ztensor, v_ztensor_alloc);
+    ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &o_pre_tfm_desc, &o_tfm_desc, &o_ztensor, o_ztensor_alloc);
+
+    ZDNN_CHECK(zdnn_transform_ztensor(&q_ztensor, q_host));
+    ZDNN_CHECK(zdnn_transform_ztensor(&k_ztensor, k_host));
+    ZDNN_CHECK(zdnn_transform_ztensor(&v_ztensor, v_host));
+
+    ggml_zdnn_pool_alloc<uint8_t> softmax_save_area(ctx->pool, ZDNN_SOFTMAX_SAVEAREA_SIZE);
+
+    zdnn_tensor_desc s_pre_tfm_desc, s_tfm_desc;
+    zdnn_ztensor     s_ztensor, s2_ztensor;
+
+    ggml_zdnn_pool_alloc<uint8_t> s_ztensor_alloc, s2_ztensor_alloc;
+
+    if (single) {
+        // the bias of K*Q is the mask row of the query, repeated for every kv head
+        zdnn_tensor_desc bias_pre_tfm_desc, bias_tfm_desc;
+        zdnn_ztensor     bias_ztensor;
+
+        ggml_zdnn_pool_alloc<uint8_t> bias_ztensor_alloc;
+
+        const zdnn_ztensor * bias = &ggml_zdnn_zero_bias(ctx, n_head_kv, n_kv);
+        if (mask != nullptr) {
+            const ggml_fp16_t * mask_row = (const ggml_fp16_t *) mask->data;
+            for (int64_t hk = 0; hk < n_head_kv; hk++) {
+                for (int64_t ik = 0; ik < n_kv; ik++) {
+                    mask_host[hk*n_kv + ik] = ggml_zdnn_fattn_mask_value(mask_row[ik]);
+                }
+            }
+
+            zdnn_init_pre_transformed_desc(ZDNN_2DS, FP32, &bias_pre_tfm_desc, n_head_kv, n_kv);
+            ZDNN_CHECK(zdnn_generate_transformed_desc(&bias_pre_tfm_desc, &bias_tfm_desc));
+            ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &bias_pre_tfm_desc, &bias_tfm_desc, &bias_ztensor, bias_ztensor_alloc);
+            ZDNN_CHECK(zdnn_transform_ztensor(&bias_ztensor, mask_host));
+            bias = &bias_ztensor;
+        }
+
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &s_pre_tfm_desc, n_head_kv, group, n_kv);
+        ZDNN_CHECK(zdnn_generate_transformed_desc(&s_pre_tfm_desc, &s_tfm_desc));
+        ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &s_pre_tfm_desc, &s_tfm_desc, &s_ztensor,  s_ztensor_alloc);
+        ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &s_pre_tfm_desc, &s_tfm_desc, &s2_ztensor, s2_ztensor_alloc);
+
+        ZDNN_CHECK(zdnn_matmul_transpose_op(&q_ztensor, &k_ztensor, bias, false, true, MATMUL_OP_ADDITION, &s_ztensor));
+        ZDNN_CHECK(zdnn_softmax(&s_ztensor, softmax_save_area.get(), SOFTMAX_ACT_NONE, &s2_ztensor));
+        ZDNN_CHECK(zdnn_matmul_op(&s2_ztensor, &v_ztensor, &ggml_zdnn_zero_bias(ctx, n_head_kv, d_v), MATMUL_OP_ADDITION, &o_ztensor));
+    } else {
+        zdnn_tensor_desc mask_pre_tfm_desc, mask_tfm_desc;
+        zdnn_ztensor     mask_ztensor;
+
+        ggml_zdnn_pool_alloc<uint8_t> mask_ztensor_alloc;
+
+        if (mask != nullptr) {
+            ggml_zdnn_parallel_for(n_threads, n_q, [&](int64_t iq_start, int64_t iq_end) {
+                for (int64_t iq = iq_start; iq < iq_end; iq++) {
+                    const ggml_fp16_t * mask_row = (const ggml_fp16_t *)((const char *) mask->data + iq*mask->nb[1]);
+                    for (int64_t ik = 0; ik < n_kv; ik++) {
+                        mask_host[iq*n_kv + ik] = ggml_zdnn_fattn_mask_value(mask_row[ik]);
+                    }
+                }
+            });
+
+            zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &mask_pre_tfm_desc, 1, n_q, n_kv);
+            ZDNN_CHECK(zdnn_generate_transformed_desc(&mask_pre_tfm_desc, &mask_tfm_desc));
+            ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &mask_pre_tfm_desc, &mask_tfm_desc, &mask_ztensor, mask_ztensor_alloc);
+            ZDNN_CHECK(zdnn_transform_ztensor(&mask_ztensor, mask_host));
+        }
+
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &s_pre_tfm_desc, 1, n_q, n_kv);
+        ZDNN_CHECK(zdnn_generate_transformed_desc(&s_pre_tfm_desc, &s_tfm_desc));
+        ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &s_pre_tfm_desc, &s_tfm_desc, &s_ztensor,  s_ztensor_alloc);
+        ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &s_pre_tfm_desc, &s_tfm_desc, &s2_ztensor, s2_ztensor_alloc);
+
+        const zdnn_ztensor & k_bias = ggml_zdnn_zero_bias(ctx, 1, n_kv);
+        const zdnn_ztensor & v_bias = ggml_zdnn_zero_bias(ctx, 1, d_v);
+
+        zdnn_tensor_desc view_pre_tfm_desc[4];
+        zdnn_tensor_desc view_tfm_desc[4];
+
+        for (int64_t h = 0; h < n_head; h++) {
+            const zdnn_ztensor q_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[0], &view_tfm_desc[0], q_ztensor, n_head,    h,         1);
+            const zdnn_ztensor k_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[1], &view_tfm_desc[1], k_ztensor, n_head_kv, h / group, 1);
+            const zdnn_ztensor v_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[2], &view_tfm_desc[2], v_ztensor, n_head_kv, h / group, 1);
+            zdnn_ztensor       o_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[3], &view_tfm_desc[3], o_ztensor, n_head,    h,         1);
+
+            s_ztensor.is_transformed  = false;
+            s2_ztensor.is_transformed = false;
+
+            ZDNN_CHECK(zdnn_matmul_transpose_op(&q_view, &k_view, &k_bias, false, true, MATMUL_OP_ADDITION, &s_ztensor));
+
+            zdnn_ztensor * scores = &s_ztensor;
+            zdnn_ztensor * probs  = &s2_ztensor;
+            if (mask != nullptr) {
+                ZDNN_CHECK(zdnn_add(&s_ztensor, &mask_ztensor, &s2_ztensor));
+                scores = &s2_ztensor;
+                probs  = &s_ztensor;
+                probs->is_transformed = false;
+            }
+
+            ZDNN_CHECK(zdnn_softmax(scores, softmax_save_area.get(), SOFTMAX_ACT_NONE, probs));
+            ZDNN_CHECK(zdnn_matmul_op(probs, &v_view, &v_bias, MATMUL_OP_ADDITION, &o_view));
+        }
+
+        o_ztensor.is_transformed = true;
+    }
+
+    ZDNN_CHECK(zdnn_transform_origtensor(&o_ztensor, o_host));
+
+    // the result is [d_v][n_head][n_q]
+    for (int64_t iq = 0; iq < n_q; iq++) {
+        for (int64_t h = 0; h < n_head; h++) {
+            memcpy((char *) dst->data + h*dst->nb[1] + iq*dst->nb[2], o_host + (h*n_q + iq)*d_v, d_v*sizeof(float));
+        }
+    }
+
+    // the result is written to dst->data directly, so the dst ztensor is stale
+    ggml_backend_zdnn_buffer * dst_extra = (ggml_backend_zdnn_buffer *) dst->extra;
+    if (dst_extra != nullptr) {
+        dst_extra->ztensor.is_transformed = false;
+    }
+}

ggml/src/ggml-zdnn/fattn.hpp

@@ -0,0 +1,10 @@
+#ifndef GGML_ZDNN_FATTN_HPP
+#define GGML_ZDNN_FATTN_HPP
+
+#include "common.hpp"
+
+bool ggml_zdnn_supports_flash_attn_ext(const ggml_backend_zdnn_device_context * ctx_dev, const ggml_tensor * op);
+
+void ggml_zdnn_flash_attn_ext(ggml_backend_zdnn_context * ctx, ggml_tensor * dst);
+
+#endif  // GGML_ZDNN_FATTN_HPP

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

@@ -3,6 +3,7 @@
 #include "ggml-backend-impl.h"
 
 #include "ggml-zdnn/common.hpp"
+#include "ggml-zdnn/fattn.hpp"
 #include "ggml-zdnn/mmf.hpp"
 #include "ggml-zdnn/mmq.hpp"
 #include "ggml-zdnn/utils.hpp"
@@ -37,6 +38,11 @@ static bool ggml_zdnn_compute_forward(
                 ggml_zdnn_compute_forward_mul_mat(ctx, dst);
             } break;
 
+        case GGML_OP_FLASH_ATTN_EXT:
+            {
+                ggml_zdnn_flash_attn_ext(ctx, dst);
+            } break;
+
         default:
             return false;
     }
@@ -120,6 +126,9 @@ static bool ggml_zdnn_supports_op(const ggml_backend_zdnn_device_context * ctx_d
                 }
             } break;
 
+        case GGML_OP_FLASH_ATTN_EXT:
+            return ggml_zdnn_supports_flash_attn_ext(ctx_dev, op);
+
         default:
             return false;
     }
@@ -205,9 +214,6 @@ static void ggml_zdnn_free(ggml_backend_zdnn_context * ctx) {
     for (const auto & it : ctx->zero_bias) {
         ZDNN_CHECK(zdnn_free_ztensor_buffer(&it.second->ztensor));
     }
-    for (void * scratch : ctx->scratch_ztensor) {
-        free(scratch);
-    }
 
     delete ctx;
 }
@@ -309,7 +315,19 @@ static void ggml_backend_zdnn_buffer_memset_tensor(ggml_backend_buffer_t buffer,
 static void ggml_backend_zdnn_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
     memcpy((char *)tensor->data + offset, data, size);
 
+    // views share the ztensor of their source, which is now stale, and empty tensors have none
+    if (tensor->view_src != nullptr) {
+        ggml_backend_zdnn_buffer * src_extra = (ggml_backend_zdnn_buffer *)tensor->view_src->extra;
+        if (src_extra != nullptr && src_extra->ztensor.is_transformed) {
+            zdnn_reset_ztensor(&src_extra->ztensor);
+        }
+        return;
+    }
+
     ggml_backend_zdnn_buffer * extra = (ggml_backend_zdnn_buffer *)tensor->extra;
+    if (extra == nullptr) {
+        return;
+    }
 
     if (ggml_is_quantized(tensor->type)) {
         GGML_ASSERT(offset == 0 && size == ggml_nbytes(tensor));
@@ -319,8 +337,8 @@ static void ggml_backend_zdnn_buffer_set_tensor(ggml_backend_buffer_t buffer, gg
 
     // Fixes the LLAMA_SET_ROWS bug
     // see: https://github.com/ggml-org/llama.cpp/issues/15414
-    // compute tensors are transformed on first use instead, as not every op reads the ztensor
-    if (tensor->buffer->usage == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+    // only weights are transformed here, others on first use, as not every op reads the ztensor and some data (masks) is not finite
+    if (tensor->buffer->usage != GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
         zdnn_reset_ztensor(&extra->ztensor);
         return;
     }

ggml/src/ggml-zdnn/mmq.cpp

@@ -6,63 +6,6 @@
 #include <future>
 #include <vector>
 
-// view of the stacks [s0, s0 + n) of a stacked ztensor, the stacks are the outermost dim of its buffer
-// the view descs must outlive the view
-static zdnn_ztensor ggml_zdnn_stack_view(zdnn_tensor_desc * pre_tfm_desc,
-                                         zdnn_tensor_desc * tfm_desc,
-                                   const zdnn_ztensor     & src,
-                                         int64_t            n_stacks,
-                                         int64_t            s0,
-                                         int64_t            n) {
-
-    GGML_ASSERT(src.buffer_size % n_stacks == 0);
-    const uint64_t stack_bytes = src.buffer_size / n_stacks;
-
-    *pre_tfm_desc = *src.pre_transformed_desc;
-    *tfm_desc     = *src.transformed_desc;
-
-    switch (pre_tfm_desc->layout) {
-        case ZDNN_3DS:
-            {
-                pre_tfm_desc->dim3 = n;
-            } break;
-        case ZDNN_2DS:
-            {
-                pre_tfm_desc->dim2 = n;
-            } break;
-        default:
-            GGML_ABORT("%s: unsupported stacked layout", __func__);
-    }
-
-    tfm_desc->dim4 = n;
-    GGML_ASSERT(zdnn_getsize_ztensor(tfm_desc) == n*stack_bytes);
-
-    zdnn_ztensor view = src;
-    view.pre_transformed_desc = pre_tfm_desc;
-    view.transformed_desc     = tfm_desc;
-    view.buffer               = (char *) src.buffer + s0*stack_bytes;
-    view.buffer_size          = n * stack_bytes;
-    return view;
-}
-
-// zero bias of n_stacks stacks, created once per shape as it never changes
-static const zdnn_ztensor & ggml_zdnn_zero_bias(ggml_backend_zdnn_context * ctx, int64_t n_stacks, int64_t ne0) {
-    std::unique_ptr<ggml_backend_zdnn_buffer> & bias = ctx->zero_bias[{ n_stacks, ne0 }];
-
-    if (!bias) {
-        bias = std::make_unique<ggml_backend_zdnn_buffer>();
-
-        std::vector<float> zeros(n_stacks * ne0, 0.0f);
-
-        zdnn_init_pre_transformed_desc(ZDNN_2DS, FP32, &bias->pre_tfm_desc, n_stacks, ne0);
-        ZDNN_CHECK(zdnn_generate_transformed_desc(&bias->pre_tfm_desc, &bias->tfm_desc));
-        ZDNN_CHECK(zdnn_init_ztensor_with_malloc(&bias->pre_tfm_desc, &bias->tfm_desc, &bias->ztensor));
-        ZDNN_CHECK(zdnn_transform_ztensor(&bias->ztensor, zeros.data()));
-    }
-
-    return bias->ztensor;
-}
-
 // the weights are stacked into groups along ne00, and each group of each row of the inputs and weights has its own scale
 // zDNN multiplies the groups in stacked calls, then the partial results of each group are scaled and summed here
 static void ggml_zdnn_mul_mat_q_grouped(ggml_backend_zdnn_context * ctx,
@@ -81,18 +24,20 @@ static void ggml_zdnn_mul_mat_q_grouped(ggml_backend_zdnn_context * ctx,
     GGML_ASSERT(n_groups * gs >= ne10 && (n_groups - 1) * gs < ne10);
 
     // the rounding errors of the inputs are stacked after the inputs, see ggml_zdnn_quantize_inputs
-    int8_t * inputs_q = ggml_zdnn_resize_scratch_max(ctx->scratch_q, 2 * n_groups * gs * ne11);
+    ggml_zdnn_pool_alloc<int8_t> inputs_q_alloc(ctx->pool, 2 * n_groups * gs * ne11);
+    int8_t * inputs_q = inputs_q_alloc.get();
     std::vector<float> inputs_scales(n_groups * ne11);
     ggml_zdnn_quantize_inputs(n_threads, inputs, gs, weights_extra->channel_scales.data(),
                             inputs_q, inputs_q + n_groups * gs * ne11, inputs_scales.data());
 
-    zdnn_tensor_desc inputs_pre_tfm_desc;
-    zdnn_tensor_desc inputs_tfm_desc;
-    zdnn_ztensor     inputs_ztensor;
+    zdnn_tensor_desc              inputs_pre_tfm_desc;
+    zdnn_tensor_desc              inputs_tfm_desc;
+    zdnn_ztensor                  inputs_ztensor;
+    ggml_zdnn_pool_alloc<uint8_t> inputs_ztensor_alloc;
 
     zdnn_init_pre_transformed_desc(ZDNN_3DS, INT8, &inputs_pre_tfm_desc, 2 * n_groups, ne11, gs);
     ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&inputs_pre_tfm_desc, QUANTIZED_INT8, &inputs_tfm_desc));
-    ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, 0);
+    ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, inputs_ztensor_alloc);
     ZDNN_CHECK(zdnn_transform_quantized_ztensor(&inputs_ztensor, false, -127, 127, inputs_q));
 
     // each group of the partial results is read back as its own 2D ztensor
@@ -116,19 +61,21 @@ static void ggml_zdnn_mul_mat_q_grouped(ggml_backend_zdnn_context * ctx,
 
     // two partial buffers, so the next chunk runs on the NNPA while this one is summed
     // each holds the results of a chunk of inputs, followed by the results of their rounding errors
-    zdnn_tensor_desc partial_pre_tfm_desc;
-    zdnn_tensor_desc partial_tfm_desc;
-    zdnn_ztensor     partial_ztensor[2];
+    zdnn_tensor_desc              partial_pre_tfm_desc;
+    zdnn_tensor_desc              partial_tfm_desc;
+    zdnn_ztensor                  partial_ztensor[2];
+    ggml_zdnn_pool_alloc<uint8_t> partial_ztensor_alloc[2];
 
     zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &partial_pre_tfm_desc, 2 * n_chunk, ne11, ne0);
     ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&partial_pre_tfm_desc, QUANTIZED_DLFLOAT16, &partial_tfm_desc));
     GGML_ASSERT(zdnn_getsize_ztensor(&partial_tfm_desc) == (uint64_t)(2 * n_chunk) * group_bytes);
     for (int32_t i = 0; i < 2; i++) {
-        ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor[i], 1 + i);
+        ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor[i], partial_ztensor_alloc[i]);
     }
 
     const float * weights_scales = weights_extra->scales.data();
-    float * partial     = ggml_zdnn_resize_scratch_max(ctx->scratch_f32, 2 * ne11 * ne0);
+    ggml_zdnn_pool_alloc<float> partial_alloc(ctx->pool, 2 * ne11 * ne0);
+    float * partial     = partial_alloc.get();
     float * partial_res = partial + ne11 * ne0;
 
     const auto sum_chunk = [&](const zdnn_ztensor & chunk, int64_t g0, int64_t n) {
@@ -262,27 +209,30 @@ void ggml_zdnn_mul_mat_q(
 
     // zDNN takes one scale per zTensor, so quantize each row of the inputs with its own scale
     // the rounding errors of the inputs are stacked after the inputs, see ggml_zdnn_quantize_inputs
-    int8_t * inputs_q = ggml_zdnn_resize_scratch_max(ctx->scratch_q, 2 * ne10 * ne11);
+    ggml_zdnn_pool_alloc<int8_t> inputs_q_alloc(ctx->pool, 2 * ne10 * ne11);
+    int8_t * inputs_q = inputs_q_alloc.get();
     std::vector<float> inputs_scales(ne11);
     ggml_zdnn_quantize_inputs(n_threads, inputs, ne10, weights_extra->channel_scales.data(),
                             inputs_q, inputs_q + ne10 * ne11, inputs_scales.data());
 
-    zdnn_tensor_desc inputs_pre_tfm_desc;
-    zdnn_tensor_desc inputs_tfm_desc;
-    zdnn_ztensor     inputs_ztensor;
+    zdnn_tensor_desc              inputs_pre_tfm_desc;
+    zdnn_tensor_desc              inputs_tfm_desc;
+    zdnn_ztensor                  inputs_ztensor;
+    ggml_zdnn_pool_alloc<uint8_t> inputs_ztensor_alloc;
 
     zdnn_init_pre_transformed_desc(ZDNN_3DS, INT8, &inputs_pre_tfm_desc, 2, ne11, ne10);
     ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&inputs_pre_tfm_desc, QUANTIZED_INT8, &inputs_tfm_desc));
-    ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, 0);
+    ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, inputs_ztensor_alloc);
     ZDNN_CHECK(zdnn_transform_quantized_ztensor(&inputs_ztensor, false, -127, 127, inputs_q));
 
-    zdnn_tensor_desc partial_pre_tfm_desc;
-    zdnn_tensor_desc partial_tfm_desc;
-    zdnn_ztensor     partial_ztensor;
+    zdnn_tensor_desc              partial_pre_tfm_desc;
+    zdnn_tensor_desc              partial_tfm_desc;
+    zdnn_ztensor                  partial_ztensor;
+    ggml_zdnn_pool_alloc<uint8_t> partial_ztensor_alloc;
 
     zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &partial_pre_tfm_desc, 2, ne11, ne0);
     ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&partial_pre_tfm_desc, QUANTIZED_DLFLOAT16, &partial_tfm_desc));
-    ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor, 1);
+    ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor, partial_ztensor_alloc);
 
     // both stacks are multiplied by the same unstacked weights in one broadcast call
     // the bias is all zeros, so it is already the pre-computed bias
@@ -300,7 +250,8 @@ void ggml_zdnn_mul_mat_q(
     const uint64_t stack_bytes = zdnn_getsize_ztensor(&stack_tfm_desc);
     GGML_ASSERT(partial_ztensor.buffer_size == 2*stack_bytes);
 
-    float * partial[2] = { (float *)output->data, ggml_zdnn_resize_scratch_max(ctx->scratch_f32, ne11 * ne0) };
+    ggml_zdnn_pool_alloc<float> partial_res_alloc(ctx->pool, ne11 * ne0);
+    float * partial[2] = { (float *)output->data, partial_res_alloc.get() };
 
     ggml_zdnn_parallel_for(n_threads, 2, [&](int64_t s_start, int64_t s_end) {
         zdnn_ztensor stack_ztensor;

ggml/src/ggml-zdnn/scratch.cpp

@@ -0,0 +1,45 @@
+#include "ggml.h"
+#include "scratch.hpp"
+
+#include <algorithm>
+#include <cstdint>
+#include <cstdlib>
+
+ggml_zdnn_pool::~ggml_zdnn_pool() {
+    for (const buffer & b : cached) {
+        std::free(b.ptr);
+    }
+}
+
+void * ggml_zdnn_pool::alloc(size_t size, size_t * actual_size) {
+    int64_t ibest = -1;
+    for (int64_t i = 0; i < (int64_t)cached.size(); i++) {
+        if (cached[i].size >= size && (ibest < 0 || cached[i].size < cached[ibest].size)) {
+            ibest = i;
+        }
+    }
+
+    if (ibest >= 0) {
+        void * ptr   = cached[ibest].ptr;
+        *actual_size = cached[ibest].size;
+        cached.erase(cached.begin() + ibest);
+        return ptr;
+    }
+
+    // every cached buffer is too small, free the largest so the pool does not grow while sizes grow with the context
+    if (!cached.empty()) {
+        const auto largest = std::max_element(cached.begin(), cached.end(), [](const buffer & a, const buffer & b) { return a.size < b.size; });
+        std::free(largest->ptr);
+        cached.erase(largest);
+    }
+
+    // 5% more, so slowly growing sizes do not allocate every time
+    *actual_size = GGML_PAD(std::max<size_t>(1, (size_t)(1.05 * size)), 4096);
+    void * ptr   = std::aligned_alloc(4096, *actual_size);
+    GGML_ASSERT(ptr != nullptr);
+    return ptr;
+}
+
+void ggml_zdnn_pool::free(void * ptr, size_t size) {
+    cached.push_back({ ptr, size });
+}

ggml/src/ggml-zdnn/scratch.hpp

@@ -0,0 +1,61 @@
+#ifndef GGML_ZDNN_SCRATCH_HPP
+#define GGML_ZDNN_SCRATCH_HPP
+
+#include "ggml.h"
+
+#include <cstddef>
+#include <vector>
+
+// 4K aligned host buffers that are cached for reuse, see ggml_cuda_pool_leg
+struct ggml_zdnn_pool {
+    struct buffer {
+        void * ptr  = nullptr;
+        size_t size = 0;
+    };
+
+    std::vector<buffer> cached;
+
+    ~ggml_zdnn_pool();
+
+    void * alloc(size_t size, size_t * actual_size);
+    void   free(void * ptr, size_t size);
+};
+
+// a buffer of the pool that goes back to it at the end of its scope, see ggml_cuda_pool_alloc
+template <typename T>
+struct ggml_zdnn_pool_alloc {
+    ggml_zdnn_pool * pool = nullptr;
+    T * ptr = nullptr;
+    size_t actual_size = 0;
+
+    ggml_zdnn_pool_alloc() = default;
+
+    ggml_zdnn_pool_alloc(ggml_zdnn_pool & pool, size_t size) {
+        alloc(pool, size);
+    }
+
+    ~ggml_zdnn_pool_alloc() {
+        if (ptr != nullptr) {
+            pool->free(ptr, actual_size);
+        }
+    }
+
+    // size is in number of elements
+    T * alloc(ggml_zdnn_pool & pool, size_t size) {
+        GGML_ASSERT(ptr == nullptr);
+        this->pool = &pool;
+        ptr = (T *) pool.alloc(size * sizeof(T), &actual_size);
+        return ptr;
+    }
+
+    T * get() {
+        return ptr;
+    }
+
+    ggml_zdnn_pool_alloc(const ggml_zdnn_pool_alloc &) = delete;
+    ggml_zdnn_pool_alloc(ggml_zdnn_pool_alloc &&) = delete;
+    ggml_zdnn_pool_alloc & operator=(const ggml_zdnn_pool_alloc &) = delete;
+    ggml_zdnn_pool_alloc & operator=(ggml_zdnn_pool_alloc &&) = delete;
+};
+
+#endif  // GGML_ZDNN_SCRATCH_HPP

gguf-py/gguf/constants.py

@@ -344,6 +344,7 @@ class Keys:
         ROUTING_BLOCK_COUNT = "{arch}.decision.routing_block_count"
         HEAD_COUNT         = "{arch}.decision.head_count"
         MAX_HEAD_TOKENS    = "{arch}.decision.max_head_tokens"
+        LABELS             = "{arch}.decision.labels"  # the text of the token read for each option, in order
         TEMPERATURE        = "{arch}.decision.temperature.{name}"  # name: "<type>" or "<type>.<n_opt bucket>"
         SCORE_BIAS         = "{arch}.decision.score_bias.{n}.{i}"  # added to the logit of level i of a score question with n levels
 
@@ -419,6 +420,8 @@ class Keys:
         PREPROC_MIN_TILES     = "clip.vision.preproc_min_tiles"
         PREPROC_MAX_TILES     = "clip.vision.preproc_max_tiles"
         PREPROC_IMAGE_SIZE    = "clip.vision.preproc_image_size"
+        PREPROC_TILE_SIZE     = "clip.vision.preproc_tile_size"
+        PREPROC_MAX_PIXELS_TOLERANCE = "clip.vision.preproc_max_pixels_tolerance"
         PATCH_SIZE            = "clip.vision.patch_size"
         EMBEDDING_LENGTH      = "clip.vision.embedding_length"
         FEED_FORWARD_LENGTH   = "clip.vision.feed_forward_length"
@@ -6337,6 +6340,7 @@ class DecisionType:
     LFM2_D1 = "lfm2-d1"  # same as openjev, the labels depend on the question type
     LFM2_D1_OMNI = "lfm2-d1-omni"  # same head as laya on a bidirectional LFM2 trunk, other prompt layout
     GRANITE_SWITCH = "granite-switch"  # same as openjev, run by the decision aLoRA of granite switch
+    H2O_LIGHTNING = "h2o-lightning"  # same as openjev, labels read after an answer prefill
 
 
 class VisionProjectorType:

gguf-py/gguf/gguf_writer.py

@@ -1401,6 +1401,9 @@ class GGUFWriter:
     def add_decision_max_head_tokens(self, value: int) -> None:
         self.add_uint32(Keys.Decision.MAX_HEAD_TOKENS.format(arch=self.arch), value)
 
+    def add_decision_labels(self, labels: Sequence[str]) -> None:
+        self.add_array(Keys.Decision.LABELS.format(arch=self.arch), labels)
+
     def add_decision_temperature(self, name: str, value: float) -> None:
         self.add_float32(Keys.Decision.TEMPERATURE.format(arch=self.arch, name=name), value)
 
@@ -1470,6 +1473,12 @@ class GGUFWriter:
     def add_vision_preproc_image_size(self, value: int) -> None:
         self.add_uint32(Keys.ClipVision.PREPROC_IMAGE_SIZE, value)
 
+    def add_vision_preproc_tile_size(self, value: int) -> None:
+        self.add_uint32(Keys.ClipVision.PREPROC_TILE_SIZE, value)
+
+    def add_vision_preproc_max_pixels_tolerance(self, value: float) -> None:
+        self.add_float32(Keys.ClipVision.PREPROC_MAX_PIXELS_TOLERANCE, value)
+
     def add_vision_projector_query_side(self, value: int) -> None:
         self.add_uint32(Keys.ClipVision.Projector.QUERY_SIDE, value)
 

src/llama-adapter.cpp

@@ -179,13 +179,13 @@ static void llama_adapter_lora_init_impl(llama_model & model, FILE * file, llama
                 ? format("%s[%s,%zu]", gguf_type_name(type), gguf_type_name(gguf_get_arr_type(gguf_ctx, i)), gguf_get_arr_n(gguf_ctx, i))
                 : gguf_type_name(type);
             const char * name = gguf_get_key(gguf_ctx, i);
-            const std::string value = gguf_kv_to_str(gguf_ctx, i);
+            const size_t MAX_VALUE_LEN = 40;
+            const std::string value = gguf_kv_to_str(gguf_ctx, i, MAX_VALUE_LEN);
 
             if (type != GGUF_TYPE_ARRAY) {
                 adapter.gguf_kv.emplace(name, value);
             }
 
-            const size_t MAX_VALUE_LEN = 40;
             std::string print_value = value.size() > MAX_VALUE_LEN ? format("%s...", value.substr(0, MAX_VALUE_LEN - 3).c_str()) : value;
             replace_all(print_value, "\n", "\\n");
 

src/llama-arch.cpp

@@ -383,6 +383,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_DECISION_BLOCK_COUNT,         "%s.decision.block_count"         },
     { LLM_KV_DECISION_ROUTING_BLOCK_COUNT, "%s.decision.routing_block_count" },
     { LLM_KV_DECISION_HEAD_COUNT,          "%s.decision.head_count"          },
+    { LLM_KV_DECISION_LABELS,              "%s.decision.labels"              },
 
     { LLM_KV_TARGET_LAYERS,         "%s.target_layers"        },
     { LLM_KV_TARGET_HIDDEN_SIZE,    "%s.target_hidden_size"   },
@@ -1109,7 +1110,7 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
     {LLM_TENSOR_MASKED_EMBD_ORDERING,       {LLM_TENSOR_LAYER_INPUT,     GGML_OP_NONE}},
     // eagle3
     {LLM_TENSOR_FC,                         {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL_MAT}},
-    {LLM_TENSOR_D2T,                        {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_GET_ROWS}},
+    {LLM_TENSOR_D2T,                        {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_SET_ROWS}},
     // dspark
     {LLM_TENSOR_DSPARK_MARKOV_W1,           {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_GET_ROWS}},
     {LLM_TENSOR_DSPARK_MARKOV_W2,           {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL_MAT}},

src/llama-arch.h

@@ -431,6 +431,7 @@ enum llm_kv {
     LLM_KV_DECISION_BLOCK_COUNT,
     LLM_KV_DECISION_ROUTING_BLOCK_COUNT,
     LLM_KV_DECISION_HEAD_COUNT,
+    LLM_KV_DECISION_LABELS,
 
     LLM_KV_TARGET_LAYERS,
     LLM_KV_TARGET_HIDDEN_SIZE,

src/llama-ext.h

@@ -122,6 +122,9 @@ enum llama_decision_order {
 // The embeddings output has one value per entry: row i is the score of option i
 LLAMA_API bool llama_batch_ext_set_decision_order(struct llama_batch_ext * batch, int32_t idx, enum llama_decision_order order);
 
+// Label of option i of a decision model ("<arch>.decision.labels"), nullptr past the last one
+LLAMA_API const char * llama_model_decision_label(const struct llama_model * model, uint32_t i);
+
 // mirrors:
 // LLAMA_API float * llama_get_embeddings(struct llama_context * ctx);
 LLAMA_API float * llama_get_embeddings_nextn(struct llama_context * ctx);

src/llama-graph.cpp

@@ -1666,6 +1666,40 @@ ggml_tensor * llm_graph_context::build_lora_mm(
     return res;
 }
 
+llm_graph_amax_src1 llm_graph_context::build_amax_src1_f16(
+          ggml_tensor * cur) const {
+    llm_graph_amax_src1 result = {
+        /*.res     =*/ cur,
+        /*.scale_s =*/ nullptr,
+    };
+
+    const int64_t n_embd = cur->ne[0];
+
+    // amax = max|cur| along the rows for each [expert, token] -> [1, n_expert, n_tokens]
+    ggml_tensor * ab = ggml_abs(ctx0, cur);
+    ggml_tensor * amax = ggml_pool_1d(ctx0, ab, GGML_OP_POOL_MAX, n_embd, n_embd, 0);
+
+    // ref: https://github.com/ggml-org/llama.cpp/pull/26223
+    constexpr float k_log2_e = 1.4426950408889634f; // 1 / ln(2)
+    constexpr float k_ln_2   = 0.6931471805599453f;
+
+    ggml_tensor * e;
+    e = ggml_scale(ctx0, amax, 1.0f/32768.0f);
+    e = ggml_log  (ctx0, e);
+    e = ggml_scale(ctx0, e, k_log2_e);
+    e = ggml_ceil (ctx0, e);
+    e = ggml_clamp(ctx0, e, 0.0f, 127.0f);
+    e = ggml_scale(ctx0, e, k_ln_2);
+    e = ggml_exp  (ctx0, e);
+
+    result = {
+        /*.res     =*/ ggml_div(ctx0, cur, e),
+        /*.scale_s =*/ e,
+    };
+
+    return result;
+}
+
 ggml_tensor * llm_graph_context::build_lora_mm_id(
           ggml_tensor * w,   // ggml_tensor * as
           ggml_tensor * cur, // ggml_tensor * b
@@ -2465,12 +2499,25 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
             GGML_ABORT("fatal error");
     }
 
-    experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s, slots); // [n_embd, n_expert_used, n_tokens]
-    if (arch == LLM_ARCH_MISTRAL4 || arch == LLM_ARCH_MELLUM) {
-        // src1 can exceed F16 range
-        ggml_prec_set_src(experts, GGML_PREC_F32, 1);
+    {
+        // TODO: this graph-level rescaling should only be applied when the corresponding backend does not support
+        //       ggml_prec_set_src(experts, GGML_PREC_BF16, 1);
+        //       ref: https://github.com/ggml-org/llama.cpp/pull/30656#issuecomment-6269597922
+        llm_graph_amax_src1 amax_src1 = { nullptr, nullptr };
+
+        if (arch == LLM_ARCH_MISTRAL4 || arch == LLM_ARCH_MELLUM) {
+            amax_src1 = build_amax_src1_f16(cur);
+            cur = amax_src1.res;
+        }
+
+        experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s, slots); // [n_embd, n_expert_used, n_tokens]
+        //ggml_prec_set_src(experts, GGML_PREC_BF16, 1); // (see TODO above)
+        cb(experts, "ffn_moe_down", il);
+
+        if (amax_src1.scale_s) {
+            experts = ggml_mul(ctx0, experts, amax_src1.scale_s);
+        }
     }
-    cb(experts, "ffn_moe_down", il);
 
     if (down_exps_s) {
         cb(experts, "ffn_moe_down_scaled", il);

src/llama-graph.h

@@ -1012,6 +1012,11 @@ struct llm_graph_qkv {
     ggml_tensor * v; // [n_embd_head, n_head_kv, n_tokens]
 };
 
+struct llm_graph_amax_src1 {
+    ggml_tensor * res;     // cur rescaled
+    ggml_tensor * scale_s; // scale to undo on the matmul result, or nullptr
+};
+
 struct llm_graph_context {
     const llm_arch arch;
 
@@ -1115,6 +1120,11 @@ struct llm_graph_context {
               ggml_tensor * w_s   = nullptr,
               ggml_tensor * slots = nullptr) const;
 
+    // rescale `cur` (the src1 of a matmul) by a power of two, computed from its maximum absolute
+    // when the values already fit in F16, the factors are exactly 1.0f
+    llm_graph_amax_src1 build_amax_src1_f16(
+              ggml_tensor * cur) const;
+
     ggml_tensor * build_norm(
              ggml_tensor * cur,
              ggml_tensor * mw,

src/llama-impl.cpp

@@ -11,7 +11,6 @@
 #include <cstdarg>
 #include <cstring>
 #include <vector>
-#include <sstream>
 
 struct llama_logger_state {
     ggml_log_callback log_callback = llama_log_callback_default;
@@ -149,7 +148,7 @@ std::string llama_format_tensor_shape(const struct ggml_tensor * t) {
     return buf;
 }
 
-static std::string gguf_data_to_str(enum gguf_type type, const void * data, int i) {
+static std::string gguf_data_to_str(enum gguf_type type, const void * data, size_t i) {
     switch (type) {
         case GGUF_TYPE_UINT8:   return std::to_string(((const uint8_t  *)data)[i]);
         case GGUF_TYPE_INT8:    return std::to_string(((const int8_t   *)data)[i]);
@@ -166,7 +165,7 @@ static std::string gguf_data_to_str(enum gguf_type type, const void * data, int
     }
 }
 
-std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i) {
+std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i, size_t max_len) {
     const enum gguf_type type = gguf_get_kv_type(ctx_gguf, i);
 
     switch (type) {
@@ -175,28 +174,38 @@ std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i) {
         case GGUF_TYPE_ARRAY:
             {
                 const enum gguf_type arr_type = gguf_get_arr_type(ctx_gguf, i);
-                int arr_n = gguf_get_arr_n(ctx_gguf, i);
+                const size_t arr_n = gguf_get_arr_n(ctx_gguf, i);
                 const void * data = arr_type == GGUF_TYPE_STRING ? nullptr : gguf_get_arr_data(ctx_gguf, i);
-                std::stringstream ss;
-                ss << "[";
-                for (int j = 0; j < arr_n; j++) {
-                    if (arr_type == GGUF_TYPE_STRING) {
-                        std::string val = gguf_get_arr_str(ctx_gguf, i, j);
-                        // escape quotes
-                        replace_all(val, "\\", "\\\\");
-                        replace_all(val, "\"", "\\\"");
-                        ss << '"' << val << '"';
-                    } else if (arr_type == GGUF_TYPE_ARRAY) {
-                        ss << "???";
-                    } else {
-                        ss << gguf_data_to_str(arr_type, data, j);
+                std::string res = "[";
+                size_t j = 0;
+                for (; j < arr_n && res.size() <= max_len; j++) {
+                    if (j > 0) {
+                        res += ", ";
                     }
-                    if (j < arr_n - 1) {
-                        ss << ", ";
+                    if (arr_type == GGUF_TYPE_STRING) {
+                        const char * s = gguf_get_arr_str(ctx_gguf, i, j);
+                        const size_t need = res.size() + 2 * strlen(s) + 2;
+                        if (res.capacity() < need) {
+                            res.reserve(std::max(need, 2 * res.capacity()));
+                        }
+                        res += '"';
+                        for (; *s; s++) {
+                            if (*s == '\\' || *s == '"') {
+                                res += '\\';
+                            }
+                            res += *s;
+                        }
+                        res += '"';
+                    } else if (arr_type == GGUF_TYPE_ARRAY) {
+                        res += "???";
+                    } else {
+                        res += gguf_data_to_str(arr_type, data, j);
                     }
                 }
-                ss << "]";
-                return ss.str();
+                if (j == arr_n) {
+                    res += "]";
+                }
+                return res;
             }
         default:
             return gguf_data_to_str(type, gguf_get_val_data(ctx_gguf, i), 0);

src/llama-impl.h

@@ -107,4 +107,5 @@ std::string format(const char * fmt, ...);
 std::string llama_format_tensor_shape(const std::vector<int64_t> & ne);
 std::string llama_format_tensor_shape(const struct ggml_tensor * t);
 
-std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i);
+// arrays are rendered only until the result is longer than max_len
+std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i, size_t max_len = SIZE_MAX);

src/llama-model-loader.cpp

@@ -807,8 +807,8 @@ llama_model_loader::llama_model_loader(
                 ? format("%s[%s,%zu]", gguf_type_name(type), gguf_type_name(gguf_get_arr_type(metadata, i)), gguf_get_arr_n(metadata, i))
                 : gguf_type_name(type);
 
-            std::string value          = gguf_kv_to_str(metadata, i);
             const size_t MAX_VALUE_LEN = 40;
+            std::string value          = gguf_kv_to_str(metadata, i, MAX_VALUE_LEN);
             if (value.size() > MAX_VALUE_LEN) {
                 value = format("%s...", value.substr(0, MAX_VALUE_LEN - 3).c_str());
             }
@@ -951,6 +951,13 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w
                 ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 512);
                 op_tensor = ggml_get_rows(ctx, w, b);
             } break;
+        case GGML_OP_SET_ROWS:
+            {
+                // w holds the row indices (e.g. d2t)
+                ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, w->ne[0]);
+                ggml_tensor * b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, w->ne[0]);
+                op_tensor = ggml_set_rows(ctx, a, b, w);
+            } break;
         case GGML_OP_MUL_MAT:
             {
                 ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, w->ne[0], 512, w->ne[2], w->ne[3]);

src/llama-model.cpp

@@ -1549,6 +1549,9 @@ void llama_model_base::load_hparams(llama_model_loader & ml) {
         hparams.n_cls_out = classifier_labels.size();
     }
 
+    // for decision models
+    ml.get_arr(LLM_KV_DECISION_LABELS, decision_labels, false);
+
     // per-arch hparams
     load_arch_hparams(ml);
 
@@ -3438,6 +3441,14 @@ const char * llama_model_cls_label(const struct llama_model * model, uint32_t i)
     return nullptr;
 }
 
+const char * llama_model_decision_label(const struct llama_model * model, uint32_t i) {
+    if (i < model->decision_labels.size()) {
+        return model->decision_labels[i].c_str();
+    }
+
+    return nullptr;
+}
+
 // deprecated
 int32_t llama_n_ctx_train(const llama_model * model) {
     return llama_model_n_ctx_train(model);

src/llama-model.h

@@ -724,6 +724,7 @@ struct llama_model {
 
     // for classifier models
     std::vector<std::string> classifier_labels;
+    std::vector<std::string> decision_labels;
 
     struct ggml_tensor * tok_embd   = nullptr;
     struct ggml_tensor * type_embd  = nullptr;

src/llama-vocab.cpp

@@ -2260,21 +2260,22 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                     throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str()));
                 }
                 const int n_merges = gguf_get_arr_n(ctx, merges_keyidx);
+                bpe_ranks.reserve(n_merges);
                 for (int i = 0; i < n_merges; i++) {
-                    const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i);
+                    const char * word = gguf_get_arr_str(ctx, merges_keyidx, i);
                     //GGML_ASSERT(unicode_cpts_from_utf8(word).size() > 0);
 
                     std::string first;
                     std::string second;
 
-                    const size_t pos = word.find(' ', 1);
+                    const char * sep = *word ? strchr(word + 1, ' ') : nullptr;
 
-                    if (pos != std::string::npos) {
-                        first  = word.substr(0, pos);
-                        second = word.substr(pos + 1);
+                    if (sep) {
+                        first.assign(word, sep - word);
+                        second.assign(sep + 1);
                     }
 
-                    bpe_ranks.emplace(std::make_pair(first, second), i);
+                    bpe_ranks.emplace(std::make_pair(std::move(first), std::move(second)), i);
                 }
             }
 
@@ -2371,20 +2372,21 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                     throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str()));
                 }
                 const int n_merges = gguf_get_arr_n(ctx, merges_keyidx);
+                bpe_ranks.reserve(n_merges);
                 for (int i = 0; i < n_merges; i++) {
-                    const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i);
+                    const char * word = gguf_get_arr_str(ctx, merges_keyidx, i);
 
                     std::string first;
                     std::string second;
 
-                    const size_t pos = word.find(' ', 1);
+                    const char * sep = *word ? strchr(word + 1, ' ') : nullptr;
 
-                    if (pos != std::string::npos) {
-                        first  = word.substr(0, pos);
-                        second = word.substr(pos + 1);
+                    if (sep) {
+                        first.assign(word, sep - word);
+                        second.assign(sep + 1);
                     }
 
-                    bpe_ranks.emplace(std::make_pair(first, second), i);
+                    bpe_ranks.emplace(std::make_pair(std::move(first), std::move(second)), i);
                 }
             }
 
@@ -2764,6 +2766,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
     }
 
     id_to_token.resize(n_tokens);
+    token_to_id.reserve(n_tokens);
 
     for (uint32_t i = 0; i < n_tokens; i++) {
         std::string word = gguf_get_arr_str(ctx, token_idx, i);
@@ -3676,13 +3679,14 @@ static void llama_unescape_whitespace(std::string & word) {
 
 static std::string llama_decode_text(const std::string & text) {
     std::string decoded_text;
+    decoded_text.reserve(text.size());
 
     const auto cpts = unicode_cpts_from_utf8(text);
     for (const auto cpt : cpts) {
-        const auto utf8 = unicode_cpt_to_utf8(cpt);
         try {
-            decoded_text += unicode_utf8_to_byte(utf8);
+            decoded_text += unicode_cpt_to_byte(cpt);
         } catch (const std::out_of_range & /*e*/) {
+            const auto utf8 = unicode_cpt_to_utf8(cpt);
             decoded_text += "[UNK_BYTE_0x";
             for (const auto c : utf8) {
                 decoded_text += format("%02x", (uint8_t) c);

src/models/kolibri1.cpp

@@ -131,11 +131,8 @@ llama_model_kolibri1::graph::graph(const llama_model & model, const llm_graph_pa
         cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
         cb(cur, "ffn_norm", il);
 
-        // The router runs in f32: the selection adds the bias to the logits, the weights
-        // are the sigmoid of the unbiased logits, so the same logits feed both.
+        // The selection adds the bias to the logits, the weights are the sigmoid of the unbiased logits, so the same logits feed both.
         ggml_tensor * logits = build_lora_mm(layer.ffn_gate_inp, cur);
-        ggml_prec_set_acc(logits, GGML_PREC_F32);
-        ggml_prec_set_src(logits, GGML_PREC_F32, 1);
         cb(logits, "ffn_moe_logits", il);
 
         ggml_tensor * selection = ggml_add(ctx0, logits, layer.ffn_exp_probs_b);

src/models/lfm-base.cpp

@@ -0,0 +1,114 @@
+#include "models.h"
+
+#include "llama-memory-recurrent.h"
+
+#include <algorithm>
+
+llm_build_lfm_base::llm_build_lfm_base(const llm_graph_params & params) : llm_graph_context(params) {}
+
+ggml_tensor * llm_build_lfm_base::build_lfm_shortconv(const llama_model & model,
+                                                      const llama_memory_recurrent_context * mctx_recr,
+                                                      llm_graph_input_rs * inp_recr,
+                                                      ggml_tensor * cur,
+                                                      int il) {
+    const uint32_t kv_head      = mctx_recr->get_head();
+    const int64_t  n_seq_tokens = ubatch.n_seq_tokens;
+    const int64_t  n_seqs       = ubatch.n_seqs;
+    GGML_ASSERT(n_seqs != 0);
+    GGML_ASSERT(ubatch.equal_seqs());
+    GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);
+
+    GGML_ASSERT(hparams.n_shortconv_l_cache > 1);
+    const uint32_t d_conv = hparams.n_shortconv_l_cache - 1;
+    const int64_t  n_embd = hparams.n_embd;
+
+    // {n_embd, n_tokens} => {n_embd, n_seq_tokens, n_seqs}
+    cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs);
+
+    auto * bcx = build_lora_mm(model.layers[il].shortconv.in_proj, cur);
+    cb(bcx, "model.layers.{}.conv.in_proj", il);
+
+    constexpr auto n_chunks = 3;
+    GGML_ASSERT(bcx->ne[0] % n_chunks == 0);
+    const auto chunk_size = bcx->ne[0] / n_chunks;
+    auto *     b          = ggml_view_3d(ctx0, bcx, chunk_size, bcx->ne[1], bcx->ne[2], bcx->nb[1], bcx->nb[2],
+                                         0 * chunk_size * ggml_element_size(bcx));
+    auto *     c          = ggml_view_3d(ctx0, bcx, chunk_size, bcx->ne[1], bcx->ne[2], bcx->nb[1], bcx->nb[2],
+                                         1 * chunk_size * ggml_element_size(bcx));
+    auto *     x          = ggml_view_3d(ctx0, bcx, chunk_size, bcx->ne[1], bcx->ne[2], bcx->nb[1], bcx->nb[2],
+                                         2 * chunk_size * ggml_element_size(bcx));
+
+    auto * bx = ggml_transpose(ctx0, ggml_mul(ctx0, b, x));
+
+    // read conv state
+    auto * conv_state = mctx_recr->get_r_l(il);
+    auto * conv_rs    = build_rs(inp_recr, conv_state, hparams.n_embd_r(), n_seqs);
+    auto * conv       = ggml_reshape_3d(ctx0, conv_rs, d_conv, n_embd, n_seqs);
+
+    // causal prepends the state, non-causal pads symmetrically for a centered window
+    if (hparams.causal_attn) {
+        bx = ggml_concat(ctx0, conv, bx, 0);
+    } else {
+        const int64_t pad = (hparams.n_shortconv_l_cache - 1) / 2;
+        auto * left = ggml_cont(ctx0,
+            ggml_view_3d(ctx0, conv, pad, n_embd, n_seqs, conv->nb[1], conv->nb[2], (d_conv - pad) * conv->nb[0]));
+        bx = ggml_pad_ext(ctx0, ggml_concat(ctx0, left, bx, 0), 0, pad, 0, 0, 0, 0, 0, 0);
+    }
+    GGML_ASSERT(bx->ne[0] > conv->ne[0]);
+
+    // write conv states: slot 0 = the final state, slot s = the state s tokens back (partial rollback)
+    const int64_t K         = hparams.causal_attn && cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1;
+    const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
+    const auto    mem_size  = mctx_recr->get_size();
+    const size_t  row_size  = ggml_row_size(conv_state->type, (int64_t) d_conv * n_embd);
+
+    for (int64_t slot = 0; slot < n_written; ++slot) {
+        auto * conv_snap = ggml_view_3d(ctx0, bx, d_conv, bx->ne[1], bx->ne[2], bx->nb[1], bx->nb[2],
+                                        (bx->ne[0] - d_conv - slot) * ggml_element_size(bx));
+        ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_snap,
+                                               ggml_view_2d(ctx0, conv_state, (int64_t) d_conv * n_embd, n_seqs,
+                                                            conv_state->nb[1],
+                                                            ((size_t) slot * mem_size + kv_head) * row_size)));
+    }
+
+    auto * conv_kernel = model.layers[il].shortconv.conv;
+    auto * conv_out    = ggml_ssm_conv(ctx0, bx, conv_kernel);
+    cb(conv_out, "model.layers.{}.conv.conv", il);
+
+    auto * y = ggml_mul(ctx0, c, conv_out);
+    y        = build_lora_mm(model.layers[il].shortconv.out_proj, y);
+    cb(y, "model.layers.{}.conv.out_proj", il);
+    // {n_embd, n_seq_tokens, n_seqs} => {n_embd, n_tokens}
+    y = ggml_reshape_2d(ctx0, y, y->ne[0], n_seq_tokens * n_seqs);
+
+    return y;
+}
+
+ggml_tensor * llm_build_lfm_base::build_lfm_dense_ffn(const llama_model & model,
+                                                      ggml_tensor * cur,
+                                                      int il) {
+    GGML_ASSERT(!model.layers[il].ffn_up_b);
+    GGML_ASSERT(!model.layers[il].ffn_gate_b);
+    GGML_ASSERT(!model.layers[il].ffn_down_b);
+    return build_ffn(cur,
+        model.layers[il].ffn_up,   NULL, NULL,
+        model.layers[il].ffn_gate, NULL, NULL,
+        model.layers[il].ffn_down, NULL, NULL,
+        NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
+}
+
+ggml_tensor * llm_build_lfm_base::build_lfm_moe_ffn(const llama_model & model,
+                                                    ggml_tensor * cur,
+                                                    int il) {
+    return build_moe_ffn(cur,
+            model.layers[il].ffn_gate_inp,
+            model.layers[il].ffn_up_exps,
+            model.layers[il].ffn_gate_exps,
+            model.layers[il].ffn_down_exps,
+            model.layers[il].ffn_exp_probs_b,
+            n_expert, n_expert_used,
+            LLM_FFN_SILU, true,
+            hparams.expert_weights_scale,
+            static_cast<llama_expert_gating_func_type>(hparams.expert_gating_func),
+            il);
+}

src/models/lfm2.cpp

@@ -147,35 +147,11 @@ std::unique_ptr<llm_graph_context> llama_model_lfm2::build_arch_graph(const llm_
 
 template <bool iswa>
 llama_model_lfm2::graph<iswa>::graph(const llama_model & model, const llm_graph_params & params) :
-    llm_graph_context(params) {
+    llm_build_lfm_base(params) {
     using inp_hybrid_type = std::conditional_t<iswa, llm_graph_input_mem_hybrid_iswa,  llm_graph_input_mem_hybrid>;
     using inp_attn_type   = std::conditional_t<iswa, llm_graph_input_attn_kv_iswa,     llm_graph_input_attn_kv>;
     using mem_hybrid_ctx  = std::conditional_t<iswa, llama_memory_hybrid_iswa_context, llama_memory_hybrid_context>;
 
-    // lambda helpers for readability
-    auto build_dense_feed_forward = [&model, this](ggml_tensor * cur, int il) -> ggml_tensor * {
-        GGML_ASSERT(!model.layers[il].ffn_up_b);
-        GGML_ASSERT(!model.layers[il].ffn_gate_b);
-        GGML_ASSERT(!model.layers[il].ffn_down_b);
-        return build_ffn(cur,
-            model.layers[il].ffn_up, NULL, NULL,
-            model.layers[il].ffn_gate, NULL, NULL,
-            model.layers[il].ffn_down, NULL, NULL,
-            NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
-    };
-    auto build_moe_feed_forward = [&model, this](ggml_tensor * cur, int il) -> ggml_tensor * {
-        return build_moe_ffn(cur,
-                model.layers[il].ffn_gate_inp,
-                model.layers[il].ffn_up_exps,
-                model.layers[il].ffn_gate_exps,
-                model.layers[il].ffn_down_exps,
-                model.layers[il].ffn_exp_probs_b,
-                n_expert, n_expert_used,
-                LLM_FFN_SILU, true,
-                hparams.expert_weights_scale,
-                static_cast<llama_expert_gating_func_type>(hparams.expert_gating_func),
-                il);
-    };
     auto build_attn_block = [&model, this](ggml_tensor *   cur,
                                            ggml_tensor *   inp_pos,
                                            inp_attn_type * inp_attn,
@@ -207,82 +183,6 @@ llama_model_lfm2::graph<iswa>::graph(const llama_model & model, const llm_graph_
 
         return cur;
     };
-    auto build_shortconv_block = [&model, this](ggml_tensor *        cur,
-                                                llm_graph_input_rs * inp_recr,
-                                                int                  il) -> ggml_tensor * {
-        const auto * mctx_cur = static_cast<const mem_hybrid_ctx *>(mctx)->get_recr();
-        const uint32_t kv_head      = mctx_cur->get_head();
-        const int64_t  n_seq_tokens = ubatch.n_seq_tokens;
-        const int64_t  n_seqs       = ubatch.n_seqs;
-        GGML_ASSERT(n_seqs != 0);
-        GGML_ASSERT(ubatch.equal_seqs());
-        GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);
-
-        GGML_ASSERT(hparams.n_shortconv_l_cache > 1);
-        const uint32_t d_conv = hparams.n_shortconv_l_cache - 1;
-
-        // {n_embd, n_tokens} => {n_embd, n_seq_tokens, n_seqs}
-        cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs);
-
-        auto * bcx = build_lora_mm(model.layers[il].shortconv.in_proj, cur);
-        cb(bcx, "model.layers.{}.conv.in_proj", il);
-
-        constexpr auto n_chunks = 3;
-        GGML_ASSERT(bcx->ne[0] % n_chunks == 0);
-        const auto chunk_size = bcx->ne[0] / n_chunks;
-        auto *     b          = ggml_view_3d(ctx0, bcx, chunk_size, bcx->ne[1], bcx->ne[2], bcx->nb[1], bcx->nb[2],
-                                             0 * chunk_size * ggml_element_size(bcx));
-        auto *     c          = ggml_view_3d(ctx0, bcx, chunk_size, bcx->ne[1], bcx->ne[2], bcx->nb[1], bcx->nb[2],
-                                             1 * chunk_size * ggml_element_size(bcx));
-        auto *     x          = ggml_view_3d(ctx0, bcx, chunk_size, bcx->ne[1], bcx->ne[2], bcx->nb[1], bcx->nb[2],
-                                             2 * chunk_size * ggml_element_size(bcx));
-
-        auto * bx = ggml_transpose(ctx0, ggml_mul(ctx0, b, x));
-
-        // read conv state
-        auto * conv_state = mctx_cur->get_r_l(il);
-        auto * conv_rs    = build_rs(inp_recr, conv_state, hparams.n_embd_r(), n_seqs);
-        auto * conv       = ggml_reshape_3d(ctx0, conv_rs, d_conv, hparams.n_embd, n_seqs);
-
-        // causal prepends the state, non-causal pads symmetrically for a centered window
-        if (hparams.causal_attn) {
-            bx = ggml_concat(ctx0, conv, bx, 0);
-        } else {
-            const int64_t pad = (hparams.n_shortconv_l_cache - 1) / 2;
-            auto * left = ggml_cont(ctx0,
-                ggml_view_3d(ctx0, conv, pad, hparams.n_embd, n_seqs, conv->nb[1], conv->nb[2], (d_conv - pad) * conv->nb[0]));
-            bx = ggml_pad_ext(ctx0, ggml_concat(ctx0, left, bx, 0), 0, pad, 0, 0, 0, 0, 0, 0);
-        }
-        GGML_ASSERT(bx->ne[0] > conv->ne[0]);
-
-        // write conv states: slot 0 = the final state, slot s = the state s tokens back (partial rollback)
-        const int64_t K         = hparams.causal_attn && cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1;
-        const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
-        const auto    mem_size  = mctx_cur->get_size();
-        const size_t  row_size  = ggml_row_size(conv_state->type, (int64_t) d_conv * n_embd);
-
-        for (int64_t slot = 0; slot < n_written; ++slot) {
-            auto * conv_snap = ggml_view_3d(ctx0, bx, d_conv, bx->ne[1], bx->ne[2], bx->nb[1], bx->nb[2],
-                                            (bx->ne[0] - d_conv - slot) * ggml_element_size(bx));
-            ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_snap,
-                                                   ggml_view_2d(ctx0, conv_state, (int64_t) d_conv * n_embd, n_seqs,
-                                                                conv_state->nb[1],
-                                                                ((size_t) slot * mem_size + kv_head) * row_size)));
-        }
-
-        auto * conv_kernel = model.layers[il].shortconv.conv;
-        auto * conv_out    = ggml_ssm_conv(ctx0, bx, conv_kernel);
-        cb(conv_out, "model.layers.{}.conv.conv", il);
-
-        auto * y = ggml_mul(ctx0, c, conv_out);
-        y        = build_lora_mm(model.layers[il].shortconv.out_proj, y);
-        cb(y, "model.layers.{}.conv.out_proj", il);
-        // {n_embd, n_seq_tokens, n_seqs} => {n_embd, n_tokens}
-        y = ggml_reshape_2d(ctx0, y, y->ne[0], n_seq_tokens * n_seqs);
-
-        return y;
-    };
-
     // actual graph construction starts here
     ggml_tensor * cur = build_inp_embd(model.tok_embd);
     cb(cur, "model.embed_tokens", -1);
@@ -308,8 +208,12 @@ llama_model_lfm2::graph<iswa>::graph(const llama_model & model, const llm_graph_
         cur             = build_norm(cur, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
         cb(cur, "model.layers.{}.operator_norm", il);
 
-        cur = hparams.is_recr(il) ? build_shortconv_block(cur, inp_hybrid->get_recr(), il) :
-                                    build_attn_block(cur, inp_pos, inp_hybrid->get_attn(), il);
+        if (hparams.is_recr(il)) {
+            const auto * mctx_recr = static_cast<const mem_hybrid_ctx *>(mctx)->get_recr();
+            cur = build_lfm_shortconv(model, mctx_recr, inp_hybrid->get_recr(), cur, il);
+        } else {
+            cur = build_attn_block(cur, inp_pos, inp_hybrid->get_attn(), il);
+        }
 
         if (il == n_layer - 1 && inp_out_ids) {
             cur      = ggml_get_rows(ctx0, cur, inp_out_ids);
@@ -322,7 +226,7 @@ llama_model_lfm2::graph<iswa>::graph(const llama_model & model, const llm_graph_
         cb(ffn_norm_out, "model.layers.{}.ffn_norm", il);
 
         ggml_tensor * ffn_out =
-            is_moe_layer ? build_moe_feed_forward(ffn_norm_out, il) : build_dense_feed_forward(ffn_norm_out, il);
+            is_moe_layer ? build_lfm_moe_ffn(model, ffn_norm_out, il) : build_lfm_dense_ffn(model, ffn_norm_out, il);
         cb(ffn_norm_out, "model.layers.{}.ffn_out", il);
 
         cur = ggml_add(ctx0, cur, ffn_out);

src/models/models.h

@@ -100,6 +100,26 @@ struct llm_build_delta_net_base : public llm_graph_context {
             int                  il);
 };
 
+struct llm_build_lfm_base : public llm_graph_context {
+    llm_build_lfm_base(const llm_graph_params & params);
+
+    virtual ~llm_build_lfm_base() = default;
+
+    ggml_tensor * build_lfm_shortconv(const llama_model & model,
+                                      const llama_memory_recurrent_context * mctx_recr,
+                                      llm_graph_input_rs * inp_recr,
+                                      ggml_tensor * cur,
+                                      int il);
+
+    ggml_tensor * build_lfm_dense_ffn(const llama_model & model,
+                                      ggml_tensor * cur,
+                                      int il);
+
+    ggml_tensor * build_lfm_moe_ffn(const llama_model & model,
+                                    ggml_tensor * cur,
+                                    int il);
+};
+
 struct llm_build_rwkv6_base : public llm_graph_context {
     const llama_model & model;
 
@@ -2264,7 +2284,7 @@ struct llama_model_lfm2 : public llama_model_base {
     void load_arch_tensors(llama_model_loader & ml) override;
 
     template <bool iswa>
-    struct graph : public llm_graph_context {
+    struct graph : public llm_build_lfm_base {
         graph(const llama_model & model, const llm_graph_params & params);
     };
 

src/models/qwen35.cpp

@@ -465,17 +465,14 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear(
     // Apply gated normalization: self.norm(core_attn_out, z)
     ggml_tensor * attn_out_norm = build_norm_gated(output, model.layers[il].ssm_norm, z_2d, il);
 
-    // Final reshape: [head_dim, n_heads, n_tokens, n_seqs] -> [n_tokens, n_seqs, n_heads * head_dim]
-    ggml_tensor * final_output = ggml_reshape_3d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens, n_seqs);
+    // Final reshape: [head_dim, n_heads, n_tokens, n_seqs] -> [n_heads * head_dim , n_tokens * n_seqs ]
+    ggml_tensor * final_output = ggml_reshape_2d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens * n_seqs);
     cb(final_output, "final_output", il);
 
     // Output projection
     cur = build_lora_mm(model.layers[il].ssm_out, final_output, model.layers[il].ssm_out_s);
     cb(cur, "linear_attn_out", il);
 
-    // Reshape back to original dimensions
-    cur = ggml_reshape_2d(ctx0, cur, n_embd, n_seq_tokens * n_seqs);
-
     return cur;
 }
 

src/models/qwen35moe.cpp

@@ -475,17 +475,13 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_attn_linear(
     // Apply gated normalization: self.norm(core_attn_out, z)
     ggml_tensor * attn_out_norm = build_norm_gated(output, model.layers[il].ssm_norm, z_2d, il);
 
-    // Final reshape: [head_dim, n_heads, n_tokens, n_seqs] -> [n_tokens, n_seqs, n_heads * head_dim]
-    ggml_tensor * final_output = ggml_reshape_3d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens, n_seqs);
-    cb(final_output, "final_output", il);
+    // Final reshape: [head_dim, n_heads, n_tokens, n_seqs] -> [n_heads * head_dim , n_tokens * n_seqs ]
+    ggml_tensor * final_output = ggml_reshape_2d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens * n_seqs);
 
     // Output projection
     cur = build_lora_mm(model.layers[il].ssm_out, final_output, model.layers[il].ssm_out_s);
     cb(cur, "linear_attn_out", il);
 
-    // Reshape back to original dimensions
-    cur = ggml_reshape_2d(ctx0, cur, n_embd, n_seq_tokens * n_seqs);
-
     return cur;
 }
 

src/models/qwen3next.cpp

@@ -542,17 +542,14 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn_linear(
     // Apply gated normalization: self.norm(core_attn_out, z)
     ggml_tensor * attn_out_norm = build_norm_gated(output, model.layers[il].ssm_norm, z_2d, il);
 
-    // Final reshape: [head_dim, n_heads, n_tokens, n_seqs] -> [n_tokens, n_seqs, n_heads * head_dim]
-    ggml_tensor * final_output = ggml_reshape_3d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens, n_seqs);
+    // Final reshape: [head_dim, n_heads, n_tokens, n_seqs] -> [n_heads * head_dim , n_tokens * n_seqs ]
+    ggml_tensor * final_output = ggml_reshape_2d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens * n_seqs);
     cb(final_output, "final_output", il);
 
     // Output projection
     cur = build_lora_mm(model.layers[il].ssm_out, final_output);
     cb(cur, "linear_attn_out", il);
 
-    // Reshape back to original dimensions
-    cur = ggml_reshape_2d(ctx0, cur, n_embd, n_seq_tokens * n_seqs);
-
     return cur;
 }
 

src/models/qwen4exp.cpp

@@ -1131,14 +1131,12 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear(
     // gated normalization, as self.norm(core_attn_out, z) in the reference
     ggml_tensor * attn_out_norm = build_norm_gated(output, model.layers[il].ssm_norm, z_2d, il);
 
-    ggml_tensor * final_output = ggml_reshape_3d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens, n_seqs);
+    ggml_tensor * final_output = ggml_reshape_2d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens * n_seqs);
     cb(final_output, "final_output", il);
 
     cur = build_lora_mm(model.layers[il].ssm_out, final_output, model.layers[il].ssm_out_s);
     cb(cur, "linear_attn_out", il);
 
-    cur = ggml_reshape_2d(ctx0, cur, n_embd, n_seq_tokens * n_seqs);
-
     return cur;
 }
 

src/unicode.cpp

@@ -202,6 +202,19 @@ static std::unordered_map<std::string, uint8_t> unicode_utf8_to_byte_map() {
     return map;
 }
 
+static std::vector<int> unicode_cpt_to_byte_map() {
+    std::vector<int> map;
+    for (int ch = 0; ch < 256; ++ch) {
+        size_t offset = 0;
+        const uint32_t cpt = unicode_cpt_from_utf8(unicode_byte_to_utf8(ch), offset);
+        if (cpt >= map.size()) {
+            map.resize(cpt + 1, -1);
+        }
+        map[cpt] = ch;
+    }
+    return map;
+}
+
 static std::vector<std::string> unicode_byte_encoding_process(const std::vector<std::string> & bpe_words) {
     std::vector<std::string> bpe_encoded_words;
     for (const auto & word : bpe_words) {
@@ -1416,6 +1429,14 @@ uint8_t unicode_utf8_to_byte(const std::string & utf8) {
     return map.at(utf8);
 }
 
+uint8_t unicode_cpt_to_byte(uint32_t cpt) {
+    static const std::vector<int> map = unicode_cpt_to_byte_map();
+    if (cpt >= map.size() || map[cpt] < 0) {
+        throw std::out_of_range("unicode_cpt_to_byte: not a byte token code point");
+    }
+    return map[cpt];
+}
+
 uint32_t unicode_tolower(uint32_t cpt) {
     // binary search
     auto it = std::lower_bound(unicode_map_lowercase.begin(), unicode_map_lowercase.end(), cpt,

src/unicode.h

@@ -104,6 +104,7 @@ unicode_cpt_flags unicode_cpt_flags_from_utf8(const std::string & utf8);
 
 std::string unicode_byte_to_utf8(uint8_t byte);
 uint8_t     unicode_utf8_to_byte(const std::string & utf8);
+uint8_t     unicode_cpt_to_byte (uint32_t cpt);
 
 uint32_t unicode_tolower(uint32_t cpt);
 

tests/CMakeLists.txt

@@ -332,6 +332,7 @@ if (NOT GGML_BACKEND_DL AND GGML_CPU)
     llama_build_and_test(test-rope.cpp LINK ggml)
     llama_build_and_test(test-col2im-1d.cpp LINK ggml)
     llama_build_and_test(test-conv-1d.cpp LINK ggml)
+    llama_build_and_test(test-gelu-quick.cpp LINK ggml)
 endif()
 
 # libmtmd

tests/test-backend-ops.cpp

@@ -2069,6 +2069,15 @@ struct test_unary : public test_case {
         return VARS_TO_STR3(type, ne_a, v);
     }
 
+    double max_nmse_err(ggml_backend_t backend) override {
+        // Metal computes f32 gelu_quick in f32, so an fp16 table in the CPU reference (NMSE about 5e-8) fails this
+        ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
+        if (op == GGML_UNARY_OP_GELU_QUICK && type == GGML_TYPE_F32 && strcmp(ggml_backend_reg_name(reg), "MTL") == 0) {
+            return 1e-10;
+        }
+        return test_case::max_nmse_err(backend);
+    }
+
     test_unary(ggml_unary_op op,
             ggml_type type = GGML_TYPE_F32,
             std::array<int64_t, 4> ne_a = {128, 2, 2, 2},
@@ -2157,6 +2166,15 @@ struct test_glu : public test_case {
         return VARS_TO_STR4(type, ne_a, v, swapped);
     }
 
+    double max_nmse_err(ggml_backend_t backend) override {
+        // Metal computes f32 geglu_quick in f32, so an fp16 table in the CPU reference (NMSE about 5e-8) fails this
+        ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
+        if (op == GGML_GLU_OP_GEGLU_QUICK && type == GGML_TYPE_F32 && strcmp(ggml_backend_reg_name(reg), "MTL") == 0) {
+            return 1e-10;
+        }
+        return test_case::max_nmse_err(backend);
+    }
+
     test_glu(ggml_glu_op op,
             ggml_type type = GGML_TYPE_F32,
             std::array<int64_t, 4> ne_a = {128, 2, 2, 2},
@@ -2202,6 +2220,15 @@ struct test_glu_split : public test_case {
         return VARS_TO_STR3(type, ne_a, v) + ",split";
     }
 
+    double max_nmse_err(ggml_backend_t backend) override {
+        // Metal computes f32 geglu_quick in f32, so an fp16 table in the CPU reference (NMSE about 5e-8) fails this
+        ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
+        if (op == GGML_GLU_OP_GEGLU_QUICK && type == GGML_TYPE_F32 && strcmp(ggml_backend_reg_name(reg), "MTL") == 0) {
+            return 1e-10;
+        }
+        return test_case::max_nmse_err(backend);
+    }
+
     test_glu_split(ggml_glu_op op,
             ggml_type type = GGML_TYPE_F32,
             std::array<int64_t, 4> ne_a = {128, 2, 2, 2},
@@ -2454,7 +2481,9 @@ struct test_get_rows : public test_case {
 };
 
 // GGML_OP_GET_ROWS on a weight also used by GGML_OP_MUL_MAT (tied embeddings)
-// the mul_mat runs first, so a backend that changes the weight layout must still return the right rows
+// that the preceding mul_mat may have put in the SYCL reorder layout. The
+// mul_mat runs first, so a backend that changes the weight layout in place
+// must still return the right rows from get_rows.
 struct test_get_rows_after_mul_mat : public test_case {
     const ggml_type type;
     const int n; // cols
@@ -2510,6 +2539,48 @@ struct test_get_rows_after_mul_mat : public test_case {
     }
 };
 
+// The persistent SYCL weight reorder installs on the first small-batch mul_mat; a
+// later >8-column mul_mat on the same weight goes through the fused-dequant GEMM /
+// conversion routes, which must decode the reordered layout (regression test for the
+// #29107 crash: the plain fused-dequant GEMM decoded standard layout unconditionally).
+struct test_mul_mat_reorder_then_big_batch : public test_case {
+    const ggml_type type;
+    const int n; // cols of the big batch
+    const int m; // rows
+
+    std::string vars() override {
+        return VARS_TO_STR3(type, n, m);
+    }
+
+    test_mul_mat_reorder_then_big_batch(ggml_type type = GGML_TYPE_IQ3_S, int n = 64, int m = 64)
+        : type(type), n(n), m(m) {}
+
+    double max_nmse_err() override {
+        return 5e-4; // same as test_mul_mat
+    }
+
+    ggml_tensor * build_graph(ggml_context * ctx) override {
+        ggml_tensor * w = ggml_new_tensor_2d(ctx, type, 1*256, m);
+        ggml_set_name(w, "w");
+
+        ggml_tensor * x1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1*256, 1);
+        ggml_set_name(x1, "x1");
+
+        ggml_tensor * x2 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1*256, n);
+        ggml_set_name(x2, "x2");
+
+        ggml_tensor * mm1 = ggml_mul_mat(ctx, w, x1);
+        ggml_set_name(mm1, "mm1");
+
+        ggml_tensor * mm2 = ggml_mul_mat(ctx, w, x2);
+        ggml_set_name(mm2, "mm2");
+
+        // mm1 is the first source of the final ADD, so the reorder-installing
+        // small-batch mul_mat is computed before the big-batch one
+        return ggml_add(ctx, ggml_repeat(ctx, mm1, mm2), mm2);
+    }
+};
+
 // GGML_OP_GET_ROWS_BACK
 struct test_get_rows_back : public test_case {
     const ggml_type type;
@@ -9885,9 +9956,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     }
 
     test_cases.emplace_back(new test_get_rows(GGML_TYPE_F32, 1, 8, 2, 1, 1, false));
-    for (ggml_type type : {GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K}) {
+    for (ggml_type type : {GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K,
+                          GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ3_S}) {
         test_cases.emplace_back(new test_get_rows_after_mul_mat(type));
     }
+    for (ggml_type type : {GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ3_S}) {
+        test_cases.emplace_back(new test_mul_mat_reorder_then_big_batch(type));
+    }
     for (ggml_type type : all_types) {
         for (int b : {1, 7}) {
             for (bool v : {false, true}) {
@@ -10659,7 +10734,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     // ne[2] > 65535 exceeds the CUDA gridDim.y limit (#27901)
     test_cases.emplace_back(new test_rms_norm_mul_rope({4, 1, 65536, 1}, 1e-6f, false, false, false, GGML_ROPE_TYPE_NORMAL));
     test_cases.emplace_back(new test_rms_norm_mul_rope({4, 1, 65536, 1}, 1e-6f, false, true,  false, GGML_ROPE_TYPE_NEOX));
-    for (int64_t d_conv : {3, 4, 9}) {
+    for (int64_t d_conv : {3, 4, 7, 9}) {
         for (int64_t d_inner: {1024, 1536, 2048}) {
             test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {d_conv, d_inner, 1, 1}, {d_conv, d_inner, 1, 1}));
             test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {2 * d_conv, d_inner, 1, 1}, {d_conv, d_inner, 1, 1}));
@@ -10672,7 +10747,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
 
     // fused ssm_conv + (optional) bias_add + silu. The bias-only graph (no silu) is intentionally
     // not tested since there's no fusion for that pattern in ggml_cuda_can_fuse.
-    for (int64_t d_conv : {3, 4, 9}) {
+    for (int64_t d_conv : {3, 4, 7, 9}) {
         for (int64_t d_inner : {1024, 1536, 2048}) {
             for (bool fuse_bias : {false, true}) {
                 // short token path (n_t <= 32)
@@ -12199,6 +12274,17 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
         }
     }
 
+    // POOL_1D max over MoE down-projection amax shapes:
+    // amax is [1, n_expert, n_tokens] after reducing the whole intermediate dim.
+    for (int64_t n_ff : { 2048, 4096, 8192 }) {
+        for (int64_t n_expert : { 8, 16 }) {
+            for (int64_t n_tokens : { 1, 512 }) {
+                test_cases.emplace_back(new test_pool1d(
+                    GGML_OP_POOL_MAX, GGML_TYPE_F32, {n_ff, n_expert, n_tokens, 1}, n_ff, n_ff, 0));
+            }
+        }
+    }
+
     // LEAKY_RELU at FFN activation width, for direct comparison with RELU
     for (int64_t n_tokens : {512, 2048}) {
         test_cases.emplace_back(new test_leaky_relu(GGML_TYPE_F32, { 17408, n_tokens, 1, 1 }, 0.1f));

tests/test-gelu-quick.cpp

@@ -0,0 +1,131 @@
+// test-gelu-quick.cpp: validate the F32 gelu_quick and geglu_quick CPU kernels against the exact formula.
+//
+// test-backend-ops uses the CPU backend as its reference, so it cannot see an error in the CPU kernels
+// themselves. This test computes gelu_quick(x) = x / (1 + exp(-1.702 x)) in double precision and compares
+// every F32 output of the CPU backend with it: dense ranges, values outside the F16 range, values near
+// the float limits, and lengths that leave a tail after the vector loop.
+
+#include "ggml.h"
+#include "ggml-cpu.h"
+
+#include <cfloat>
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+#include <vector>
+
+// Same coefficient as the CPU kernels
+static const double GELU_QUICK_COEF_REF = (double) 1.702f;
+
+static double gelu_quick_ref(double x) {
+    return x / (1.0 + std::exp(-GELU_QUICK_COEF_REF * x));
+}
+
+// Deterministic LCG mapped to [-1, 1]
+static uint64_t g_rng = 0x12345678ULL;
+static float frand(void) {
+    g_rng = g_rng * 6364136223846793005ULL + 1442695040888963407ULL;
+    return (float)((g_rng >> 33) & 0xffffff) / (float)0x800000 - 1.0f;
+}
+
+static std::vector<float> make_inputs(void) {
+    std::vector<float> x;
+    for (int i = 0; i <= 200000; i++) {
+        x.push_back(-10.0f + 20.0f * (float) i / 200000.0f);
+    }
+    for (int i = 0; i <= 2000; i++) {
+        x.push_back(-150.0f + 300.0f * (float) i / 2000.0f);
+    }
+    // outside the F16 range, near the float limits, and subnormal; both signs
+    const float edge[] = { 65504.0f, 65520.0f, 1e5f, 1e10f, 1e30f, FLT_MAX, 1e-30f, FLT_MIN, FLT_MIN / 16, 0.0f };
+    for (float v : edge) {
+        x.push_back(v);
+        x.push_back(-v);
+    }
+    return x;
+}
+
+static bool matches(float y, double ref) {
+    if (std::isnan(ref)) {
+        return std::isnan(y);
+    }
+    if (std::isinf(ref)) {
+        return (double) y == ref;
+    }
+    return std::fabs((double) y - ref) <= 1e-6 + 1e-6 * std::fabs(ref);
+}
+
+// Runs gelu_quick and geglu_quick on x[0, n) and returns the number of outputs that differ from the reference
+static int check(const float * x, const float * g, int64_t n) {
+    struct ggml_init_params params = {
+        /* .mem_size   = */ (size_t) 64 << 20,
+        /* .mem_base   = */ NULL,
+        /* .no_alloc   = */ false,
+    };
+    struct ggml_context * ctx = ggml_init(params);
+
+    struct ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n);
+    struct ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n);
+    memcpy(a->data, x, n * sizeof(float));
+    memcpy(b->data, g, n * sizeof(float));
+
+    struct ggml_tensor * y_gelu  = ggml_gelu_quick(ctx, a);
+    struct ggml_tensor * y_geglu = ggml_geglu_quick_split(ctx, a, b);
+
+    struct ggml_cgraph * gf = ggml_new_graph(ctx);
+    ggml_build_forward_expand(gf, y_gelu);
+    ggml_build_forward_expand(gf, y_geglu);
+    ggml_graph_compute_with_ctx(ctx, gf, 1);
+
+    int    bad     = 0;
+    double max_err = 0.0;
+    for (int64_t i = 0; i < n; i++) {
+        const double r0 = gelu_quick_ref(x[i]);
+        const double r1 = r0 * g[i];
+        const float  y0 = ((const float *) y_gelu->data)[i];
+        const float  y1 = ((const float *) y_geglu->data)[i];
+        if (std::isfinite(r0) && std::isfinite(y0)) {
+            max_err = std::fmax(max_err, std::fabs((double) y0 - r0) / (1.0 + std::fabs(r0)));
+        }
+        if (!matches(y0, r0)) {
+            if (bad < 3) {
+                printf("  gelu_quick(%g) = %g, expected %g\n", x[i], y0, r0);
+            }
+            bad++;
+        }
+        if (!matches(y1, r1)) {
+            if (bad < 3) {
+                printf("  geglu_quick(%g, %g) = %g, expected %g\n", x[i], g[i], y1, r1);
+            }
+            bad++;
+        }
+    }
+    printf("gelu_quick n=%6d: max err %.2e, %d mismatches %s\n", (int) n, max_err, bad, bad == 0 ? "OK" : "FAIL");
+
+    ggml_free(ctx);
+    return bad;
+}
+
+int main(void) {
+    const std::vector<float> x = make_inputs();
+    std::vector<float> g(x.size());
+    for (float & v : g) {
+        v = frand();
+    }
+
+    int fails = 0;
+
+    // every input in one call
+    fails += check(x.data(), g.data(), (int64_t) x.size()) != 0;
+
+    // short lengths end on the edge values and leave a tail after the vector loop
+    const int64_t lengths[] = { 1, 2, 3, 5, 7, 8, 9, 15, 16, 17, 23, 31, 33 };
+    for (int64_t n : lengths) {
+        const int64_t off = (int64_t) x.size() - n;
+        fails += check(x.data() + off, g.data() + off, n) != 0;
+    }
+
+    printf(fails == 0 ? "all gelu_quick checks passed\n" : "%d gelu_quick checks FAILED\n", fails);
+    return fails == 0 ? 0 : 1;
+}

tests/test-mtmd-impl.cpp

@@ -91,11 +91,19 @@ MAKE_TEST(test_image_preprocessor_dyn_size) {
     }
 }
 
+static std::string image_size_to_str(const clip_image_size & size) {
+    return std::to_string(size.width) + "x" + std::to_string(size.height);
+}
+
 MAKE_TEST(test_image_preprocessor_lfm2) {
     clip_hparams hparams;
     hparams.patch_size = 16;
     hparams.n_merge = 2;
     hparams.set_limit_image_tokens(64, 256);
+    hparams.preproc_min_tiles            = 2;
+    hparams.preproc_max_tiles            = 10;
+    hparams.preproc_tile_size            = 512;
+    hparams.preproc_max_pixels_tolerance = 2.0f;
 
     // { image size, expected tiling }
     const std::vector<std::pair<clip_image_size, bool>> cases = {
@@ -116,6 +124,38 @@ MAKE_TEST(test_image_preprocessor_lfm2) {
             std::string(expected ? "tiled" : "single"),
             std::string(actual   ? "tiled" : "single"));
     }
+
+    // { image size, expected size from HF smart_resize }
+    const std::vector<std::pair<clip_image_size, clip_image_size>> resize_cases = {
+        { {  400, 600 }, { 384, 608 } }, // 400 / 32 = 12.5 rounds half to even
+        { {  480, 336 }, { 480, 320 } },
+        { {   64, 169 }, { 160, 416 } }, // 169 * beta / 32 is exactly 13 in double
+        { { 1024, 768 }, { 576, 416 } },
+    };
+
+    for (const auto & [size, expected] : resize_cases) {
+        t.assert_equal("smart_resize for " + image_size_to_str(size), image_size_to_str(expected),
+            image_size_to_str(mtmd_image_preprocessor_lfm2::smart_resize(hparams, size)));
+    }
+
+    // { image size, expected grid from HF }
+    const std::vector<std::pair<clip_image_size, clip_image_size>> grid_cases = {
+        { {  720, 1600 }, { 2, 4 } }, // 0.45 is an exact tie between 2/4 and 2/5 in double, not in float
+        { {  420, 1440 }, { 1, 3 } },
+        { {  255, 2160 }, { 1, 8 } },
+        { { 1024,  768 }, { 3, 2 } },
+    };
+
+    for (const auto & [size, expected] : grid_cases) {
+        t.assert_equal("grid for " + image_size_to_str(size), image_size_to_str(expected),
+            image_size_to_str(mtmd_image_preprocessor_lfm2::get_grid_layout(hparams, size.height, size.width)));
+    }
+
+    // do_image_splitting=false
+    hparams.preproc_min_tiles = 1;
+    hparams.preproc_max_tiles = 1;
+    t.assert_equal("tiling for 1024x977 with max_tiles=1", std::string("single"),
+            std::string(mtmd_image_preprocessor_lfm2::should_tile(hparams, { 1024, 977 }) ? "tiled" : "single"));
 }
 
 //

tools/mtmd/clip-impl.h

@@ -56,6 +56,8 @@
 #define KEY_PREPROC_MIN_TILES       "clip.vision.preproc_min_tiles"
 #define KEY_PREPROC_MAX_TILES       "clip.vision.preproc_max_tiles"
 #define KEY_PREPROC_IMAGE_SIZE      "clip.vision.preproc_image_size"
+#define KEY_PREPROC_TILE_SIZE       "clip.vision.preproc_tile_size"
+#define KEY_PREPROC_MAX_PIXELS_TOLERANCE "clip.vision.preproc_max_pixels_tolerance"
 #define KEY_PATCH_SIZE              "clip.vision.patch_size"
 #define KEY_IMAGE_MEAN              "clip.vision.image_mean"
 #define KEY_IMAGE_STD               "clip.vision.image_std"

tools/mtmd/clip-model.h

@@ -66,6 +66,7 @@ struct clip_hparams {
     resize_algo image_resize_algo = RESIZE_ALGO_BICUBIC;
     pad_style image_resize_pad = PAD_CEIL; // padding style when resizing
     bool image_round_half_even = false;    // smart_resize rounds .5 to even, like Python round() in the HF processor
+    bool image_floor_before_align = false;
     std::array<uint8_t, 3> image_pad_color = {0, 0, 0};
 
     // (preprocessor) for llava-uhd style models
@@ -73,7 +74,8 @@ struct clip_hparams {
     int32_t preproc_min_tiles = 0;
     int32_t preproc_max_tiles = 0;
     int32_t max_slice_nums = 9; // llava-uhd slice cap; per-model, carried in the GGUF
-    int32_t preproc_tile_size = 0; // local tile size (deepseek-ocr)
+    int32_t preproc_tile_size = 0; // local tile size (deepseek-ocr, lfm2)
+    float preproc_max_pixels_tolerance = 1.0f;
     resize_algo image_resize_algo_rf = RESIZE_ALGO_BICUBIC;
     resize_algo image_resize_algo_ov = RESIZE_ALGO_BICUBIC;
     pad_style image_pad_rf = PAD_CEIL;  // padding style for the refined image (e.g. llava-1.6)

tools/mtmd/clip.cpp

@@ -1593,12 +1593,20 @@ struct clip_model_loader {
                         }
                         hparams.image_resize_algo_rf = hparams.image_resize_algo;
                         hparams.image_resize_algo_ov = hparams.image_resize_algo;
-                        // the tiles stretch the image to the grid (d1-omni vision.py)
-                        if (model.proj_type == PROJECTOR_TYPE_D1OMNI_V) {
-                            hparams.image_pad_rf = PAD_NONE;
-                        }
+                        // the tiles stretch the image to the grid (Lfm2VlImageProcessor, d1-omni vision.py)
+                        hparams.image_pad_rf = PAD_NONE;
                         get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge, false);
+                        // defaults for GGUFs converted before tiling params were written
                         // ref: https://huggingface.co/LiquidAI/LFM2.5-VL-1.6B/blob/main/processor_config.json
+                        hparams.preproc_min_tiles            = 2;
+                        hparams.preproc_max_tiles            = 10;
+                        hparams.preproc_tile_size            = 512;
+                        hparams.preproc_max_pixels_tolerance = 2.0f;
+                        get_u32(KEY_PREPROC_MIN_TILES, hparams.preproc_min_tiles, false);
+                        get_u32(KEY_PREPROC_MAX_TILES, hparams.preproc_max_tiles, false);
+                        get_u32(KEY_PREPROC_TILE_SIZE, hparams.preproc_tile_size, false);
+                        get_f32(KEY_PREPROC_MAX_PIXELS_TOLERANCE, hparams.preproc_max_pixels_tolerance, false);
+                        GGML_ASSERT(hparams.preproc_min_tiles >= 1 && hparams.preproc_min_tiles <= hparams.preproc_max_tiles);
                         hparams.set_limit_image_tokens(64, 256);
                     } break;
                 case PROJECTOR_TYPE_PHI4:
@@ -1617,7 +1625,23 @@ struct clip_model_loader {
                         hparams.image_resize_algo = RESIZE_ALGO_BICUBIC;
                         hparams.rope_theta = 10000.0f;
                         get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false);
-                        hparams.set_limit_image_tokens(8, 1024);
+                        get_u32(KEY_PREPROC_IMAGE_SIZE, hparams.image_longest_edge, false);
+                        if (hparams.image_longest_edge > 0) {
+                            const int64_t patch_size = int64_t(hparams.patch_size) * hparams.n_merge;
+                            if (patch_size <= 0) {
+                                throw std::runtime_error("Pixtral merged patch size must be positive");
+                            }
+                            const int64_t n_tokens_per_side =
+                                (hparams.image_longest_edge + patch_size - 1) / patch_size;
+                            if (n_tokens_per_side * patch_size > std::sqrt(std::numeric_limits<int>::max())) {
+                                throw std::runtime_error("Pixtral preprocessing image size is too large");
+                            }
+                            hparams.image_resize_pad         = PAD_NONE;
+                            hparams.image_floor_before_align = true;
+                            hparams.set_limit_image_tokens(0, n_tokens_per_side * n_tokens_per_side);
+                        } else {
+                            hparams.set_limit_image_tokens(8, 1024);
+                        }
                         hparams.set_warmup_n_tokens(256); // avoid OOM on warmup
                     } break;
                 case PROJECTOR_TYPE_LIGHTONOCR:
@@ -1725,6 +1749,8 @@ struct clip_model_loader {
                 case PROJECTOR_TYPE_GEMMA4UV:
                     {
                         hparams.rope_theta = 100.0f;
+                        hparams.ffn_op = FFN_GELU;
+                        log_ffn_op = "gelu";
                         hparams.n_merge = 3; // pooling_kernel_size
                         hparams.image_resize_algo = RESIZE_ALGO_BICUBIC;
                         get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge, false);
@@ -4252,8 +4278,9 @@ struct clip_init_result clip_init(const char * fname, struct clip_context_params
 struct clip_cap clip_get_cap(const char * fname) {
     clip_cap res;
     clip_model_loader loader(fname, /* skip_tensors= */ true);
-    res.has_vision = loader.has_vision;
-    res.has_audio  = loader.has_audio;
+    res.has_vision    = loader.has_vision;
+    res.has_audio     = loader.has_audio;
+    res.has_gen_audio = loader.has_gen_audio;
     return res;
 }
 

tools/mtmd/clip.h

@@ -148,5 +148,6 @@ std::map<ggml_backend_dev_t, size_t> clip_get_mem_usage(const struct clip_ctx *
 struct clip_cap {
     bool has_vision;
     bool has_audio;
+    bool has_gen_audio;
 };
 struct clip_cap clip_get_cap(const char * fname);

tools/mtmd/mtmd-helper-gen.cpp

@@ -50,29 +50,38 @@ static llama_token find_special_token(const llama_vocab * vocab, const std::stri
     return LLAMA_TOKEN_NULL;
 }
 
+static void put_bytes(std::vector<char> & buf, const void * p, size_t n) {
+    const char * c = (const char *) p;
+    buf.insert(buf.end(), c, c + n);
+}
+
+// data_sz == UINT32_MAX writes the "unknown length" sentinel (streaming), same as ffmpeg does on a pipe
+static void write_wav16_header(std::vector<char> & buf, uint32_t data_sz, int32_t rate) {
+    const uint32_t riff_sz   = data_sz == UINT32_MAX ? UINT32_MAX : 36 + data_sz;
+    const uint32_t fmt_sz    = 16, byte_rate = (uint32_t) rate * 2;
+    const uint16_t fmt = 1, ch = 1, align = 2, bits = 16;
+    const uint32_t rate32    = (uint32_t) rate;
+    put_bytes(buf, "RIFF", 4); put_bytes(buf, &riff_sz, 4); put_bytes(buf, "WAVE", 4);
+    put_bytes(buf, "fmt ", 4); put_bytes(buf, &fmt_sz, 4);
+    put_bytes(buf, &fmt, 2); put_bytes(buf, &ch, 2); put_bytes(buf, &rate32, 4);
+    put_bytes(buf, &byte_rate, 4); put_bytes(buf, &align, 2); put_bytes(buf, &bits, 2);
+    put_bytes(buf, "data", 4); put_bytes(buf, &data_sz, 4);
+}
+
+static void append_wav16_pcm(std::vector<char> & buf, const float * pcm, size_t n) {
+    for (size_t i = 0; i < n; i++) {
+        int16_t s = (int16_t) (std::max(-1.0f, std::min(1.0f, pcm[i])) * 32767.0f);
+        put_bytes(buf, &s, 2);
+    }
+}
+
 static bool write_wav16(std::vector<char> & buf, const std::vector<float> & pcm, int32_t rate) {
     // RIFF chunk sizes are 32-bit; refuse to emit a file with a truncated header
     if (pcm.size() > ((size_t) UINT32_MAX - 36) / 2) {
         return false;
     }
-    const uint32_t data_sz   = (uint32_t) (pcm.size() * 2);
-    const uint32_t riff_sz   = 36 + data_sz;
-    const uint32_t fmt_sz    = 16, byte_rate = (uint32_t) rate * 2;
-    const uint16_t fmt = 1, ch = 1, align = 2, bits = 16;
-    const uint32_t rate32    = (uint32_t) rate;
-    auto put = [&](const void * p, size_t n) {
-        const char * c = (const char *) p;
-        buf.insert(buf.end(), c, c + n);
-    };
-    put("RIFF", 4); put(&riff_sz, 4); put("WAVE", 4);
-    put("fmt ", 4); put(&fmt_sz, 4);
-    put(&fmt, 2); put(&ch, 2); put(&rate32, 4);
-    put(&byte_rate, 4); put(&align, 2); put(&bits, 2);
-    put("data", 4); put(&data_sz, 4);
-    for (float v : pcm) {
-        int16_t s = (int16_t) (std::max(-1.0f, std::min(1.0f, v)) * 32767.0f);
-        put(&s, 2);
-    }
+    write_wav16_header(buf, (uint32_t) (pcm.size() * 2), rate);
+    append_wav16_pcm(buf, pcm.data(), pcm.size());
     return true;
 }
 
@@ -92,6 +101,8 @@ public:
     // set out_stop on end-of-speech, h_state_out must be null if no frame is generated
     virtual int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) = 0;
     virtual int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) = 0;
+    // forces any buffered codes through code2wav now, regardless of window_frames
+    virtual int32_t flush() = 0;
 
 protected:
     llama_context * lctx;
@@ -121,6 +132,9 @@ public:
         prompt_batch.reset();
         n_prompt = 0;
         prompt_pos = 0;
+        stream = false;
+        pcm_sent = 0;
+        wav_header_sent = false;
     }
 
     int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override {
@@ -208,6 +222,7 @@ public:
         top_p = inp->top_p > 0 ? inp->top_p : def.top_p;
         seed  = inp->seed;
         out_type = inp->out_type;
+        stream = inp->stream;
 
         // the prompt above holds the whole text stream up to tts_eos, so every generated
         // frame adds tts_pad on top of the codes embedding
@@ -301,31 +316,60 @@ public:
     }
 
     int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) override {
-        if (!flush_gen_wav()) {
-            return 1;
+        *out_sample_rate = info.sample_rate;
+
+        if (!stream) {
+            // one-shot call: force out whatever's left, regardless of window_frames
+            if (!flush_gen_wav()) {
+                return 1;
+            }
+            if (out_n_samples) {
+                *out_n_samples = (int64_t) audio_pcm.size();
+            }
+            if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
+                *out_data     = (const char *) audio_pcm.data();
+                *out_data_len = audio_pcm.size() * sizeof(float);
+                return 0;
+            }
+            out_buf.clear();
+            if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) {
+                LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n");
+                return 1;
+            }
+            *out_data     = out_buf.data();
+            *out_data_len = out_buf.size();
+            return 0;
         }
 
-        *out_sample_rate = info.sample_rate;
+        // streaming: only return audio produced since the previous call
+        const size_t n_new = audio_pcm.size() - pcm_sent;
         if (out_n_samples) {
-            *out_n_samples = (int64_t) audio_pcm.size();
+            *out_n_samples = (int64_t) n_new;
         }
 
         if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
-            *out_data     = (const char *) audio_pcm.data();
-            *out_data_len = audio_pcm.size() * sizeof(float);
+            *out_data     = (const char *) (audio_pcm.data() + pcm_sent);
+            *out_data_len = n_new * sizeof(float);
+            pcm_sent = audio_pcm.size();
             return 0;
         }
 
         out_buf.clear();
-        if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) {
-            LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n");
-            return 1;
+        if (!wav_header_sent) {
+            write_wav16_header(out_buf, UINT32_MAX, info.sample_rate);
+            wav_header_sent = true;
         }
+        append_wav16_pcm(out_buf, audio_pcm.data() + pcm_sent, n_new);
+        pcm_sent = audio_pcm.size();
         *out_data     = out_buf.data();
         *out_data_len = out_buf.size();
         return 0;
     }
 
+    int32_t flush() override {
+        return flush_gen_wav() ? 0 : 1;
+    }
+
 private:
     bool ensure_cache() {
         if (specials_ok) {
@@ -369,7 +413,7 @@ private:
             LOG_ERR("mtmd_helper_gen_audio: mmproj has no speaker/audio encoder\n");
             return false;
         }
-        const std::string  marker = mtmd_default_marker();
+        const std::string  marker = mtmd_get_marker(mctx);
         mtmd_input_text     text{ marker.c_str(), marker.size(), false, true };
         mtmd_input_chunks * chunks = mtmd_input_chunks_init();
         const mtmd_bitmap * bptr = bitmap;
@@ -455,6 +499,9 @@ private:
     std::vector<float> h_state_buf;
     mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
     std::vector<char> out_buf;
+    bool   stream          = false;
+    size_t pcm_sent        = 0; // samples already returned by get_output()
+    bool   wav_header_sent = false;
 };
 
 // settings that only live in the reference's per-pack yaml, not in the checkpoint
@@ -506,6 +553,9 @@ public:
         chunk_idx = 0;
         n_voice_pos = 0;
         chunk_budget = 0;
+        stream = false;
+        pcm_sent = 0;
+        wav_header_sent = false;
     }
 
     int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override {
@@ -575,6 +625,7 @@ public:
 
         seed     = inp->seed;
         out_type = inp->out_type;
+        stream   = inp->stream;
 
         return 0;
     }
@@ -662,31 +713,60 @@ public:
     }
 
     int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) override {
-        if (!flush_gen_wav()) {
-            return 1;
+        *out_sample_rate = info.sample_rate;
+
+        if (!stream) {
+            // one-shot call: force out whatever's left, regardless of window_frames
+            if (!flush_gen_wav()) {
+                return 1;
+            }
+            if (out_n_samples) {
+                *out_n_samples = (int64_t) audio_pcm.size();
+            }
+            if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
+                *out_data     = (const char *) audio_pcm.data();
+                *out_data_len = audio_pcm.size() * sizeof(float);
+                return 0;
+            }
+            out_buf.clear();
+            if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) {
+                LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n");
+                return 1;
+            }
+            *out_data     = out_buf.data();
+            *out_data_len = out_buf.size();
+            return 0;
         }
 
-        *out_sample_rate = info.sample_rate;
+        // streaming: only return audio produced since the previous call
+        const size_t n_new = audio_pcm.size() - pcm_sent;
         if (out_n_samples) {
-            *out_n_samples = (int64_t) audio_pcm.size();
+            *out_n_samples = (int64_t) n_new;
         }
 
         if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
-            *out_data     = (const char *) audio_pcm.data();
-            *out_data_len = audio_pcm.size() * sizeof(float);
+            *out_data     = (const char *) (audio_pcm.data() + pcm_sent);
+            *out_data_len = n_new * sizeof(float);
+            pcm_sent = audio_pcm.size();
             return 0;
         }
 
         out_buf.clear();
-        if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) {
-            LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n");
-            return 1;
+        if (!wav_header_sent) {
+            write_wav16_header(out_buf, UINT32_MAX, info.sample_rate);
+            wav_header_sent = true;
         }
+        append_wav16_pcm(out_buf, audio_pcm.data() + pcm_sent, n_new);
+        pcm_sent = audio_pcm.size();
         *out_data     = out_buf.data();
         *out_data_len = out_buf.size();
         return 0;
     }
 
+    int32_t flush() override {
+        return flush_gen_wav() ? 0 : 1;
+    }
+
 private:
     bool ensure_cache() {
         if (specials_ok) {
@@ -907,7 +987,7 @@ private:
             LOG_ERR("mtmd_helper_gen_audio: mmproj has no voice encoder\n");
             return false;
         }
-        const std::string  marker = mtmd_default_marker();
+        const std::string  marker = mtmd_get_marker(mctx);
         mtmd_input_text     text{ marker.c_str(), marker.size(), false, true };
         mtmd_input_chunks * chunks = mtmd_input_chunks_init();
         const mtmd_bitmap * bptr = bitmap;
@@ -989,6 +1069,9 @@ private:
     std::vector<float>   h_state_buf;
     mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
     std::vector<char> out_buf;
+    bool   stream          = false;
+    size_t pcm_sent        = 0; // samples already returned by get_output()
+    bool   wav_header_sent = false;
 };
 
 static std::unique_ptr<mtmd_gen_audio_pipeline> make_pipeline(llama_context * lctx, mtmd_context * mctx) {
@@ -1017,11 +1100,20 @@ void mtmd_helper_gen_audio_free(mtmd_helper_gen_audio * ctx) {
 }
 
 void mtmd_helper_gen_audio_reset(mtmd_helper_gen_audio * ctx) {
-    if (ctx->pipeline) {
+    if (ctx && ctx->pipeline) {
         ctx->pipeline->reset();
     }
 }
 
+struct mtmd_helper_gen_audio_inp mtmd_helper_gen_audio_inp_default(void) {
+    mtmd_helper_gen_audio_inp inp{};
+    inp.top_k    = 50;
+    inp.top_p    = 1.0f;
+    inp.seed     = UINT32_MAX; // random
+    inp.out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
+    return inp;
+}
+
 int32_t mtmd_helper_gen_audio_set_input(mtmd_helper_gen_audio * ctx, const mtmd_helper_gen_audio_inp * inp) {
     if (!ctx->pipeline) {
         LOG_ERR("mtmd_helper_gen_audio: unsupported or missing gen-audio pipeline\n");
@@ -1058,3 +1150,10 @@ int32_t mtmd_helper_gen_audio_get_output(mtmd_helper_gen_audio * ctx, int32_t *
     }
     return ctx->pipeline->get_output(out_sample_rate, out_data, out_data_len, out_n_samples);
 }
+
+int32_t mtmd_helper_gen_audio_flush(mtmd_helper_gen_audio * ctx) {
+    if (!ctx->pipeline) {
+        return 1;
+    }
+    return ctx->pipeline->flush();
+}

tools/mtmd/mtmd-helper.h

@@ -204,6 +204,7 @@ enum mtmd_helper_gen_audio_outtype {
     MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV, // WAV PCM 16-bit LE, mono
 };
 struct mtmd_helper_gen_audio_inp {
+    bool stream; // if true, output() must be called after each step_gen()
     llama_seq_id seq_id;
 
     const char * prompt;
@@ -219,6 +220,8 @@ struct mtmd_helper_gen_audio_inp {
     enum mtmd_helper_gen_audio_outtype out_type;
 };
 
+MTMD_API struct mtmd_helper_gen_audio_inp mtmd_helper_gen_audio_inp_default(void);
+
 MTMD_API mtmd_helper_gen_audio * mtmd_helper_gen_audio_init(
                                     struct llama_context * lctx,
                                     struct mtmd_context * mctx);
@@ -250,6 +253,8 @@ MTMD_API int32_t mtmd_helper_gen_audio_step_gen(
 
 // out_data valid until next get_output() or reset() call
 // out_n_samples (optional, can be NULL) receives the number of generated PCM samples
+// if inp->stream is true: returns only audio produced since the previous call, and
+// *out_data_len == 0 whenever a full window_frames batch hasn't accumulated yet
 MTMD_API int32_t mtmd_helper_gen_audio_get_output(
                         mtmd_helper_gen_audio * ctx,
                         int32_t * out_sample_rate,
@@ -257,6 +262,10 @@ MTMD_API int32_t mtmd_helper_gen_audio_get_output(
                         size_t * out_data_len,
                         int64_t * out_n_samples);
 
+// forces any buffered codes through code2wav now, regardless of window_frames;
+// call once when generation has ended, before the last get_output() in stream mode
+MTMD_API int32_t mtmd_helper_gen_audio_flush(mtmd_helper_gen_audio * ctx);
+
 #ifdef __cplusplus
 } // extern "C"
 #endif
@@ -283,8 +292,41 @@ struct mtmd_helper_gen_audio_deleter {
 };
 using gen_audio_ptr = std::unique_ptr<mtmd_helper_gen_audio, mtmd_helper_gen_audio_deleter>;
 struct gen_audio {
+
+    // sub-struct, RAII wrapper for mtmd_helper_gen_audio_inp
+    struct inp {
+        mtmd_helper_gen_audio_inp data = mtmd_helper_gen_audio_inp_default();
+        std::string prompt_str;
+        std::string lang_str;
+        mtmd::bitmap_ptr speaker_ref_ptr;
+
+        inp()                             = default;
+        inp(inp &&)                       = default;
+        inp & operator=(inp &&)           = default;
+        inp(const inp &)                  = delete;
+        inp & operator=(const inp &)      = delete;
+
+        void set_prompt     (std::string p)        { prompt_str = std::move(p); }
+        void set_lang       (std::string l)        { lang_str   = std::move(l); }
+        void set_speaker_ref(mtmd::bitmap_ptr bmp) { speaker_ref_ptr = std::move(bmp); }
+
+        // pointers are only valid as long as *this is alive
+        const mtmd_helper_gen_audio_inp * get() {
+            data.prompt      = prompt_str.c_str();
+            data.prompt_len  = prompt_str.size();
+            data.lang        = lang_str.empty() ? nullptr : lang_str.c_str();
+            data.speaker_ref = speaker_ref_ptr.get();
+            return &data;
+        }
+    };
+
     gen_audio_ptr ctx;
-    gen_audio(struct llama_context * lctx, struct mtmd_context * mctx) : ctx(mtmd_helper_gen_audio_init(lctx, mctx)) {}
+    void init(struct llama_context * lctx, struct mtmd_context * mctx) {
+        ctx.reset(mtmd_helper_gen_audio_init(lctx, mctx));
+    }
+    bool valid() const {
+        return ctx.get() != nullptr;
+    }
     void reset() {
         mtmd_helper_gen_audio_reset(ctx.get());
     }
@@ -300,6 +342,9 @@ struct gen_audio {
     int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples = nullptr) {
         return mtmd_helper_gen_audio_get_output(ctx.get(), out_sample_rate, out_data, out_data_len, out_n_samples);
     }
+    int32_t flush() {
+        return mtmd_helper_gen_audio_flush(ctx.get());
+    }
 };
 
 } // namespace mtmd_helper

tools/mtmd/mtmd-image.cpp

@@ -873,16 +873,28 @@ mtmd_image_preproc_out mtmd_image_preprocessor_glm5v::preprocess(const clip_imag
 
 clip_image_size mtmd_image_preprocessor_longest_edge::get_target_size(const clip_hparams & hparams, const clip_image_size & original_size) {
     GGML_ASSERT(hparams.image_longest_edge > 0);
+    clip_image_size size = original_size;
+    if (hparams.image_floor_before_align) {
+        const double ratio = double(std::max(size.width, size.height)) / hparams.image_longest_edge;
+        if (ratio > 1.0) {
+            size.width  = static_cast<int>(std::floor(size.width / ratio));
+            size.height = static_cast<int>(std::floor(size.height / ratio));
+            if (size.width <= 0 || size.height <= 0) {
+                throw std::runtime_error("image dimensions must be positive after resizing");
+            }
+        }
+    }
+
     // the original pixtral model doesn't have n_merge
     const int cur_merge = hparams.n_merge == 0 ? 1 : hparams.n_merge;
     return img_tool::calc_size_preserved_ratio(
-        original_size,
+        size,
         {
             /* align_size   */ hparams.patch_size * cur_merge,
             /* min_pixels   */ std::max(0, hparams.image_min_pixels),
             /* max_pixels   */ std::max(0, hparams.image_max_pixels),
             /* longest_edge */ std::min(hparams.image_longest_edge,
-                                        std::max(original_size.width, original_size.height)),
+                                        std::max(size.width, size.height)),
         });
 }
 
@@ -947,18 +959,41 @@ bool mtmd_image_preprocessor_lfm2::should_tile(
     const int h_bar = std::max(hparams.patch_size, round_by_factor(original_size.height));
     const int w_bar = std::max(hparams.patch_size, round_by_factor(original_size.width));
 
-    return static_cast<double>(h_bar) * static_cast<double>(w_bar) >
-           static_cast<double>(hparams.image_max_pixels) * max_pixels_tolerance;
+    return hparams.preproc_max_tiles > 1 &&
+           static_cast<double>(h_bar) * static_cast<double>(w_bar) >
+           static_cast<double>(hparams.image_max_pixels) * hparams.preproc_max_pixels_tolerance;
+}
+
+// same as Lfm2VlImageProcessor.smart_resize: Python's round half to even, double precision
+clip_image_size mtmd_image_preprocessor_lfm2::smart_resize(
+        const clip_hparams & hparams,
+        const clip_image_size & original_size) {
+    const int    align_size = hparams.patch_size * hparams.n_merge;
+    const double height     = original_size.height;
+    const double width      = original_size.width;
+
+    int h_bar = std::max(align_size, static_cast<int>(std::nearbyint(height / align_size)) * align_size);
+    int w_bar = std::max(align_size, static_cast<int>(std::nearbyint(width  / align_size)) * align_size);
+
+    if (static_cast<double>(h_bar) * w_bar > hparams.image_max_pixels) {
+        const double beta = std::sqrt(height * width / hparams.image_max_pixels);
+        h_bar = std::max(align_size, static_cast<int>(std::floor(height / beta / align_size)) * align_size);
+        w_bar = std::max(align_size, static_cast<int>(std::floor(width  / beta / align_size)) * align_size);
+    } else if (static_cast<double>(h_bar) * w_bar < hparams.image_min_pixels) {
+        const double beta = std::sqrt(hparams.image_min_pixels / (height * width));
+        h_bar = static_cast<int>(std::ceil(height * beta / align_size)) * align_size;
+        w_bar = static_cast<int>(std::ceil(width  * beta / align_size)) * align_size;
+    }
+
+    return {w_bar, h_bar};
 }
 
 mtmd_image_preprocessor_llava_uhd::slice_instructions mtmd_image_preprocessor_lfm2::get_slice_instructions(const clip_image_size & original_size) const {
     mtmd_image_preprocessor_llava_uhd::slice_instructions inst;
-    const int align_size = hparams.patch_size * hparams.n_merge;
-    inst.overview_size = img_tool::calc_size_preserved_ratio(
-                            original_size,
-                            { align_size, hparams.image_min_pixels, hparams.image_max_pixels, 0 });
+    inst.overview_size = smart_resize(hparams, original_size);
 
     const bool needs_tiling = should_tile(hparams, original_size);
+    const int  tile_size    = hparams.preproc_tile_size;
 
     if (!needs_tiling) {
         inst.refined_size = clip_image_size{0, 0};
@@ -966,7 +1001,7 @@ mtmd_image_preprocessor_llava_uhd::slice_instructions mtmd_image_preprocessor_lf
         return inst;
     }
 
-    const clip_image_size grid = get_grid_layout(original_size.height, original_size.width);
+    const clip_image_size grid = get_grid_layout(hparams, original_size.height, original_size.width);
 
     inst.grid_size    = grid;
     inst.refined_size = clip_image_size{tile_size * grid.width, tile_size * grid.height};
@@ -994,23 +1029,26 @@ mtmd_image_preprocessor_llava_uhd::slice_instructions mtmd_image_preprocessor_lf
     return inst;
 }
 
+// double precision as in Python, float breaks exact ties such as 720/1600 between 2/4 and 2/5
 clip_image_size mtmd_image_preprocessor_lfm2::find_closest_aspect_ratio(
-        float aspect_ratio,
+        const clip_hparams & hparams,
+        double aspect_ratio,
         const std::vector<clip_image_size> & target_ratios,
-        int width, int height) const {
-    float best_ratio_diff = std::numeric_limits<float>::max();
+        int width, int height) {
+    double best_ratio_diff = std::numeric_limits<double>::max();
     clip_image_size best_ratio = {1, 1};
-    const float area = static_cast<float>(width * height);
+    const double area = static_cast<double>(width) * height;
 
     for (const auto & ratio : target_ratios) {
-        const float target_aspect_ratio = static_cast<float>(ratio.width) / ratio.height;
-        const float ratio_diff = std::abs(aspect_ratio - target_aspect_ratio);
+        const double target_aspect_ratio = static_cast<double>(ratio.width) / ratio.height;
+        const double ratio_diff = std::abs(aspect_ratio - target_aspect_ratio);
         if (ratio_diff < best_ratio_diff) {
             best_ratio_diff = ratio_diff;
             best_ratio = ratio;
         } else if (ratio_diff == best_ratio_diff) {
-            const float target_area = static_cast<float>(tile_size * tile_size * ratio.width * ratio.height);
-            if (area > 0.5f * target_area) {
+            const int    tile_size   = hparams.preproc_tile_size;
+            const double target_area = static_cast<double>(tile_size * tile_size * ratio.width * ratio.height);
+            if (area > 0.5 * target_area) {
                 best_ratio = ratio;
             }
         }
@@ -1018,7 +1056,9 @@ clip_image_size mtmd_image_preprocessor_lfm2::find_closest_aspect_ratio(
     return best_ratio;
 }
 
-std::vector<clip_image_size> mtmd_image_preprocessor_lfm2::get_target_ratios() const {
+std::vector<clip_image_size> mtmd_image_preprocessor_lfm2::get_target_ratios(const clip_hparams & hparams) {
+    const int min_tiles = hparams.preproc_min_tiles;
+    const int max_tiles = hparams.preproc_max_tiles;
     std::vector<clip_image_size> ratios;
     for (int n = min_tiles; n <= max_tiles; n++) {
         for (int w = 1; w <= n; w++) {
@@ -1044,10 +1084,10 @@ std::vector<clip_image_size> mtmd_image_preprocessor_lfm2::get_target_ratios() c
     return ratios;
 }
 
-clip_image_size mtmd_image_preprocessor_lfm2::get_grid_layout(int height, int width) const {
-    const float aspect_ratio = static_cast<float>(width) / height;
-    const auto ratios = get_target_ratios();
-    return find_closest_aspect_ratio(aspect_ratio, ratios, width, height);
+clip_image_size mtmd_image_preprocessor_lfm2::get_grid_layout(const clip_hparams & hparams, int height, int width) {
+    const double aspect_ratio = static_cast<double>(width) / height;
+    const auto ratios = get_target_ratios(hparams);
+    return find_closest_aspect_ratio(hparams, aspect_ratio, ratios, width, height);
 }
 
 //

tools/mtmd/mtmd-image.h

@@ -191,27 +191,24 @@ protected:
 };
 
 // custom llava-uhd slicing logic for LFM2
+// tiling params (min/max tiles, tile size, pixels tolerance) come from hparams
 // ref: https://github.com/huggingface/transformers/blob/v5.1.0/src/transformers/models/lfm2_vl/image_processing_lfm2_vl_fast.py
 struct mtmd_image_preprocessor_lfm2 : mtmd_image_preprocessor_llava_uhd {
-    // ref: https://huggingface.co/LiquidAI/LFM2.5-VL-1.6B/blob/main/processor_config.json
-    static constexpr int   min_tiles            = 2;
-    static constexpr int   max_tiles            = 10;
-    static constexpr float max_pixels_tolerance = 2.0f;
-    static constexpr int   tile_size            = 512;
-
     using mtmd_image_preprocessor_llava_uhd::mtmd_image_preprocessor_llava_uhd;
     mtmd_image_preproc_out preprocess(const clip_image_u8 & img) const override;
     slice_instructions get_slice_instructions(const clip_image_size & original_size) const override;
 
     static bool should_tile(const clip_hparams & hparams, const clip_image_size & original_size);
+    static clip_image_size smart_resize(const clip_hparams & hparams, const clip_image_size & original_size);
+    static clip_image_size get_grid_layout(const clip_hparams & hparams, int height, int width);
 
 private:
-    clip_image_size find_closest_aspect_ratio(
-            float aspect_ratio,
+    static clip_image_size find_closest_aspect_ratio(
+            const clip_hparams & hparams,
+            double aspect_ratio,
             const std::vector<clip_image_size> & target_ratios,
-            int width, int height) const;
-    std::vector<clip_image_size> get_target_ratios() const;
-    clip_image_size get_grid_layout(int height, int width) const;
+            int width, int height);
+    static std::vector<clip_image_size> get_target_ratios(const clip_hparams & hparams);
 };
 
 struct mtmd_image_preprocessor_idefics3 : mtmd_image_preprocessor_llava_uhd {

tools/mtmd/mtmd.cpp

@@ -804,7 +804,11 @@ struct mtmd_context {
                 {
                     // https://github.com/huggingface/transformers/blob/1cd110c6cb6a6237614130c470e9a902dbc1a4bd/docs/source/en/model_doc/pixtral.md
                     img_end = "[IMG_END]";
-                    image_preproc = std::make_unique<mtmd_image_preprocessor_dyn_size>(ctx_v);
+                    if (clip_get_hparams(ctx_v)->image_longest_edge > 0) {
+                        image_preproc = std::make_unique<mtmd_image_preprocessor_longest_edge>(ctx_v);
+                    } else {
+                        image_preproc = std::make_unique<mtmd_image_preprocessor_dyn_size>(ctx_v);
+                    }
                 } break;
             case PROJECTOR_TYPE_PHI4:
                 {
@@ -2842,10 +2846,11 @@ struct mtmd_caps mtmd_get_cap_from_file(const char * fname) {
         mtmd_caps cap;
         cap.inp_audio  = tmp.has_audio;
         cap.inp_vision = tmp.has_vision;
+        cap.gen_audio  = tmp.has_gen_audio;
         return cap;
     } catch (const std::exception & e) {
         LOG_ERR("%s: failed to get capabilities from file '%s': %s\n", __func__, fname, e.what());
-        return mtmd_caps{ false, false };
+        return mtmd_caps{ false, false, false };
     }
 }
 

tools/mtmd/mtmd.h

@@ -360,6 +360,7 @@ MTMD_API void mtmd_log_set(ggml_log_callback log_callback, void * user_data);
 struct mtmd_caps {
     bool inp_vision;
     bool inp_audio;
+    bool gen_audio;
 };
 MTMD_API struct mtmd_caps mtmd_get_cap_from_file(const char * mmproj_fname);
 

tools/server/README.md

@@ -802,6 +802,50 @@ curl http://127.0.0.1:8012/v1/rerank \
     }' | jq
 ```
 
+### POST `/tts`: Generate speech audio from text
+
+Returns raw audio bytes (`audio/wav` by default) rather than JSON. For more info, see [tts/README.md](../tts/README.md)
+
+*Options:*
+
+`input`: The text to speak (alias: `prompt`).
+
+`lang`: Language code for the utterance (model-dependent, e.g. `en`, `zh`). Optional.
+
+`speaker_ref_b64`: Base64-encoded reference audio to clone the speaker's voice. Optional. Alternatively, upload the reference audio as a `speaker_ref` file field via `multipart/form-data` (see example below) - if both are provided, the uploaded file takes precedence.
+
+`top_k`, `top_p`: Sampling params for the acoustic code predictor. Optional, model-dependent defaults apply.
+
+`repeat_penalty`: Repetition penalty applied to the backbone sampler over the whole generation (default `1.05`). Without this, the backbone can loop and re-generate the same utterance.
+
+`seed`: RNG seed for both the backbone sampler and the codec/vocoder (`-1` = random, the default). Set to a fixed value for reproducible output.
+
+`n_predict`: Max number of audio frames to generate. Defaults to `512`; generation normally stops earlier once the model emits an end-of-speech token.
+
+`response_format`: `wav` (default) or `pcm` (raw `float32` little-endian mono samples, no header; the response `Content-Type` carries the sample rate, e.g. `audio/pcm;rate=24000;encoding=float;bits=32`).
+
+`stream`: If `true`, the response is streamed as audio becomes available instead of waiting for the full generation to finish. WAV streaming writes an RFC-noncompliant header with an unknown (`0xFFFFFFFF`) size field, since the final length isn't known up front; most players and decoders (ffmpeg, VLC, ...) handle this by reading until EOF.
+
+Note: it's highly recommended to always provide a speaker reference voice; otherwise, the model's performance may be degraded.
+
+*Examples:*
+
+```shell
+curl -X POST http://127.0.0.1:9931/tts \
+    -H "Content-Type: application/json" \
+    -d '{"input": "Hello, this is a test."}' \
+    -o output.wav
+```
+
+With a speaker reference uploaded as a file (`multipart/form-data`), instead of base64-encoding it into the JSON body:
+
+```shell
+curl -X POST http://127.0.0.1:9931/tts \
+    -F "input=Hello, this is a test." \
+    -F "speaker_ref=@/path/to/speaker-reference.wav;type=audio/wav" \
+    -o output.wav
+```
+
 ### POST `/infill`: For code infilling.
 
 Takes a prefix and a suffix and returns the predicted completion as stream.
@@ -1262,11 +1306,12 @@ Each object in `data` has an `architecture` object. It has two string arrays:
 - `input_modalities` lists what the model can read. It always has `text`, plus each media type that the model supports.
 - `output_modalities` lists what the model can produce.
 
-One output value is special:
+Some output values are special:
 
 | Value | Meaning |
 |---|---|
 | `decisions` | The model is a native decision model. Serve it with [`/v1/systemone`](#post-v1systemone-typesafe-compatible-system-one-api). |
+| `audio` | The model can generate speech. Serve it with [`/tts`](#post-tts-generate-speech-audio-from-text). A TTS-only model has `["audio"]` without `text`. |
 
 A language model that classifies with prompts does not get `decisions`. Only native decision models do.
 
@@ -1491,6 +1536,31 @@ Acts on an in-flight completion identified by its `id` (the `id` field streamed
 
 Returns a JSON object with a boolean `success` field, and an optional `message` field describing the reason when `success` is `false`.
 
+### WebSocket `/v1/realtime`: Realtime transcription
+
+Transcribes streamed audio, for models with a streaming audio encoder. Compatible with OAI `/v1/realtime` schema.
+
+Audio must be PCM16, mono, at the model sample rate. The sample rate is returned in `session.created` as `audio.input.format.rate`. Each `input_audio_buffer.append` carries a base64-encoded chunk of PCM16 samples.
+
+*Client events:*
+
+- `session.update`: accepted for compatibility, the server replies with `session.updated`
+- `input_audio_buffer.append`: `audio` holds base64 PCM16 samples
+- `input_audio_buffer.commit`: flushes the audio received so far, so the words still held back by the model's delay are produced. The server replies with `input_audio_buffer.committed` and starts a new item. With `"final": true`, the input ends instead and the last item is completed once the stream drains.
+- `input_audio_buffer.clear`: accepted for compatibility, audio already sent cannot be discarded
+
+Closing the websocket also ends the input.
+
+*Server events:*
+
+- `session.created` and `session.updated`: carry the `session` object
+- `conversation.item.input_audio_transcription.delta`: `delta` holds the new text for `item_id`
+- `conversation.item.input_audio_transcription.completed`: `transcript` holds the full text of `item_id`
+- `input_audio_buffer.committed`: `item_id` of the committed item, `previous_item_id` of the item committed before it (`null` for the first one)
+- `error`: `error.message` describes the problem
+
+Note that words held back by the model's delay at the time of a commit are attributed to the next item.
+
 ### POST `/v1/responses`: OpenAI-compatible Responses API
 
 *Options:*
@@ -1725,13 +1795,13 @@ Follows the [TypeSafe API](https://docs.typesafe.ai/api), streaming is not suppo
 
 The questions of a request are answered independently, an answer does not depend on the other questions. The exception is clef: it reads all the questions in one prompt and decides them jointly.
 
-The number of options of a `choice` question is limited by the model, for example: 52 for openjev, 255 for granite-switch, laya, clef, pplx-decider, vllm-decision2, vllm-d3, lfm2-d1 and lfm2-d1-omni. For laya, long questions and options are truncated to the token budget the model was trained with.
+The number of options of a `choice` question is limited by the model, for example: 52 for openjev, 255 for granite-switch, laya, clef, pplx-decider, vllm-decision2, vllm-d3, lfm2-d1, lfm2-d1-omni and h2o-lightning. For laya, long questions and options are truncated to the token budget the model was trained with.
 
-For laya, clef, vllm-d3 and lfm2-d1-omni, the whole prompt is evaluated in one batch: it must fit in `--ubatch-size`. An lfm2-d1-omni prompt is cut to 16384 tokens. For vllm-decision2, the options of a question and the end of the prompt must fit in one batch. A server that runs clef only serves this endpoint, text generation is not available. A decision model that is also a text model, such as granite-switch, serves both and reports `["text", "decisions"]`.
+For laya, clef, vllm-d3 and lfm2-d1-omni, the whole prompt is evaluated in one batch: it must fit in `--ubatch-size`. An lfm2-d1-omni prompt is cut to 16384 tokens. For vllm-decision2, the options of a question and the end of the prompt must fit in one batch. An h2o-lightning state longer than 32000 tokens keeps its first and last 16000 tokens. A server that runs clef only serves this endpoint, text generation is not available. A decision model that is also a text model, such as granite-switch, serves both and reports `["text", "decisions"]`.
 
 *Image input:*
 
-Image input needs a model that supports it (for example: openjev, clef, pplx-decider, vllm-d3, lfm2-d1, lfm2-d1-omni) and its multimodal projector, see `--mmproj`.
+Image input needs a model that supports it (for example: openjev, clef, pplx-decider, vllm-d3, lfm2-d1, lfm2-d1-omni, h2o-lightning) and its multimodal projector, see `--mmproj`.
 
 Images can be given in two ways, and both can be used in the same request:
 
@@ -1760,6 +1830,8 @@ All the images are placed before the state in the prompt, the ones from `files`
 
 The probabilities are scaled with the temperatures stored in the model file. They are not guaranteed to be calibrated for your data.
 
+For h2o-lightning, `score` questions can have 1 to 255 levels, and the `instructions` and `type` of a question can be left out (`choice` by default; `"instructions": null` is written as `null`).
+
 *Examples:*
 
 ```shell
@@ -1990,6 +2062,8 @@ For **GET** endpoints (`/props`, `/metrics`, etc.) The router uses the `model` q
 GET /props?model=ggml-org%2Fgemma-3-4b-it-GGUF%3AQ4_K_M
 ```
 
+For the **WebSocket** endpoint (`/v1/realtime`) The router also uses the `model` query parameter.
+
 By default, the model will be loaded automatically if it's not loaded. To disable this, add `--no-models-autoload` when starting the server. Additionally, you can include `?autoload=true|false` in the query param to control this behavior per-request.
 
 ### GET `/models`: List available models
@@ -2025,7 +2099,7 @@ Note:
     - If a model is not running, it will be added or updated according to the source
 2. When the model is loaded, the info from `/v1/models` is forwarded to router's `/v1/models`. This includes metadata about the model and the runtime instance.
 
-Each object in `data` has the same `architecture` object as [`GET /v1/models`](#get-v1models-openai-compatible-model-info-api) of a direct server. The server computes both arrays offline. It does not load the model, download files, or run inference. `output_modalities` comes from the GGUF metadata. `input_modalities` comes from the projector file. A native decision model shows `decisions` before its first load, after unload, and while it sleeps:
+Each object in `data` has the same `architecture` object as [`GET /v1/models`](#get-v1models-openai-compatible-model-info-api) of a direct server. The server computes both arrays offline. It does not load the model, download files, or run inference. `output_modalities` comes from the GGUF metadata, plus `audio` when the projector file can generate speech. `input_modalities` comes from the projector file. A native decision model shows `decisions` before its first load, after unload, and while it sleeps:
 
 ```json
 {

tools/server/server-chat.cpp

@@ -590,7 +590,7 @@ json server_chat_convert_anthropic_to_oai(const json & body) {
     }
 
     // Pass through common params
-    for (const auto & key : {"temperature", "top_p", "top_k", "stream", "chat_template_kwargs"}) {
+    for (const auto & key : {"temperature", "top_p", "top_k", "stream", "chat_template_kwargs", "id_slot"}) {
         if (body.contains(key)) {
             oai_body[key] = body.at(key);
         }

tools/server/server-common.cpp

@@ -202,7 +202,8 @@ const char * get_media_marker() {
 // model output modalities
 //
 
-std::vector<std::string> server_model_output_modalities(common_decision_type decision_type) {
+std::vector<std::string> server_model_output_modalities(common_decision_type decision_type, bool has_cap_chat, bool has_cap_tts) {
+    std::vector<std::string> out;
     switch (decision_type) {
         case COMMON_DECISION_TYPE_OPENJEV:
         case COMMON_DECISION_TYPE_LEV:
@@ -215,13 +216,28 @@ std::vector<std::string> server_model_output_modalities(common_decision_type dec
         case COMMON_DECISION_TYPE_VLLM_D3:
         case COMMON_DECISION_TYPE_LFM2_D1:
         case COMMON_DECISION_TYPE_LFM2_D1_OMNI:
-            return {"decisions"};
+        case COMMON_DECISION_TYPE_H2O_LIGHTNING:
+            out = {"decisions"};
+            break;
         case COMMON_DECISION_TYPE_GRANITE_SWITCH:
-            return {"text", "decisions"};
+            out = {"text", "decisions"};
+            break;
         default:
-            // fallback when there is no decision type or the metadata is bad
-            return {"text"};
+            // a tts-only model cannot generate text
+            if (has_cap_chat || !has_cap_tts) {
+                // fallback when there is no decision type or the metadata is bad
+                out = {"text"};
+            }
+            break;
     }
+    if (has_cap_tts) {
+        out.push_back("audio");
+    }
+    return out;
+}
+
+std::vector<std::string> server_model_output_modalities(const llama_model * model, bool has_cap_chat, bool has_cap_tts) {
+    return server_model_output_modalities(common_get_decision_type(model), has_cap_chat, has_cap_tts);
 }
 
 json server_model_architecture_json(

tools/server/server-common.h

@@ -114,7 +114,10 @@ const char * get_media_marker();
 //
 
 // output modalities for architecture.output_modalities in GET /models
-std::vector<std::string> server_model_output_modalities(common_decision_type decision_type);
+// has_cap_chat: false for a tts-only model, has_cap_tts: the model can generate audio
+std::vector<std::string> server_model_output_modalities(common_decision_type decision_type, bool has_cap_chat, bool has_cap_tts);
+// same, with the decision type read from a loaded model
+std::vector<std::string> server_model_output_modalities(const llama_model * model, bool has_cap_chat, bool has_cap_tts);
 
 // architecture object of GET /models; shared by the direct server and the router
 json server_model_architecture_json(

tools/server/server-context.cpp

@@ -1,4 +1,7 @@
 #include "server-context.h"
+#include <atomic>
+#include <thread>
+#include "base64.hpp"
 #include "server-chat.h"
 #include "server-common.h"
 #include "server-decision.h"
@@ -17,6 +20,7 @@
 #include "speculative.h"
 #include "mtmd.h"
 #include "mtmd-helper.h"
+#include "base64.hpp"
 
 #include <algorithm>
 #include <cstddef>
@@ -40,6 +44,14 @@
 constexpr int HTTP_POLLING_SECONDS = 1;
 
 static common_speculative_output_limits server_output_limits(const common_params & params) {
+    if (!params.mmproj.path.empty()) {
+        const auto mcaps = mtmd_get_cap_from_file(params.mmproj.path.c_str());
+        if (mcaps.gen_audio) {
+            // some TTS models decode with embeddings output for the whole batch
+            return { params.n_batch, params.n_batch };
+        }
+    }
+
     if (params.embedding ||
             (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != LLAMA_POOLING_TYPE_NONE)) {
         return { params.n_batch, 1 };
@@ -254,6 +266,27 @@ struct server_slot {
     mtmd_context * mctx = nullptr;
     mtmd::batch_ptr mbatch = nullptr;
 
+    struct tts_ctx {
+        mtmd_helper::gen_audio ctx;
+        std::vector<float> h_state_prompt; // the first h_state after step_prompt done
+        const float * h_state;
+        llama_token   sampled;
+        int32_t       n_decoded;
+        bool          stop;
+        bool is_supported() const {
+            return ctx.valid();
+        }
+        void reset() {
+            ctx.reset();
+            h_state_prompt.clear();
+            h_state   = nullptr;
+            sampled   = LLAMA_TOKEN_NULL;
+            n_decoded = 0;
+            stop      = false;
+        }
+    };
+    tts_ctx tts;
+
     // speculative decoding
     common_speculative * spec;
 
@@ -360,8 +393,13 @@ struct server_slot {
 
     // for TTS models, this is the embd generated from prev step, decode this to generate next hidden state
     // corresponding to one token position (size = n_embd)
+    // for realtime transcription, these are the rows from the last audio stream step (size = n_rows * n_embd)
     std::vector<float> inp_embd;
 
+    // realtime transcription, see SERVER_TASK_TYPE_TRANSCRIBE_RT
+    mtmd::audio_stream_ptr audio_stream;
+    bool audio_wait = false; // the stream needs more audio before it can produce rows
+
     server_slot_stats stats;
 
     // accepted tokens per draft position
@@ -413,6 +451,12 @@ struct server_slot {
 
         // clear multimodal state
         mbatch.reset();
+
+        tts.reset();
+
+        audio_stream.reset();
+        audio_wait = false;
+        inp_embd.clear();
     }
 
     void init_sampler() const {
@@ -542,16 +586,28 @@ struct server_slot {
     // add sampled token of this slot to the batch, optionally add the speculative draft tokens if any
     void handle_last_sampled_token(server_batch & batch) {
         bool add_ok = true;
+        // the audio stream has no token before the first step, track the position with a dummy token
+        const llama_token tok_track = sampled == LLAMA_TOKEN_NULL ? 0 : sampled;
         if (spec_draft.empty()) {
             // no speculative decoding
             i_batch = batch.size();
 
             if (!inp_embd.empty()) {
-                // text embeddings broadcast the same position across the M-RoPE sections
-                const llama_pos p = prompt.tokens.pos_next();
-                const llama_pos pos[GGML_MROPE_SECTIONS] = { p, p, p, 0 };
-                GGML_ASSERT((int32_t) inp_embd.size() == batch.n_embd);
-                add_ok &= batch.add_embd(id, inp_embd.data(), pos, true, false);
+                const int32_t n_rows = (int32_t) inp_embd.size() / batch.n_embd;
+                GGML_ASSERT((int32_t) inp_embd.size() == n_rows * batch.n_embd);
+                i_batch = batch.size() + n_rows - 1;
+                for (int32_t i = 0; i < n_rows; i++) {
+                    // text embeddings broadcast the same position across the M-RoPE sections
+                    const llama_pos p = prompt.tokens.pos_next();
+                    const llama_pos pos[GGML_MROPE_SECTIONS] = { p, p, p, 0 };
+                    add_ok &= batch.add_embd(id, inp_embd.data() + (size_t) i * batch.n_embd, pos, i == n_rows - 1, false);
+                    if (i < n_rows - 1) {
+                        prompt.tokens.push_back(tok_track);
+                    }
+                }
+                if (audio_stream) {
+                    inp_embd.clear(); // rows are consumed, the next step produces new ones
+                }
             } else {
                 add_ok &= batch.add(id, sampled, prompt.tokens.pos_next(), true, false);
             }
@@ -579,7 +635,7 @@ struct server_slot {
 
         GGML_ASSERT(add_ok && "batch must be large enough to hold the sampled and draft tokens");
 
-        prompt.tokens.push_back(sampled);
+        prompt.tokens.push_back(tok_track);
         prompt.tokens.insert(spec_draft);
     }
 
@@ -595,7 +651,12 @@ struct server_slot {
 
             // do not keep context of the child slots - the parent's context is enough
             // nor of the stateless slots - no task ever reuses it
-            if (task->is_child() || task->is_stateless()) {
+            if (task->is_child() || task->is_stateless() || audio_stream) {
+                prompt_clear();
+            }
+
+            // mtmd_helper always re-eval the whole prompt
+            if (task->type == SERVER_TASK_TYPE_TTS) {
                 prompt_clear();
             }
 
@@ -952,6 +1013,14 @@ public:
     mtmd_helper_init_opt init_opt = mtmd_helper_init_opt_default();
     const llama_vocab * vocab = nullptr;
 
+    bool has_cap_tts() const {
+        return mctx != nullptr && mtmd_gen_audio_get_info(mctx).type != MTMD_GEN_AUDIO_TYPE_NONE;
+    }
+
+    bool has_cap_chat() const {
+        return mctx == nullptr || mtmd_helper_model_can_chat(ctx_tgt, mctx);
+    }
+
     server_decision_context decision;
 
     server_queue    queue_tasks;
@@ -1180,7 +1249,8 @@ private:
         }
 
         // optionally fit mmproj memory usage
-        if (has_mmproj && params_base.fit_params) {
+        // on resume params_base already holds the margin from the first load
+        if (has_mmproj && params_base.fit_params && !is_resume) {
             const auto & mmproj_mem = mmproj_usage.backend_mem_usage;
             const int64_t t_elapsed = mmproj_usage_t_us;
             if (!mmproj_mem.empty()) {
@@ -1439,6 +1509,10 @@ private:
             slot.mctx                   = mctx;
             slot.prompt.tokens.has_mtmd = mctx != nullptr;
 
+            if (has_cap_tts()) {
+                slot.tts.ctx.init(ctx_tgt, mctx);
+            }
+
             SLT_TRC(slot, "new slot, n_ctx = %d\n", slot.n_ctx);
 
             slot.callback_on_release = [this](int id_slot) {
@@ -1920,6 +1994,20 @@ private:
 
         SLT_DBG(slot, "launching slot : %s\n", safe_json_to_str(slot.to_json()).c_str());
 
+        if (task.type == SERVER_TASK_TYPE_TTS) {
+            GGML_ASSERT(has_cap_tts()); // should already checked in route handler
+            if (!slot.tts.is_supported()) {
+                slot.tts.ctx.init(ctx_tgt, slot.mctx);
+            }
+            // mtmd_helper always re-eval the whole prompt
+            slot.prompt_clear();
+            task.tts_inp.data.seq_id = slot.id;
+            if (slot.tts.ctx.set_input(task.tts_inp.get()) != 0) {
+                send_error(task, "failed to process TTS prompt", ERROR_TYPE_SERVER);
+                return false;
+            }
+        }
+
         // initialize samplers
         if (task.need_sampling()) {
             try {
@@ -1937,6 +2025,9 @@ private:
             // TODO: getting pre sampling logits is not yet supported with backend sampling
             use_backend_sampling &= !need_pre_sample_logits;
 
+            // TODO: check verify if this actually works with TTS
+            use_backend_sampling &= task.type != SERVER_TASK_TYPE_TTS;
+
             // TODO: tmp until backend sampling is fully implemented
             if (use_backend_sampling) {
                 llama_set_sampler(ctx_tgt, slot.id, common_sampler_get(slot.smpl.get()));
@@ -1962,9 +2053,30 @@ private:
 
         slot.task = std::make_unique<const server_task>(std::move(task));
 
-        slot.state = slot.task->is_child()
-            ? SLOT_STATE_WAIT_OTHER // wait for the parent to process prompt
-            : SLOT_STATE_STARTED;
+        if (slot.task->type == SERVER_TASK_TYPE_TTS) {
+            slot.state = SLOT_STATE_PROCESSING_PROMPT;
+        } else {
+            slot.state = slot.task->is_child()
+                ? SLOT_STATE_WAIT_OTHER // wait for the parent to process prompt
+                : SLOT_STATE_STARTED;
+        }
+
+        // no prompt, the rows come from the audio stream, which audio controls reach from now on
+        if (slot.task->type == SERVER_TASK_TYPE_TRANSCRIBE_RT) {
+            slot.audio_stream.reset(mtmd_audio_stream_init(mctx, mtmd_audio_stream_params_default()));
+            if (!slot.audio_stream) {
+                send_error(slot, "failed to initialize the audio stream", ERROR_TYPE_SERVER);
+                slot.release();
+                return false;
+            }
+            slot.prompt_clear();
+            slot.init_sampler();
+            slot.sampled = LLAMA_TOKEN_NULL;
+            slot.i_batch = -1;
+            slot.stats.update_prompt_start();
+            slot.state = SLOT_STATE_GENERATING;
+            send_partial_response(slot, {}, false, true);
+        }
 
         // reset server kill-switch counter
         n_empty_consecutive = 0;
@@ -2227,6 +2339,18 @@ private:
         queue_results.send(std::move(res));
     }
 
+    void send_tts_result(server_slot & slot, int32_t sample_rate, const char * data, size_t data_len, bool final) {
+        auto res = std::make_unique<server_task_result_tts>();
+
+        res->id          = slot.task->id;
+        res->index       = slot.task->index;
+        res->sample_rate = sample_rate;
+        res->audio.assign(data, data_len);
+        res->final       = final;
+
+        queue_results.send(std::move(res));
+    }
+
     void send_final_response(server_slot & slot) {
         auto res = std::make_unique<server_task_result_cmpl_final>();
 
@@ -2828,8 +2952,10 @@ private:
         switch (task.type) {
             case SERVER_TASK_TYPE_COMPLETION:
             case SERVER_TASK_TYPE_INFILL:
+            case SERVER_TASK_TYPE_TRANSCRIBE_RT:
             case SERVER_TASK_TYPE_EMBEDDING:
             case SERVER_TASK_TYPE_RERANK:
+            case SERVER_TASK_TYPE_TTS:
             case SERVER_TASK_TYPE_DECISION:
                 {
                     // special case: if input is provided via CLI, tokenize it first
@@ -2915,15 +3041,15 @@ private:
 
                     server_slot * slot = get_slot_by_cmpl_id(task.params.control_cmpl_id);
                     if (slot == nullptr) {
-                        SRV_WRN("control %s on unknown completion id=%s, no live slot\n",
-                                task.params.control_action.c_str(), task.params.control_cmpl_id.c_str());
+                        SRV_WRN("control %d on unknown completion id=%s, no live slot\n",
+                                (int) task.params.control_action, task.params.control_cmpl_id.c_str());
                         res->success = false;
                         res->message = "no active completion for this id";
                         queue_results.send(std::move(res));
                         break;
                     }
 
-                    if (task.params.control_action == "reasoning_end") {
+                    if (task.params.control_action == SERVER_CONTROL_ACTION_REASONING_END) {
                         // the budget sampler only exists when reasoning control was armed
                         if (!slot->task->params.sampling.reasoning_control) {
                             res->success = false;
@@ -2934,6 +3060,22 @@ private:
                         // act on the live slot mid generation, never defer
                         common_sampler_reasoning_budget_force(slot->smpl.get());
                         res->success = true;
+                    } else if (task.params.control_action == SERVER_CONTROL_ACTION_AUDIO_APPEND
+                            || task.params.control_action == SERVER_CONTROL_ACTION_AUDIO_COMMIT
+                            || task.params.control_action == SERVER_CONTROL_ACTION_AUDIO_END) {
+                        if (!slot->audio_stream) {
+                            res->success = false;
+                            res->message = "not a realtime transcription slot";
+                        } else {
+                            auto * stream = slot->audio_stream.get();
+                            if (task.params.control_action == SERVER_CONTROL_ACTION_AUDIO_APPEND) {
+                                mtmd_audio_stream_push(stream, task.params.control_audio.data(), task.params.control_audio.size());
+                            } else {
+                                mtmd_audio_stream_flush(stream, task.params.control_action == SERVER_CONTROL_ACTION_AUDIO_END);
+                            }
+                            slot->audio_wait = false;
+                            res->success = true;
+                        }
                     } else {
                         res->success = false;
                         res->message = "unknown control action";
@@ -3253,6 +3395,39 @@ private:
     };
 #endif
 
+    // fill slot.inp_embd with the next rows of the audio stream
+    // returns false if the slot has nothing to decode now
+    bool slot_audio_step(server_slot & slot) {
+        if (!slot.inp_embd.empty()) {
+            return true; // rows from a previous step, not decoded yet
+        }
+
+        const int32_t n = mtmd_audio_stream_step(slot.audio_stream.get(), slot.sampled);
+        if (n > 0) {
+            const float * rows = mtmd_audio_stream_get_output_embd(slot.audio_stream.get());
+            slot.inp_embd.assign(rows, rows + (size_t) n * batch.n_embd);
+            slot.audio_wait = false;
+            return true;
+        }
+        if (n == 0) {
+            slot.audio_wait = true;
+            return false;
+        }
+        if (n == -2) {
+            send_error(slot, "audio stream error", ERROR_TYPE_SERVER);
+            slot.release();
+            return false;
+        }
+
+        // end of stream
+        slot.stop           = STOP_TYPE_EOS;
+        slot.has_next_token = false;
+        slot.print_timings();
+        send_final_response(slot);
+        slot.release();
+        return false;
+    }
+
     void update_slots() {
 #ifdef DEBUG_TIMINGS
         static int64_t t_prev = 0;
@@ -3268,11 +3443,12 @@ private:
 #endif
 
         // check if all slots are idle
+        // a realtime slot waiting for audio has no work, the next audio control task wakes the loop
         {
             bool all_idle = true;
 
             for (auto & slot : slots) {
-                if (slot.is_processing()) {
+                if (slot.is_processing() && !slot.audio_wait) {
                     all_idle = false;
                     break;
                 }
@@ -3305,6 +3481,14 @@ private:
             return;
         }
 
+        // note: TTS slots bypass the shared batch entirely
+        try {
+            process_tts_slots();
+        } catch (const std::exception & e) {
+            SRV_ERR("process_tts_slots() failed: %s\n", e.what());
+            abort_all_slots("process_tts_slots() failed: " + std::string(e.what()));
+        }
+
         GGML_ASSERT(batch.slot_batched || batch.size() == 0);
 
         if (batch.slot_batched) {
@@ -3367,10 +3551,89 @@ private:
         }
     }
 
+    void process_tts_slots() {
+        iterate(slots, [&](server_slot & slot) {
+            if (!slot.is_processing() || slot.task->type != SERVER_TASK_TYPE_TTS) {
+                return;
+            }
+
+            llama_set_embeddings(ctx_tgt, true);
+
+            if (slot.state == SLOT_STATE_PROCESSING_PROMPT) {
+                const int32_t ret = slot.tts.ctx.step_prompt(llama_n_batch(ctx_tgt));
+                if (ret < 0) {
+                    send_error(slot, "TTS prompt processing failed", ERROR_TYPE_SERVER);
+                    slot.release();
+                } else if (ret == 0) {
+                    // done prompt, do sample and save h_state_prompt
+                    slot.tts.sampled = common_sampler_sample(slot.smpl.get(), ctx_tgt, -1);
+                    common_sampler_accept(slot.smpl.get(), slot.tts.sampled, true);
+                    const float * h_embd = llama_get_embeddings_ith(ctx_tgt, -1);
+                    slot.tts.h_state_prompt.assign(h_embd, h_embd + llama_model_n_embd(model_tgt));
+                    slot.tts.h_state = slot.tts.h_state_prompt.data();
+                    slot.state = SLOT_STATE_GENERATING;
+                }
+                return;
+            }
+
+            const int32_t n_predict = slot.n_predict_max > 0 ? slot.n_predict_max : 512;
+            if (slot.tts.stop || slot.tts.n_decoded >= n_predict) {
+                int32_t      sample_rate = 0;
+                const char * data        = nullptr;
+                size_t       data_len    = 0;
+                // generation truly ends here: force out any sub-window remainder still buffered
+                if (slot.tts.ctx.flush() != 0 || slot.tts.ctx.get_output(&sample_rate, &data, &data_len) != 0) {
+                    send_error(slot, "failed to finalize TTS output", ERROR_TYPE_SERVER);
+                } else {
+                    send_tts_result(slot, sample_rate, data, data_len, true);
+                }
+                slot.release();
+                return;
+            }
+
+            const float * h_state_next = nullptr;
+            bool          stop         = false;
+            if (slot.tts.ctx.step_gen(slot.tts.sampled, slot.tts.h_state, &h_state_next, &stop) != 0) {
+                send_error(slot, "TTS generation failed", ERROR_TYPE_SERVER);
+                slot.release();
+                return;
+            }
+            if (h_state_next == nullptr) {
+                // end-of-speech without a new frame (e.g. pocket-tts eos head)
+                slot.tts.stop = true;
+                return;
+            }
+            slot.tts.h_state = h_state_next;
+            slot.tts.n_decoded++;
+            slot.tts.stop = stop;
+
+            if (!stop) {
+                slot.tts.sampled = common_sampler_sample(slot.smpl.get(), ctx_tgt, -1);
+                common_sampler_accept(slot.smpl.get(), slot.tts.sampled, true);
+            }
+
+            if (slot.task->params.stream) {
+                int32_t      sample_rate = 0;
+                const char * data        = nullptr;
+                size_t       data_len    = 0;
+                if (slot.tts.ctx.get_output(&sample_rate, &data, &data_len) != 0) {
+                    send_error(slot, "TTS streaming output failed", ERROR_TYPE_SERVER);
+                    slot.release();
+                } else if (data_len > 0) {
+                    send_tts_result(slot, sample_rate, data, data_len, false);
+                }
+            }
+        });
+    }
+
     void pre_decode() {
         // apply context-shift if needed
         // TODO: simplify and improve
         iterate(slots, [&](server_slot & slot) {
+            if (slot.task && slot.task->type == SERVER_TASK_TYPE_TTS) {
+                // TTS slots drive their own decode loop in process_tts_slots(), never enter the shared batch
+                return;
+            }
             if (slot.state == SLOT_STATE_GENERATING && slot.prompt.n_tokens() + 1 >= slot.n_ctx) {
                 if (!params_base.ctx_shift) {
                     // this check is redundant (for good)
@@ -3443,10 +3706,14 @@ private:
 
         // determine which slots are generating and drafting
         iterate(slots, [&](server_slot & slot) {
-            if (slot.state != SLOT_STATE_GENERATING) {
+            if (slot.state != SLOT_STATE_GENERATING || slot.task->type == SERVER_TASK_TYPE_TTS) {
                 return;
             }
 
+            if (slot.audio_stream && !slot_audio_step(slot)) {
+                return; // waiting for audio, or the stream ended
+            }
+
             // check if we can batch this slot with the previous one
             if (!slot_batched) {
                 slot_batched = &slot;
@@ -3585,7 +3852,7 @@ private:
                     return; // batch is full, skip remaining slots
                 }
 
-                if (!slot.is_processing()) {
+                if (!slot.is_processing() || slot.task->type == SERVER_TASK_TYPE_TTS) {
                     return;
                 }
 
@@ -3600,6 +3867,11 @@ private:
                     return;
                 }
 
+                // no prompt, the rows come from the audio stream
+                if (slot.task->type == SERVER_TASK_TYPE_TRANSCRIBE_RT) {
+                    return;
+                }
+
                 // this slot still has a prompt to be processed
                 if (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_STARTED) {
                     const auto & input_tokens = slot.task->tokens;
@@ -4521,11 +4793,14 @@ private:
 
             GGML_ASSERT(n_draft > 0);
 
+            // batch indices of the draft tokens, used to get the token probabilities below
+            auto spec_i_batch = std::move(slot.spec_i_batch);
+
             // verify and try to accept the draft
             {
                 common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get()));
 
-                GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1);
+                GGML_ASSERT(spec_i_batch.size() == n_draft + 1);
                 GGML_ASSERT(slot.spec_draft_q.empty() || (slot.spec_draft_q.size() == slot.spec_draft.size()));
                 const auto & synth_probs = common_speculative_get_synth_probs(spec.get());
 
@@ -4536,16 +4811,15 @@ private:
                 if (!synth_probs.empty()) {
                     // synthetic acceptance replaces verification entirely, so it comes first
                     accepted = server_sample_and_accept_synth(
-                            slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft,
+                            slot.smpl.get(), slot.ctx_tgt, spec_i_batch, slot.spec_draft,
                             synth_probs, slot.spec_synth_rng, slot.spec_is_replay);
                 } else if (slot.spec_is_replay && slot.use_spec_rejection()) {
-                    accepted = server_accept_replay(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft);
+                    accepted = server_accept_replay(slot.smpl.get(), slot.ctx_tgt, spec_i_batch, slot.spec_draft);
                 } else if (use_rejection) {
-                    accepted = common_sampler_sample_and_accept_n_rejection(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft, slot.spec_draft_q);
+                    accepted = common_sampler_sample_and_accept_n_rejection(slot.smpl.get(), slot.ctx_tgt, spec_i_batch, slot.spec_draft, slot.spec_draft_q);
                 } else {
-                    accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft);
+                    accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, spec_i_batch, slot.spec_draft);
                 }
-                slot.spec_i_batch.clear();
 
                 GGML_ASSERT(accepted.size() >= 1);
 
@@ -4632,7 +4906,10 @@ private:
                 result.text_to_send = common_token_to_piece(slot.ctx_tgt, result.tok, accept_special_token(slot, result.tok));
                 result.prob         = 1.0f; // set later
 
-                // TODO: set result.probs
+                // post_sampling_probs is not supported with speculative decoding
+                if (slot.task->params.sampling.n_probs > 0 && !slot.task->params.post_sampling_probs) {
+                    populate_token_probs(slot, result, false, params_base.special, spec_i_batch[i]);
+                }
 
                 slot.stats.n_gen += 1;
 
@@ -4820,11 +5097,14 @@ server_context_meta server_context::get_meta() const {
         /* model_aliases           */ impl->model_aliases,
         /* model_tags              */ impl->model_tags,
         /* model_path              */ impl->params_base.model.path,
-        /* model_output_modalities */ server_model_output_modalities(common_get_decision_type(impl->model_tgt)),
+        /* model_output_modalities */ server_model_output_modalities(impl->model_tgt, impl->has_cap_chat(), impl->has_cap_tts()),
         /* has_mtmd                */ impl->mctx != nullptr,
         /* has_inp_image           */ impl->chat_params.allow_image,
         /* has_inp_audio           */ impl->chat_params.allow_audio,
         /* has_inp_video           */ impl->chat_params.allow_video,
+        /* has_cap_chat            */ impl->has_cap_chat(),
+        /* has_cap_tts             */ impl->has_cap_tts(),
+        /* has_audio_stream        */ impl->mctx != nullptr && mtmd_support_audio_stream(impl->mctx),
         /* json_ui_settings        */ impl->json_ui_settings,
         /* slot_n_ctx              */ impl->n_ctx_slot(),
         /* pooling_type            */ llama_pooling_type(impl->ctx_tgt),
@@ -4899,6 +5179,11 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
 
     res->set_req(&req); // will also set spipe if needed
 
+    if (!ctx_server.has_cap_chat()) {
+        res->error(format_error_response("this server does not support chat/completions", ERROR_TYPE_NOT_SUPPORTED));
+        return res;
+    }
+
     int32_t sse_ping_interval = params.sse_ping_interval;
 
     try {
@@ -4992,10 +5277,20 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
                 res->ok(arr[0]);
             } else if (res_type == TASK_RESPONSE_TYPE_OAI_CHAT || res_type == TASK_RESPONSE_TYPE_OAI_CMPL) {
                 // if multiple results in OAI format, we need to re-format them
+                json & usage   = arr[0]["usage"];
                 json & choices = arr[0]["choices"];
                 for (size_t i = 1; i < arr.size(); i++) {
                     choices.push_back(std::move(arr[i]["choices"][0]));
                 }
+
+                // OAI counts generated tokens across all choices
+                int64_t n_completion_tokens = 0;
+                for (auto & item : arr) {
+                    n_completion_tokens += item["usage"]["completion_tokens"].get<int64_t>();
+                }
+                usage["completion_tokens"] = n_completion_tokens;
+                usage["total_tokens"]      = usage["prompt_tokens"].get<int64_t>() + n_completion_tokens;
+
                 res->ok(arr[0]);
             } else {
                 // multi-results, non-OAI compat
@@ -5605,7 +5900,7 @@ void server_routes::init_routes() {
             server_task task(SERVER_TASK_TYPE_CONTROL);
             task.id              = rd.get_new_id();
             task.params.control_cmpl_id = cmpl_id;
-            task.params.control_action  = action;
+            task.params.control_action  = SERVER_CONTROL_ACTION_REASONING_END;
             rd.post_task(std::move(task));
         }
 
@@ -5680,6 +5975,197 @@ void server_routes::init_routes() {
             session);
     };
 
+    // realtime transcription over websocket, modeled after the OpenAI realtime API
+    // audio is PCM16 mono at the model sample rate, sent as base64 in input_audio_buffer.append events
+    // input_audio_buffer.commit with "final": true ends the input, closing the socket also ends it
+    this->ws_realtime = [this](const server_http_req &, server_http_context::ws_conn & conn) {
+        auto send_event = [&conn](json ev) {
+            ev["event_id"] = "event_" + random_string();
+            conn.send(safe_json_to_str(ev));
+        };
+        auto send_error_event = [&send_event](const std::string & message, const std::string & type) {
+            send_event(json {
+                {"type", "error"},
+                {"error", {
+                    {"type",    type},
+                    {"code",    nullptr},
+                    {"message", message},
+                    {"param",   nullptr},
+                }},
+            });
+        };
+
+        if (!meta->has_audio_stream) {
+            send_error_event("The current model does not support realtime transcription.", "invalid_request_error");
+            conn.close();
+            return;
+        }
+
+        const std::string session_id = "sess_" + random_string();
+        const json session_config = json {
+            {"id",     session_id},
+            {"object", "realtime.transcription_session"},
+            {"type",   "transcription"},
+            {"audio",  {{"input", {
+                {"format",          {{"type", "audio/pcm"}, {"rate", mtmd_get_audio_sample_rate(ctx_server.mctx)}}},
+                {"transcription",   {{"model", meta->model_name}}},
+                {"turn_detection",  nullptr},
+                {"noise_reduction", nullptr},
+            }}}},
+        };
+
+        auto res = create_response();
+        auto & rd = res->rd;
+        {
+            server_task task(SERVER_TASK_TYPE_TRANSCRIBE_RT);
+            task.id = rd.get_new_id();
+            task.params.sampling          = params.sampling;
+            task.params.stream            = true;
+            task.params.oaicompat_cmpl_id = session_id;
+            task.params.oaicompat_model   = meta->model_name;
+            rd.post_task(std::move(task));
+        }
+
+        // the slot answers once its audio stream is live, the session starts there
+        {
+            auto result = rd.next([&conn]() { return !conn.is_open(); });
+            if (!result || result->is_error()) {
+                if (result) {
+                    send_error_event(json_value(result->to_json(), "message", std::string("server error")), "server_error");
+                }
+                conn.close();
+                return;
+            }
+        }
+        send_event(json {{"type", "session.created"}, {"session", session_config}});
+
+        // one item per committed audio buffer, deltas are attributed to the current item
+        std::mutex  mutex_item;
+        std::string item_id = "item_" + random_string();
+        std::string item_prev; // the committed item before item_id, empty for the first one
+        std::string item_text;
+
+        // called with mutex_item held
+        auto send_completed = [&]() {
+            send_event(json {
+                {"type",          "conversation.item.input_audio_transcription.completed"},
+                {"item_id",       item_id},
+                {"content_index", 0},
+                {"transcript",    item_text},
+            });
+            item_text.clear();
+        };
+
+        auto post_control = [&](server_control_action action, std::vector<float> && audio) {
+            server_response_reader ctl(queue_tasks, queue_results, HTTP_POLLING_SECONDS);
+            server_task task(SERVER_TASK_TYPE_CONTROL);
+            task.id                     = ctl.get_new_id();
+            task.params.control_action  = action;
+            task.params.control_cmpl_id = session_id;
+            task.params.control_audio   = std::move(audio);
+            ctl.post_task(std::move(task));
+            auto result = ctl.next([&conn]() { return !conn.is_open(); });
+            auto * control = result ? dynamic_cast<server_task_result_control *>(result.get()) : nullptr;
+            return control && control->success;
+        };
+
+        // client events are read on this thread, the transcript is sent from the thread below
+        std::atomic<bool> input_ended{false};
+        std::thread reader([&]() {
+            std::string msg;
+            while (conn.read(msg) != 0) {
+                json ev;
+                try {
+                    ev = json::parse(msg);
+                } catch (const std::exception &) {
+                    send_error_event("invalid JSON", "invalid_request_error");
+                    continue;
+                }
+                const std::string type = json_value(ev, "type", std::string());
+                if (type == "session.update" || type == "transcription_session.update") {
+                    send_event(json {{"type", "session.updated"}, {"session", session_config}});
+                } else if (type == "input_audio_buffer.append") {
+                    std::string pcm;
+                    try {
+                        pcm = base64::decode(json_value(ev, "audio", std::string()));
+                    } catch (const std::exception &) {
+                        send_error_event("invalid base64 audio", "invalid_request_error");
+                        continue;
+                    }
+                    std::vector<float> samples(pcm.size() / 2);
+                    for (size_t i = 0; i < samples.size(); i++) {
+                        int16_t v;
+                        memcpy(&v, pcm.data() + 2 * i, sizeof(v));
+                        samples[i] = v / 32768.0f;
+                    }
+                    if (!post_control(SERVER_CONTROL_ACTION_AUDIO_APPEND, std::move(samples))) {
+                        break;
+                    }
+                } else if (type == "input_audio_buffer.commit") {
+                    // "final": true ends the input, the last item is completed when the stream drains
+                    if (json_value(ev, "final", false)) {
+                        input_ended = true;
+                        post_control(SERVER_CONTROL_ACTION_AUDIO_END, {});
+                        continue;
+                    }
+                    if (!post_control(SERVER_CONTROL_ACTION_AUDIO_COMMIT, {})) {
+                        break;
+                    }
+                    // the current item is the committed one, its transcript completes it and the next item opens
+                    std::lock_guard<std::mutex> lock(mutex_item);
+                    send_event(json {
+                        {"type",             "input_audio_buffer.committed"},
+                        {"previous_item_id", item_prev.empty() ? json(nullptr) : json(item_prev)},
+                        {"item_id",          item_id},
+                    });
+                    send_completed();
+                    item_prev = item_id;
+                    item_id   = "item_" + random_string();
+                } else if (type == "input_audio_buffer.clear") {
+                    send_event(json {{"type", "input_audio_buffer.cleared"}});
+                } else {
+                    send_error_event("unknown event type: " + type, "invalid_request_error");
+                }
+            }
+            if (!input_ended.exchange(true)) {
+                post_control(SERVER_CONTROL_ACTION_AUDIO_END, {});
+            }
+            // the result loop waits on the socket, nobody reads it from here on
+            conn.close();
+        });
+
+        while (true) {
+            auto result = rd.next([&conn]() { return !conn.is_open(); });
+            if (!result) {
+                break; // connection closed
+            }
+            if (result->is_error()) {
+                send_error_event(json_value(result->to_json(), "message", std::string("server error")), "server_error");
+                break;
+            }
+            if (auto * partial = dynamic_cast<server_task_result_cmpl_partial *>(result.get())) {
+                if (!partial->content.text.empty()) {
+                    std::lock_guard<std::mutex> lock(mutex_item);
+                    item_text += partial->content.text;
+                    send_event(json {
+                        {"type",          "conversation.item.input_audio_transcription.delta"},
+                        {"item_id",       item_id},
+                        {"content_index", 0},
+                        {"delta",         partial->content.text},
+                    });
+                }
+            } else if (result->is_stop()) {
+                std::lock_guard<std::mutex> lock(mutex_item);
+                send_completed();
+                break;
+            }
+        }
+
+        input_ended = true;
+        conn.close();
+        reader.join();
+    };
+
     this->post_anthropic_messages = [this](const server_http_req & req) {
         auto res = create_response();
         std::vector<raw_buffer> files;
@@ -5877,6 +6363,171 @@ void server_routes::init_routes() {
         return res;
     };
 
+    this->post_tts = [this](const server_http_req & req) {
+        auto res = create_response();
+        res->set_req(&req); // will also set spipe if needed
+
+        if (!ctx_server.has_cap_tts()) {
+            res->error(format_error_response("this server does not support audio generation", ERROR_TYPE_NOT_SUPPORTED));
+            return res;
+        }
+
+        const auto info = mtmd_gen_audio_get_info(ctx_server.mctx);
+
+        json body = json::parse(req.body);
+
+        // multipart form fields arrive as strings, correct the types of the numeric and boolean params
+        for (const char * k : { "top_k", "n_predict", "max_tokens", "seed" }) {
+            if (body.contains(k) && body[k].is_string()) {
+                body[k] = std::stoll(body[k].get<std::string>());
+            }
+        }
+        for (const char * k : { "top_p", "repeat_penalty" }) {
+            if (body.contains(k) && body[k].is_string()) {
+                body[k] = std::stof(body[k].get<std::string>());
+            }
+        }
+        if (body.contains("stream") && body["stream"].is_string()) {
+            body["stream"] = body["stream"] == "true";
+        }
+
+        std::string prompt = json_value(body, "input", json_value(body, "prompt", std::string()));
+        if (prompt.empty()) {
+            res->error(format_error_response("\"input\" must be a non-empty string", ERROR_TYPE_INVALID_REQUEST));
+            return res;
+        }
+
+        const std::string response_format = json_value(body, "response_format", std::string("wav"));
+
+        server_task task(SERVER_TASK_TYPE_TTS);
+        server_schema::eval_tts_schema(params, task, body);
+        const bool stream = task.params.stream;
+        task.tts_inp.set_prompt(prompt);
+        task.tts_inp.set_lang(json_value(body, "lang", std::string()));
+        task.tts_inp.data.stream    = stream;
+        task.tts_inp.data.out_type  = response_format == "pcm"
+            ? MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM
+            : MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
+        // note: -1 is clamped to 0 by the sampler, it will disable the penalty
+        task.params.sampling.penalty_last_n = task.params.n_predict > 0 ? task.params.n_predict : 512;
+        task.params.sampling.seed = task.tts_inp.data.seed; // same seed for the codec/vocoder RNG
+        if (task.tts_inp.data.top_k > 0) {
+            task.params.sampling.top_k = task.tts_inp.data.top_k;
+        }
+        if (task.tts_inp.data.top_p > 0) {
+            task.params.sampling.top_p = task.tts_inp.data.top_p;
+        }
+
+        // speaker reference: either an uploaded form file ("speaker_ref") or a base64 JSON field ("speaker_ref_b64")
+        const unsigned char * speaker_ref_data = nullptr;
+        size_t speaker_ref_len  = 0;
+        std::string speaker_ref_b64_decoded;
+
+        auto speaker_ref_file = req.files.find("speaker_ref");
+        if (speaker_ref_file != req.files.end()) {
+            speaker_ref_data = speaker_ref_file->second.data.data();
+            speaker_ref_len  = speaker_ref_file->second.data.size();
+        } else {
+            std::string speaker_ref_b64 = json_value(body, "speaker_ref_b64", std::string());
+            if (!speaker_ref_b64.empty()) {
+                try {
+                    speaker_ref_b64_decoded = base64::decode(speaker_ref_b64);
+                } catch (const std::exception &) {
+                    res->error(format_error_response("\"speaker_ref_b64\" is not valid base64", ERROR_TYPE_INVALID_REQUEST));
+                    return res;
+                }
+                speaker_ref_data = (const unsigned char *) speaker_ref_b64_decoded.data();
+                speaker_ref_len  = speaker_ref_b64_decoded.size();
+            }
+        }
+
+        if (speaker_ref_len > 0) {
+            auto wrapper = mtmd_helper_bitmap_init_from_buf(ctx_server.mctx, speaker_ref_data, speaker_ref_len, false, ctx_server.init_opt);
+            if (!wrapper.bitmap) {
+                res->error(format_error_response("failed to decode \"speaker_ref\"", ERROR_TYPE_INVALID_REQUEST));
+                return res;
+            }
+            task.tts_inp.set_speaker_ref(mtmd::bitmap_ptr(wrapper.bitmap));
+        } else {
+            SRV_WRN("%s", "no speaker reference provided, the model may behave randomly\n");
+        }
+
+        auto & rd = res->rd;
+        task.id = rd.get_new_id();
+        rd.post_task(std::move(task));
+
+        // raw float32 LE mono samples; audio/L16 would imply 16-bit big-endian (RFC 2586)
+        const std::string content_type = response_format == "pcm"
+            ? "audio/pcm;rate=" + std::to_string(info.sample_rate) + ";encoding=float;bits=32"
+            : "audio/wav";
+
+        if (!stream) {
+            auto result = rd.next(req.should_stop);
+            if (!result) {
+                GGML_ASSERT(req.should_stop());
+                return res; // connection is closed
+            }
+            if (result->is_error()) {
+                res->error(result->to_json());
+                return res;
+            }
+            auto * tts_res = dynamic_cast<server_task_result_tts *>(result.get());
+            GGML_ASSERT(tts_res != nullptr);
+            res->status       = 200;
+            res->content_type = content_type;
+            res->data         = std::move(tts_res->audio);
+            return res;
+        } else {
+            auto first_result = rd.next(req.should_stop);
+            if (!first_result) {
+                GGML_ASSERT(req.should_stop());
+                return res; // connection is closed
+            }
+            if (first_result->is_error()) {
+                res->error(first_result->to_json());
+                return res;
+            }
+            auto * first_tts_res = dynamic_cast<server_task_result_tts *>(first_result.get());
+            GGML_ASSERT(first_tts_res != nullptr);
+
+            res->status       = 200;
+            res->content_type = content_type;
+            res->data          = std::move(first_tts_res->audio);
+            bool is_done       = first_tts_res->final;
+
+            res->set_next([res_this = res.get(), is_done](std::string & output) mutable -> bool {
+                // flush buffered audio before is_done - the first result can already be final
+                if (!res_this->data.empty()) {
+                    output = std::move(res_this->data);
+                    res_this->data.clear();
+                    return true;
+                }
+                if (is_done) {
+                    return false;
+                }
+                if (res_this->should_stop()) {
+                    return false;
+                }
+
+                server_response_reader & rd = res_this->rd;
+                if (!rd.has_next()) {
+                    return false;
+                }
+                auto result = rd.next([&res_this]() { return res_this->should_stop(); });
+                if (!result || result->is_error()) {
+                    return false;
+                }
+                auto * tts_res = dynamic_cast<server_task_result_tts *>(result.get());
+                GGML_ASSERT(tts_res != nullptr);
+                output  = std::move(tts_res->audio);
+                is_done = tts_res->final;
+                return true;
+            });
+        }
+
+        return res;
+    };
+
     this->post_systemone = [this](const server_http_req & req) {
         auto res = create_response();
         const auto & decision = ctx_server.decision;
@@ -5889,7 +6540,7 @@ void server_routes::init_routes() {
         const auto questions = decision.parse_questions(body);
 
         std::vector<raw_buffer> files;
-        const json state = decision.parse_state(body, files);
+        const json state = decision.prepare_state(decision.parse_state(body, files));
         if (!files.empty() && (!decision.can_use_images() || !meta->has_inp_image)) {
             res->error(format_error_response("This server does not support image input for decisions. For a model that supports it, start it with `--mmproj`", ERROR_TYPE_NOT_SUPPORTED));
             return res;

tools/server/server-context.h

@@ -24,6 +24,9 @@ struct server_context_meta {
     bool has_inp_image;
     bool has_inp_audio;
     bool has_inp_video;
+    bool has_cap_chat;
+    bool has_cap_tts;
+    bool has_audio_stream; // realtime transcription, see mtmd_support_audio_stream()
     json json_ui_settings;
     int slot_n_ctx;
     enum llama_pooling_type pooling_type;
@@ -144,6 +147,7 @@ struct server_routes {
     server_http_context::handler_t post_responses_oai;
     server_http_context::handler_t post_responses_tok_oai;
     server_http_context::handler_t post_transcriptions_oai;
+    server_http_context::ws_handler_t ws_realtime;
     server_http_context::handler_t post_anthropic_messages;
     server_http_context::handler_t post_anthropic_count_tokens;
     server_http_context::handler_t post_apply_template;
@@ -153,6 +157,7 @@ struct server_routes {
     server_http_context::handler_t post_embeddings;
     server_http_context::handler_t post_embeddings_oai;
     server_http_context::handler_t post_rerank;
+    server_http_context::handler_t post_tts;
     server_http_context::handler_t post_systemone;
     server_http_context::handler_t get_lora_adapters;
     server_http_context::handler_t post_lora_adapters;