opencl: write KV-cache views in set_tensor so a state restore is not dropped (#30471)

* hex-mmid: vectorize skipped row zeroing and added zeroing to fused MM_ID_NX

* hex-mmid: remove scalar div

.github/workflows/build-apple.yml

@@ -107,62 +107,6 @@ jobs:
           # ref: https://github.com/ggml-org/llama.cpp/pull/19802#issuecomment-4013704023
           ctest -L main -E "test-llama-archs|test-llama-context|test-speculative-rejection-sampling" --verbose --timeout 900
 
-  macos-latest-x64:
-    runs-on: macos-15-intel
-
-    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: ccache-buckets-restore
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
-        with:
-          key: apple-x64
-          folder: llama.cpp
-          hf_bucket: ggml-org/cache
-
-      - name: Build
-        id: cmake_build
-        run: |
-          sysctl -a
-          # Metal is disabled due to intermittent failures with Github runners not having a GPU:
-          # https://github.com/ggml-org/llama.cpp/actions/runs/8635935781/job/23674807267#step:5:2313
-          cmake -B build \
-            -DCMAKE_BUILD_RPATH="@loader_path" \
-            -DLLAMA_FATAL_WARNINGS=ON \
-            -DLLAMA_BUILD_BORINGSSL=ON \
-            -DGGML_METAL=OFF \
-            -DGGML_RPC=ON \
-            -DCMAKE_OSX_DEPLOYMENT_TARGET=13.3
-          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: apple-x64
-          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
-
   macos-latest-ios-xcode:
     runs-on: macos-latest
 

ci/run.sh

@@ -460,10 +460,10 @@ function gg_run_qwen3_0_6b {
     #       https://github.com/ggml-org/llama.cpp/actions/runs/39639985732
     if [ -z ${GG_BUILD_WEBGPU} ]; then
         (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa off --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
-        (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa on  --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
         (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa off                ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
-        (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa on                 ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
     fi
+    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa on  --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa on                 ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
 
     check_ppl "f16"  "$(cat $OUT/${ci}-tg-f16.log  | grep "^\[1\]")"
     if [ -z ${GG_BUILD_NO_BF16} ]; then

common/common.cpp

@@ -742,50 +742,56 @@ static inline bool glob_class_match(const char c, const char * pattern, const ch
 }
 
 // simple glob: * matches non-/ chars, ** matches anything including /, [] matches character class
-static inline bool glob_match(const char * pattern, const char * str) {
-    if (*pattern == '\0') {
-        return *str == '\0';
-    }
-    if (pattern[0] == '*' && pattern[1] == '*') {
-        const char * p = pattern + 2;
-        if (glob_match(p, str)) return true;
-        if (*str != '\0') return glob_match(pattern, str + 1);
+// only wildcards recurse, too many wildcards: no match
+static bool glob_match(const char * pattern, const char * str, int depth = 0) {
+    if (depth >= COMMON_RECURSIVE_LIMIT) {
         return false;
     }
-    if (*pattern == '*') {
-        const char * p = pattern + 1;
-        for (; *str != '\0' && *str != '/'; str++) {
-            if (glob_match(p, str)) return true;
-        }
-        return glob_match(p, str);
-    }
-    if (*pattern == '?' && *str != '\0' && *str != '/') {
-        return glob_match(pattern + 1, str + 1);
-    }
-    if (*pattern == '[') {
-        const char * class_end = pattern + 1;
-        // If first character after '[' is ']' or '-', treat it as literal
-        if (*class_end == ']' || *class_end == '-') {
-            class_end++;
-        }
-        while (*class_end != '\0' && *class_end != ']') {
-            class_end++;
-        }
-        if (*class_end == ']') {
-            if (*str == '\0') return false;
-            bool matched = glob_class_match(*str, pattern + 1, class_end);
-            return matched && glob_match(class_end + 1, str + 1);
-        } else {
-            if (*str == '[') {
-                return glob_match(pattern + 1, str + 1);
+    while (*pattern != '\0') {
+        if (pattern[0] == '*' && pattern[1] == '*') {
+            const char * p = pattern + 2;
+            for (;; str++) {
+                if (glob_match(p, str, depth + 1)) return true;
+                if (*str == '\0') return false;
             }
-            return false;
         }
+        if (*pattern == '*') {
+            const char * p = pattern + 1;
+            for (; *str != '\0' && *str != '/'; str++) {
+                if (glob_match(p, str, depth + 1)) return true;
+            }
+            return glob_match(p, str, depth + 1);
+        }
+        if (*pattern == '?' && *str != '\0' && *str != '/') {
+            pattern++;
+            str++;
+            continue;
+        }
+        if (*pattern == '[') {
+            const char * class_end = pattern + 1;
+            // If first character after '[' is ']' or '-', treat it as literal
+            if (*class_end == ']' || *class_end == '-') {
+                class_end++;
+            }
+            while (*class_end != '\0' && *class_end != ']') {
+                class_end++;
+            }
+            if (*class_end == ']') {
+                if (*str == '\0' || !glob_class_match(*str, pattern + 1, class_end)) return false;
+                pattern = class_end + 1;
+                str++;
+                continue;
+            }
+            if (*str != '[') return false;
+            pattern++;
+            str++;
+            continue;
+        }
+        if (*pattern != *str) return false;
+        pattern++;
+        str++;
     }
-    if (*pattern == *str) {
-        return glob_match(pattern + 1, str + 1);
-    }
-    return false;
+    return *str == '\0';
 }
 
 bool glob_match(const std::string & pattern, const std::string & str) {

common/common.h

@@ -31,6 +31,8 @@
 #define DIRECTORY_SEPARATOR '/'
 #endif // _WIN32
 
+#define COMMON_RECURSIVE_LIMIT 128
+
 #define COM_DBG(fmt, ...) LOG_DBG("cmn  %12.*s: " fmt, 12, __func__, __VA_ARGS__)
 #define COM_TRC(fmt, ...) LOG_TRC("cmn  %12.*s: " fmt, 12, __func__, __VA_ARGS__)
 #define COM_INF(fmt, ...) LOG_INF("cmn  %12.*s: " fmt, 12, __func__, __VA_ARGS__)

common/jinja/parser.cpp

@@ -21,6 +21,17 @@ static bool is_type(const statement_ptr & ptr) {
 class parser {
     const std::vector<token> & tokens;
     size_t current = 0;
+    int depth = 0;
+
+    struct depth_guard {
+        parser & p;
+        depth_guard(parser & p) : p(p) {
+            if (++p.depth > JINJA_RECURSIVE_LIMIT) {
+                throw parser_exception("Parser Error: template is nested more than " + std::to_string(JINJA_RECURSIVE_LIMIT) + " levels deep", p.source, p.peek().pos);
+            }
+        }
+        ~depth_guard() { p.depth--; }
+    };
 
     std::string source; // for error reporting
 
@@ -94,6 +105,7 @@ private:
     }
 
     statement_ptr parse_any() {
+        depth_guard guard(*this);
         size_t start_pos = current;
         switch (peek().t) {
             case token::comment:
@@ -254,6 +266,7 @@ private:
             size_t pos0 = current;
             ++current; // consume {%
             ++current; // consume 'elif'
+            depth_guard guard(*this);
             alternate.push_back(parse_if_statement(pos0)); // nested If
         } else if (is_statement({"else"})) {
             ++current; // consume {%
@@ -330,6 +343,7 @@ private:
     }
 
     statement_ptr parse_if_expression() {
+        depth_guard guard(*this);
         auto a = parse_logical_or_expression();
         if (is_identifier("if")) {
             // Ternary expression
@@ -373,6 +387,7 @@ private:
     statement_ptr parse_logical_negation_expression() {
         // Try parse unary operators
         if (is_identifier("not")) {
+            depth_guard guard(*this);
             size_t start_pos = current;
             auto op = next();
             return mk_stmt<unary_expression>(start_pos, op, parse_logical_negation_expression());
@@ -458,6 +473,7 @@ private:
 
     statement_ptr parse_unary_expression() {
         if (is(token::unary_operator)) {
+            depth_guard guard(*this);
             size_t start_pos = current;
             auto op = next();
             return mk_stmt<unary_expression>(start_pos, op, parse_unary_expression());
@@ -477,9 +493,11 @@ private:
         size_t start_pos = current;
         auto expr = mk_stmt<call_expression>(start_pos, std::move(callee), parse_args());
         auto member = parse_member_expression(std::move(expr)); // foo.x().y
-        return is(token::open_paren)
-            ? parse_call_expression(std::move(member)) // foo.x()()
-            : std::move(member);
+        if (!is(token::open_paren)) {
+            return member;
+        }
+        depth_guard guard(*this);
+        return parse_call_expression(std::move(member)); // foo.x()()
     }
 
     statements parse_args() {

common/jinja/runtime.cpp

@@ -3,6 +3,7 @@
 #include "value.h"
 #include "utils.h"
 
+#include <exception>
 #include <string>
 #include <vector>
 #include <memory>
@@ -52,32 +53,39 @@ static void ensure_key_type_allowed(const value & val) {
 
 // execute with error handling
 value statement::execute(context & ctx) const {
+    std::exception_ptr pending;
+    std::string message;
     try {
         return execute_impl(ctx);
     } catch (const continue_statement::signal & /* ex */) {
-        throw;
+        pending = std::current_exception();
     } catch (const break_statement::signal & /* ex */) {
-        throw;
+        pending = std::current_exception();
     } catch (const rethrown_exception & /* ex */) {
-        throw;
+        pending = std::current_exception();
     } catch (const not_implemented_exception & /* ex */) {
-        throw;
+        pending = std::current_exception();
     } catch (const std::exception & e) {
         const std::string & source = *ctx.src;
         if (source.empty()) {
             std::ostringstream oss;
             oss << "\nError executing " << type() << " at position " << pos << ": " << e.what();
-            throw rethrown_exception(oss.str());
+            message = oss.str();
         } else {
             std::ostringstream oss;
             oss << "\n------------\n";
             oss << "While executing " << type() << " at " << get_line_col(source, pos) << " in source:\n";
             oss << peak_source(source, pos) << "\n";
             oss << "Error: " << e.what();
-            // throw as another exception to avoid repeated formatting
-            throw rethrown_exception(oss.str());
+            message = oss.str();
         }
     }
+    // throw outside the catch: MSVC keeps the unwound frames on the stack while a catch block runs
+    if (pending) {
+        std::rethrow_exception(pending);
+    }
+    // throw as another exception to avoid repeated formatting
+    throw rethrown_exception(message);
 }
 
 value identifier::execute_impl(context & ctx) const {
@@ -747,7 +755,11 @@ value macro_statement::execute_impl(context & ctx) const {
     std::string name = cast_stmt<identifier>(this->name)->val;
 
     const func_handler func = [this, name](const func_args & args) -> value {
+        if (args.ctx.macro_depth >= JINJA_RECURSIVE_LIMIT) {
+            throw std::runtime_error("macro calls are nested more than " + std::to_string(JINJA_RECURSIVE_LIMIT) + " levels deep");
+        }
         context macro_ctx(args.ctx); // new scope for macro execution
+        macro_ctx.macro_depth = args.ctx.macro_depth + 1;
 
         bind_parameters(name, this->args, args, macro_ctx);
 
@@ -778,7 +790,11 @@ value call_statement::execute_impl(context & ctx) const {
     context caller_ctx(ctx); // new scope for caller execution
 
     const func_handler func = [this, caller_ctx = std::move(caller_ctx)](const func_args & args) -> value {
+        if (args.ctx.macro_depth >= JINJA_RECURSIVE_LIMIT) {
+            throw std::runtime_error("macro calls are nested more than " + std::to_string(JINJA_RECURSIVE_LIMIT) + " levels deep");
+        }
         context block_ctx(caller_ctx); // new scope for block execution
+        block_ctx.macro_depth = args.ctx.macro_depth + 1;
 
         bind_parameters("caller", this->caller_args, args, block_ctx);
 

common/jinja/runtime.h

@@ -58,6 +58,8 @@ struct context {
 
     bool is_get_stats = false; // whether to collect stats
 
+    int macro_depth = 0;
+
     visitor_fn visitor;
 
     // src is optional, used for error reporting
@@ -78,6 +80,7 @@ struct context {
         src(parent.src),
         current_time(parent.current_time),
         is_get_stats(parent.is_get_stats),
+        macro_depth(parent.macro_depth),
         // inherit variables (for example, when entering a new scope)
         env(mk_val<value_object>(*parent.env)) {}
 
@@ -710,7 +713,7 @@ struct rethrown_exception : public std::exception {
 
 //////////////////////
 
-static void gather_string_parts_recursive(const value & val, value_string & parts) {
+static void gather_string_parts_recursive(const value & val, value_string & parts, int depth = 0) {
     // TODO: probably allow print value_none as "None" string? currently this breaks some templates
     if (is_val<value_string>(val)) {
         const auto & str_val = cast_val<value_string>(val)->val_str;
@@ -718,10 +721,12 @@ static void gather_string_parts_recursive(const value & val, value_string & part
     } else if (is_val<value_int>(val) || is_val<value_float>(val) || is_val<value_bool>(val)) {
         std::string str_val = val->as_string().str();
         parts->val_str.append(str_val);
+    } else if (is_val<value_array>(val) && depth >= JINJA_RECURSIVE_LIMIT) {
+        throw std::runtime_error("value is nested more than " + std::to_string(JINJA_RECURSIVE_LIMIT) + " levels deep");
     } else if (is_val<value_array>(val)) {
         auto items = cast_val<value_array>(val)->as_array();
         for (const auto & item : items) {
-            gather_string_parts_recursive(item, parts);
+            gather_string_parts_recursive(item, parts, depth + 1);
         }
     }
 }

common/jinja/utils.h

@@ -6,6 +6,8 @@
 #include <cstdint>
 #include <cstring>
 
+#define JINJA_RECURSIVE_LIMIT 128
+
 namespace jinja {
 
 static void string_replace_all(std::string & s, const std::string & search, const std::string & replace) {

common/jinja/value.cpp

@@ -1441,8 +1441,6 @@ const func_builtins & value_undefined_t::get_builtins() const {
 
 
 // the runtime walks values with recursion, a deeper input overflows the stack
-static constexpr int FROM_JSON_MAX_DEPTH = 128;
-
 static value from_json(const common_json & j, bool mark_input, int depth = 0) {
     if (j.is_null()) {
         return mk_val<value_none>();
@@ -1458,8 +1456,8 @@ static value from_json(const common_json & j, bool mark_input, int depth = 0) {
             str->mark_input();
         }
         return str;
-    } else if ((j.is_array() || j.is_object()) && depth >= FROM_JSON_MAX_DEPTH) {
-        throw std::invalid_argument("JSON is nested more than " + std::to_string(FROM_JSON_MAX_DEPTH) + " levels deep");
+    } else if ((j.is_array() || j.is_object()) && depth >= JINJA_RECURSIVE_LIMIT) {
+        throw std::invalid_argument("JSON is nested more than " + std::to_string(JINJA_RECURSIVE_LIMIT) + " levels deep");
     } else if (j.is_array()) {
         auto arr = mk_val<value_array>();
         for (const auto & item : j) {
@@ -1664,6 +1662,8 @@ static void value_to_json_internal(std::ostringstream & oss, const value & val,
             }
         }
         oss << "\"";
+    } else if ((is_val<value_array>(val) || is_val<value_object>(val)) && curr_lvl >= JINJA_RECURSIVE_LIMIT) {
+        throw std::runtime_error("value is nested more than " + std::to_string(JINJA_RECURSIVE_LIMIT) + " levels deep");
     } else if (is_val<value_array>(val)) {
         const auto & arr = val->as_array();
         oss << "[";
@@ -1729,17 +1729,20 @@ std::string value_to_string_repr(const value & val) {
 }
 
 // stats utility
-void value_t::stats_t::mark_used(value & val, bool deep) {
+void value_t::stats_t::mark_used(value & val, bool deep, int depth) {
     val->stats.used = true;
+    if (deep && (is_val<value_array>(val) || is_val<value_object>(val)) && depth >= JINJA_RECURSIVE_LIMIT) {
+        throw std::runtime_error("value is nested more than " + std::to_string(JINJA_RECURSIVE_LIMIT) + " levels deep");
+    }
     if (deep) {
         if (is_val<value_array>(val)) {
             for (auto & item : val->val_arr) {
-                mark_used(item, deep);
+                mark_used(item, deep, depth + 1);
             }
         } else if (is_val<value_object>(val)) {
             for (auto & pair : val->val_obj) {
-                mark_used(pair.first, deep);
-                mark_used(pair.second, deep);
+                mark_used(pair.first, deep, depth + 1);
+                mark_used(pair.second, deep, depth + 1);
             }
         }
     }

common/jinja/value.h

@@ -119,7 +119,7 @@ struct value_t {
         // ops can be builtin calls or operators: "array_access", "object_access"
         std::set<std::string> ops;
         // utility to recursively mark value and its children as used
-        static void mark_used(value & val, bool deep = false);
+        static void mark_used(value & val, bool deep = false, int depth = 0);
     } stats;
 
     value_t() = default;

common/json-schema-to-grammar.cpp

@@ -272,15 +272,12 @@ static bool is_reserved_name(const std::string & name) {
     return RESERVED_NAMES.find(name) != RESERVED_NAMES.end();
 }
 
-static std::regex INVALID_RULE_CHARS_RE("[^a-zA-Z0-9-]+");
 static std::regex GRAMMAR_LITERAL_ESCAPE_RE("[\r\n\"\\\\]");
 static std::regex GRAMMAR_RANGE_LITERAL_ESCAPE_RE("[\r\n\"\\]\\-\\\\]");
 static std::unordered_map<char, std::string> GRAMMAR_LITERAL_ESCAPES = {
     {'\r', "\\r"}, {'\n', "\\n"}, {'"', "\\\""}, {'-', "\\-"}, {']', "\\]"}, {'\\', "\\\\"}
 };
 
-static const int MAX_PATTERN_DEPTH = 100;
-
 static std::unordered_set<char> NON_LITERAL_SET = {'|', '.', '(', ')', '[', ']', '{', '}', '*', '+', '?', '^', '$'};
 static std::unordered_set<char> ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS = {'^', '$', '.', '[', ']', '(', ')', '|', '{', '}', '*', '+', '?'};
 
@@ -312,6 +309,22 @@ static std::string format_literal(const std::string & literal) {
 
 std::string gbnf_format_literal(const std::string & literal) { return format_literal(literal); }
 
+// same as std::regex_replace with [^a-zA-Z0-9-]+ -> "-", no std::regex: libstdc++ recurses per char
+std::string gbnf_format_rule_name(const std::string & name) {
+    std::string out;
+    bool in_invalid = false;
+    for (char c : name) {
+        if ((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '-') {
+            out += c;
+            in_invalid = false;
+        } else if (!in_invalid) {
+            out += '-';
+            in_invalid = true;
+        }
+    }
+    return out;
+}
+
 static size_t gbnf_escape_length(const std::string & pattern, size_t pos) {
     if (pos + 1 >= pattern.length() || pattern[pos] != '\\') {
         return 0;
@@ -347,6 +360,17 @@ private:
     std::unordered_set<std::string> _refs_being_resolved;
     std::vector<std::string> _errors;
     std::vector<std::string> _warnings;
+    int _depth = 0;
+
+    struct depth_guard {
+        int & depth;
+        depth_guard(int & depth) : depth(depth) {
+            if (++depth > COMMON_RECURSIVE_LIMIT) {
+                throw std::invalid_argument("JSON schema is nested more than " + std::to_string(COMMON_RECURSIVE_LIMIT) + " levels deep");
+            }
+        }
+        ~depth_guard() { depth--; }
+    };
 
     template <typename T>
     static const T & as(const common_chat_schema & node) {
@@ -354,7 +378,7 @@ private:
     }
 
     std::string _add_rule(const std::string & name, const std::string & rule) {
-        std::string esc_name = regex_replace(name, INVALID_RULE_CHARS_RE, "-");
+        std::string esc_name = gbnf_format_rule_name(name);
         if (_rules.find(esc_name) == _rules.end() || _rules[esc_name] == rule) {
             _rules[esc_name] = rule;
             return esc_name;
@@ -482,7 +506,7 @@ private:
                         }
                     }
                     paren_depth++;
-                    if (paren_depth > MAX_PATTERN_DEPTH) {
+                    if (paren_depth > COMMON_RECURSIVE_LIMIT) {
                         throw unsupported_pattern("pattern nesting too deep");
                     }
                     seq.emplace_back("(" + to_rule(transform()) + ")", false);
@@ -644,7 +668,11 @@ private:
         std::string char_rule = _add_primitive("char", PRIMITIVE_RULES.at("char"));
         std::ostringstream out;
         out << "[\"] ( ";
-        std::function<void(size_t)> visit = [&](size_t idx) {
+        // one level per char of the longest name
+        std::function<void(size_t, int)> visit = [&](size_t idx, int depth) {
+            if (depth >= COMMON_RECURSIVE_LIMIT) {
+                throw std::invalid_argument("JSON schema property name is longer than " + std::to_string(COMMON_RECURSIVE_LIMIT) + " chars");
+            }
             const auto & node = trie.nodes[idx];
             std::string rejects;
             auto first = true;
@@ -659,7 +687,7 @@ private:
                 out << "[" << c << "]";
                 if (!trie.nodes[child].children.empty()) {
                     out << " (";
-                    visit(child);
+                    visit(child, depth + 1);
                     out << ")";
                 } else {
                     out << " " << char_rule << "+";
@@ -669,7 +697,7 @@ private:
                 out << " | [^\"" << rejects << "] " << char_rule << "*";
             }
         };
-        visit(0);
+        visit(0, 0);
 
         out << " )";
         if (trie.nodes[0].pattern < 0) {
@@ -682,8 +710,7 @@ private:
     std::string _resolve_ref(const common_chat_schema_ref & schema) {
         auto it = schema.ref.find('#');
         std::string ref_fragment = it != std::string::npos ? schema.ref.substr(it + 1) : schema.ref;
-        static const std::regex nonalphanumeric_regex(R"([^a-zA-Z0-9-]+)");
-        std::string ref_name = "ref" + std::regex_replace(ref_fragment, nonalphanumeric_regex, "-");
+        std::string ref_name = "ref" + gbnf_format_rule_name(ref_fragment);
         if (_rules.find(ref_name) == _rules.end() && _refs_being_resolved.find(schema.ref) == _refs_being_resolved.end()) {
             if (!schema.target) {
                 _errors.push_back("Unresolved $ref " + schema.ref);
@@ -833,10 +860,13 @@ public:
         std::unordered_set<std::string> required;
         std::vector<std::pair<std::string, const common_chat_schema *>> properties;
         std::map<std::string, size_t> enum_values;
-        std::function<void(const common_chat_schema &, bool)> add_component = [&](const common_chat_schema & comp, bool is_required) {
+        std::function<void(const common_chat_schema &, bool, int)> add_component = [&](const common_chat_schema & comp, bool is_required, int depth) {
+            if (depth >= COMMON_RECURSIVE_LIMIT) {
+                throw std::invalid_argument("JSON schema is nested more than " + std::to_string(COMMON_RECURSIVE_LIMIT) + " levels deep");
+            }
             if (comp.kind() == common_chat_schema::KIND_REF) {
                 if (const auto * target = as<common_chat_schema_ref>(comp).target) {
-                    add_component(*target, is_required);
+                    add_component(*target, is_required, depth + 1);
                 }
             } else if (comp.kind() == common_chat_schema::KIND_OBJECT) {
                 for (const auto & prop : as<common_chat_schema_object>(comp).properties) {
@@ -854,10 +884,10 @@ public:
         for (const auto & child : schema.children) {
             if (child->kind() == common_chat_schema::KIND_ANY_OF) {
                 for (const auto & alt : as<common_chat_schema_any_of>(*child).children) {
-                    add_component(*alt, false);
+                    add_component(*alt, false, 0);
                 }
             } else {
-                add_component(*child, true);
+                add_component(*child, true, 0);
             }
         }
         if (!enum_values.empty()) {
@@ -875,6 +905,8 @@ public:
     }
 
     std::string visit(const common_chat_schema & schema, const std::string & name) {
+        // the $ref targets are visited too, a $ref chain is not bounded by the builder
+        depth_guard guard(_depth);
         std::string rule_name = is_reserved_name(name) ? name + "-" : name.empty() ? "root" : name;
         std::string sub_name  = name + (name.empty() ? "" : "-");
 

common/json-schema-to-grammar.h

@@ -19,5 +19,6 @@ struct common_grammar_options {
 };
 
 std::string gbnf_format_literal(const std::string & literal);
+std::string gbnf_format_rule_name(const std::string & name);
 
 std::string build_grammar(const std::function<void(const common_grammar_builder &)> & cb, const common_grammar_options & options = {});

common/json-schema.cpp

@@ -9,6 +9,18 @@
 #include <utility>
 #include <vector>
 
+namespace {
+struct depth_guard {
+    int & depth;
+    depth_guard(int & depth) : depth(depth) {
+        if (++depth > COMMON_RECURSIVE_LIMIT) {
+            throw std::invalid_argument("JSON schema is nested more than " + std::to_string(COMMON_RECURSIVE_LIMIT) + " levels deep");
+        }
+    }
+    ~depth_guard() { depth--; }
+};
+} // namespace
+
 class common_chat_schema_builder {
     const common_json &      root_;
     common_chat_schema_document & doc_;
@@ -224,18 +236,11 @@ class common_chat_schema_builder {
         return node;
     }
 
-    // the builder and the code that walks the built schema recurse once per level, a deeper schema overflows the stack
-    static constexpr int max_depth = 128;
-    int                  depth_    = 0;
+    int depth_ = 0;
 
     common_chat_schema_ptr build_node(const common_json & schema, const std::string & path) {
-        if (depth_ >= max_depth) {
-            throw std::invalid_argument("JSON schema is nested more than " + std::to_string(max_depth) + " levels deep");
-        }
-        depth_++;
-        auto node = build_node_impl(schema, path);
-        depth_--;
-        return node;
+        depth_guard guard(depth_);
+        return build_node_impl(schema, path);
     }
 
     common_chat_schema_ptr build_node_impl(const common_json & schema, const std::string & path) {
@@ -379,7 +384,9 @@ static common_chat_schema::value_type json_type(const common_json & value) {
     return common_chat_schema::TYPE_OBJECT;
 }
 
-static common_chat_schema::type_set value_types_impl(const common_chat_schema & s, std::unordered_set<const common_chat_schema *> & visited) {
+// depth counts the $ref, anyOf and allOf hops, a $ref chain is not bounded by the builder
+static common_chat_schema::type_set value_types_impl(const common_chat_schema & s, std::unordered_set<const common_chat_schema *> & visited, int & depth) {
+    depth_guard guard(depth);
     switch (s.kind()) {
         case common_chat_schema::KIND_ANY:
             return common_chat_schema::type_set::all();
@@ -413,21 +420,21 @@ static common_chat_schema::type_set value_types_impl(const common_chat_schema &
                 // a cycle contributes no type, to be safe
                 return {};
             }
-            auto types = value_types_impl(*target, visited);
+            auto types = value_types_impl(*target, visited, depth);
             visited.erase(target);
             return types;
         }
         case common_chat_schema::KIND_ANY_OF: {
             common_chat_schema::type_set types;
             for (const auto & child : static_cast<const common_chat_schema_any_of &>(s).children) {
-                types |= value_types_impl(*child, visited);
+                types |= value_types_impl(*child, visited, depth);
             }
             return types;
         }
         case common_chat_schema::KIND_ALL_OF: {
             auto types = common_chat_schema::type_set::all();
             for (const auto & child : static_cast<const common_chat_schema_all_of &>(s).children) {
-                types &= value_types_impl(*child, visited);
+                types &= value_types_impl(*child, visited, depth);
             }
             return types;
         }
@@ -437,10 +444,12 @@ static common_chat_schema::type_set value_types_impl(const common_chat_schema &
 
 common_chat_schema::type_set common_chat_schema::value_types() const {
     std::unordered_set<const common_chat_schema *> visited;
-    return value_types_impl(*this, visited);
+    int depth = 0;
+    return value_types_impl(*this, visited, depth);
 }
 
-static bool may_be_string_impl(const common_chat_schema & s, std::unordered_set<const common_chat_schema *> & visited) {
+static bool may_be_string_impl(const common_chat_schema & s, std::unordered_set<const common_chat_schema *> & visited, int & depth) {
+    depth_guard guard(depth);
     switch (s.kind()) {
         case common_chat_schema::KIND_STRING:
             return true;
@@ -459,13 +468,13 @@ static bool may_be_string_impl(const common_chat_schema & s, std::unordered_set<
             if (!target || !visited.insert(target).second) {
                 return false;
             }
-            bool result = may_be_string_impl(*target, visited);
+            bool result = may_be_string_impl(*target, visited, depth);
             visited.erase(target);
             return result;
         }
         case common_chat_schema::KIND_ANY_OF:
             for (const auto & child : static_cast<const common_chat_schema_any_of &>(s).children) {
-                if (may_be_string_impl(*child, visited)) {
+                if (may_be_string_impl(*child, visited, depth)) {
                     return true;
                 }
             }
@@ -477,7 +486,7 @@ static bool may_be_string_impl(const common_chat_schema & s, std::unordered_set<
                 if (child->kind() == common_chat_schema::KIND_ANY) {
                     continue;
                 }
-                if (!may_be_string_impl(*child, visited)) {
+                if (!may_be_string_impl(*child, visited, depth)) {
                     return false;
                 }
                 any_string = true;
@@ -491,7 +500,8 @@ static bool may_be_string_impl(const common_chat_schema & s, std::unordered_set<
 
 bool common_chat_schema::may_be_string() const {
     std::unordered_set<const common_chat_schema *> visited;
-    return may_be_string_impl(*this, visited);
+    int depth = 0;
+    return may_be_string_impl(*this, visited, depth);
 }
 
 const char * common_chat_schema::kind_name(node_kind kind) {

common/peg-parser.cpp

@@ -10,10 +10,12 @@
 #include <initializer_list>
 #include <map>
 #include <memory>
-#include <regex>
 #include <set>
 #include <stdexcept>
 
+// The parser can run with a lower limit than COMMON_RECURSIVE_LIMIT
+static constexpr int RULE_RECURSIVE_LIMIT = 64;
+
 // Trick to catch missing branches
 template <typename T>
 inline constexpr bool is_always_false_v = false;
@@ -186,13 +188,17 @@ std::string common_peg_ast_node::sanitized_text() const {
 }
 
 void common_peg_ast_arena::visit(common_peg_ast_id id, const common_peg_ast_visitor & visitor) const {
-    if (id == COMMON_PEG_INVALID_AST_ID) {
-        return;
-    }
-    const auto & node = get(id);
-    visitor(node);
-    for (const auto & child : node.children) {
-        visit(child, visitor);
+    // no recursion: the AST is several nodes deep per level of the input
+    std::vector<common_peg_ast_id> stack = { id };
+    while (!stack.empty()) {
+        const auto cur = stack.back();
+        stack.pop_back();
+        if (cur == COMMON_PEG_INVALID_AST_ID) {
+            continue;
+        }
+        const auto & node = get(cur);
+        visitor(node);
+        stack.insert(stack.end(), node.children.rbegin(), node.children.rend());
     }
 }
 
@@ -736,8 +742,15 @@ struct parser_executor {
     }
 
     common_peg_parse_result operator()(const common_peg_rule_parser & p) {
+        // only rules recurse, a recursive rule is entered once per input level
+        if (ctx.rule_depth[p.child] >= RULE_RECURSIVE_LIMIT) {
+            throw common_peg_recursion_error("input is nested more than " + std::to_string(RULE_RECURSIVE_LIMIT) + " levels deep");
+        }
+
         // Parse the child
+        ctx.rule_depth[p.child]++;
         auto result = arena.parse(p.child, ctx, start_pos);
+        ctx.rule_depth[p.child]--;
 
         if (!result.fail()) {
             std::string_view text;
@@ -819,6 +832,10 @@ common_peg_parse_result common_peg_arena::parse(common_peg_parse_context & ctx,
     if (root_ == COMMON_PEG_INVALID_PARSER_ID) {
         throw std::runtime_error("No root parser set");
     }
+
+    // Initialize rule depth for all parsers
+    ctx.rule_depth.assign(parsers_.size(), 0);
+
     return parse(root_, ctx, start);
 }
 
@@ -1069,11 +1086,6 @@ common_peg_parser operator|(const std::string & str, const common_peg_parser & p
     return operator|(str.c_str(), p);
 }
 
-static std::string rule_name(const std::string & name) {
-    static const std::regex invalid_rule_chars_re("[^a-zA-Z0-9-]+");
-    return std::regex_replace(name, invalid_rule_chars_re, "-");
-}
-
 common_peg_parser_builder::common_peg_parser_builder() {}
 
 common_peg_parser common_peg_parser_builder::sequence(const std::vector<common_peg_parser_id> & parsers) {
@@ -1155,14 +1167,14 @@ common_peg_parser common_peg_parser_builder::schema(const common_peg_parser & p,
 }
 
 common_peg_parser common_peg_parser_builder::rule(const std::string & name, const common_peg_parser & p, bool trigger) {
-    auto clean_name = rule_name(name);
+    auto clean_name = gbnf_format_rule_name(name);
     auto rule_id = arena_.add_parser(common_peg_rule_parser{clean_name, p.id(), trigger});
     arena_.add_rule(clean_name, rule_id);
     return ref(clean_name);
 }
 
 common_peg_parser common_peg_parser_builder::rule(const std::string & name, const std::function<common_peg_parser()> & builder_fn, bool trigger) {
-    auto clean_name = rule_name(name);
+    auto clean_name = gbnf_format_rule_name(name);
     if (arena_.has_rule(clean_name)) {
         return ref(clean_name);
     }

common/peg-parser.h

@@ -6,6 +6,7 @@
 
 #include <memory>
 #include <set>
+#include <stdexcept>
 #include <unordered_map>
 #include <string>
 #include <string_view>
@@ -182,6 +183,10 @@ inline common_peg_parse_flags operator~(common_peg_parse_flags a) {
     return static_cast<common_peg_parse_flags>(~int(a));
 }
 
+struct common_peg_recursion_error : std::invalid_argument {
+    using std::invalid_argument::invalid_argument;
+};
+
 struct common_peg_parse_context {
     std::string input;               // [h,   e,  l,  l,  o,  _,  w,  o,  r,  l,  d]
     std::vector<llama_token> tokens; // [id, -1, -1, -1, -1, id, -1, -1, -1, -1, -1]
@@ -189,6 +194,7 @@ struct common_peg_parse_context {
     common_peg_ast_arena ast;
 
     int parse_depth;
+    std::vector<int> rule_depth; // index: rule child id
 
     common_peg_parse_context(common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE)
         : flags(flags), parse_depth(0) {}

conversion/deepseek.py

@@ -803,9 +803,6 @@ class DeepseekV4Model(TextModel):
         self._dsv4_mxfp4_generated = True
         return ()
 
-    def _format_dsv4_tensor_name(self, key: gguf.MODEL_TENSOR, bid: int | None, suffix: str = ".weight") -> str:
-        return self.format_tensor_name(key, bid, suffix)
-
     def _map_dsv4_tensor_name(self, name: str, bid: int | None) -> tuple[gguf.MODEL_TENSOR, str]:
         root_map: dict[str, tuple[gguf.MODEL_TENSOR, str]] = {
             "embed.weight": (gguf.MODEL_TENSOR.TOKEN_EMBD, ".weight"),
@@ -879,17 +876,17 @@ class DeepseekV4Model(TextModel):
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         if re.match(r"layers\.\d+\.ffn\.experts\.\d+\.w[123]\.(weight|scale)$", name):
-            return []
+            return
 
         # hash layers route text tokens via tid2eid and image tokens via bias_vl; gate.bias is unused
         if name.endswith(".ffn.gate.bias") and bid is not None and bid < self.hparams["num_hash_layers"]:
-            return []
+            return
 
         tensor_key, suffix = self._map_dsv4_tensor_name(name, bid)
         if tensor_key == gguf.MODEL_TENSOR.FFN_GATE_TID2EID:
-            return []
+            return
 
-        return [(self._format_dsv4_tensor_name(tensor_key, bid, suffix), data_torch)]
+        yield self.format_tensor_name(tensor_key, bid, suffix), data_torch
 
     def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
         del bid  # unused

conversion/deepseek41.py

@@ -145,12 +145,6 @@ class DeepseekV41Model(DeepseekV4Model):
         # num_hidden_layers lives in text_config, which is merged only after indexing; MTP stages are skipped anyway
         return TextModel.index_tensors(self, remote_hf_model_id=remote_hf_model_id)
 
-    @classmethod
-    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
-        if item[0].startswith("vision."):
-            return None
-        return super().filter_tensors(item)
-
     def set_gguf_parameters(self):
         super().set_gguf_parameters()
         hp = self.hparams
@@ -282,12 +276,13 @@ class DeepseekV41Model(DeepseekV4Model):
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         if name.endswith(".engram.k_weight"):
-            return []  # folded into engram_gate
+            return  # folded into engram_gate
         if name.endswith(".engram.q_weight"):
             # Engram gate weight = q_weight * k_weight (bf16 * bf16 is exact in f32)
             k = self.model_tensors[name.removesuffix("q_weight") + "k_weight"]
-            return [(self.format_tensor_name(gguf.MODEL_TENSOR.ENGRAM_GATE, bid), data_torch.float() * k().float())]
-        return super().modify_tensors(data_torch, name, bid)
+            yield self.format_tensor_name(gguf.MODEL_TENSOR.ENGRAM_GATE, bid), data_torch.float() * k().float()
+            return
+        yield from super().modify_tensors(data_torch, name, bid)
 
     def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
         if new_name.endswith((".engram_gate.weight", ".ffn_gate_inp.weight")) or name in self._dsv4_f32_tensors:
@@ -411,13 +406,15 @@ class DeepseekV41DSparkModel(DeepseekV41Model):
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         if name.endswith(".ffn.gate.bias_vl"):
-            return []
+            return
         if name in self._DSPARK_ROOT:
-            return [(self.format_tensor_name(self._DSPARK_ROOT[name], None), data_torch)]
+            yield self.format_tensor_name(self._DSPARK_ROOT[name], None), data_torch
+            return
         m = re.match(r"layers\.\d+\.(main_proj\.weight|main_norm\.weight|norm\.weight)$", name)
         if m is not None:
-            return [(self.format_tensor_name(self._DSPARK_ROOT[m.group(1)], None), data_torch)]
-        return super().modify_tensors(data_torch, name, bid)
+            yield self.format_tensor_name(self._DSPARK_ROOT[m.group(1)], None), data_torch
+            return
+        yield from super().modify_tensors(data_torch, name, bid)
 
     def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
         if new_name == "conf_proj.weight":

conversion/qwen4exp.py

@@ -41,6 +41,8 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
         self._ple_shards: dict[int, str] = {}
         self._ple_row_dim: int | None = None
         self._mtp_fc: dict[str, Tensor] = {}
+        # halves of the merged weights, keyed by (bid, merged tensor), see _merge_shared_input
+        self._merge_parts: dict[tuple[int, gguf.MODEL_TENSOR], dict[int, Tensor]] = {}
 
     @classmethod
     def filter_tensors(cls, item):
@@ -131,44 +133,79 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
         # int64 hash constants must stay exact; 1-D tensors force F32, so use KV
         if name.endswith("ple_embedding.layer_multipliers"):
             self._ple_multipliers = [int(x) for x in data_torch.tolist()]
-            return []
+            return
         if name.endswith("ple_embedding.ngram_heads_offsets"):
             self._ple_head_offsets = [int(x) for x in data_torch.tolist()]
-            return []
+            return
         if name.endswith("ple_embedding.ngram_heads_vocab_sizes"):
             self._ple_head_vocab_sizes = [int(x) for x in data_torch.tolist()]
-            return []
+            return
 
         if ".ngram_embedding.shard_" in name:
-            return self._place_ple_shard(data_torch, name)
+            yield from self._place_ple_shard(data_torch, name)
+            return
 
         # eh_proj([e ; h_s]) = fc_embedding(e) + fc_hidden(h_s) for every hc stream s
         if name.endswith((".nextn_fc_embedding.weight", ".nextn_fc_hidden.weight")):
             self._mtp_fc[name.rsplit(".", 2)[1]] = data_torch
             if len(self._mtp_fc) < 2:
-                return []
+                return
             eh = torch.cat([self._mtp_fc.pop("nextn_fc_embedding"), self._mtp_fc.pop("nextn_fc_hidden")], dim=1)
-            return [(self.format_tensor_name(gguf.MODEL_TENSOR.NEXTN_EH_PROJ, bid, ".weight"), eh)]
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.NEXTN_EH_PROJ, bid, ".weight"), eh)
+            return
 
         # one projection feeds indexer q and k; split it, as minimax-m3 does
         if ".indexer.index_qk_proj.weight" in name:
             n_q = self.hparams["indexer_n_heads"] * self.hparams["indexer_head_dim"]
-            q = data_torch[:n_q]
-            k = data_torch[n_q:]
-            return [
-                (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_PROJ, bid, ".weight"), q),
-                (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_K_PROJ, bid, ".weight"), k),
-            ]
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_PROJ, bid, ".weight"), data_torch[:n_q])
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_K_PROJ, bid, ".weight"), data_torch[n_q:])
+            return
 
         # Gemma zero-centred gammas the inherited norm.weight rule misses
         if name.endswith((".ple.norm_key.weight", ".ple.norm_query.weight", ".ple.norm_conv.weight",
                           ".indexer.q_layernorm.weight", ".indexer.k_layernorm.weight")):
-            return [(self.map_tensor_name(name), data_torch + 1)]
+            yield (self.map_tensor_name(name), data_torch + 1)
+            return
 
         if name.endswith(".ple.conv1d.weight"):
-            return [(self.map_tensor_name(name), data_torch.squeeze())]
+            yield (self.map_tensor_name(name), data_torch.squeeze())
+            return
 
-        return super().modify_tensors(data_torch, name, bid)
+        for new_name, data in super().modify_tensors(data_torch, name, bid):
+            yield from self._merge_shared_input(new_name, data, bid)
+
+    # matmuls that read the same input, written as one weight [a | b] so the runtime does one matmul:
+    # (merged tensor, (a, b), linear-attention layers only)
+    _SHARED_INPUT_MERGES = (
+        (gguf.MODEL_TENSOR.HC_ATTN_DOWN_INJECT, (gguf.MODEL_TENSOR.HC_ATTN_DOWN, gguf.MODEL_TENSOR.HC_ATTN_INJECT), False),
+        (gguf.MODEL_TENSOR.HC_FFN_DOWN_INJECT,  (gguf.MODEL_TENSOR.HC_FFN_DOWN,  gguf.MODEL_TENSOR.HC_FFN_INJECT),  False),
+        (gguf.MODEL_TENSOR.ATTN_QKV_GATE,       (gguf.MODEL_TENSOR.ATTN_QKV,     gguf.MODEL_TENSOR.ATTN_GATE),      True),
+        (gguf.MODEL_TENSOR.SSM_BETA_ALPHA_CAT,  (gguf.MODEL_TENSOR.SSM_BETA,     gguf.MODEL_TENSOR.SSM_ALPHA),      True),
+    )
+
+    # caches one half of a merged weight until the other half arrives, then yields [a | b]
+    def _merge_shared_input(self, new_name: str, data: Tensor, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # a pre-quantized checkpoint carries per-tensor scales, which one merged weight cannot hold
+        if bid is None or self.hparams.get("quantization_config"):
+            yield (new_name, data)
+            return
+
+        layer_types = self.hparams.get("layer_types") or []
+        is_linear = bid < len(layer_types) and layer_types[bid] == "linear_attention"
+
+        match = next(((merged, parts.index(part)) for merged, parts, linear_only in self._SHARED_INPUT_MERGES
+                      for part in parts
+                      if (is_linear or not linear_only) and self.match_model_tensor_name(new_name, part, bid)), None)
+        if match is None:
+            yield (new_name, data)
+            return
+
+        merged, idx = match
+        halves = self._merge_parts.setdefault((bid, merged), {})
+        halves[idx] = data
+        if len(halves) == 2:
+            del self._merge_parts[(bid, merged)]
+            yield (self.format_tensor_name(merged, bid), torch.cat([halves[0], halves[1]], dim=0))
 
     # the shards concatenate into a tensor of well over 100 GB
     # use LazyChunkedTensor here, a single shard resident at a time
@@ -214,6 +251,8 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
         super().prepare_tensors()
         if self._mtp_fc:
             raise ValueError(f"MTP projection missing its other half: {sorted(self._mtp_fc)}")
+        if self._merge_parts:
+            raise ValueError(f"merged weights missing a half: {sorted((bid, t.name) for bid, t in self._merge_parts)}")
         n_parts = self.hparams.get("split_ngram_parts", 0)
         if self._ple_shards and len(self._ple_shards) != n_parts:
             raise ValueError(

ggml/src/ggml-cpu/arch/arm/quants.c

@@ -287,6 +287,9 @@ void ggml_vec_dot_q2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi
 
     sumf = vaddvq_f32(sumv);
 #else
+    UNUSED(nb);
+    UNUSED(x);
+    UNUSED(y);
     ggml_vec_dot_q2_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc);
     return;
 #endif

ggml/src/ggml-cpu/ops.cpp

@@ -1914,213 +1914,87 @@ void ggml_compute_forward_repeat_back(
 }
 
 // ggml_compute_forward_concat
-static void ggml_compute_forward_concat_any(
-    const ggml_compute_params * params,
-    ggml_tensor * dst) {
 
-    const ggml_tensor * src0 = dst->src[0];
-    const ggml_tensor * src1 = dst->src[1];
+// rows of dst copied together, so that a strided source is read along neighboring rows
+static constexpr int64_t GGML_CONCAT_TILE = 32;
 
-    const int ith = params->ith;
-    const int nth = params->nth;
-
-    GGML_TENSOR_BINARY_OP_LOCALS
-
-    const int32_t dim = ggml_get_op_params_i32(dst, 0);
-
-    GGML_ASSERT(dim >= 0 && dim < 4);
-    GGML_ASSERT(ggml_is_contiguous_rows(src0));
-    GGML_ASSERT(ggml_is_contiguous_rows(src1));
-
-    int64_t o[4] = {0, 0, 0, 0};
-
-    if (dim == 0) {
-        GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0);
-        GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0);
-
-        o[dim] = src0->ne[dim]/ggml_blck_size(src0->type);
-    } else {
-        o[dim] = src0->ne[dim];
-    }
-
-    // Region 1: copy rows from src0
-    for (int i3 = 0; i3 < ne03; i3++) {
-        for (int i2 = ith; i2 < ne02; i2 += nth) {
-            for (int i1 = 0; i1 < ne01; i1++) {
-                const char * x = (const char *) src0->data + i1*nb01 + i2*nb02 + i3*nb03;
-                      char * y = (      char *) dst->data  + i1*nb1  + i2*nb2  + i3*nb3;
-                memcpy(y, x, ggml_row_size(src0->type, ne00));
-            }
-        }
-    }
-
-    // Region 2: copy rows from src1, offset into dst by o[]
-    for (int i3 = 0; i3 < ne13; i3++) {
-        for (int i2 = ith; i2 < ne12; i2 += nth) {
-            for (int i1 = 0; i1 < ne11; i1++) {
-                const char * x = (const char *) src1->data + i1*nb11         + i2*nb12         + i3*nb13;
-                      char * y = (      char *)  dst->data + (i1 + o[1])*nb1 + (i2 + o[2])*nb2 + (i3 + o[3])*nb3 + o[0]*nb0;
-                memcpy(y, x, ggml_row_size(src1->type, ne10));
-            }
+// copy n_rows rows of n elements of es bytes from a source with element stride nb0, element by element
+template <size_t es>
+static void ggml_concat_copy_strided(char * y, size_t nby, const char * x, size_t nbx, size_t nb0, int64_t n, int64_t n_rows) {
+    for (int64_t i0 = 0; i0 < n; i0++) {
+        for (int64_t r = 0; r < n_rows; r++) {
+            memcpy(y + r*nby + i0*es, x + r*nbx + i0*nb0, es);
         }
     }
 }
 
-static void ggml_compute_forward_concat_i8(
-    const ggml_compute_params * params,
-    ggml_tensor * dst) {
+// copy n_rows rows of n elements of src into dst, rows are nbx bytes apart in src and nby bytes apart in dst
+static void ggml_concat_copy_rows(char * y, size_t nby, const ggml_tensor * src, const char * x, size_t nbx, int64_t n, int64_t n_rows) {
+    const size_t es = ggml_type_size(src->type);
 
-    const ggml_tensor * src0 = dst->src[0];
-    const ggml_tensor * src1 = dst->src[1];
-
-    GGML_ASSERT(ggml_type_size(src0->type) == sizeof(int8_t));
-
-    const int ith = params->ith;
-    const int nth = params->nth;
-
-    GGML_TENSOR_BINARY_OP_LOCALS
-
-    const int32_t dim = ggml_get_op_params_i32(dst, 0);
-
-    GGML_ASSERT(dim >= 0 && dim < 4);
-
-    int64_t o[4] = {0, 0, 0, 0};
-    o[dim] = src0->ne[dim];
-
-    const int8_t * x;
-
-    // TODO: smarter multi-theading
-    for (int i3 = 0; i3 < ne3; i3++) {
-        for (int i2 = ith; i2 < ne2; i2 += nth) {
-            for (int i1 = 0; i1 < ne1; i1++) {
-                for (int i0 = 0; i0 < ne0; i0++) {
-                    if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) {
-                        x = (const int8_t *) ((const char *)src0->data + (i0       )*nb00 + (i1       )*nb01 + (i2       )*nb02 + (i3       )*nb03);
-                    } else {
-                        x = (const int8_t *) ((const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13);
-                    }
-
-                    int8_t * y = (int8_t *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3);
-
-                    *y = *x;
-                }
-            }
-        }
-    }
-}
-
-static void ggml_compute_forward_concat_f16(
-    const ggml_compute_params * params,
-    ggml_tensor * dst) {
-
-    const ggml_tensor * src0 = dst->src[0];
-    const ggml_tensor * src1 = dst->src[1];
-
-    GGML_ASSERT(ggml_type_size(src0->type) == sizeof(ggml_fp16_t));
-
-    const int ith = params->ith;
-    const int nth = params->nth;
-
-    GGML_TENSOR_BINARY_OP_LOCALS
-
-    const int32_t dim = ggml_get_op_params_i32(dst, 0);
-
-    GGML_ASSERT(dim >= 0 && dim < 4);
-
-    int64_t o[4] = {0, 0, 0, 0};
-    o[dim] = src0->ne[dim];
-
-    const ggml_fp16_t * x;
-
-    // TODO: smarter multi-theading
-    for (int i3 = 0; i3 < ne3; i3++) {
-        for (int i2 = ith; i2 < ne2; i2 += nth) {
-            for (int i1 = 0; i1 < ne1; i1++) {
-                for (int i0 = 0; i0 < ne0; i0++) {
-                    if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) {
-                        x = (const ggml_fp16_t *) ((const char *)src0->data + (i0       )*nb00 + (i1       )*nb01 + (i2       )*nb02 + (i3       )*nb03);
-                    } else {
-                        x = (const ggml_fp16_t *) ((const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13);
-                    }
-
-                    ggml_fp16_t * y = (ggml_fp16_t *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3);
-
-                    *y = *x;
-                }
-            }
-        }
-    }
-}
-
-static void ggml_compute_forward_concat_f32(
-    const ggml_compute_params * params,
-    ggml_tensor * dst) {
-
-    const ggml_tensor * src0 = dst->src[0];
-    const ggml_tensor * src1 = dst->src[1];
-
-    GGML_ASSERT(ggml_type_size(src0->type) == sizeof(float));
-
-    const int ith = params->ith;
-    const int nth = params->nth;
-
-    GGML_TENSOR_BINARY_OP_LOCALS
-
-    const int32_t dim = ggml_get_op_params_i32(dst, 0);
-
-    GGML_ASSERT(dim >= 0 && dim < 4);
-
-    int64_t o[4] = {0, 0, 0, 0};
-    o[dim] = src0->ne[dim];
-
-    const float * x;
-
-    // TODO: smarter multi-theading
-    for (int i3 = 0; i3 < ne3; i3++) {
-        for (int i2 = ith; i2 < ne2; i2 += nth) {
-            for (int i1 = 0; i1 < ne1; i1++) {
-                for (int i0 = 0; i0 < ne0; i0++) {
-                    if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) {
-                        x = (const float *) ((const char *)src0->data + (i0       )*nb00 + (i1       )*nb01 + (i2       )*nb02 + (i3       )*nb03);
-                    } else {
-                        x = (const float *) ((const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13);
-                    }
-
-                    float * y = (float *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3);
-
-                    *y = *x;
-                }
-            }
+    if (src->nb[0] == es) {
+        const size_t row_size = ggml_row_size(src->type, n);
+        for (int64_t r = 0; r < n_rows; r++) {
+            memcpy(y + r*nby, x + r*nbx, row_size);
         }
+        return;
+    }
+
+    switch (es) {
+        case 1: ggml_concat_copy_strided<1>(y, nby, x, nbx, src->nb[0], n, n_rows); break;
+        case 2: ggml_concat_copy_strided<2>(y, nby, x, nbx, src->nb[0], n, n_rows); break;
+        case 4: ggml_concat_copy_strided<4>(y, nby, x, nbx, src->nb[0], n, n_rows); break;
+        case 8: ggml_concat_copy_strided<8>(y, nby, x, nbx, src->nb[0], n, n_rows); break;
+        default: GGML_ABORT("unsupported element size");
     }
 }
 
+// each dst row comes from src0 or src1, or for dim 0 from the row of src0 followed by the row of src1
 void ggml_compute_forward_concat(
     const ggml_compute_params * params,
     ggml_tensor * dst) {
 
     const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
 
-    switch (src0->type) {
-        case GGML_TYPE_F16:
-        case GGML_TYPE_BF16:
-        case GGML_TYPE_I16:
-            {
-                ggml_compute_forward_concat_f16(params, dst);
-            } break;
-        case GGML_TYPE_I8:
-            {
-                ggml_compute_forward_concat_i8(params, dst);
-            } break;
-        case GGML_TYPE_F32:
-        case GGML_TYPE_I32:
-            {
-                ggml_compute_forward_concat_f32(params, dst);
-            } break;
-        default:
-            {
-                ggml_compute_forward_concat_any(params, dst);
-            }
+    GGML_TENSOR_BINARY_OP_LOCALS
+
+    const int32_t dim = ggml_get_op_params_i32(dst, 0);
+
+    GGML_ASSERT(dim >= 0 && dim < 4);
+    GGML_ASSERT(nb0 == ggml_type_size(dst->type));
+    GGML_ASSERT(!ggml_is_quantized(dst->type) || (nb00 == nb0 && nb10 == nb0));
+    GGML_ASSERT(dim != 0 || (ne00 % ggml_blck_size(dst->type) == 0 && ne10 % ggml_blck_size(dst->type) == 0));
+
+    const size_t y1 = dim == 0 ? ggml_row_size(dst->type, ne00) : 0;
+
+    const int64_t nr = ne1*ne2*ne3;
+    const int64_t dr = (nr + params->nth - 1)/params->nth;
+
+    const int64_t ir0 = dr*params->ith;
+    const int64_t ir1 = MIN(ir0 + dr, nr);
+
+    for (int64_t ir = ir0; ir < ir1; ) {
+        int64_t i[4] = { 0, ir % ne1, (ir/ne1) % ne2, ir/(ne1*ne2) };
+
+        // a tile of consecutive rows sharing i2, i3 and their source
+        int64_t nt = MIN(MIN(ir1 - ir, ne1 - i[1]), GGML_CONCAT_TILE);
+        if (dim == 1 && i[1] < ne01) {
+            nt = MIN(nt, ne01 - i[1]);
+        }
+
+        char * y = (char *) dst->data + i[1]*nb1 + i[2]*nb2 + i[3]*nb3;
+
+        if (dim == 0 || i[dim] < src0->ne[dim]) {
+            ggml_concat_copy_rows(y, nb1, src0, (const char *) src0->data + i[1]*nb01 + i[2]*nb02 + i[3]*nb03, nb01, ne00, nt);
+        }
+
+        if (dim == 0 || i[dim] >= src0->ne[dim]) {
+            i[dim] -= src0->ne[dim]; // i[0] does not address rows
+            ggml_concat_copy_rows(y + y1, nb1, src1, (const char *) src1->data + i[1]*nb11 + i[2]*nb12 + i[3]*nb13, nb11, ne10, nt);
+        }
+
+        ir += nt;
     }
 }
 

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

@@ -214,52 +214,36 @@ static cudaError_t ggml_cuda_device_malloc(void ** ptr, size_t size, int device)
 }
 
 #if defined(GGML_USE_HIP)
-static int ggml_cuda_parse_id(char devName[]) {
+static int ggml_cuda_parse_id(const char * devName) {
     // A list of possible Target IDs can be found under the rocclr/clr repo in device.cpp
     // these values are not stable so this is susceptible to breakage
     // https://github.com/ROCm/clr/blob/amd-staging/rocclr/device/device.cpp
-    int archMajor = 0x0;
-    int archMinor = 0x0;
-    int archNum = GGML_CUDA_CC_OFFSET_AMD;
-    int archLen = strlen(devName);
-    char archName[archLen + 1];
+    // Extract the base architecture model ("gfx908:sramecc+:xnack-" -> "908")
+    std::string archName = devName;
+    archName.erase(0, 3);
+    archName = archName.substr(0, archName.find(':'));
 
-    // strip leading 'gfx' while copying into our buffer
-    if (archLen > 3) {
-        strcpy(archName, &devName[3]);
-        archLen -= 3;
-    }
+    const auto hex = [](const std::string & s) { return (int) strtoul(s.c_str(), nullptr, 16); };
 
-    // trim trailing :xnack- or :sramecc- statuses
-    archLen = strcspn(archName, ":");
-    archName[archLen] = '\0';
+    int cc = 0;
 
     // tease out the version information
-    if (archLen > 8) {
+    if (archName.size() > 8) {
         // versions labeled generic use '-' as delimiter
-        // strip the trailing "-generic" then iterate through what remains
-        if ((strstr(archName, "-generic"))) {
-            archName[archLen - 8] = '\0';
-            char * pch;
-            if ((pch = strtok(archName, "-"))) {
-                archMajor = (int)strtoul(pch, 0, 16);
-                if ((pch = strtok(NULL, "-"))) {
-                    archMinor = 0x10 * (int)strtoul(pch, 0, 16);
-                }
+        if (archName.compare(archName.size() - 8, 8, "-generic") == 0) {
+            archName.resize(archName.size() - 8);
+            const size_t sep = archName.find('-');
+            cc = 0x100 * hex(archName.substr(0, sep));
+            if (sep != std::string::npos) {
+                cc += 0x10 * hex(archName.substr(sep + 1));
             }
         }
-    } else if (archLen >= 3) {
-        // last two digits should be the minor * 0x10 + stepping
-        archMinor = (int)strtoul(&archName[archLen - 2], 0, 16);
-        archName[archLen - 2] = '\0';
-
-        // only the major version remains
-        archMajor = (int)strtoul(archName, 0, 16);
+    } else if (archName.size() >= 3) {
+        // ordinary IDs are the version in hex
+        cc = hex(archName);
     }
-    archNum += archMajor * 0x100;
-    archNum += archMinor;
 
-    return archNum;
+    return GGML_CUDA_CC_OFFSET_AMD + cc;
 }
 #endif // defined(GGML_USE_HIP)
 

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

@@ -16,6 +16,7 @@
 #include "hvx-dump.h"
 #include "hvx-arith.h"
 #include "hvx-reduce.h"
+#include "hvx-copy.h"
 
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
@@ -1413,7 +1414,7 @@ static void hvx_mv_id_tiled(unsigned int nth, unsigned int ith, void * data) {
 
     for (uint32_t ie1 = 0; ie1 < n_aids; ++ie1) {  // for each expert
         const int32_t eid = *(const int32_t *) ((const uint8_t *) src2->data + ie1 * src2->nb[0]);
-        if (eid == -1) {
+        if (eid < 0) {
             continue;
         }
         assert(eid >= 0 && eid < (int32_t) n_ids);
@@ -4111,7 +4112,7 @@ static int hvx_mm_matmul_id_nx(
     return HTP_STATUS_OK;
 }
 
-static inline void scan_expert_ids_n(
+static __attribute__((noinline)) uint32_t scan_expert_ids_n(
     const struct htp_tensor * ids,
     const uint32_t n_ids,
     uint32_t n_as,
@@ -4121,15 +4122,17 @@ static inline void scan_expert_ids_n(
 ) {
     const size_t ids_nb1 = ids->nb[1];
     const uint8_t * ids_data = (const uint8_t *) ids->data;
+    uint32_t n_skipped = 0;
 
     for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
         const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
         for (uint32_t id = 0; id < n_ids; ++id) {
             const int32_t i02 = row_ptr[id];
-            if (i02 == -1) {
+            if (i02 < 0) {
+                n_skipped++;
                 continue;
             }
-            assert(i02 >= 0 && i02 < n_as);
+            assert(i02 >= 0 && i02 < (int32_t) n_as);
 
             if (matrix_rows) {
                 matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
@@ -4137,31 +4140,10 @@ static inline void scan_expert_ids_n(
             counts[i02] += 1;
         }
     }
+    return n_skipped;
 }
 
-// the matmuls only write the rows that an expert owns, so the rows of the skipped slots are zeroed here
-static inline void zero_skipped_dst_rows(
-    const struct htp_tensor * ids,
-    uint32_t n_ids,
-    const struct htp_tensor * dst,
-    uint32_t ne0,
-    size_t nb1,
-    size_t nb2
-) {
-    const uint8_t * ids_data = (const uint8_t *) ids->data;
-
-    for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
-        for (uint32_t id = 0; id < n_ids; ++id) {
-            const int32_t i02 = *(const int32_t *) (ids_data + iid1 * ids->nb[1] + id * ids->nb[0]);
-            if (i02 != -1) {
-                continue;
-            }
-            memset((uint8_t *) (dst->data + id * nb1 + iid1 * nb2), 0, ne0 * sizeof(float));
-        }
-    }
-}
-
-static inline void scan_expert_ids(
+static __attribute__((noinline)) uint32_t scan_expert_ids(
     const struct htp_tensor * ids,
     uint32_t n_ids,
     uint32_t n_as,
@@ -4173,23 +4155,25 @@ static inline void scan_expert_ids(
 
     if (ids_nb0 == 4) {
         switch (n_ids) {
-            case 8:  scan_expert_ids_n(ids, 8,     n_as, counts, matrix_rows, mapping_stride); break;
-            case 4:  scan_expert_ids_n(ids, 4,     n_as, counts, matrix_rows, mapping_stride); break;
-            case 2:  scan_expert_ids_n(ids, 2,     n_as, counts, matrix_rows, mapping_stride); break;
-            default: scan_expert_ids_n(ids, n_ids, n_as, counts, matrix_rows, mapping_stride); break;
+            case 8:  return scan_expert_ids_n(ids, 8,     n_as, counts, matrix_rows, mapping_stride);
+            case 4:  return scan_expert_ids_n(ids, 4,     n_as, counts, matrix_rows, mapping_stride);
+            case 2:  return scan_expert_ids_n(ids, 2,     n_as, counts, matrix_rows, mapping_stride);
+            default: return scan_expert_ids_n(ids, n_ids, n_as, counts, matrix_rows, mapping_stride);
         }
     } else {
         // Strided fallback
         const size_t ids_nb1 = ids->nb[1];
         const uint8_t * ids_data = (const uint8_t *) ids->data;
+        uint32_t n_skipped = 0;
         for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
             const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
             for (uint32_t id = 0; id < n_ids; ++id) {
                 const int32_t i02 = *(const int32_t *) ((const uint8_t *) row_ptr + id * ids_nb0);
-                if (i02 == -1) {
+                if (i02 < 0) {
+                    n_skipped++;
                     continue;
                 }
-                assert(i02 >= 0 && i02 < n_as);
+                assert(i02 >= 0 && i02 < (int32_t) n_as);
 
                 if (matrix_rows) {
                     matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
@@ -4197,9 +4181,124 @@ static inline void scan_expert_ids(
                 counts[i02] += 1;
             }
         }
+        return n_skipped;
     }
 }
 
+static __attribute__((noinline)) uint32_t count_skipped_ids_single(const struct htp_tensor * ids, uint32_t n_ids) {
+    const uint8_t * ids_data = (const uint8_t *) ids->data;
+    const size_t ids_nb0 = ids->nb[0];
+    uint32_t n_skipped = 0;
+    for (uint32_t id = 0; id < n_ids; ++id) {
+        const int32_t i02 = *(const int32_t *) (ids_data + id * ids_nb0);
+        if (i02 < 0) {
+            n_skipped++;
+        }
+    }
+    return n_skipped;
+}
+
+// Zero rows for skipped slots (ids < 0). Matmuls only write rows that an expert owns.
+static void zero_skipped_dst_rows(
+    struct htp_ops_context * octx,
+    const struct htp_tensor * ids,
+    uint32_t n_ids,
+    const struct htp_tensor * dst
+) {
+    const uint32_t elem_size = (dst->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float);
+    uint32_t col_start = 0;
+    uint32_t col_count = dst->ne[0];
+
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_can_row_partition(dst, elem_size);
+        const uint32_t chunk_elems = (dst->type == HTP_TYPE_F16) ? 64 : 32;
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
+            dst->ne[0], can_split ? chunk_elems : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        col_start = range.start;
+        col_count = range.count;
+    }
+
+    if (col_count == 0) {
+        return;
+    }
+
+    const uint8_t * ids_data = (const uint8_t *) ids->data;
+    const size_t ids_nb0 = ids->nb[0];
+    const size_t ids_nb1 = ids->nb[1];
+    const size_t dst_nb1 = dst->nb[1];
+    const size_t dst_nb2 = dst->nb[2];
+    const uint32_t col_offset_bytes = col_start * elem_size;
+
+    for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
+        const uint8_t * ids_row = ids_data + iid1 * ids_nb1;
+        uint8_t * dst_row_group = (uint8_t *) dst->data + iid1 * dst_nb2 + col_offset_bytes;
+        for (uint32_t id = 0; id < n_ids; ++id) {
+            const int32_t i02 = *(const int32_t *) (ids_row + id * ids_nb0);
+            if (i02 < 0) {
+                hvx_splat_u(dst_row_group + id * dst_nb1, Q6_V_vzero(), col_count, elem_size);
+            }
+        }
+    }
+}
+
+static int htp_mmid_prepare_mapping(
+    struct htp_ops_context * octx,
+    const struct htp_tensor * ids,
+    uint32_t n_ids,
+    uint32_t n_as,
+    uint32_t act_nrows,
+    uint8_t ** out_mapping_buf,
+    uint32_t ** out_matrix_row_counts,
+    struct mmid_row_mapping ** out_matrix_rows,
+    uint32_t * out_mapping_stride,
+    uint32_t * out_n_skipped
+) {
+    uint8_t  * mapping_buf       = octx->ctx->ddr_spad_base;
+    uint32_t   mapping_stride    = 1;
+    uint32_t * matrix_row_counts = (uint32_t *) mapping_buf;
+    struct mmid_row_mapping * matrix_rows = NULL;
+    uint32_t n_skipped = 0;
+
+    if (act_nrows > 1) {
+        const size_t matrix_row_counts_size = n_as * sizeof(uint32_t);
+        assert(octx->ctx->ddr_spad_size >= matrix_row_counts_size);
+
+        hex_l2fetch_block((const void *) ids->data, ids->ne[1] * ids->nb[1]);
+
+        hvx_splat_u(matrix_row_counts, Q6_V_vzero(), n_as, sizeof(uint32_t));
+        n_skipped = scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, NULL, 0);
+
+        uint32_t max_count = hvx_reduce_max_i32((const uint8_t *) matrix_row_counts, n_as);
+        mapping_stride = max_count > 0 ? max_count : 1;
+
+        size_t matrix_row_map_size  = n_as * mapping_stride * sizeof(struct mmid_row_mapping);
+        const size_t total_map_size = matrix_row_counts_size + matrix_row_map_size;
+
+        if (total_map_size > octx->ctx->ddr_spad_size) {
+            mapping_buf = memalign(128, total_map_size);
+            if (!mapping_buf) {
+                return HTP_STATUS_INTERNAL_ERR;
+            }
+        }
+
+        matrix_row_counts = (uint32_t *) mapping_buf;
+        matrix_rows       = (struct mmid_row_mapping *) (mapping_buf + matrix_row_counts_size);
+
+        hvx_splat_u(matrix_row_counts, Q6_V_vzero(), n_as, sizeof(uint32_t));
+        scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride);
+    } else {
+        n_skipped = count_skipped_ids_single(ids, n_ids);
+    }
+
+    *out_mapping_buf       = mapping_buf;
+    *out_matrix_row_counts = matrix_row_counts;
+    *out_matrix_rows       = matrix_rows;
+    *out_mapping_stride    = mapping_stride;
+    *out_n_skipped         = n_skipped;
+
+    return HTP_STATUS_OK;
+}
+
 int op_matmul_id(struct htp_ops_context * octx) {
     htp_matmul_tensors_preamble;
 
@@ -4235,41 +4334,22 @@ int op_matmul_id(struct htp_ops_context * octx) {
     const int n_ids = ids->ne[0];  // n_expert_used
     const int n_as  = ne02;        // n_expert
 
-    uint8_t  * mapping_buf       = octx->ctx->ddr_spad_base;
-    uint32_t   mapping_stride    = 1;
-    uint32_t * matrix_row_counts = (uint32_t *) mapping_buf;
+    uint8_t  * mapping_buf       = NULL;
+    uint32_t * matrix_row_counts = NULL;
     struct mmid_row_mapping * matrix_rows = NULL;
+    uint32_t   mapping_stride    = 1;
+    uint32_t   n_skipped         = 0;
 
-    if (act_nrows > 1) {
-        const size_t matrix_row_counts_size = n_as * sizeof(uint32_t);
-        assert(octx->ctx->ddr_spad_size >= matrix_row_counts_size);
-
-        hex_l2fetch_block((const void *) ids->data, ids->ne[1] * ids->nb[1]);
-
-        memset(matrix_row_counts, 0, matrix_row_counts_size);
-        scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, NULL, 0);
-
-        uint32_t max_count = hvx_reduce_max_i32((const uint8_t *) matrix_row_counts, n_as);
-        mapping_stride = max_count > 0 ? max_count : 1;
-
-        size_t matrix_row_map_size  = n_as * mapping_stride * sizeof(struct mmid_row_mapping);
-        const size_t total_map_size = matrix_row_counts_size + matrix_row_map_size;
-
-        if (total_map_size > octx->ctx->ddr_spad_size) {
-            mapping_buf = memalign(128, total_map_size);
-            if (!mapping_buf) {
-                return HTP_STATUS_INTERNAL_ERR;
-            }
-        }
-
-        matrix_row_counts = (uint32_t *) mapping_buf;
-        matrix_rows       = (struct mmid_row_mapping *) (mapping_buf + matrix_row_counts_size);
-
-        memset(matrix_row_counts, 0, n_as * sizeof(uint32_t));
-        scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride);
+    int map_status = htp_mmid_prepare_mapping(octx, ids, n_ids, n_as, act_nrows,
+                                              &mapping_buf, &matrix_row_counts, &matrix_rows,
+                                              &mapping_stride, &n_skipped);
+    if (map_status != HTP_STATUS_OK) {
+        return map_status;
     }
 
-    zero_skipped_dst_rows(ids, n_ids, dst, ne0, nb1, nb2);
+    if (n_skipped > 0) {
+        zero_skipped_dst_rows(octx, ids, n_ids, dst);
+    }
 
     mmctx->matrix_row_counts    = matrix_row_counts;
     mmctx->matrix_rows          = matrix_rows;
@@ -4360,38 +4440,26 @@ int op_matmul_id_nx(struct htp_ops_context * octx) {
     const int n_ids = ids->ne[0];
     const int n_as  = src0->ne[2];
 
-    uint8_t  * mapping_buf       = octx->ctx->ddr_spad_base;
-    uint32_t   mapping_stride    = 1;
-    uint32_t * matrix_row_counts = (uint32_t *) mapping_buf;
+    uint8_t  * mapping_buf       = NULL;
+    uint32_t * matrix_row_counts = NULL;
     struct mmid_row_mapping * matrix_rows = NULL;
+    uint32_t   mapping_stride    = 1;
+    uint32_t   n_skipped         = 0;
 
-    if (act_nrows > 1) {
-        const size_t matrix_row_counts_size = n_as * sizeof(uint32_t);
-        assert(octx->ctx->ddr_spad_size >= matrix_row_counts_size);
+    int map_status = htp_mmid_prepare_mapping(octx, ids, n_ids, n_as, act_nrows,
+                                              &mapping_buf, &matrix_row_counts, &matrix_rows,
+                                              &mapping_stride, &n_skipped);
+    if (map_status != HTP_STATUS_OK) {
+        return map_status;
+    }
 
-        hex_l2fetch_block((const void *) ids->data, ids->ne[1] * ids->nb[1]);
-
-        memset(matrix_row_counts, 0, matrix_row_counts_size);
-        scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, NULL, 0);
-
-        uint32_t max_count = hvx_reduce_max_i32((const uint8_t *) matrix_row_counts, n_as);
-        mapping_stride = max_count > 0 ? max_count : 1;
-
-        size_t matrix_row_map_size  = n_as * mapping_stride * sizeof(struct mmid_row_mapping);
-        const size_t total_map_size = matrix_row_counts_size + matrix_row_map_size;
-
-        if (total_map_size > octx->ctx->ddr_spad_size) {
-            mapping_buf = memalign(128, total_map_size);
-            if (!mapping_buf) {
-                return HTP_STATUS_INTERNAL_ERR;
+    if (n_skipped > 0) {
+        for (uint32_t p = 0; p < n_weights; ++p) {
+            const struct htp_tensor * restrict d = octx->dsts[p];
+            if (d) {
+                zero_skipped_dst_rows(octx, ids, n_ids, d);
             }
         }
-
-        matrix_row_counts = (uint32_t *) mapping_buf;
-        matrix_rows       = (struct mmid_row_mapping *) (mapping_buf + matrix_row_counts_size);
-
-        memset(matrix_row_counts, 0, n_as * sizeof(uint32_t));
-        scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride);
     }
 
     mmctx->matrix_row_counts    = matrix_row_counts;

ggml/src/ggml-metal/kernels/mul_mm.metal

@@ -323,10 +323,10 @@ kernel void kernel_mul_mm(
         // if no bounds checks on the output are needed, we can directly write to device memory
         device float * C = (device float *) dst +
             (r0 + 32*(sgitg &  1)) + \
-            (r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0;
+            (uint64_t) (r1 + 16*(sgitg >> 1))*args.ne0 + (uint64_t) im*args.ne1*args.ne0;
 
         for (short i = 0; i < 8; i++) {
-            simdgroup_store(mc[i], C + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false);
+            simdgroup_store(mc[i], C + 8*(i%4) + (uint64_t) 8*args.ne0*(i/4), args.ne0, 0, false);
         }
     } else {
         // block is smaller than 64x32, we should avoid writing data outside of the matrix
@@ -342,7 +342,7 @@ kernel void kernel_mul_mm(
 
         if (sgitg == 0) {
             for (int j = tiitg; j < nr1; j += NR1) {
-                device float  * D  = (device float  *) dst + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0;
+                device float  * D  = (device float  *) dst + r0 + (uint64_t) (r1 + j)*args.ne0 + (uint64_t) im*args.ne1*args.ne0;
                 device float4 * D4 = (device float4 *) D;
 
                 threadgroup float  * C  = temp_str + (j*NR0);
@@ -633,7 +633,7 @@ kernel void kernel_mul_mm_id(
     const int id = ids_i32[idb + lr1];
 
     const short i11 = (id % args.ne20) % args.ne11;
-    const short i12 = (id / args.ne20);
+    const int   i12 = (id / args.ne20);
 
     device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
 
@@ -866,9 +866,9 @@ kernel void kernel_mul_mm_id(
         const int id = ids_i32[idb + j];
 
         const short ide = id % args.ne20;
-        const short idt = id / args.ne20;
+        const int   idt = id / args.ne20;
 
-        device float  * D  = (device float  *) dst + r0 + ide*args.ne0 + idt*args.ne1*args.ne0;
+        device float  * D  = (device float  *) dst + r0 + (uint64_t) ide*args.ne0 + (uint64_t) idt*args.ne1*args.ne0;
         device float4 * D4 = (device float4 *) D;
 
         threadgroup float  * C  = (threadgroup float  *) shmem + j*NR0;

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

@@ -10901,7 +10901,16 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer,
     // conversion kernel.
     if (tensor->type == GGML_TYPE_Q4_0) {
         // Views can't SoA-ify here — parent owns the layout (see q8_0 guard).
+        // A view of an AoS tensor (the KV cache, written by a state restore)
+        // is written in place.
         if (tensor->view_src != nullptr || !ggml_is_contiguous(tensor)) {
+            if (!ggml_cl_is_q4_0_soa(tensor)) {
+                ggml_tensor_extra_cl * extra_aos = (ggml_tensor_extra_cl *) tensor->extra;
+                CL_CHECK(clEnqueueWriteBuffer(
+                    queue, extra_aos->data_device, CL_TRUE,
+                    extra_aos->offset + tensor->view_offs + offset,
+                    size, data, 0, NULL, NULL));
+            }
             return;
         }
         // Tensors should have been preallocated, therefore they should
@@ -11605,7 +11614,16 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer,
     }
     if (tensor->type == GGML_TYPE_Q8_0) {
         // Views share the parent's buffer; parent owns SoA conversion.
+        // A view of an AoS tensor (the KV cache, written by a state restore)
+        // is written in place.
         if (tensor->view_src != nullptr || !ggml_is_contiguous(tensor)) {
+            if (!ggml_cl_is_q8_0_soa(tensor)) {
+                ggml_tensor_extra_cl * extra_aos = (ggml_tensor_extra_cl *) tensor->extra;
+                CL_CHECK(clEnqueueWriteBuffer(
+                    queue, extra_aos->data_device, CL_TRUE,
+                    extra_aos->offset + tensor->view_offs + offset,
+                    size, data, 0, NULL, NULL));
+            }
             return;
         }
 
@@ -12440,7 +12458,7 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer,
         GGML_ASSERT(extra);
 
         cl_ulong n_elements = size / sizeof(ggml_fp16_t);
-        cl_ulong off_dst = (extra->offset + offset) / sizeof(ggml_fp16_t);
+        cl_ulong off_dst = (extra->offset + tensor->view_offs + offset) / sizeof(ggml_fp16_t);
 
         cl_int err;
         cl_mem data_device = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
@@ -12468,8 +12486,9 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer,
     ggml_tensor_extra_cl * extra = (ggml_tensor_extra_cl *) tensor->extra;
     GGML_ASSERT(extra);
 
+    // a view reuses its parent's extra, whose offset does not include view_offs
     CL_CHECK(clEnqueueWriteBuffer(
-        queue, extra->data_device, CL_TRUE, extra->offset + offset,
+        queue, extra->data_device, CL_TRUE, extra->offset + tensor->view_offs + offset,
         size, data, 0, NULL, NULL));
 
     GGML_UNUSED(buffer);

ggml/src/ggml-sycl/concat.cpp

@@ -196,7 +196,8 @@ void concat_impl_sycl(ggml_backend_sycl_context & ctx, ggml_tensor *dst) {
     }
 }
 
-static void concat_impl_q4_0_sycl(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+template <typename block_t, int qk, ggml_type qtype>
+static void concat_impl_q_sycl(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
     const ggml_tensor * src0   = dst->src[0];
     const ggml_tensor * src1   = dst->src[1];
@@ -204,25 +205,25 @@ static void concat_impl_q4_0_sycl(ggml_backend_sycl_context & ctx, ggml_tensor *
 
     const int32_t dim = ((int32_t *) dst->op_params)[0];
 
-    GGML_ASSERT(src0->type == GGML_TYPE_Q4_0);
-    GGML_ASSERT(src1->type == GGML_TYPE_Q4_0);
-    GGML_ASSERT(dst->type == GGML_TYPE_Q4_0);
-    GGML_ASSERT(src0->ne[0] % QK4_0 == 0);
-    GGML_ASSERT(src1->ne[0] % QK4_0 == 0);
-    GGML_ASSERT(dst->ne[0] % QK4_0 == 0);
+    GGML_ASSERT(src0->type == qtype);
+    GGML_ASSERT(src1->type == qtype);
+    GGML_ASSERT(dst->type == qtype);
+    GGML_ASSERT(src0->ne[0] % qk == 0);
+    GGML_ASSERT(src1->ne[0] % qk == 0);
+    GGML_ASSERT(dst->ne[0] % qk == 0);
 
-    const int ne00_blk = src0->ne[0] / QK4_0;
-    const int ne0_blk  = dst->ne[0] / QK4_0;
+    const int ne00_blk = src0->ne[0] / qk;
+    const int ne0_blk  = dst->ne[0] / qk;
 
     if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) {
-        const block_q4_0 * src0_d = (const block_q4_0 *) src0->data;
-        const block_q4_0 * src1_d = (const block_q4_0 *) src1->data;
-        block_q4_0 * dst_d        = (block_q4_0 *) dst->data;
-        const size_t type_size    = sizeof(block_q4_0);
+        const block_t * src0_d = (const block_t *) src0->data;
+        const block_t * src1_d = (const block_t *) src1->data;
+        block_t * dst_d        = (block_t *) dst->data;
+        const size_t type_size = sizeof(block_t);
 
         if (dim != 3) {
             for (int i3 = 0; i3 < dst->ne[3]; i3++) {
-                concat_T_sycl<block_q4_0>(
+                concat_T_sycl<block_t>(
                     src0_d + i3 * (src0->nb[3] / type_size),
                     src1_d + i3 * (src1->nb[3] / type_size),
                     dst_d + i3 * (dst->nb[3] / type_size),
@@ -237,223 +238,12 @@ static void concat_impl_q4_0_sycl(ggml_backend_sycl_context & ctx, ggml_tensor *
             SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy((char *) dst_d + size0, src1_d, size1)));
         }
     } else {
-        concat_T_sycl_non_cont<block_q4_0>(
+        concat_T_sycl_non_cont<block_t>(
             stream, (const char *) src0->data, (const char *) src1->data,
             (char *) dst->data,
             ne00_blk, src0->ne[1], src0->ne[2], src0->ne[3],
             src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3],
-            src1->ne[0] / QK4_0, src1->ne[1], src1->ne[2], src1->ne[3],
-            src1->nb[0], src1->nb[1], src1->nb[2], src1->nb[3],
-            ne0_blk, dst->ne[1], dst->ne[2], dst->ne[3],
-            dst->nb[0], dst->nb[1], dst->nb[2], dst->nb[3], dim);
-    }
-}
-
-static void concat_impl_q4_1_sycl(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
-    const ggml_tensor * src0   = dst->src[0];
-    const ggml_tensor * src1   = dst->src[1];
-    queue_ptr           stream = ctx.stream();
-
-    const int32_t dim = ((int32_t *) dst->op_params)[0];
-
-    GGML_ASSERT(src0->type == GGML_TYPE_Q4_1);
-    GGML_ASSERT(src1->type == GGML_TYPE_Q4_1);
-    GGML_ASSERT(dst->type == GGML_TYPE_Q4_1);
-    GGML_ASSERT(src0->ne[0] % QK4_1 == 0);
-    GGML_ASSERT(src1->ne[0] % QK4_1 == 0);
-    GGML_ASSERT(dst->ne[0] % QK4_1 == 0);
-
-    const int ne00_blk = src0->ne[0] / QK4_1;
-    const int ne0_blk  = dst->ne[0] / QK4_1;
-
-    if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) {
-        const block_q4_1 * src0_d = (const block_q4_1 *) src0->data;
-        const block_q4_1 * src1_d = (const block_q4_1 *) src1->data;
-        block_q4_1 * dst_d        = (block_q4_1 *) dst->data;
-        const size_t type_size    = sizeof(block_q4_1);
-
-        if (dim != 3) {
-            for (int i3 = 0; i3 < dst->ne[3]; i3++) {
-                concat_T_sycl<block_q4_1>(
-                    src0_d + i3 * (src0->nb[3] / type_size),
-                    src1_d + i3 * (src1->nb[3] / type_size),
-                    dst_d + i3 * (dst->nb[3] / type_size),
-                    ne00_blk, src0->ne[1], src0->ne[2], ne0_blk,
-                    dst->ne[1], dst->ne[2], dim, stream);
-            }
-        } else {
-            const size_t size0 = ggml_nbytes(src0);
-            const size_t size1 = ggml_nbytes(src1);
-
-            SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy(dst_d, src0_d, size0)));
-            SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy((char *) dst_d + size0, src1_d, size1)));
-        }
-    } else {
-        concat_T_sycl_non_cont<block_q4_1>(
-            stream, (const char *) src0->data, (const char *) src1->data,
-            (char *) dst->data,
-            ne00_blk, src0->ne[1], src0->ne[2], src0->ne[3],
-            src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3],
-            src1->ne[0] / QK4_1, src1->ne[1], src1->ne[2], src1->ne[3],
-            src1->nb[0], src1->nb[1], src1->nb[2], src1->nb[3],
-            ne0_blk, dst->ne[1], dst->ne[2], dst->ne[3],
-            dst->nb[0], dst->nb[1], dst->nb[2], dst->nb[3], dim);
-    }
-}
-
-static void concat_impl_q5_0_sycl(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
-    const ggml_tensor * src0   = dst->src[0];
-    const ggml_tensor * src1   = dst->src[1];
-    queue_ptr           stream = ctx.stream();
-
-    const int32_t dim = ((int32_t *) dst->op_params)[0];
-
-    GGML_ASSERT(src0->type == GGML_TYPE_Q5_0);
-    GGML_ASSERT(src1->type == GGML_TYPE_Q5_0);
-    GGML_ASSERT(dst->type == GGML_TYPE_Q5_0);
-    GGML_ASSERT(src0->ne[0] % QK5_0 == 0);
-    GGML_ASSERT(src1->ne[0] % QK5_0 == 0);
-    GGML_ASSERT(dst->ne[0] % QK5_0 == 0);
-
-    const int ne00_blk = src0->ne[0] / QK5_0;
-    const int ne0_blk  = dst->ne[0] / QK5_0;
-
-    if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) {
-        const block_q5_0 * src0_d = (const block_q5_0 *) src0->data;
-        const block_q5_0 * src1_d = (const block_q5_0 *) src1->data;
-        block_q5_0 * dst_d        = (block_q5_0 *) dst->data;
-        const size_t type_size    = sizeof(block_q5_0);
-
-        if (dim != 3) {
-            for (int i3 = 0; i3 < dst->ne[3]; i3++) {
-                concat_T_sycl<block_q5_0>(
-                    src0_d + i3 * (src0->nb[3] / type_size),
-                    src1_d + i3 * (src1->nb[3] / type_size),
-                    dst_d + i3 * (dst->nb[3] / type_size),
-                    ne00_blk, src0->ne[1], src0->ne[2], ne0_blk,
-                    dst->ne[1], dst->ne[2], dim, stream);
-            }
-        } else {
-            const size_t size0 = ggml_nbytes(src0);
-            const size_t size1 = ggml_nbytes(src1);
-
-            SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy(dst_d, src0_d, size0)));
-            SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy((char *) dst_d + size0, src1_d, size1)));
-        }
-    } else {
-        concat_T_sycl_non_cont<block_q5_0>(
-            stream, (const char *) src0->data, (const char *) src1->data,
-            (char *) dst->data,
-            ne00_blk, src0->ne[1], src0->ne[2], src0->ne[3],
-            src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3],
-            src1->ne[0] / QK5_0, src1->ne[1], src1->ne[2], src1->ne[3],
-            src1->nb[0], src1->nb[1], src1->nb[2], src1->nb[3],
-            ne0_blk, dst->ne[1], dst->ne[2], dst->ne[3],
-            dst->nb[0], dst->nb[1], dst->nb[2], dst->nb[3], dim);
-    }
-}
-
-static void concat_impl_q5_1_sycl(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
-    const ggml_tensor * src0   = dst->src[0];
-    const ggml_tensor * src1   = dst->src[1];
-    queue_ptr           stream = ctx.stream();
-
-    const int32_t dim = ((int32_t *) dst->op_params)[0];
-
-    GGML_ASSERT(src0->type == GGML_TYPE_Q5_1);
-    GGML_ASSERT(src1->type == GGML_TYPE_Q5_1);
-    GGML_ASSERT(dst->type == GGML_TYPE_Q5_1);
-    GGML_ASSERT(src0->ne[0] % QK5_1 == 0);
-    GGML_ASSERT(src1->ne[0] % QK5_1 == 0);
-    GGML_ASSERT(dst->ne[0] % QK5_1 == 0);
-
-    const int ne00_blk = src0->ne[0] / QK5_1;
-    const int ne0_blk  = dst->ne[0] / QK5_1;
-
-    if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) {
-        const block_q5_1 * src0_d = (const block_q5_1 *) src0->data;
-        const block_q5_1 * src1_d = (const block_q5_1 *) src1->data;
-        block_q5_1 * dst_d        = (block_q5_1 *) dst->data;
-        const size_t type_size    = sizeof(block_q5_1);
-
-        if (dim != 3) {
-            for (int i3 = 0; i3 < dst->ne[3]; i3++) {
-                concat_T_sycl<block_q5_1>(
-                    src0_d + i3 * (src0->nb[3] / type_size),
-                    src1_d + i3 * (src1->nb[3] / type_size),
-                    dst_d + i3 * (dst->nb[3] / type_size),
-                    ne00_blk, src0->ne[1], src0->ne[2], ne0_blk,
-                    dst->ne[1], dst->ne[2], dim, stream);
-            }
-        } else {
-            const size_t size0 = ggml_nbytes(src0);
-            const size_t size1 = ggml_nbytes(src1);
-
-            SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy(dst_d, src0_d, size0)));
-            SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy((char *) dst_d + size0, src1_d, size1)));
-        }
-    } else {
-        concat_T_sycl_non_cont<block_q5_1>(
-            stream, (const char *) src0->data, (const char *) src1->data,
-            (char *) dst->data,
-            ne00_blk, src0->ne[1], src0->ne[2], src0->ne[3],
-            src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3],
-            src1->ne[0] / QK5_1, src1->ne[1], src1->ne[2], src1->ne[3],
-            src1->nb[0], src1->nb[1], src1->nb[2], src1->nb[3],
-            ne0_blk, dst->ne[1], dst->ne[2], dst->ne[3],
-            dst->nb[0], dst->nb[1], dst->nb[2], dst->nb[3], dim);
-    }
-}
-
-static void concat_impl_q8_0_sycl(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
-    const ggml_tensor * src0   = dst->src[0];
-    const ggml_tensor * src1   = dst->src[1];
-    queue_ptr           stream = ctx.stream();
-
-    const int32_t dim = ((int32_t *) dst->op_params)[0];
-
-    GGML_ASSERT(src0->type == GGML_TYPE_Q8_0);
-    GGML_ASSERT(src1->type == GGML_TYPE_Q8_0);
-    GGML_ASSERT(dst->type == GGML_TYPE_Q8_0);
-    GGML_ASSERT(src0->ne[0] % QK8_0 == 0);
-    GGML_ASSERT(src1->ne[0] % QK8_0 == 0);
-    GGML_ASSERT(dst->ne[0] % QK8_0 == 0);
-
-    const int ne00_blk = src0->ne[0] / QK8_0;
-    const int ne0_blk  = dst->ne[0] / QK8_0;
-
-    if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) {
-        const block_q8_0 * src0_d = (const block_q8_0 *) src0->data;
-        const block_q8_0 * src1_d = (const block_q8_0 *) src1->data;
-        block_q8_0 * dst_d        = (block_q8_0 *) dst->data;
-        const size_t type_size    = sizeof(block_q8_0);
-
-        if (dim != 3) {
-            for (int i3 = 0; i3 < dst->ne[3]; i3++) {
-                concat_T_sycl<block_q8_0>(
-                    src0_d + i3 * (src0->nb[3] / type_size),
-                    src1_d + i3 * (src1->nb[3] / type_size),
-                    dst_d + i3 * (dst->nb[3] / type_size),
-                    ne00_blk, src0->ne[1], src0->ne[2], ne0_blk,
-                    dst->ne[1], dst->ne[2], dim, stream);
-            }
-        } else {
-            const size_t size0 = ggml_nbytes(src0);
-            const size_t size1 = ggml_nbytes(src1);
-            SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy(dst_d, src0_d, size0)));
-            SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy((char *) dst_d + size0, src1_d, size1)));
-        }
-    } else {
-        concat_T_sycl_non_cont<block_q8_0>(
-            stream, (const char *) src0->data, (const char *) src1->data,
-            (char *) dst->data,
-            ne00_blk, src0->ne[1], src0->ne[2], src0->ne[3],
-            src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3],
-            src1->ne[0] / QK8_0, src1->ne[1], src1->ne[2], src1->ne[3],
+            src1->ne[0] / qk, src1->ne[1], src1->ne[2], src1->ne[3],
             src1->nb[0], src1->nb[1], src1->nb[2], src1->nb[3],
             ne0_blk, dst->ne[1], dst->ne[2], dst->ne[3],
             dst->nb[0], dst->nb[1], dst->nb[2], dst->nb[3], dim);
@@ -487,19 +277,19 @@ void ggml_sycl_op_concat(ggml_backend_sycl_context & ctx, ggml_tensor *dst) {
         concat_impl_sycl<int8_t>(ctx, dst);
         break;
     case GGML_TYPE_Q4_0:
-        concat_impl_q4_0_sycl(ctx, dst);
+        concat_impl_q_sycl<block_q4_0, QK4_0, GGML_TYPE_Q4_0>(ctx, dst);
         break;
     case GGML_TYPE_Q4_1:
-        concat_impl_q4_1_sycl(ctx, dst);
+        concat_impl_q_sycl<block_q4_1, QK4_1, GGML_TYPE_Q4_1>(ctx, dst);
         break;
     case GGML_TYPE_Q5_0:
-        concat_impl_q5_0_sycl(ctx, dst);
+        concat_impl_q_sycl<block_q5_0, QK5_0, GGML_TYPE_Q5_0>(ctx, dst);
         break;
     case GGML_TYPE_Q5_1:
-        concat_impl_q5_1_sycl(ctx, dst);
+        concat_impl_q_sycl<block_q5_1, QK5_1, GGML_TYPE_Q5_1>(ctx, dst);
         break;
     case GGML_TYPE_Q8_0:
-        concat_impl_q8_0_sycl(ctx, dst);
+        concat_impl_q_sycl<block_q8_0, QK8_0, GGML_TYPE_Q8_0>(ctx, dst);
         break;
     default:
         fprintf(stderr, "%s: unsupported types: dst: %s\n", __func__, ggml_type_name(dst->type));

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

@@ -842,6 +842,19 @@ catch (sycl::exception const &exc) {
   std::exit(1);
 }
 
+// allocates the pinned upload staging on first use; returns false if staging is disabled or the allocation failed
+static bool ggml_sycl_buffer_init_staging(ggml_backend_sycl_buffer_context * ctx) {
+    const int n_slots = g_ggml_sycl_upload_staging_slots;
+    if (n_slots > 0 && ctx->staging.data == nullptr) {
+        ctx->staging.data = sycl::malloc_host(n_slots * ctx->staging_slot_size, *ctx->stream);
+        if (ctx->staging.data != nullptr) {
+            ctx->staging.events.resize(n_slots);
+            ctx->staging.submitted.assign(n_slots, false);
+        }
+    }
+    return ctx->staging.data != nullptr;
+}
+
 static void ggml_backend_sycl_buffer_set_tensor(ggml_backend_buffer_t buffer,
                                                 ggml_tensor *tensor,
                                                 const void *data, size_t offset,
@@ -854,15 +867,7 @@ static void ggml_backend_sycl_buffer_set_tensor(ggml_backend_buffer_t buffer,
 
     // copy through pinned memory so the device never reads mmap()ed pages directly
     // chunks pipeline on the in-order compute queue, so no drain per tensor is needed
-    const int n_slots = g_ggml_sycl_upload_staging_slots;
-    if (n_slots > 0 && ctx->staging.data == nullptr) {
-        ctx->staging.data = sycl::malloc_host(n_slots * ctx->staging_slot_size, *ctx->stream);
-        if (ctx->staging.data != nullptr) {
-            ctx->staging.events.resize(n_slots);
-            ctx->staging.submitted.assign(n_slots, false);
-        }
-    }
-    if (ctx->staging.data != nullptr) {
+    if (ggml_sycl_buffer_init_staging(ctx)) {
         queue_ptr    stream    = ctx->stream;
         char *       dst       = (char *) tensor->data + offset;
         const char * src       = (const char *) data;
@@ -904,6 +909,49 @@ catch (sycl::exception const &exc) {
   std::exit(1);
 }
 
+// Copy n_copies rows in one call. -sm tensor uploads split weights this way; the generic fallback does one set_tensor per row.
+static void ggml_backend_sycl_buffer_set_tensor_2d(ggml_backend_buffer_t buffer, ggml_tensor * tensor,
+                                                   const void * data, size_t offset, size_t size,
+                                                   size_t n_copies, size_t stride_tensor,
+                                                   size_t stride_data) try {
+    ggml_backend_sycl_buffer_context * ctx = (ggml_backend_sycl_buffer_context *) buffer->context;
+    const char * src = (const char *) data;
+    ggml_sycl_set_device(ctx->device);
+
+    if (stride_tensor != size || size > ctx->staging_slot_size || !ggml_sycl_buffer_init_staging(ctx)) {
+        // dst rows not contiguous, row larger than a staging slot, or no staging: same as the generic fallback
+        for (size_t i = 0; i < n_copies; ++i) {
+            ggml_backend_sycl_buffer_set_tensor(buffer, tensor, src + i * stride_data, offset + i * stride_tensor, size);
+        }
+        return;
+    }
+
+    // dst rows are contiguous: pack as many rows as fit into a staging slot and copy each slot with one memcpy
+    // a slot is only waited on when it is reused, so packing the next slot overlaps with the previous copies
+    queue_ptr    stream        = ctx->stream;
+    char *       dst           = (char *) tensor->data + offset;
+    const size_t rows_per_slot = ctx->staging_slot_size / size;
+    for (size_t i0 = 0; i0 < n_copies; i0 += rows_per_slot) {
+        const size_t n    = std::min(rows_per_slot, n_copies - i0);
+        const int    slot = ctx->staging.next;
+        ctx->staging.next = (ctx->staging.next + 1) % (int) ctx->staging.submitted.size();
+        if (ctx->staging.submitted[slot]) {
+            ctx->staging.events[slot].wait_and_throw();
+        }
+        char * stage = (char *) ctx->staging.data + slot * ctx->staging_slot_size;
+        for (size_t i = 0; i < n; ++i) {
+            memcpy(stage + i * size, src + (i0 + i) * stride_data, size);
+        }
+        ctx->staging.events[slot]    = stream->memcpy(dst + i0 * size, stage, n * size);
+        ctx->staging.submitted[slot] = true;
+    }
+}
+catch (sycl::exception const &exc) {
+  std::cerr << exc.what() << "Exception caught at file:" << __FILE__
+            << ", line:" << __LINE__ << std::endl;
+  std::exit(1);
+}
+
 static void ggml_backend_sycl_buffer_get_tensor(ggml_backend_buffer_t buffer,
                                                 const ggml_tensor *tensor,
                                                 void *data, size_t offset,
@@ -1122,7 +1170,7 @@ static const ggml_backend_buffer_i ggml_backend_sycl_buffer_interface = {
     /* .memset_tensor   = */ ggml_backend_sycl_buffer_memset_tensor,
     /* .set_tensor      = */ ggml_backend_sycl_buffer_set_tensor,
     /* .get_tensor      = */ ggml_backend_sycl_buffer_get_tensor,
-    /* .set_tensor_2d   = */ NULL,
+    /* .set_tensor_2d   = */ ggml_backend_sycl_buffer_set_tensor_2d,
     /* .get_tensor_2d   = */ NULL,
     /* .cpy_tensor      = */ ggml_backend_sycl_buffer_cpy_tensor,
     /* .clear           = */ ggml_backend_sycl_buffer_clear,
@@ -4414,7 +4462,15 @@ static inline void * sycl_ext_malloc_device(dpct::queue_ptr stream, size_t size)
     bool use_async = g_ggml_sycl_use_async_mem_op;
 #if defined(GGML_SYCL_GRAPH) && SYCL_EXT_ONEAPI_ASYNC_MEMORY_ALLOC
     if (use_async) {
-        void * ptr = syclex::async_malloc(*stream, sycl::usm::alloc::device, size);
+        // async_malloc throws on out-of-memory; return nullptr like the sync path so callers can fall back.
+        void * ptr = nullptr;
+        try {
+            ptr = syclex::async_malloc(*stream, sycl::usm::alloc::device, size);
+        } catch (const sycl::exception & e) {
+            GGML_LOG_WARN("%s: async_malloc of %zu bytes failed: %s\n", __func__, size, e.what());
+            ggml_sycl_memtrace_fail(GGML_SYCL_MEM_ASYNC, size);
+            return nullptr;
+        }
         ggml_sycl_memtrace_add(GGML_SYCL_MEM_ASYNC, ptr, size);
         return ptr;
     }
@@ -4481,307 +4537,317 @@ private:
     bool host_fallback = false;
 };
 
-static bool reorder_qw_q4_0(uint8_t * data_device, const int ncols, const int nrows, size_t size, size_t offset,
-                            dpct::queue_ptr stream) {
-    sycl_reorder_temp_buffer tmp(stream, size);
+// Reorder a 2D tensor of quantized blocks in place into [head of every block][tail of every block] without a
+// full-size temporary copy. The head is the contiguous `head_bytes` at `head_offset` inside each block. Heads only
+// move towards the front, so they are packed in chunks that never overwrite unread blocks. The tail overlaps
+// blocks that are not packed yet, so `write_tail(x, ib, tail, nblocks)` first writes every block's remaining fields
+// (in their final layout, relative to the start of the tail) into a side buffer that is copied to the end afterwards.
+template <typename block_t, typename tail_fn_t>
+static bool reorder_qw_planar(const char * name, uint8_t * data_device, size_t size, size_t head_offset,
+                              size_t head_bytes, tail_fn_t write_tail, dpct::queue_ptr stream) {
+    GGML_ASSERT(size % sizeof(block_t) == 0);
+    GGML_ASSERT(head_offset + head_bytes <= sizeof(block_t));
+
+    const int64_t nblocks    = size / sizeof(block_t);
+    const size_t  tail_bytes = (sizeof(block_t) - head_bytes) * nblocks;
+
+    sycl_reorder_temp_buffer tail_tmp(stream, tail_bytes);
+    if (!tail_tmp) {
+        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", name, tail_bytes);
+        return false;
+    }
+    uint8_t * tail_buf = static_cast<uint8_t *>(tail_tmp.ptr);
+
+    const size_t  max_chunk_bytes = 32u << 20;
+    const int64_t chunk_blocks    = std::max<int64_t>(1, std::min<int64_t>(nblocks, max_chunk_bytes / sizeof(block_t)));
+    sycl_reorder_temp_buffer tmp(stream, chunk_blocks * sizeof(block_t));
     if (!tmp) {
-        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, size);
+        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", name,
+                      (size_t) chunk_blocks * sizeof(block_t));
         return false;
     }
     uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
 
-    sycl::event copy_event;
-    SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, size)));
+    sycl::event event = stream->parallel_for(nblocks, [=](auto i) {
+        const int64_t ib = i;
+        write_tail(((const block_t *) data_device)[ib], ib, tail_buf, nblocks);
+    });
     if (!g_ggml_sycl_use_async_mem_op) {
-        copy_event.wait();
+        event.wait_and_throw();
     }
 
-    GGML_ASSERT((size % sizeof(block_q4_0) == 0));
-    GGML_ASSERT((offset % sizeof(block_q4_0) == 0));
-    int offset_blks = offset / sizeof(block_q4_0);
-    auto qs_ptr      = data_device + offset_blks * QK4_0 / 2;
-    auto d_ptr = (sycl::half*)(qs_ptr + ncols * nrows / 2) + offset_blks;
+    for (int64_t b0 = 0; b0 < nblocks; b0 += chunk_blocks) {
+        const int64_t n_chunk = std::min(chunk_blocks, nblocks - b0);
 
-    auto reorder_event = stream->parallel_for(
-        size / sizeof(block_q4_0),
-            [=](auto i) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-            const block_q4_0* x = (const block_q4_0*)tmp_buf;
-            const int ib = i;
+        SYCL_CHECK(CHECK_TRY_ERROR(event = stream->memcpy(tmp_buf, data_device + b0 * sizeof(block_t),
+                                                          n_chunk * sizeof(block_t))));
+        if (!g_ggml_sycl_use_async_mem_op) {
+            event.wait();
+        }
 
-            for (int j = 0; j < QK4_0/2; j ++)
-            {
-                *(qs_ptr + ib * QK4_0 / 2 + j) = x[ib].qs[j];
+        uint8_t * head_ptr = data_device + b0 * head_bytes;
+        event = stream->parallel_for(n_chunk, [=](auto i) {
+            const int64_t   ib  = i;
+            const uint8_t * src = tmp_buf + ib * sizeof(block_t) + head_offset;
+            for (size_t j = 0; j < head_bytes; ++j) {
+                head_ptr[ib * head_bytes + j] = src[j];
             }
-            *(d_ptr + ib) = x[ib].d;
         });
+        if (!g_ggml_sycl_use_async_mem_op) {
+            event.wait_and_throw();
+        }
+    }
+
+    SYCL_CHECK(CHECK_TRY_ERROR(event = stream->memcpy(data_device + head_bytes * nblocks, tail_buf, tail_bytes)));
     if (!g_ggml_sycl_use_async_mem_op) {
-        reorder_event.wait_and_throw();
+        event.wait();
     }
     return true;
 }
 
+static bool reorder_qw_q4_0(uint8_t * data_device, const int ncols, const int nrows, size_t size, size_t offset,
+                            dpct::queue_ptr stream) {
+    GGML_ASSERT(offset == 0);
+    GGML_ASSERT((size_t) ncols * nrows / 2 == size / sizeof(block_q4_0) * (QK4_0 / 2));
+    return reorder_qw_planar<block_q4_0>(__func__, data_device, size, offsetof(block_q4_0, qs), QK4_0 / 2,
+        [](const block_q4_0 & x, int64_t ib, uint8_t * tail, int64_t) {
+            ((sycl::half *) tail)[ib] = x.d;
+        }, stream);
+}
+
 static bool reorder_qw_q8_0(uint8_t * data_device, const int ncols, const int nrows, size_t size, size_t offset,
                             dpct::queue_ptr stream) {
-    sycl_reorder_temp_buffer tmp(stream, size);
-    if (!tmp) {
-        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, size);
-        return false;
-    }
-    uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
-
-    sycl::event copy_event;
-    SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, size)));
-    if (!g_ggml_sycl_use_async_mem_op) {
-        copy_event.wait();
-    }
-
-    GGML_ASSERT((size % sizeof(block_q8_0) == 0));
-    GGML_ASSERT((offset % sizeof(block_q8_0) == 0));
-    int offset_blks = offset / sizeof(block_q8_0);
-    auto qs_ptr = data_device + offset_blks * QK8_0;
-    auto d_ptr = (sycl::half*)(qs_ptr + ncols * nrows) + offset_blks;
-
-    auto reorder_event = stream->parallel_for(
-        size / sizeof(block_q8_0),
-            [=](auto i) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-            const block_q8_0* x = (const block_q8_0*)tmp_buf;
-            const int ib = i;
-
-            for (int j = 0; j < QK8_0; j++)
-            {
-                *((int8_t*)qs_ptr + ib * QK8_0 + j) = x[ib].qs[j];
-            }
-            *(d_ptr + ib) = x[ib].d;
-        });
-    if (!g_ggml_sycl_use_async_mem_op) {
-        reorder_event.wait_and_throw();
-    }
-    return true;
-}
-
-// Reorder each Q8_0 expert slice into a self-contained [qs][d] SoA layout.
-static bool reorder_qw_q8_0_moe(uint8_t * data_device, size_t expert_bytes, int64_t n_expert, dpct::queue_ptr stream) {
-    GGML_ASSERT(expert_bytes % sizeof(block_q8_0) == 0);
-    const int    blocks_per_expert = (int) (expert_bytes / sizeof(block_q8_0));
-    const size_t total_bytes       = expert_bytes * (size_t) n_expert;
-
-    sycl_reorder_temp_buffer tmp(stream, total_bytes);
-    if (!tmp) {
-        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, total_bytes);
-        return false;
-    }
-    uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
-
-    sycl::event copy_event;
-    SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, total_bytes)));
-    if (!g_ggml_sycl_use_async_mem_op) {
-        copy_event.wait();
-    }
-
-    const int total_blocks = blocks_per_expert * (int) n_expert;
-    auto reorder_event = stream->parallel_for(total_blocks, [=](auto gb_) {
-        const int          gb   = gb_;
-        const int          e    = gb / blocks_per_expert;
-        const int          ib   = gb % blocks_per_expert;
-        const block_q8_0 * x    = (const block_q8_0 *) (tmp_buf + (size_t) e * expert_bytes);
-        uint8_t *          base = data_device + (size_t) e * expert_bytes;
-
-        auto * qs_ptr = base;
-        auto * d_ptr  = (sycl::half *) (qs_ptr + (size_t) blocks_per_expert * QK8_0);
-
-        for (int j = 0; j < QK8_0; ++j) {
-            *((int8_t *) qs_ptr + ib * QK8_0 + j) = x[ib].qs[j];
-        }
-        d_ptr[ib] = x[ib].d;
-    });
-    if (!g_ggml_sycl_use_async_mem_op) {
-        reorder_event.wait_and_throw();
-    }
-    return true;
+    GGML_ASSERT(offset == 0);
+    GGML_ASSERT((size_t) ncols * nrows == size / sizeof(block_q8_0) * QK8_0);
+    return reorder_qw_planar<block_q8_0>(__func__, data_device, size, offsetof(block_q8_0, qs), QK8_0,
+        [](const block_q8_0 & x, int64_t ib, uint8_t * tail, int64_t) {
+            ((sycl::half *) tail)[ib] = x.d;
+        }, stream);
 }
 
 static bool reorder_qw_q4_k(uint8_t * data_device, size_t size, size_t offset, dpct::queue_ptr stream) {
-    GGML_ASSERT(size % sizeof(block_q4_K) == 0);
-    GGML_ASSERT(offset % sizeof(block_q4_K) == 0);
+    GGML_ASSERT(offset == 0);
+    return reorder_qw_planar<block_q4_K>(__func__, data_device, size, offsetof(block_q4_K, qs), QK_K / 2,
+        [](const block_q4_K & x, int64_t ib, uint8_t * tail, int64_t nblocks) {
+            uint8_t * scales = tail;
+            for (int j = 0; j < K_SCALE_SIZE; ++j) {
+                scales[ib * K_SCALE_SIZE + j] = x.scales[j];
+            }
+            ((sycl::half2 *) (scales + K_SCALE_SIZE * nblocks))[ib] = x.dm;
+        }, stream);
+}
 
-    const int nblocks = size / sizeof(block_q4_K);
+// Reorder each Q8_0 expert slice into a self-contained [qs][d] SoA layout, a few experts at a time.
+static bool reorder_qw_q8_0_moe(uint8_t * data_device, size_t expert_bytes, int64_t n_expert, dpct::queue_ptr stream) {
+    GGML_ASSERT(expert_bytes % sizeof(block_q8_0) == 0);
+    const int     blocks_per_expert = (int) (expert_bytes / sizeof(block_q8_0));
+    const size_t  max_chunk_bytes   = 32u << 20;
+    const int64_t chunk_experts     = std::max<int64_t>(1, std::min<int64_t>(n_expert, (int64_t) (max_chunk_bytes / expert_bytes)));
 
-    sycl_reorder_temp_buffer tmp(stream, size);
+    sycl_reorder_temp_buffer tmp(stream, (size_t) chunk_experts * expert_bytes);
     if (!tmp) {
-        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, size);
+        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__,
+                      (size_t) chunk_experts * expert_bytes);
         return false;
     }
     uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
 
-    sycl::event copy_event;
-    SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, size)));
-    if (!g_ggml_sycl_use_async_mem_op) {
-        copy_event.wait();
-    }
+    for (int64_t e0 = 0; e0 < n_expert; e0 += chunk_experts) {
+        const int64_t n_chunk = std::min(chunk_experts, n_expert - e0);
+        uint8_t *     chunk   = data_device + (size_t) e0 * expert_bytes;
 
-    auto * qs_ptr     = data_device;
-    auto * scales_ptr = qs_ptr + QK_K / 2 * nblocks;
-    auto * dm_ptr     = (sycl::half2 *) (scales_ptr + K_SCALE_SIZE * nblocks);
-
-    auto reorder_event = stream->parallel_for(nblocks, [=](auto i) {
-        const block_q4_K * x  = (const block_q4_K *) tmp_buf;
-        const int          ib = i;
-
-        for (int j = 0; j < QK_K / 2; ++j) {
-            qs_ptr[ib * (QK_K / 2) + j] = x[ib].qs[j];
+        sycl::event copy_event;
+        SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, chunk, (size_t) n_chunk * expert_bytes)));
+        if (!g_ggml_sycl_use_async_mem_op) {
+            copy_event.wait();
         }
 
-        for (int j = 0; j < K_SCALE_SIZE; ++j) {
-            scales_ptr[ib * K_SCALE_SIZE + j] = x[ib].scales[j];
-        }
+        const int total_blocks = blocks_per_expert * (int) n_chunk;
+        auto reorder_event = stream->parallel_for(total_blocks, [=](auto gb_) {
+            const int          gb   = gb_;
+            const int          e    = gb / blocks_per_expert;
+            const int          ib   = gb % blocks_per_expert;
+            const block_q8_0 * x    = (const block_q8_0 *) (tmp_buf + (size_t) e * expert_bytes);
+            uint8_t *          base = chunk + (size_t) e * expert_bytes;
 
-        dm_ptr[ib] = x[ib].dm;
-    });
-    if (!g_ggml_sycl_use_async_mem_op) {
-        reorder_event.wait_and_throw();
+            auto * qs_ptr = base;
+            auto * d_ptr  = (sycl::half *) (qs_ptr + (size_t) blocks_per_expert * QK8_0);
+
+            for (int j = 0; j < QK8_0; ++j) {
+                *((int8_t *) qs_ptr + ib * QK8_0 + j) = x[ib].qs[j];
+            }
+            d_ptr[ib] = x[ib].d;
+        });
+        if (!g_ggml_sycl_use_async_mem_op) {
+            reorder_event.wait_and_throw();
+        }
     }
     return true;
 }
 
 // Reorder each expert slice into a self-contained SoA layout.
+// Like the MXFP4 version below, this goes a few experts at a time through a small temporary:
+// a whole-tensor temporary (GBs for large MoE models) can exceed the VRAM left on a nearly full card.
 static bool reorder_qw_q4_k_moe(uint8_t * data_device, size_t expert_bytes, int64_t n_expert, dpct::queue_ptr stream) {
     GGML_ASSERT(expert_bytes % sizeof(block_q4_K) == 0);
-    const int    blocks_per_expert = (int) (expert_bytes / sizeof(block_q4_K));
-    const size_t total_bytes       = expert_bytes * (size_t) n_expert;
+    const int     blocks_per_expert = (int) (expert_bytes / sizeof(block_q4_K));
+    const size_t  max_chunk_bytes   = 32u << 20;
+    const int64_t chunk_experts     = std::max<int64_t>(1, std::min<int64_t>(n_expert, (int64_t) (max_chunk_bytes / expert_bytes)));
 
-    sycl_reorder_temp_buffer tmp(stream, total_bytes);
+    sycl_reorder_temp_buffer tmp(stream, (size_t) chunk_experts * expert_bytes);
     if (!tmp) {
-        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, total_bytes);
+        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__,
+                      (size_t) chunk_experts * expert_bytes);
         return false;
     }
     uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
 
-    sycl::event copy_event;
-    SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, total_bytes)));
-    if (!g_ggml_sycl_use_async_mem_op) {
-        copy_event.wait();
-    }
+    // the queue is in-order: each chunk's copy into tmp_buf waits for the previous chunk's kernel
+    for (int64_t e0 = 0; e0 < n_expert; e0 += chunk_experts) {
+        const int64_t n_chunk = std::min(chunk_experts, n_expert - e0);
+        uint8_t *     chunk   = data_device + (size_t) e0 * expert_bytes;
 
-    const int total_blocks = blocks_per_expert * (int) n_expert;
-    auto reorder_event = stream->parallel_for(total_blocks, [=](auto gb_) {
-        const int          gb   = gb_;
-        const int          e    = gb / blocks_per_expert;
-        const int          ib   = gb % blocks_per_expert;
-        const block_q4_K * x    = (const block_q4_K *) (tmp_buf + (size_t) e * expert_bytes);
-        uint8_t *          base = data_device + (size_t) e * expert_bytes;
-
-        auto * qs_ptr     = base;
-        auto * scales_ptr = qs_ptr + QK_K / 2 * blocks_per_expert;
-        auto * dm_ptr     = (sycl::half2 *) (scales_ptr + K_SCALE_SIZE * blocks_per_expert);
-
-        for (int j = 0; j < QK_K / 2; ++j) {
-            qs_ptr[ib * (QK_K / 2) + j] = x[ib].qs[j];
+        sycl::event copy_event;
+        SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, chunk, (size_t) n_chunk * expert_bytes)));
+        if (!g_ggml_sycl_use_async_mem_op) {
+            copy_event.wait();
         }
-        for (int j = 0; j < K_SCALE_SIZE; ++j) {
-            scales_ptr[ib * K_SCALE_SIZE + j] = x[ib].scales[j];
+
+        const int total_blocks = blocks_per_expert * (int) n_chunk;
+        auto reorder_event = stream->parallel_for(total_blocks, [=](auto gb_) {
+            const int          gb   = gb_;
+            const int          e    = gb / blocks_per_expert;
+            const int          ib   = gb % blocks_per_expert;
+            const block_q4_K * x    = (const block_q4_K *) (tmp_buf + (size_t) e * expert_bytes);
+            uint8_t *          base = chunk + (size_t) e * expert_bytes;
+
+            auto * qs_ptr     = base;
+            auto * scales_ptr = qs_ptr + QK_K / 2 * blocks_per_expert;
+            auto * dm_ptr     = (sycl::half2 *) (scales_ptr + K_SCALE_SIZE * blocks_per_expert);
+
+            for (int j = 0; j < QK_K / 2; ++j) {
+                qs_ptr[ib * (QK_K / 2) + j] = x[ib].qs[j];
+            }
+            for (int j = 0; j < K_SCALE_SIZE; ++j) {
+                scales_ptr[ib * K_SCALE_SIZE + j] = x[ib].scales[j];
+            }
+            dm_ptr[ib] = x[ib].dm;
+        });
+        if (!g_ggml_sycl_use_async_mem_op) {
+            reorder_event.wait_and_throw();
         }
-        dm_ptr[ib] = x[ib].dm;
-    });
-    if (!g_ggml_sycl_use_async_mem_op) {
-        reorder_event.wait_and_throw();
     }
     return true;
 }
 
-// Reorder each Q5_K expert slice into [qs][qh][scales][dm].
+// Reorder each Q5_K expert slice into [qs][qh][scales][dm], a few experts at a time (see Q4_K above).
 static bool reorder_qw_q5_k_moe(uint8_t * data_device, size_t expert_bytes, int64_t n_expert, dpct::queue_ptr stream) {
     GGML_ASSERT(expert_bytes % sizeof(block_q5_K) == 0);
-    const int    blocks_per_expert = (int) (expert_bytes / sizeof(block_q5_K));
-    const size_t total_bytes       = expert_bytes * (size_t) n_expert;
+    const int     blocks_per_expert = (int) (expert_bytes / sizeof(block_q5_K));
+    const size_t  max_chunk_bytes   = 32u << 20;
+    const int64_t chunk_experts     = std::max<int64_t>(1, std::min<int64_t>(n_expert, (int64_t) (max_chunk_bytes / expert_bytes)));
 
-    sycl_reorder_temp_buffer tmp(stream, total_bytes);
+    sycl_reorder_temp_buffer tmp(stream, (size_t) chunk_experts * expert_bytes);
     if (!tmp) {
-        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, total_bytes);
+        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__,
+                      (size_t) chunk_experts * expert_bytes);
         return false;
     }
     uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
 
-    sycl::event copy_event;
-    SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, total_bytes)));
-    if (!g_ggml_sycl_use_async_mem_op) {
-        copy_event.wait();
-    }
+    for (int64_t e0 = 0; e0 < n_expert; e0 += chunk_experts) {
+        const int64_t n_chunk = std::min(chunk_experts, n_expert - e0);
+        uint8_t *     chunk   = data_device + (size_t) e0 * expert_bytes;
 
-    const int total_blocks = blocks_per_expert * (int) n_expert;
-    auto reorder_event = stream->parallel_for(total_blocks, [=](auto gb_) {
-        const int          gb   = gb_;
-        const int          e    = gb / blocks_per_expert;
-        const int          ib   = gb % blocks_per_expert;
-        const block_q5_K * x    = (const block_q5_K *) (tmp_buf + (size_t) e * expert_bytes);
-        uint8_t *          base = data_device + (size_t) e * expert_bytes;
+        sycl::event copy_event;
+        SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, chunk, (size_t) n_chunk * expert_bytes)));
+        if (!g_ggml_sycl_use_async_mem_op) {
+            copy_event.wait();
+        }
 
-        auto * qs_ptr     = base;
-        auto * qh_ptr     = qs_ptr + (QK_K / 2) * blocks_per_expert;
-        auto * scales_ptr = qh_ptr + (QK_K / 8) * blocks_per_expert;
-        auto * dm_ptr     = (sycl::half2 *) (scales_ptr + K_SCALE_SIZE * blocks_per_expert);
+        const int total_blocks = blocks_per_expert * (int) n_chunk;
+        auto reorder_event = stream->parallel_for(total_blocks, [=](auto gb_) {
+            const int          gb   = gb_;
+            const int          e    = gb / blocks_per_expert;
+            const int          ib   = gb % blocks_per_expert;
+            const block_q5_K * x    = (const block_q5_K *) (tmp_buf + (size_t) e * expert_bytes);
+            uint8_t *          base = chunk + (size_t) e * expert_bytes;
 
-        for (int j = 0; j < QK_K / 2; ++j) {
-            qs_ptr[ib * (QK_K / 2) + j] = x[ib].qs[j];
+            auto * qs_ptr     = base;
+            auto * qh_ptr     = qs_ptr + (QK_K / 2) * blocks_per_expert;
+            auto * scales_ptr = qh_ptr + (QK_K / 8) * blocks_per_expert;
+            auto * dm_ptr     = (sycl::half2 *) (scales_ptr + K_SCALE_SIZE * blocks_per_expert);
+
+            for (int j = 0; j < QK_K / 2; ++j) {
+                qs_ptr[ib * (QK_K / 2) + j] = x[ib].qs[j];
+            }
+            for (int j = 0; j < QK_K / 8; ++j) {
+                qh_ptr[ib * (QK_K / 8) + j] = x[ib].qh[j];
+            }
+            for (int j = 0; j < K_SCALE_SIZE; ++j) {
+                scales_ptr[ib * K_SCALE_SIZE + j] = x[ib].scales[j];
+            }
+            dm_ptr[ib] = x[ib].dm;
+        });
+        if (!g_ggml_sycl_use_async_mem_op) {
+            reorder_event.wait_and_throw();
         }
-        for (int j = 0; j < QK_K / 8; ++j) {
-            qh_ptr[ib * (QK_K / 8) + j] = x[ib].qh[j];
-        }
-        for (int j = 0; j < K_SCALE_SIZE; ++j) {
-            scales_ptr[ib * K_SCALE_SIZE + j] = x[ib].scales[j];
-        }
-        dm_ptr[ib] = x[ib].dm;
-    });
-    if (!g_ggml_sycl_use_async_mem_op) {
-        reorder_event.wait_and_throw();
     }
     return true;
 }
 
-// Reorder each Q6_K expert slice into [ql][qh][scales][d].
+// Reorder each Q6_K expert slice into [ql][qh][scales][d], a few experts at a time (see Q4_K above).
 static bool reorder_qw_q6_k_moe(uint8_t * data_device, size_t expert_bytes, int64_t n_expert, dpct::queue_ptr stream) {
     GGML_ASSERT(expert_bytes % sizeof(block_q6_K) == 0);
-    const int    blocks_per_expert = (int) (expert_bytes / sizeof(block_q6_K));
-    const size_t total_bytes       = expert_bytes * (size_t) n_expert;
+    const int     blocks_per_expert = (int) (expert_bytes / sizeof(block_q6_K));
+    const size_t  max_chunk_bytes   = 32u << 20;
+    const int64_t chunk_experts     = std::max<int64_t>(1, std::min<int64_t>(n_expert, (int64_t) (max_chunk_bytes / expert_bytes)));
 
-    sycl_reorder_temp_buffer tmp(stream, total_bytes);
+    sycl_reorder_temp_buffer tmp(stream, (size_t) chunk_experts * expert_bytes);
     if (!tmp) {
-        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, total_bytes);
+        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__,
+                      (size_t) chunk_experts * expert_bytes);
         return false;
     }
     uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
 
-    sycl::event copy_event;
-    SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, total_bytes)));
-    if (!g_ggml_sycl_use_async_mem_op) {
-        copy_event.wait();
-    }
+    for (int64_t e0 = 0; e0 < n_expert; e0 += chunk_experts) {
+        const int64_t n_chunk = std::min(chunk_experts, n_expert - e0);
+        uint8_t *     chunk   = data_device + (size_t) e0 * expert_bytes;
 
-    const int total_blocks = blocks_per_expert * (int) n_expert;
-    auto reorder_event = stream->parallel_for(total_blocks, [=](auto gb_) {
-        const int          gb   = gb_;
-        const int          e    = gb / blocks_per_expert;
-        const int          ib   = gb % blocks_per_expert;
-        const block_q6_K * x    = (const block_q6_K *) (tmp_buf + (size_t) e * expert_bytes);
-        uint8_t *          base = data_device + (size_t) e * expert_bytes;
+        sycl::event copy_event;
+        SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, chunk, (size_t) n_chunk * expert_bytes)));
+        if (!g_ggml_sycl_use_async_mem_op) {
+            copy_event.wait();
+        }
 
-        auto * ql_ptr     = base;
-        auto * qh_ptr     = ql_ptr + (QK_K / 2) * blocks_per_expert;
-        auto * scales_ptr = qh_ptr + (QK_K / 4) * blocks_per_expert;
-        auto * d_ptr      = (sycl::half *) (scales_ptr + (QK_K / 16) * blocks_per_expert);
+        const int total_blocks = blocks_per_expert * (int) n_chunk;
+        auto reorder_event = stream->parallel_for(total_blocks, [=](auto gb_) {
+            const int          gb   = gb_;
+            const int          e    = gb / blocks_per_expert;
+            const int          ib   = gb % blocks_per_expert;
+            const block_q6_K * x    = (const block_q6_K *) (tmp_buf + (size_t) e * expert_bytes);
+            uint8_t *          base = chunk + (size_t) e * expert_bytes;
 
-        for (int j = 0; j < QK_K / 2; ++j) {
-            ql_ptr[ib * (QK_K / 2) + j] = x[ib].ql[j];
+            auto * ql_ptr     = base;
+            auto * qh_ptr     = ql_ptr + (QK_K / 2) * blocks_per_expert;
+            auto * scales_ptr = qh_ptr + (QK_K / 4) * blocks_per_expert;
+            auto * d_ptr      = (sycl::half *) (scales_ptr + (QK_K / 16) * blocks_per_expert);
+
+            for (int j = 0; j < QK_K / 2; ++j) {
+                ql_ptr[ib * (QK_K / 2) + j] = x[ib].ql[j];
+            }
+            for (int j = 0; j < QK_K / 4; ++j) {
+                qh_ptr[ib * (QK_K / 4) + j] = x[ib].qh[j];
+            }
+            for (int j = 0; j < QK_K / 16; ++j) {
+                scales_ptr[ib * (QK_K / 16) + j] = x[ib].scales[j];
+            }
+            d_ptr[ib] = x[ib].d;
+        });
+        if (!g_ggml_sycl_use_async_mem_op) {
+            reorder_event.wait_and_throw();
         }
-        for (int j = 0; j < QK_K / 4; ++j) {
-            qh_ptr[ib * (QK_K / 4) + j] = x[ib].qh[j];
-        }
-        for (int j = 0; j < QK_K / 16; ++j) {
-            scales_ptr[ib * (QK_K / 16) + j] = x[ib].scales[j];
-        }
-        d_ptr[ib] = x[ib].d;
-    });
-    if (!g_ggml_sycl_use_async_mem_op) {
-        reorder_event.wait_and_throw();
     }
     return true;
 }
@@ -4839,278 +4905,89 @@ static bool reorder_qw_mxfp4_moe(uint8_t * data_device, size_t expert_bytes, int
 }
 
 static bool reorder_qw_q2_k(uint8_t * data_device, size_t size, size_t offset, dpct::queue_ptr stream) {
-    GGML_ASSERT(size % sizeof(block_q2_K) == 0);
-    GGML_ASSERT(offset % sizeof(block_q2_K) == 0);
-
-    const int nblocks = size / sizeof(block_q2_K);
-
-    sycl_reorder_temp_buffer tmp(stream, size);
-    if (!tmp) {
-        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, size);
-        return false;
-    }
-    uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
-
-    sycl::event copy_event;
-    SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, size)));
-    if (!g_ggml_sycl_use_async_mem_op) {
-        copy_event.wait();
-    }
-
-    auto *        qs_ptr     = data_device;
-    auto *        scales_ptr = qs_ptr + (QK_K / 4) * nblocks;
-    sycl::half2 * dm_ptr     = (sycl::half2 *) (scales_ptr + (QK_K / 16) * nblocks);
-
-    auto reorder_event = stream->parallel_for(nblocks, [=](auto i) {
-        const block_q2_K * x  = (const block_q2_K *) tmp_buf;
-        const int          ib = i;
-
-        for (int j = 0; j < QK_K / 4; ++j) {
-            qs_ptr[ib * (QK_K / 4) + j] = x[ib].qs[j];
-        }
-
-        for (int j = 0; j < QK_K / 16; ++j) {
-            scales_ptr[ib * (QK_K / 16) + j] = x[ib].scales[j];
-        }
-
-        dm_ptr[ib] = x[ib].dm;
-    });
-    if (!g_ggml_sycl_use_async_mem_op) {
-        reorder_event.wait_and_throw();
-    }
-    return true;
+    GGML_ASSERT(offset == 0);
+    return reorder_qw_planar<block_q2_K>(__func__, data_device, size, offsetof(block_q2_K, qs), QK_K / 4,
+        [](const block_q2_K & x, int64_t ib, uint8_t * tail, int64_t nblocks) {
+            uint8_t * scales = tail;
+            for (int j = 0; j < QK_K / 16; ++j) {
+                scales[ib * (QK_K / 16) + j] = x.scales[j];
+            }
+            ((sycl::half2 *) (scales + (QK_K / 16) * nblocks))[ib] = x.dm;
+        }, stream);
 }
 
 static bool reorder_qw_q3_k(uint8_t * data_device, size_t size, size_t offset, dpct::queue_ptr stream) {
-    GGML_ASSERT(size % sizeof(block_q3_K) == 0);
-    GGML_ASSERT(offset % sizeof(block_q3_K) == 0);
-
-    const int nblocks = size / sizeof(block_q3_K);
-
-    sycl_reorder_temp_buffer tmp(stream, size);
-    if (!tmp) {
-        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, size);
-        return false;
-    }
-    uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
-
-    sycl::event copy_event;
-    SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, size)));
-    if (!g_ggml_sycl_use_async_mem_op) {
-        copy_event.wait();
-    }
-
-    auto *       qs_ptr     = data_device;
-    auto *       hmask_ptr  = qs_ptr + (QK_K / 4) * nblocks;
-    auto *       scales_ptr = hmask_ptr + (QK_K / 8) * nblocks;
-    sycl::half * d_ptr      = (sycl::half *) (scales_ptr + 12 * nblocks);
-
-    auto reorder_event = stream->parallel_for(nblocks, [=](auto i) {
-        const block_q3_K * x  = (const block_q3_K *) tmp_buf;
-        const int          ib = i;
-
-        for (int j = 0; j < QK_K / 4; ++j) {
-            qs_ptr[ib * (QK_K / 4) + j] = x[ib].qs[j];
-        }
-
-        for (int j = 0; j < QK_K / 8; ++j) {
-            hmask_ptr[ib * (QK_K / 8) + j] = x[ib].hmask[j];
-        }
-
-        for (int j = 0; j < 12; ++j) {
-            scales_ptr[ib * 12 + j] = x[ib].scales[j];
-        }
-
-        d_ptr[ib] = x[ib].d;
-    });
-    if (!g_ggml_sycl_use_async_mem_op) {
-        reorder_event.wait_and_throw();
-    }
-    return true;
+    GGML_ASSERT(offset == 0);
+    return reorder_qw_planar<block_q3_K>(__func__, data_device, size, offsetof(block_q3_K, qs), QK_K / 4,
+        [](const block_q3_K & x, int64_t ib, uint8_t * tail, int64_t nblocks) {
+            uint8_t * hmask  = tail;
+            uint8_t * scales = hmask + (QK_K / 8) * nblocks;
+            for (int j = 0; j < QK_K / 8; ++j) {
+                hmask[ib * (QK_K / 8) + j] = x.hmask[j];
+            }
+            for (int j = 0; j < 12; ++j) {
+                scales[ib * 12 + j] = x.scales[j];
+            }
+            ((sycl::half *) (scales + 12 * nblocks))[ib] = x.d;
+        }, stream);
 }
 
 static bool reorder_qw_q5_k(uint8_t * data_device, size_t size, size_t offset, dpct::queue_ptr stream) {
-    GGML_ASSERT(size % sizeof(block_q5_K) == 0);
-    GGML_ASSERT(offset % sizeof(block_q5_K) == 0);
-
-    const int nblocks = size / sizeof(block_q5_K);
-
-    sycl_reorder_temp_buffer tmp(stream, size);
-    if (!tmp) {
-        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, size);
-        return false;
-    }
-    uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
-
-    sycl::event copy_event;
-    SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, size)));
-    if (!g_ggml_sycl_use_async_mem_op) {
-        copy_event.wait();
-    }
-
-    auto * qs_ptr     = data_device;
-    auto * qh_ptr     = qs_ptr + (QK_K / 2) * nblocks;
-    auto * scales_ptr = qh_ptr + (QK_K / 8) * nblocks;
-    auto * dm_ptr     = (sycl::half2 *) (scales_ptr + K_SCALE_SIZE * nblocks);
-
-    auto reorder_event = stream->parallel_for(nblocks, [=](auto i) {
-        const block_q5_K * x  = (const block_q5_K *) tmp_buf;
-        const int          ib = i;
-
-        for (int j = 0; j < QK_K / 2; ++j) {
-            qs_ptr[ib * (QK_K / 2) + j] = x[ib].qs[j];
-        }
-
-        for (int j = 0; j < QK_K / 8; ++j) {
-            qh_ptr[ib * (QK_K / 8) + j] = x[ib].qh[j];
-        }
-
-        for (int j = 0; j < K_SCALE_SIZE; ++j) {
-            scales_ptr[ib * K_SCALE_SIZE + j] = x[ib].scales[j];
-        }
-
-        dm_ptr[ib] = x[ib].dm;
-    });
-    if (!g_ggml_sycl_use_async_mem_op) {
-        reorder_event.wait_and_throw();
-    }
-    return true;
+    GGML_ASSERT(offset == 0);
+    return reorder_qw_planar<block_q5_K>(__func__, data_device, size, offsetof(block_q5_K, qs), QK_K / 2,
+        [](const block_q5_K & x, int64_t ib, uint8_t * tail, int64_t nblocks) {
+            uint8_t * qh     = tail;
+            uint8_t * scales = qh + (QK_K / 8) * nblocks;
+            for (int j = 0; j < QK_K / 8; ++j) {
+                qh[ib * (QK_K / 8) + j] = x.qh[j];
+            }
+            for (int j = 0; j < K_SCALE_SIZE; ++j) {
+                scales[ib * K_SCALE_SIZE + j] = x.scales[j];
+            }
+            ((sycl::half2 *) (scales + K_SCALE_SIZE * nblocks))[ib] = x.dm;
+        }, stream);
 }
 
 static bool reorder_qw_q6_k(uint8_t * data_device, size_t size, size_t offset, dpct::queue_ptr stream) {
-    GGML_ASSERT(size % sizeof(block_q6_K) == 0);
-    GGML_ASSERT(offset % sizeof(block_q6_K) == 0);
-
-    const int nblocks = size / sizeof(block_q6_K);
-
-    sycl_reorder_temp_buffer tmp(stream, size);
-    if (!tmp) {
-        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, size);
-        return false;
-    }
-    uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
-
-    sycl::event copy_event;
-    SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, size)));
-    if (!g_ggml_sycl_use_async_mem_op) {
-        copy_event.wait();
-    }
-
-    auto *       ql_ptr     = data_device;
-    auto *       qh_ptr     = ql_ptr + (QK_K / 2) * nblocks;
-    auto *       scales_ptr = qh_ptr + (QK_K / 4) * nblocks;
-    sycl::half * dm_ptr     = (sycl::half *) (scales_ptr + (QK_K / 16) * nblocks);
-
-    auto reorder_event = stream->parallel_for(nblocks, [=](auto i) {
-        const block_q6_K * x  = (const block_q6_K *) tmp_buf;
-        const int          ib = i;
-
-        const uint8_t * ql              = x[ib].ql;
-        const uint8_t * qh              = x[ib].qh;
-        uint8_t *       base_ql_ptr     = ql_ptr + (QK_K / 2) * ib;
-        uint8_t *       base_qh_ptr     = qh_ptr + (QK_K / 4) * ib;
-        uint8_t *       base_scales_ptr = scales_ptr + (QK_K / 16) * ib;
-
-        for (int j = 0; j < QK_K / 2; ++j) {
-            base_ql_ptr[j] = ql[j];
-        }
-        for (int j = 0; j < QK_K / 4; ++j) {
-            base_qh_ptr[j] = qh[j];
-        }
-
-        for (int j = 0; j < QK_K / 16; ++j) {
-            base_scales_ptr[j] = x[ib].scales[j];
-        }
-
-        dm_ptr[ib] = x[ib].d;
-    });
-    if (!g_ggml_sycl_use_async_mem_op) {
-        reorder_event.wait_and_throw();
-    }
-    return true;
+    GGML_ASSERT(offset == 0);
+    return reorder_qw_planar<block_q6_K>(__func__, data_device, size, offsetof(block_q6_K, ql), QK_K / 2,
+        [](const block_q6_K & x, int64_t ib, uint8_t * tail, int64_t nblocks) {
+            uint8_t * qh     = tail;
+            uint8_t * scales = qh + (QK_K / 4) * nblocks;
+            for (int j = 0; j < QK_K / 4; ++j) {
+                qh[ib * (QK_K / 4) + j] = x.qh[j];
+            }
+            for (int j = 0; j < QK_K / 16; ++j) {
+                scales[ib * (QK_K / 16) + j] = x.scales[j];
+            }
+            ((sycl::half *) (scales + (QK_K / 16) * nblocks))[ib] = x.d;
+        }, stream);
 }
 
 static bool reorder_qw_iq3_s(uint8_t * data_device, size_t size, dpct::queue_ptr stream) {
-    GGML_ASSERT(size % sizeof(block_iq3_s) == 0);
-
-    const int nblocks = size / sizeof(block_iq3_s);
-    sycl_reorder_temp_buffer tmp(stream, size);
-    if (!tmp) {
-        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, size);
-        return false;
-    }
-    uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
-
-    sycl::event copy_event;
-    SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, size)));
-    if (!g_ggml_sycl_use_async_mem_op) {
-        copy_event.wait();
-    }
-
-    auto * qs_ptr       = data_device;
-    auto * qh_ptr       = qs_ptr + (QK_K / 4) * nblocks;
-    auto * signs_ptr    = qh_ptr + (QK_K / 32) * nblocks;
-    auto * metadata_ptr = signs_ptr + (QK_K / 8) * nblocks;
-
-    auto reorder_event = stream->parallel_for(nblocks, [=](auto i) {
-        const block_iq3_s * x = (const block_iq3_s *) tmp_buf;
-        const int ib = i;
-
-        for (int j = 0; j < QK_K / 4; ++j) {
-            qs_ptr[ib * (QK_K / 4) + j] = x[ib].qs[j];
-        }
-        for (int j = 0; j < QK_K / 32; ++j) {
-            qh_ptr[ib * (QK_K / 32) + j] = x[ib].qh[j];
-        }
-        for (int j = 0; j < QK_K / 8; ++j) {
-            signs_ptr[ib * (QK_K / 8) + j] = x[ib].signs[j];
-        }
-
-        uint8_t * metadata = metadata_ptr + ib * (sizeof(ggml_half) + IQ3S_N_SCALE);
-        *reinterpret_cast<ggml_half *>(metadata) = x[ib].d;
-        for (int j = 0; j < IQ3S_N_SCALE; ++j) {
-            metadata[sizeof(ggml_half) + j] = x[ib].scales[j];
-        }
-    });
-    if (!g_ggml_sycl_use_async_mem_op) {
-        reorder_event.wait_and_throw();
-    }
-    return true;
+    return reorder_qw_planar<block_iq3_s>(__func__, data_device, size, offsetof(block_iq3_s, qs), QK_K / 4,
+        [](const block_iq3_s & x, int64_t ib, uint8_t * tail, int64_t nblocks) {
+            uint8_t * qh       = tail;
+            uint8_t * signs    = qh + (QK_K / 32) * nblocks;
+            uint8_t * metadata = signs + (QK_K / 8) * nblocks + ib * (sizeof(ggml_half) + IQ3S_N_SCALE);
+            for (int j = 0; j < QK_K / 32; ++j) {
+                qh[ib * (QK_K / 32) + j] = x.qh[j];
+            }
+            for (int j = 0; j < QK_K / 8; ++j) {
+                signs[ib * (QK_K / 8) + j] = x.signs[j];
+            }
+            *reinterpret_cast<ggml_half *>(metadata) = x.d;
+            for (int j = 0; j < IQ3S_N_SCALE; ++j) {
+                metadata[sizeof(ggml_half) + j] = x.scales[j];
+            }
+        }, stream);
 }
 
 static bool reorder_qw_iq3_xxs(uint8_t * data_device, size_t size, dpct::queue_ptr stream) {
-    GGML_ASSERT(size % sizeof(block_iq3_xxs) == 0);
-
-    const int nblocks = size / sizeof(block_iq3_xxs);
-    sycl_reorder_temp_buffer tmp(stream, size);
-    if (!tmp) {
-        GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, size);
-        return false;
-    }
-    uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
-
-    sycl::event copy_event;
-    SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, size)));
-    if (!g_ggml_sycl_use_async_mem_op) {
-        copy_event.wait();
-    }
-
-    auto * qs_ptr = data_device;
-    auto * d_ptr  = reinterpret_cast<ggml_half *>(qs_ptr + (3 * QK_K / 8) * nblocks);
-
-    auto reorder_event = stream->parallel_for(nblocks, [=](auto i) {
-        const block_iq3_xxs * x = (const block_iq3_xxs *) tmp_buf;
-        const int ib = i;
-
-        for (int j = 0; j < 3 * QK_K / 8; ++j) {
-            qs_ptr[ib * (3 * QK_K / 8) + j] = x[ib].qs[j];
-        }
-        d_ptr[ib] = x[ib].d;
-    });
-    if (!g_ggml_sycl_use_async_mem_op) {
-        reorder_event.wait_and_throw();
-    }
-    return true;
+    return reorder_qw_planar<block_iq3_xxs>(__func__, data_device, size, offsetof(block_iq3_xxs, qs), 3 * QK_K / 8,
+        [](const block_iq3_xxs & x, int64_t ib, uint8_t * tail, int64_t) {
+            reinterpret_cast<ggml_half *>(tail)[ib] = x.d;
+        }, stream);
 }
 
 static bool reorder_qw(const ggml_tensor * src0, dpct::queue_ptr stream) {

ggml/src/ggml-vulkan/ggml-vulkan-types.h

@@ -848,18 +848,18 @@ struct vk_device_struct {
     vk_pipeline pipeline_set_rows_i64[2][GGML_TYPE_COUNT];
     vk_pipeline pipeline_norm_f32;
     vk_pipeline pipeline_group_norm_f32;
-    vk_pipeline pipeline_rms_norm_f32;
-    vk_pipeline pipeline_rms_norm_mul_f32;
-    vk_pipeline pipeline_rms_norm_mul_add_f32;
-    vk_pipeline pipeline_rms_norm_mul_add_mul_f32;
+    std::array<vk_pipeline, 3> pipeline_rms_norm_f32;
+    std::array<vk_pipeline, 3> pipeline_rms_norm_mul_f32;
+    std::array<vk_pipeline, 3> pipeline_rms_norm_mul_add_f32;
+    std::array<vk_pipeline, 3> pipeline_rms_norm_mul_add_mul_f32;
     vk_pipeline pipeline_rms_norm_mul_add_partials_f32;
     vk_pipeline pipeline_rms_norm_mul_add_mul_partials_f32;
     vk_pipeline pipeline_rms_norm_set_rows_f32_f32;
     vk_pipeline pipeline_rms_norm_set_rows_f32_f16;
     vk_pipeline pipeline_rms_norm_partials_f32;
     vk_pipeline pipeline_rms_norm_mul_partials_f32;
-    vk_pipeline pipeline_rms_norm_mul_rope_f32_f32;
-    vk_pipeline pipeline_rms_norm_mul_rope_f32_f16;
+    std::array<vk_pipeline, 3> pipeline_rms_norm_mul_rope_f32_f32;
+    std::array<vk_pipeline, 3> pipeline_rms_norm_mul_rope_f32_f16;
     vk_pipeline pipeline_rms_norm_back_f32;
     vk_pipeline pipeline_l2_norm_f32;
 

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

@@ -3256,21 +3256,34 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     ggml_vk_create_pipeline(device, device->pipeline_norm_f32, "norm_f32", norm_f32_len, norm_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_group_norm_f32, "group_norm_f32", group_norm_f32_len, group_norm_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1);
 
-    ggml_vk_create_pipeline(device, device->pipeline_rms_norm_f32, "rms_norm_f32", rms_norm_f32_len, rms_norm_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true);
-    ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_f32, "rms_norm_mul_f32", rms_norm_f32_len, rms_norm_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1}, 1, true);
-    ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_f32, "rms_norm_mul_add_f32", rms_norm_mul_add_f32_len, rms_norm_mul_add_f32_data, "main", 5, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 0}, 1, true);
-    ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_mul_f32, "rms_norm_mul_add_mul_f32", rms_norm_mul_add_f32_len, rms_norm_mul_add_f32_data, "main", 5, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 1}, 1, true);
+    const bool rms_subgroups = (device->vendor_id == VK_VENDOR_ID_INTEL || device->vendor_id == VK_VENDOR_ID_NVIDIA || device->vendor_id == VK_VENDOR_ID_AMD) && device->subgroup_arithmetic && device->subgroup_require_full_support;
+    const uint32_t rms_set_rows_block_size = 128;
+#define RMS_SHADER(name) (rms_subgroups ? name##_subgroup_len : name##_len), (rms_subgroups ? name##_subgroup_data : name##_data)
+    // 128/256/512 block sizes, picked by row length
+    for (uint32_t i = 0; i < 3; ++i) {
+        const uint32_t block_size = 128u << i;
+        const std::string suffix = "_" + std::to_string(block_size);
+        ggml_vk_create_pipeline2(device, device->pipeline_rms_norm_f32[i], "rms_norm_f32" + suffix, RMS_SHADER(rms_norm_f32), "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0, 0, block_size}, 1, true, rms_subgroups);
+        ggml_vk_create_pipeline2(device, device->pipeline_rms_norm_mul_f32[i], "rms_norm_mul_f32" + suffix, RMS_SHADER(rms_norm_f32), "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 0, block_size}, 1, true, rms_subgroups);
+        ggml_vk_create_pipeline2(device, device->pipeline_rms_norm_mul_add_f32[i], "rms_norm_mul_add_f32" + suffix, RMS_SHADER(rms_norm_mul_add_f32), "main", 5, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 0, block_size}, 1, true, rms_subgroups);
+        ggml_vk_create_pipeline2(device, device->pipeline_rms_norm_mul_add_mul_f32[i], "rms_norm_mul_add_mul_f32" + suffix, RMS_SHADER(rms_norm_mul_add_f32), "main", 5, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 1, block_size}, 1, true, rms_subgroups);
+    }
     ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_partials_f32, "rms_norm_mul_add_partials_f32", rms_norm_mul_add_partials_f32_len, rms_norm_mul_add_partials_f32_data, "main", 6, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 0}, 1, true);
     ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_mul_partials_f32, "rms_norm_mul_add_mul_partials_f32", rms_norm_mul_add_partials_f32_len, rms_norm_mul_add_partials_f32_data, "main", 6, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 1}, 1, true);
-    ggml_vk_create_pipeline(device, device->pipeline_rms_norm_set_rows_f32_f32, "rms_norm_set_rows_f32_f32", rms_norm_set_rows_f32_f32_len, rms_norm_set_rows_f32_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true);
-    ggml_vk_create_pipeline(device, device->pipeline_rms_norm_set_rows_f32_f16, "rms_norm_set_rows_f32_f16", rms_norm_set_rows_f32_f16_len, rms_norm_set_rows_f32_f16_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true);
+    ggml_vk_create_pipeline(device, device->pipeline_rms_norm_set_rows_f32_f32, "rms_norm_set_rows_f32_f32", RMS_SHADER(rms_norm_set_rows_f32_f32), "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0, 0, rms_set_rows_block_size}, 1, true, rms_subgroups);
+    ggml_vk_create_pipeline(device, device->pipeline_rms_norm_set_rows_f32_f16, "rms_norm_set_rows_f32_f16", RMS_SHADER(rms_norm_set_rows_f32_f16), "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0, 0, rms_set_rows_block_size}, 1, true, rms_subgroups);
     ggml_vk_create_pipeline(device, device->pipeline_rms_norm_partials_f32, "rms_norm_partials_f32", rms_norm_partials_f32_len, rms_norm_partials_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true);
     ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_partials_f32, "rms_norm_mul_partials_f32", rms_norm_partials_f32_len, rms_norm_partials_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1}, 1, true);
 
     if (sizeof(vk_op_rms_norm_mul_rope_push_constants) <= device->properties.limits.maxPushConstantsSize) {
-        ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_rope_f32_f32, "rms_norm_mul_rope_f32_f32", rms_norm_mul_rope_f32_f32_len, rms_norm_mul_rope_f32_f32_data, "main", 7, sizeof(vk_op_rms_norm_mul_rope_push_constants), {1, 1, 1}, {0, 1}, 1, true);
-        ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_rope_f32_f16, "rms_norm_mul_rope_f32_f16", rms_norm_mul_rope_f32_f16_len, rms_norm_mul_rope_f32_f16_data, "main", 7, sizeof(vk_op_rms_norm_mul_rope_push_constants), {1, 1, 1}, {0, 1}, 1, true);
+        for (uint32_t i = 0; i < 3; ++i) {
+            const uint32_t block_size = 128u << i;
+            const std::string suffix = "_" + std::to_string(block_size);
+            ggml_vk_create_pipeline2(device, device->pipeline_rms_norm_mul_rope_f32_f32[i], "rms_norm_mul_rope_f32_f32" + suffix, RMS_SHADER(rms_norm_mul_rope_f32_f32), "main", 7, sizeof(vk_op_rms_norm_mul_rope_push_constants), {1, 1, 1}, {0, 1, 0, block_size}, 1, true, rms_subgroups);
+            ggml_vk_create_pipeline2(device, device->pipeline_rms_norm_mul_rope_f32_f16[i], "rms_norm_mul_rope_f32_f16" + suffix, RMS_SHADER(rms_norm_mul_rope_f32_f16), "main", 7, sizeof(vk_op_rms_norm_mul_rope_push_constants), {1, 1, 1}, {0, 1, 0, block_size}, 1, true, rms_subgroups);
+        }
     }
+#undef RMS_SHADER
 
     ggml_vk_create_pipeline(device, device->pipeline_rms_norm_back_f32, "rms_norm_back_f32", rms_norm_back_f32_len, rms_norm_back_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_l2_norm_f32, "l2_norm_f32", l2_norm_f32_len, l2_norm_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1);
@@ -8740,6 +8753,16 @@ static vk_conv_shapes ggml_vk_conv3d_select_shape(ggml_backend_vk_context * ctx,
     return ggml_vk_conv_select_shape(ctx, K, NPQ);
 }
 
+static vk_pipeline ggml_vk_rms_norm_pipeline(const std::array<vk_pipeline, 3> & pipelines, int64_t ncols) {
+    if (pipelines[0] && ncols <= 128) {
+        return pipelines[0];
+    }
+    if (pipelines[1] && ncols <= 256) {
+        return pipelines[1];
+    }
+    return pipelines[2];
+}
+
 static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * dst, ggml_op op) {
     switch (op) {
     case GGML_OP_GET_ROWS:
@@ -8980,7 +9003,7 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
             if (ctx->do_add_rms_partials) {
                 return ctx->fused_rms_norm_mode == RMS_NORM_MUL ? ctx->device->pipeline_rms_norm_mul_partials_f32 : ctx->device->pipeline_rms_norm_partials_f32;
             }
-            return ctx->fused_rms_norm_mode == RMS_NORM_MUL ? ctx->device->pipeline_rms_norm_mul_f32 : ctx->device->pipeline_rms_norm_f32;
+            return ggml_vk_rms_norm_pipeline(ctx->fused_rms_norm_mode == RMS_NORM_MUL ? ctx->device->pipeline_rms_norm_mul_f32 : ctx->device->pipeline_rms_norm_f32, src0->ne[0]);
         }
         return nullptr;
     case GGML_OP_RMS_NORM_BACK:
@@ -10942,8 +10965,8 @@ void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const s
             pipeline = do_post_multiply ?
                 ctx->device->pipeline_rms_norm_mul_add_mul_partials_f32 : ctx->device->pipeline_rms_norm_mul_add_partials_f32;
         } else {
-            pipeline = do_post_multiply ?
-                ctx->device->pipeline_rms_norm_mul_add_mul_f32 : ctx->device->pipeline_rms_norm_mul_add_f32;
+            pipeline = ggml_vk_rms_norm_pipeline(do_post_multiply ?
+                ctx->device->pipeline_rms_norm_mul_add_mul_f32 : ctx->device->pipeline_rms_norm_mul_add_f32, src0->ne[0]);
         }
         ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
         if (ctx->do_add_rms_partials) {
@@ -11016,7 +11039,8 @@ void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const s
         pc.bin = bin;
         pc.rope = ggml_vk_make_rope_constants(rope, rope->src[0], tensors[4] != nullptr, false, set_rows_stride);
 
-        vk_pipeline pipeline = tensors[5]->type == GGML_TYPE_F16 ? ctx->device->pipeline_rms_norm_mul_rope_f32_f16 : ctx->device->pipeline_rms_norm_mul_rope_f32_f32;
+        const auto & pipelines = tensors[5]->type == GGML_TYPE_F16 ? ctx->device->pipeline_rms_norm_mul_rope_f32_f16 : ctx->device->pipeline_rms_norm_mul_rope_f32_f32;
+        vk_pipeline pipeline = ggml_vk_rms_norm_pipeline(pipelines, src0->ne[0]);
 
         ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
 

ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp

@@ -3,6 +3,15 @@
 #include "generic_binary_head.glsl"
 #include "types.glsl"
 
+#extension GL_EXT_control_flow_attributes : enable
+#if USE_SUBGROUP_ADD
+#extension GL_KHR_shader_subgroup_basic : require
+#extension GL_KHR_shader_subgroup_arithmetic : require
+#endif
+
+layout (constant_id = 3) const int BLOCK_SIZE = 512;
+layout(local_size_x_id = 3, local_size_y = 1, local_size_z = 1) in;
+
 #if RMS_NORM_ROPE_FUSION
 
 layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
@@ -11,7 +20,8 @@ layout (binding = 1) readonly buffer B {B_TYPE data_b[];};
 // data is passed from rms_norm -> rope through shared memory.
 // rms_norm calls this data_d, rope calls this rope_data_a.
 // Binding 2 is not used
-shared FLOAT_TYPE rope_data_a[1024];
+// host picks a block size such that ncols <= 2*BLOCK_SIZE
+shared FLOAT_TYPE rope_data_a[2*BLOCK_SIZE];
 #define data_d rope_data_a
 
 layout (binding = 3) readonly buffer R_Y {int rope_data_pos[];};
@@ -38,16 +48,11 @@ layout (binding = 3) readonly buffer I {uvec2 data_i[];};
 
 #endif
 
-#extension GL_EXT_control_flow_attributes : enable
-#define BLOCK_SIZE 512
-
 layout (constant_id = 1) const bool do_multiply = false;
 #if RMS_NORM_ADD_FUSION
 layout (constant_id = 2) const bool do_post_multiply = false;
 #endif
 
-layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
-
 shared FLOAT_TYPE sumsh[BLOCK_SIZE];
 
 void rms_norm(uint num_iters) {
@@ -88,17 +93,38 @@ void rms_norm(uint num_iters) {
                 sum += xi * xi;
             }
 
-            sumsh[tid] = sum;
-            // sum up partial sums and write back result
-            barrier();
-            [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
-                if (tid < s) {
-                    sum += sumsh[tid + s];
-                    sumsh[tid] = sum;
+            #if USE_SUBGROUP_ADD
+                sum = subgroupAdd(sum);
+                if (subgroupElect()) {
+                    sumsh[gl_SubgroupID] = sum;
                 }
                 barrier();
-            }
-            sum = sumsh[0];
+                sum = FLOAT_TYPE(0.0f);
+                if(gl_NumSubgroups <= 4) { // fast path to skip the 2nd subgroupAdd
+                    [[unroll]] for(uint i=0; i < gl_NumSubgroups; ++i) {
+                        sum += sumsh[i];
+                    }
+                }
+                else {
+                    // Each subgroup reads the partials, so one workgroup barrier is enough.
+                    for (uint i = gl_SubgroupInvocationID; i < gl_NumSubgroups; i += gl_SubgroupSize) {
+                        sum += sumsh[i];
+                    }
+                    sum = subgroupAdd(sum);
+                }
+            #else
+                sumsh[tid] = sum;
+                // sum up partial sums and write back result
+                barrier();
+                [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+                    if (tid < s) {
+                        sum += sumsh[tid + s];
+                        sumsh[tid] = sum;
+                    }
+                    barrier();
+                }
+                sum = sumsh[0];
+            #endif
 
             const FLOAT_TYPE mean = sum / FLOAT_TYPE(ncols);
             const FLOAT_TYPE scale = inversesqrt(mean + FLOAT_TYPE(p.param1));

ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp

@@ -853,14 +853,18 @@ void process_shaders() {
     // Norms
     string_to_spv("norm_f32", "norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
     string_to_spv("group_norm_f32", "group_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
-    string_to_spv("rms_norm_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
-    string_to_spv("rms_norm_mul_add_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"RMS_NORM_ADD_FUSION", "1"}}));
+    for (bool use_subgroups : {false, true}) {
+        const std::string suffix = use_subgroups ? "_subgroup" : "";
+        const auto rms_dict = merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD", use_subgroups ? "1" : "0"}});
+        string_to_spv("rms_norm_f32" + suffix, "rms_norm.comp", rms_dict);
+        string_to_spv("rms_norm_mul_add_f32" + suffix, "rms_norm.comp", merge_maps(rms_dict, {{"RMS_NORM_ADD_FUSION", "1"}}));
+        string_to_spv("rms_norm_set_rows_f32_f32" + suffix, "rms_norm.comp", merge_maps(rms_dict, {{"RMS_NORM_SET_ROWS_FUSION", "1"}}));
+        string_to_spv("rms_norm_set_rows_f32_f16" + suffix, "rms_norm.comp", merge_maps({{"D_TYPE", "float16_t"}, {"RMS_NORM_SET_ROWS_FUSION", "1"}}, rms_dict)); // keep rms_dict last so 'merge_maps' overwrites the value for key 'D_TYPE'
+        string_to_spv("rms_norm_mul_rope_f32_f32" + suffix, "rms_norm.comp", merge_maps(rms_dict, {{"ROPE_D_TYPE", "float"}, {"RMS_NORM_ROPE_FUSION", "1"}}));
+        string_to_spv("rms_norm_mul_rope_f32_f16" + suffix, "rms_norm.comp", merge_maps(rms_dict, {{"ROPE_D_TYPE", "float16_t"}, {"RMS_NORM_ROPE_FUSION", "1"}}));
+    }
     string_to_spv("rms_norm_mul_add_partials_f32", "rms_norm_partials.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"RMS_NORM_ADD_FUSION", "1"}}));
-    string_to_spv("rms_norm_set_rows_f32_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"RMS_NORM_SET_ROWS_FUSION", "1"}}));
-    string_to_spv("rms_norm_set_rows_f32_f16", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float16_t"}, {"RMS_NORM_SET_ROWS_FUSION", "1"}}));
     string_to_spv("rms_norm_partials_f32", "rms_norm_partials.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
-    string_to_spv("rms_norm_mul_rope_f32_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"ROPE_D_TYPE", "float"}, {"RMS_NORM_ROPE_FUSION", "1"}}));
-    string_to_spv("rms_norm_mul_rope_f32_f16", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"ROPE_D_TYPE", "float16_t"}, {"RMS_NORM_ROPE_FUSION", "1"}}));
     string_to_spv("rms_norm_back_f32", "rms_norm_back.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
     string_to_spv("l2_norm_f32", "l2_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
 

ggml/src/ggml-webgpu/ggml-webgpu.cpp

@@ -670,6 +670,8 @@ static void ggml_backend_webgpu_buffer_memset(webgpu_global_context & ctx,
                                               uint32_t                value,
                                               size_t                  offset,
                                               size_t                  size) {
+    std::lock_guard<std::recursive_mutex> lock(ctx->mutex);
+
     std::vector<uint32_t>             params  = { (uint32_t) offset, (uint32_t) size, value };
     std::vector<wgpu::BindGroupEntry> entries = { ggml_webgpu_make_bind_group_entry(0, buf, 0, buf.GetSize()) };
     size_t                            bytes_per_wg =

gguf-py/gguf/constants.py

@@ -920,6 +920,10 @@ class MODEL_TENSOR(IntEnum):
     HC_FFN_DOWN          = auto() # qwen4exp
     HC_FFN_UP            = auto() # qwen4exp
     HC_FFN_INJECT        = auto() # qwen4exp
+    HC_ATTN_DOWN_INJECT  = auto() # qwen4exp, [hc_attn_down | hc_attn_inject]
+    HC_FFN_DOWN_INJECT   = auto() # qwen4exp, [hc_ffn_down | hc_ffn_inject]
+    ATTN_QKV_GATE        = auto() # qwen4exp, [attn_qkv | attn_gate]
+    SSM_BETA_ALPHA_CAT   = auto() # qwen4exp, [ssm_beta | ssm_alpha], unlike the grouped qwen3next ssm_ba
     PLE_KEY              = auto() # qwen4exp
     PLE_VALUE            = auto() # qwen4exp
     PLE_NORM_KEY         = auto() # qwen4exp
@@ -1746,6 +1750,10 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
     MODEL_TENSOR.HC_FFN_DOWN:               "blk.{bid}.hc_ffn_down",          # qwen4exp
     MODEL_TENSOR.HC_FFN_UP:                 "blk.{bid}.hc_ffn_up",            # qwen4exp
     MODEL_TENSOR.HC_FFN_INJECT:             "blk.{bid}.hc_ffn_inject",        # qwen4exp
+    MODEL_TENSOR.HC_ATTN_DOWN_INJECT:       "blk.{bid}.hc_attn_down_inject",  # qwen4exp
+    MODEL_TENSOR.HC_FFN_DOWN_INJECT:        "blk.{bid}.hc_ffn_down_inject",   # qwen4exp
+    MODEL_TENSOR.ATTN_QKV_GATE:             "blk.{bid}.attn_qkv_gate",        # qwen4exp
+    MODEL_TENSOR.SSM_BETA_ALPHA_CAT:        "blk.{bid}.ssm_beta_alpha",       # qwen4exp
     MODEL_TENSOR.PLE_KEY:                   "blk.{bid}.ple_key",              # qwen4exp
     MODEL_TENSOR.PLE_VALUE:                 "blk.{bid}.ple_value",            # qwen4exp
     MODEL_TENSOR.PLE_NORM_KEY:              "blk.{bid}.ple_norm_key",         # qwen4exp
@@ -3309,6 +3317,10 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.HC_FFN_DOWN,
         MODEL_TENSOR.HC_FFN_UP,
         MODEL_TENSOR.HC_FFN_INJECT,
+        MODEL_TENSOR.HC_ATTN_DOWN_INJECT,
+        MODEL_TENSOR.HC_FFN_DOWN_INJECT,
+        MODEL_TENSOR.ATTN_QKV_GATE,
+        MODEL_TENSOR.SSM_BETA_ALPHA_CAT,
         # full attention layers: ATTN_Q holds [q|gate] interleaved per head
         MODEL_TENSOR.ATTN_Q,
         MODEL_TENSOR.ATTN_Q_NORM,

include/llama.h

@@ -1564,6 +1564,10 @@ extern "C" {
                                float   penalty_freq,     // must be finite, 0.0 = disabled
                                float   penalty_present); // must be finite, 0.0 = disabled
 
+    /// maximum number of distinct DRY sequence breakers, duplicates not counted
+    /// processing each breaker scans the entire vocabulary, so the init cost is O(num_breakers x n_vocab)
+    #define LLAMA_DRY_MAX_SEQ_BREAKERS 64
+
     ///  @details DRY sampler, designed by p-e-w, as described in: https://github.com/oobabooga/text-generation-webui/pull/5677, porting Koboldcpp implementation authored by pi6am: https://github.com/LostRuins/koboldcpp/pull/982
     LLAMA_API struct llama_sampler * llama_sampler_init_dry(
             const struct llama_vocab *  vocab,

skills/code-review/SKILL.md

@@ -58,6 +58,7 @@ Mandatory on every review; any finding here is **blocking**. Rule of thumb: GGUF
 - **Reused/reserved buffers:** recheck bounds after a buffer is shrunk or reused; watch `reserve()` then index-by-assumed-size, and header fields read before their length is checked.
 - **Server JSON ints:** clamp client-supplied integers (token/discard counts, offsets) to non-negative and an upper bound before they reach index/pointer arithmetic.
 - **RPC-deserialized fields:** treat every field (type/buffer/data/ne/nb/op_params) as hostile - validate before use. Null/zero buffers skipping validation, attacker data pointers, out-of-range type indices, and negative strides sign-extending past a corner-only assert all give arbitrary read/write.
+- **Recursion on untrusted input:** a function that recurses once per level (or per char) of request JSON, model output, a grammar, a template or a regex pattern overflows the stack - the HTTP threads have 512 KB on macOS, so a few hundred levels are enough. Every such function needs its own depth bound, checked in the function itself, not a pre-walk at a caller (the next caller skips it): libcommon uses `COMMON_RECURSIVE_LIMIT` (`common/common.h`), jinja `JINJA_RECURSIVE_LIMIT` (`common/jinja/utils.h`), the grammar `GRAMMAR_RECURSIVE_LIMIT` (`src/llama-grammar.h`). Count the levels of the input, not the frames, and throw `std::invalid_argument` for request input so the server answers 400. Look for recursion the code does not spell out: mutual recursion, recursive `std::function` lambdas, visitor/virtual dispatch, `std::regex` (libstdc++ recurses per char of the input), and per-char helpers that recurse where a loop would do.
 - **Lifetime/UAF:** flag stored raw pointers to caller/temporary storage, cached pointers to buffers a later free releases, async ops whose source may drop before completion, and structures not invalidated on free/realloc. Null-check conditionally-built or "not required" tensors before dereferencing.
 
 ## Approach and design (when a new component/infra is introduced)

src/llama-arch.cpp

@@ -1,5 +1,6 @@
 #include "llama-arch.h"
 
+#include "llama-ext.h"
 #include "llama-impl.h"
 
 #include <map>
@@ -635,6 +636,10 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
     { LLM_TENSOR_HC_FFN_DOWN,                            "blk.%d.hc_ffn_down" },
     { LLM_TENSOR_HC_FFN_UP,                              "blk.%d.hc_ffn_up" },
     { LLM_TENSOR_HC_FFN_INJECT,                          "blk.%d.hc_ffn_inject" },
+    { LLM_TENSOR_HC_ATTN_DOWN_INJECT,                    "blk.%d.hc_attn_down_inject" },
+    { LLM_TENSOR_HC_FFN_DOWN_INJECT,                     "blk.%d.hc_ffn_down_inject" },
+    { LLM_TENSOR_ATTN_QKV_GATE,                          "blk.%d.attn_qkv_gate" },
+    { LLM_TENSOR_SSM_BETA_ALPHA_CAT,                     "blk.%d.ssm_beta_alpha" },
     { LLM_TENSOR_PLE_KEY,                                "blk.%d.ple_key" },
     { LLM_TENSOR_PLE_VALUE,                              "blk.%d.ple_value" },
     { LLM_TENSOR_PLE_NORM_KEY,                           "blk.%d.ple_norm_key" },
@@ -907,6 +912,10 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
     {LLM_TENSOR_HC_FFN_DOWN,                {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_HC_FFN_UP,                  {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_HC_FFN_INJECT,              {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_HC_ATTN_DOWN_INJECT,        {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_HC_FFN_DOWN_INJECT,         {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_ATTN_QKV_GATE,              {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_SSM_BETA_ALPHA_CAT,         {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_PLE_KEY,                    {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_PLE_VALUE,                  {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_PLE_NORM_KEY,               {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
@@ -1348,3 +1357,43 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
             return true;
     }
 }
+
+bool llama_arch_supports_save(llm_arch arch) {
+    switch (arch) {
+        case LLM_ARCH_SENSENOVA_U1:
+        case LLM_ARCH_GEMMA3N:
+        case LLM_ARCH_BITNET:
+        case LLM_ARCH_T5:
+        case LLM_ARCH_APERTUS:
+        case LLM_ARCH_STEP35:
+        case LLM_ARCH_CLEF: // the head tensors are not saved
+        case LLM_ARCH_VLLM_DECISION2:
+            return false;
+        default:
+            return true;
+    }
+}
+
+std::vector<llm_arch> llama_arch_all() {
+    return llm_arch_all();
+}
+
+const char * llama_arch_name(llm_arch arch) {
+    return llm_arch_name(arch);
+}
+
+llm_arch llama_arch_from_string(const std::string & name) {
+    return llm_arch_from_string(name);
+}
+
+std::string llama_arch_kv_name(llm_arch arch, llm_kv key) {
+    return LLM_KV(arch)(key);
+}
+
+bool llama_arch_supports_sm_tensor(llm_arch arch) {
+    return llm_arch_supports_sm_tensor(arch);
+}
+
+bool llama_arch_supports_mixed_batch(llm_arch arch) {
+    return llm_arch_supports_mixed_batch(arch);
+}

src/llama-arch.h

@@ -679,6 +679,10 @@ enum llm_tensor {
     LLM_TENSOR_HC_FFN_DOWN,    // qwen4exp
     LLM_TENSOR_HC_FFN_UP,      // qwen4exp
     LLM_TENSOR_HC_FFN_INJECT,  // qwen4exp
+    LLM_TENSOR_HC_ATTN_DOWN_INJECT, // qwen4exp
+    LLM_TENSOR_HC_FFN_DOWN_INJECT,  // qwen4exp
+    LLM_TENSOR_ATTN_QKV_GATE,       // qwen4exp
+    LLM_TENSOR_SSM_BETA_ALPHA_CAT,  // qwen4exp
     LLM_TENSOR_PLE_KEY,        // qwen4exp
     LLM_TENSOR_PLE_VALUE,      // qwen4exp
     LLM_TENSOR_PLE_NORM_KEY,   // qwen4exp

src/llama-context.cpp

@@ -489,6 +489,7 @@ llama_context::llama_context(
             /*.swa_full  =*/ params.swa_full,
             /*.ctx_type  =*/ cparams.ctx_type,
             /*.mem_other =*/ llama_get_memory(cparams.ctx_other),
+            /*.no_alloc  =*/ model.no_alloc(),
         };
 
         memory.reset(model.create_memory(params_mem, cparams));
@@ -559,7 +560,7 @@ llama_context::llama_context(
         }
 
         if (cparams.moe_cache_size > 0) {
-            moe_cache = std::make_unique<llama_moe_cache>(model, backend_ptrs, backend_buft, cparams.moe_cache_size);
+            moe_cache = std::make_unique<llama_moe_cache>(model, model.no_alloc(), backend_ptrs, backend_buft, cparams.moe_cache_size);
         }
 
         sched_reserve();
@@ -587,7 +588,7 @@ llama_context::~llama_context() {
     synchronize();
 
     // when training, ggml_opt allocates extra buffers through the scheduler, so the sizes no longer match the expectation
-    if (!model.hparams.no_alloc && !opt_ctx) {
+    if (!model.no_alloc() && !opt_ctx) {
         for (size_t i = 0; i < backend_ptrs.size(); ++i) {
             ggml_backend_t             backend = backend_ptrs[i];
             ggml_backend_buffer_type_t buft    = backend_buft[i];
@@ -788,7 +789,7 @@ void llama_context::sched_reserve() {
     // reserve pp (prompt processing) graph first so that buffers are only allocated once
     {
         auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(),
-                model.hparams.no_alloc, model.hparams.no_alloc ? backend_buf_exp_size.data() : nullptr);
+                model.no_alloc(), model.no_alloc() ? backend_buf_exp_size.data() : nullptr);
         if (!gf) {
             if (cparams.pipeline_parallel) {
                 LLAMA_LOG_WARN("%s: compute buffer allocation failed, retrying without pipeline parallelism\n", __func__);
@@ -815,7 +816,7 @@ void llama_context::sched_reserve() {
         n_splits_tg = n_splits_pp;
         n_nodes_tg  = n_nodes_pp;
     } else {
-        auto * gf = graph_reserve(n_seqs, n_seqs, n_seqs, mctx.get(), model.hparams.no_alloc);
+        auto * gf = graph_reserve(n_seqs, n_seqs, n_seqs, mctx.get(), model.no_alloc());
         if (!gf) {
             throw std::runtime_error("failed to allocate compute tg buffers");
         }
@@ -837,10 +838,10 @@ void llama_context::sched_reserve() {
                 // [TAG_RESERVE_DIAG_DECAY]
                 // the `inp_diag_decay` tensor size scales with `n_seq_tokens^2` which
                 // makes `n_seqs == 1` use more memory for the compute graph compared to `n_seqs > 1`
-                gf = graph_reserve(n_tokens, 1,      n_outputs_pp, mctx.get(), model.hparams.no_alloc);
+                gf = graph_reserve(n_tokens, 1,      n_outputs_pp, mctx.get(), model.no_alloc());
                 break;
             default:
-                gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.hparams.no_alloc);
+                gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.no_alloc());
         };
 
         if (!gf) {
@@ -851,7 +852,7 @@ void llama_context::sched_reserve() {
     for (size_t i = 0; i < backend_ptrs.size(); ++i) {
         ggml_backend_t             backend = backend_ptrs[i];
         ggml_backend_buffer_type_t buft    = backend_buft[i];
-        if (!model.hparams.no_alloc) {
+        if (!model.no_alloc()) {
             backend_buf_exp_size[i] = ggml_backend_sched_get_buffer_size(sched.get(), backend);
         }
         if (backend_buf_exp_size[i] > 1) {
@@ -3730,7 +3731,7 @@ llama_memory_breakdown llama_context::memory_breakdown() const {
             ret[buft].context += size;
         }
     }
-    if (model.hparams.no_alloc) {
+    if (model.no_alloc()) {
         for (size_t i = 0; i < backends.size(); ++i) {
             ggml_backend_t             backend = backends[i].get();
             ggml_backend_buffer_type_t buft    = ggml_backend_sched_get_buffer_type(sched.get(), backend);

src/llama-ext.h

@@ -5,11 +5,36 @@
 // try as much as possible to not include this header in the rest of the codebase
 
 #include "llama.h"
+#include "llama-arch.h"
 
 #include <cstdint>
 #include <map>
+#include <string>
 #include <vector>
 
+//
+// llama architecture queries
+//
+
+// all the LLM architectures known to the library
+LLAMA_API std::vector<llm_arch> llama_arch_all();
+
+// the canonical name of an architecture (e.g. "llama"), or "unknown"
+LLAMA_API const char * llama_arch_name(llm_arch arch);
+
+// the architecture with the given canonical name, or LLM_ARCH_UNKNOWN
+LLAMA_API llm_arch llama_arch_from_string(const std::string & name);
+
+// the arch-resolved GGUF key of a KV id, e.g. LLM_KV_CONTEXT_LENGTH -> "llama.context_length"
+LLAMA_API std::string llama_arch_kv_name(llm_arch arch, llm_kv key);
+
+LLAMA_API bool llama_arch_supports_sm_tensor(llm_arch arch);
+
+LLAMA_API bool llama_arch_supports_mixed_batch(llm_arch arch);
+
+// FIXME temporary function for better error messages
+LLAMA_API bool llama_arch_supports_save(llm_arch arch);
+
 // Reserve a new compute graph. It is valid until the next call to llama_graph_reserve.
 LLAMA_API struct ggml_cgraph * llama_graph_reserve(
         struct llama_context * ctx,

src/llama-grammar.cpp

@@ -8,9 +8,11 @@
 #include <algorithm>
 #include <cstdint>
 #include <set>
+#include <numeric>
 #include <stdexcept>
 
 #define MAX_REPETITION_THRESHOLD 2000
+#define GRAMMAR_RECURSIVE_LIMIT 256
 //
 // helpers
 //
@@ -440,13 +442,14 @@ const char * llama_grammar_parser::parse_alternates(
         const char        * src,
         const std::string & rule_name,
         uint32_t            rule_id,
-        bool                is_nested) {
+        bool                is_nested,
+        uint32_t            depth) {
     llama_grammar_rule rule;
-    const char * pos = parse_sequence(src, rule_name, rule, is_nested);
+    const char * pos = parse_sequence(src, rule_name, rule, is_nested, depth);
     while (*pos == '|') {
         rule.push_back({LLAMA_GRETYPE_ALT, 0});
         pos = parse_space(pos + 1, true);
-        pos = parse_sequence(pos, rule_name, rule, is_nested);
+        pos = parse_sequence(pos, rule_name, rule, is_nested, depth);
     }
     rule.push_back({LLAMA_GRETYPE_END, 0});
     add_rule(rule_id, rule);
@@ -457,7 +460,8 @@ const char * llama_grammar_parser::parse_sequence(
         const char         * src,
         const std::string  & rule_name,
         llama_grammar_rule & rule,
-        bool               is_nested) {
+        bool               is_nested,
+        uint32_t           depth) {
     size_t last_sym_start = rule.size();
     const char * pos = src;
     uint64_t n_prev_rules = 1;
@@ -595,11 +599,14 @@ const char * llama_grammar_parser::parse_sequence(
             n_prev_rules = 1;
             rule.push_back({LLAMA_GRETYPE_RULE_REF, ref_rule_id});
         } else if (*pos == '(') { // grouping
+            if (depth >= GRAMMAR_RECURSIVE_LIMIT) {
+                throw std::runtime_error("grammar recursion depth limit exceeded (max " + std::to_string(GRAMMAR_RECURSIVE_LIMIT) + "), please reduce nesting depth");
+            }
             // parse nested alternates into synthesized rule
             pos = parse_space(pos + 1, true);
             uint32_t n_rules_before = symbol_ids.size();
             uint32_t sub_rule_id = generate_symbol_id(rule_name);
-            pos = parse_alternates(pos, rule_name, sub_rule_id, true);
+            pos = parse_alternates(pos, rule_name, sub_rule_id, true, depth + 1);
             n_prev_rules = std::max(1u, (uint32_t)symbol_ids.size() - n_rules_before);
             last_sym_start = rule.size();
             // output reference to synthesized rule
@@ -941,34 +948,19 @@ static void llama_grammar_advance_stack(
     }
 }
 
-static llama_grammar_candidates llama_grammar_reject_candidates(
-        const llama_grammar_rules      & rules,
-        const llama_grammar_stacks     & stacks,
-        const llama_grammar_candidates & candidates) {
-    GGML_ASSERT(!stacks.empty()); // REVIEW
-
-    if (candidates.empty()) {
-        return {};
-    }
-
-    auto rejects = llama_grammar_reject_candidates_for_stack(rules, stacks.front(), candidates);
-
-    for (size_t i = 1, size = stacks.size(); i < size; ++i) {
-        rejects = llama_grammar_reject_candidates_for_stack(rules, stacks[i], rejects);
-    }
-
-    return rejects;
-}
-
 static bool llama_grammar_detect_left_recursion(
         const llama_grammar_rules & rules,
         size_t rule_index,
         std::vector<bool> * rules_visited,
         std::vector<bool> * rules_in_progress,
-        std::vector<bool> * rules_may_be_empty) {
+        std::vector<bool> * rules_may_be_empty,
+        uint32_t depth = 0) {
     if ((*rules_in_progress)[rule_index]) {
         return true;
     }
+    if (depth > GRAMMAR_RECURSIVE_LIMIT) {
+        throw std::runtime_error("grammar rules are chained more than " + std::to_string(GRAMMAR_RECURSIVE_LIMIT) + " levels deep");
+    }
 
     (*rules_in_progress)[rule_index] = true;
 
@@ -994,7 +986,7 @@ static bool llama_grammar_detect_left_recursion(
     bool recurse_into_nonterminal = true;
     for (size_t i = 0; i < rule.size(); i++) {
         if (rule[i].type == LLAMA_GRETYPE_RULE_REF && recurse_into_nonterminal) {
-            if (llama_grammar_detect_left_recursion(rules, (size_t)rule[i].value, rules_visited, rules_in_progress, rules_may_be_empty)) {
+            if (llama_grammar_detect_left_recursion(rules, (size_t)rule[i].value, rules_visited, rules_in_progress, rules_may_be_empty, depth + 1)) {
                 return true;
             }
             if (!((*rules_may_be_empty)[(size_t)rule[i].value])) {
@@ -1058,56 +1050,63 @@ void llama_grammar_accept(struct llama_grammar * grammar, uint32_t chr) {
     grammar->stacks = std::move(stacks_new);
 }
 
-llama_grammar_candidates llama_grammar_reject_candidates_for_stack(
+// checks the code point at depth of each candidate (an index into the input) against the stack
+// returns true when some match: they go to next, and the stacks that follow to next_stacks
+static bool llama_grammar_reject_step(
         const llama_grammar_rules      & rules,
         const llama_grammar_stack      & stack,
-        const llama_grammar_candidates & candidates) {
-
-    llama_grammar_candidates rejects;
-    rejects.reserve(candidates.size());
+        const llama_grammar_candidates & candidates,
+        size_t                           depth,
+        const std::vector<uint32_t>    & cur,
+              std::vector<uint32_t>    & rejects,
+              std::vector<uint32_t>    & next,
+              llama_grammar_stacks     & next_stacks) {
+    rejects.reserve(cur.size());
 
     if (stack.empty()) {
-        for (const auto & tok : candidates) {
-            if (*tok.code_points != 0 || tok.partial_utf8.n_remain != 0) {
-                rejects.push_back(tok);
+        for (const uint32_t i : cur) {
+            const auto & tok = candidates[i];
+            if (tok.code_points[depth] != 0 || tok.partial_utf8.n_remain != 0) {
+                rejects.push_back(i);
             }
         }
-        return rejects;
+        return false;
     }
 
     const llama_grammar_element * stack_pos = stack.back();
 
     // if the top of the stack is a token rule, then we only need to check the token id
     if (stack_pos->type == LLAMA_GRETYPE_TOKEN || stack_pos->type == LLAMA_GRETYPE_TOKEN_NOT) {
-        for (const auto & tok : candidates) {
-            if (*tok.code_points == 0) {
+        for (const uint32_t i : cur) {
+            const auto & tok = candidates[i];
+            if (tok.code_points[depth] == 0) {
                 // reached the end of a token consumed by char rules, reject iff it ended
                 // in a partial response
                 if (tok.partial_utf8.n_remain != 0) {
-                    rejects.push_back(tok);
+                    rejects.push_back(i);
                 }
             } else if (!llama_grammar_match_token(stack_pos, tok.id)) {
-                rejects.push_back(tok);
+                rejects.push_back(i);
             }
         }
-        return rejects;
+        return false;
     }
 
-    llama_grammar_candidates next_candidates;
-    next_candidates.reserve(candidates.size());
+    next.reserve(cur.size());
 
-    for (const auto & tok : candidates) {
-        if (*tok.code_points == 0) {
+    for (const uint32_t i : cur) {
+        const auto & tok = candidates[i];
+        if (tok.code_points[depth] == 0) {
             // reached end of full codepoints in token, reject iff it ended in a partial sequence
             // that cannot satisfy this position in grammar
             if (tok.partial_utf8.n_remain != 0 &&
                     !llama_grammar_match_partial_char(stack_pos, tok.partial_utf8)) {
-                rejects.push_back(tok);
+                rejects.push_back(i);
             }
-        } else if (llama_grammar_match_char(stack_pos, *tok.code_points).first) {
-            next_candidates.push_back({ tok.index, tok.code_points + 1, tok.partial_utf8, tok.id });
+        } else if (llama_grammar_match_char(stack_pos, tok.code_points[depth]).first) {
+            next.push_back(i);
         } else {
-            rejects.push_back(tok);
+            rejects.push_back(i);
         }
     }
 
@@ -1118,17 +1117,80 @@ llama_grammar_candidates llama_grammar_reject_candidates_for_stack(
     if (!llama_grammar_is_end_of_sequence(stack_pos_after)) {
         stack_after.push_back(stack_pos_after);
     }
-    llama_grammar_stacks next_stacks;
     llama_grammar_advance_stack(rules, stack_after, next_stacks);
 
-    auto next_rejects = llama_grammar_reject_candidates(rules, next_stacks, next_candidates);
-    for (const auto & tok : next_rejects) {
-        rejects.push_back({ tok.index, tok.code_points - 1, tok.partial_utf8, tok.id });
+    return true;
+}
+
+// the candidates that no stack accepts, each stack narrows the rejects of the previous one
+// no recursion: a level per code point, a token from the GGUF can be long
+static llama_grammar_candidates llama_grammar_reject_candidates(
+        const llama_grammar_rules      & rules,
+        const llama_grammar_stacks     & stacks,
+        const llama_grammar_candidates & candidates) {
+    GGML_ASSERT(!stacks.empty()); // REVIEW
+
+    // a level checks the candidates that reach its code point against its stacks in turn
+    struct level {
+        llama_grammar_stacks  stacks;
+        size_t                i;       // the stack in progress
+        std::vector<uint32_t> cur;     // the candidates rejected by stacks[0, i)
+        std::vector<uint32_t> rejects; // the rejects of stacks[i] at this code point
+    };
+
+    std::vector<uint32_t> all(candidates.size());
+    std::iota(all.begin(), all.end(), 0u);
+
+    std::vector<level> levels;
+    levels.push_back({ stacks, 0, std::move(all), {} });
+
+    while (true) {
+        level & lv = levels.back();
+
+        if (lv.i == lv.stacks.size() || lv.cur.empty()) {
+            if (levels.size() == 1) {
+                break;
+            }
+            // the rejects at the next code point belong to stacks[i] of the parent level
+            level & parent = levels[levels.size() - 2];
+            parent.rejects.insert(parent.rejects.end(), lv.cur.begin(), lv.cur.end());
+            levels.pop_back();
+            parent.cur = std::move(parent.rejects);
+            parent.rejects = {};
+            parent.i++;
+            continue;
+        }
+
+        std::vector<uint32_t> rejects;
+        std::vector<uint32_t> next;
+        llama_grammar_stacks  next_stacks;
+
+        if (!llama_grammar_reject_step(rules, lv.stacks[lv.i], candidates, levels.size() - 1, lv.cur, rejects, next, next_stacks)) {
+            lv.cur = std::move(rejects);
+            lv.i++;
+            continue;
+        }
+
+        GGML_ASSERT(!next_stacks.empty());
+        lv.rejects = std::move(rejects);
+        levels.push_back({ std::move(next_stacks), 0, std::move(next), {} });
     }
 
+    llama_grammar_candidates rejects;
+    rejects.reserve(levels[0].cur.size());
+    for (const uint32_t i : levels[0].cur) {
+        rejects.push_back(candidates[i]);
+    }
     return rejects;
 }
 
+llama_grammar_candidates llama_grammar_reject_candidates_for_stack(
+        const llama_grammar_rules      & rules,
+        const llama_grammar_stack      & stack,
+        const llama_grammar_candidates & candidates) {
+    return llama_grammar_reject_candidates(rules, llama_grammar_stacks{ stack }, candidates);
+}
+
 ////////////////////
 
 struct llama_grammar * llama_grammar_init_impl(
@@ -1163,14 +1225,19 @@ struct llama_grammar * llama_grammar_init_impl(
     std::vector<bool> rules_visited(n_rules);
     std::vector<bool> rules_in_progress(n_rules);
     std::vector<bool> rules_may_be_empty(n_rules);
-    for (size_t i = 0; i < n_rules; i++) {
-        if (rules_visited[i]) {
-            continue;
-        }
-        if (llama_grammar_detect_left_recursion(vec_rules, i, &rules_visited, &rules_in_progress, &rules_may_be_empty)) {
-            LLAMA_LOG_ERROR("unsupported grammar, left recursion detected for nonterminal at index %zu", i);
-            return nullptr;
+    try {
+        for (size_t i = 0; i < n_rules; i++) {
+            if (rules_visited[i]) {
+                continue;
+            }
+            if (llama_grammar_detect_left_recursion(vec_rules, i, &rules_visited, &rules_in_progress, &rules_may_be_empty)) {
+                LLAMA_LOG_ERROR("unsupported grammar, left recursion detected for nonterminal at index %zu", i);
+                return nullptr;
+            }
         }
+    } catch (const std::exception & err) {
+        LLAMA_LOG_ERROR("unsupported grammar: %s\n", err.what());
+        return nullptr;
     }
 
     // loop over alternates of start rule to build initial stacks
@@ -1256,14 +1323,19 @@ struct llama_grammar * llama_grammar_init_impl(
     std::vector<bool> rules_visited(n_rules);
     std::vector<bool> rules_in_progress(n_rules);
     std::vector<bool> rules_may_be_empty(n_rules);
-    for (size_t i = 0; i < n_rules; i++) {
-        if (rules_visited[i]) {
-            continue;
-        }
-        if (llama_grammar_detect_left_recursion(vec_rules, i, &rules_visited, &rules_in_progress, &rules_may_be_empty)) {
-            LLAMA_LOG_ERROR("unsupported grammar, left recursion detected for nonterminal at index %zu\n", i);
-            return nullptr;
+    try {
+        for (size_t i = 0; i < n_rules; i++) {
+            if (rules_visited[i]) {
+                continue;
+            }
+            if (llama_grammar_detect_left_recursion(vec_rules, i, &rules_visited, &rules_in_progress, &rules_may_be_empty)) {
+                LLAMA_LOG_ERROR("unsupported grammar, left recursion detected for nonterminal at index %zu\n", i);
+                return nullptr;
+            }
         }
+    } catch (const std::exception & err) {
+        LLAMA_LOG_ERROR("unsupported grammar: %s\n", err.what());
+        return nullptr;
     }
 
     // loop over alternates of start rule to build initial stacks

src/llama-grammar.h

@@ -102,13 +102,15 @@ struct llama_grammar_parser {
             const char        * src,
             const std::string & rule_name,
             uint32_t            rule_id,
-            bool                is_nested);
+            bool                is_nested,
+            uint32_t            depth = 0);
 
     const char * parse_sequence(
             const char         * src,
             const std::string  & rule_name,
             llama_grammar_rule & rule,
-            bool               is_nested);
+            bool               is_nested,
+            uint32_t           depth = 0);
 
     const char * parse_rule(const char * src);
 

src/llama-hparams.h

@@ -61,7 +61,6 @@ struct llama_hparams {
     //       for example: n_embd_out() vs n_embd_out_impl
 
     bool vocab_only;
-    bool no_alloc; // TODO: no_alloc is not a model hyperparameter; move it out of llama_hparams
     bool rope_finetuned;
     bool use_par_res;
     bool swin_norm;
@@ -358,6 +357,8 @@ struct llama_hparams {
     llama_token ple_embd_token   = LLAMA_TOKEN_NULL;
     // hash constant sets, stored back to back: 1 shared by all PLE layers (qwen4exp), or 1 per PLE layer in layer order (deepseek41)
     uint32_t ple_n_modules       = 1;
+    // deepseek41 Engram: the token that pads the n-gram history, an index into llama_model::engram_token_map
+    uint32_t engram_pad_id       = 0;
     // the file lists PLE layer indices, so this is never a per-layer gguf array and can hold one bit per layer
     std::bitset<LLAMA_MAX_LAYERS> is_ple_impl;
     // the hash multipliers reach ~2e13 and have to stay 64-bit

src/llama-kv-cache-dsa-iswa.cpp

@@ -13,6 +13,7 @@
 
 llama_kv_cache_dsa_iswa::llama_kv_cache_dsa_iswa(
         const llama_model & model,
+                     bool   no_alloc,
                 ggml_type   type_k,
                 ggml_type   type_v,
                      bool   v_trans,
@@ -63,14 +64,14 @@ llama_kv_cache_dsa_iswa::llama_kv_cache_dsa_iswa(
     LLAMA_LOG_INFO("%s: creating DSA KV cache, size = %u cells\n", __func__, size_dsa);
 
     kv_dsa = std::make_unique<llama_kv_cache_dsa>(
-            model, type_k, type_v,
+            model, no_alloc, type_k, type_v,
             v_trans, offload, unified, size_dsa, n_seq_max, n_pad,
             0, LLAMA_SWA_TYPE_NONE, filter_dsa, filter_lid, reuse);
 
     LLAMA_LOG_INFO("%s: creating SWA KV cache, size = %u cells\n", __func__, size_swa);
 
     kv_swa = std::make_unique<llama_kv_cache>(
-            model, hparams, type_k, type_v,
+            model, no_alloc, hparams, type_k, type_v,
             v_trans, offload, unified, size_swa, n_seq_max, n_pad,
             hparams.n_swa, hparams.swa_type, nullptr, filter_swa, reuse, nullptr);
 }

src/llama-kv-cache-dsa-iswa.h

@@ -14,6 +14,7 @@ class llama_kv_cache_dsa_iswa : public llama_memory_i {
 public:
     llama_kv_cache_dsa_iswa(
             const llama_model & model,
+                         bool   no_alloc,
                     ggml_type   type_k,
                     ggml_type   type_v,
                          bool   v_trans,

src/llama-kv-cache-dsa.cpp

@@ -13,6 +13,7 @@
 
 llama_kv_cache_dsa::llama_kv_cache_dsa(
         const llama_model & model,
+                     bool   no_alloc,
                 ggml_type   type_k,
                 ggml_type   type_v,
                      bool   v_trans,
@@ -31,7 +32,7 @@ llama_kv_cache_dsa::llama_kv_cache_dsa(
     LLAMA_LOG_INFO("%s: creating main KV cache, size = %u cells\n", __func__, kv_size);
 
     kv_mla = std::make_unique<llama_kv_cache>(
-            model, model.hparams, type_k, type_v,
+            model, no_alloc, model.hparams, type_k, type_v,
             v_trans, offload, unified, kv_size, n_seq_max, n_pad,
             n_swa, swa_type, nullptr, filter_mla, reuse, nullptr);
 
@@ -48,7 +49,7 @@ llama_kv_cache_dsa::llama_kv_cache_dsa(
     LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);
 
     kv_lid = std::make_unique<llama_kv_cache>(
-            model, hparams_lid, type_k, type_v,
+            model, no_alloc, hparams_lid, type_k, type_v,
             v_trans, offload, unified, kv_size, n_seq_max, n_pad,
             n_swa, swa_type, nullptr, filter_lid, reuse, nullptr);
 }

src/llama-kv-cache-dsa.h

@@ -16,6 +16,7 @@ class llama_kv_cache_dsa : public llama_memory_i {
 public:
     llama_kv_cache_dsa(
             const llama_model & model,
+                         bool   no_alloc,
                     ggml_type   type_k,
                     ggml_type   type_v,
                          bool   v_trans,

src/llama-kv-cache-dsv4.cpp

@@ -428,8 +428,11 @@ static void dsv4_state_read_k_cache(
         dsv4_state_read_tensor_streams(io, kv->get_k_storage(il), kv_size, n_rows_ref, s0, ns);
     }
 
-    // refill the mirrors from their source rows just read
-    std::vector<uint8_t> buf;
+    if (n_rows_ref == 0) {
+        return;
+    }
+
+    // refill the mirrors from their sources; a mirror gets the same rows as its source, so a whole copy is exact
     for (uint32_t il : kv->get_layer_ids()) {
         const int32_t src = kv_src[il];
         if (src == (int32_t) il) {
@@ -437,20 +440,11 @@ static void dsv4_state_read_k_cache(
         }
         ggml_tensor * t_src = kv->get_k_storage(src);
         ggml_tensor * t_dst = kv->get_k_storage(il);
-        GGML_ASSERT(t_src->type == t_dst->type && t_src->ne[0] == t_dst->ne[0]);
+        if (t_src->type != t_dst->type || !ggml_are_same_shape(t_src, t_dst)) {
+            throw std::runtime_error("DSV4 K-cache mirror layout mismatch");
+        }
 
-        const size_t row_size      = ggml_row_size(t_src->type, t_src->ne[0]);
-        const size_t stream_stride = (size_t) kv_size*row_size;
-        const size_t size          = (size_t) n_rows_ref*row_size;
-        if (size == 0) {
-            continue;
-        }
-        buf.resize(size);
-        for (uint32_t s = 0; s < ns; ++s) {
-            const size_t offset = (size_t) (s0 + s)*stream_stride;
-            ggml_backend_tensor_get(t_src, buf.data(), offset, size);
-            ggml_backend_tensor_set(t_dst, buf.data(), offset, size);
-        }
+        ggml_backend_tensor_copy(t_src, t_dst);
     }
 }
 
@@ -937,21 +931,24 @@ static void dsv4_make_k_only(llama_hparams & hparams) {
 
 llama_dsv4_comp_state::llama_dsv4_comp_state(
         const llama_model & model,
-                bool        offload,
-                bool        unified,
-            uint32_t        n_seq_max,
-            uint32_t        ratio,
-            uint32_t        state_size,
-            uint32_t        n_embd_state,
-            uint32_t        n_rs_seq,
-        const char    * name,
+                     bool   no_alloc,
+                     bool   offload,
+                     bool   unified,
+                 uint32_t   n_seq_max,
+                 uint32_t   ratio,
+                 uint32_t   state_size,
+                 uint32_t   n_embd_state,
+                 uint32_t   n_rs_seq,
+               const char * name,
         const llama_memory_i::layer_filter_cb & filter) :
-    hparams(model.hparams),
     ratio(ratio),
     state_size(state_size),
     n_embd_state(n_embd_state),
     n_stream(unified ? 1 : n_seq_max),
-    n_rs_seq(n_rs_seq) {
+    n_rs_seq(n_rs_seq),
+    ctxs_bufs(no_alloc) {
+    const llama_hparams & hparams = model.hparams;
+
     struct ggml_backend_buft_comparator {
         bool operator()(const ggml_backend_buffer_type_t & lhs, const ggml_backend_buffer_type_t & rhs) const {
             return strcmp(ggml_backend_buft_name(lhs), ggml_backend_buft_name(rhs)) < 0;
@@ -1026,25 +1023,10 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
     }
 
     for (auto & [buft, ctx] : ctx_map) {
-        ggml_backend_buffer_t buf;
-        if (hparams.no_alloc) {
-            buf = ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0); // dummy buffer
-            for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) {
-                t->buffer = buf; // set dummy buffer for DSV4 state so that the backend scheduler won't try to allocate it
-            }
-        } else {
-            buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer
-        }
-        if (!buf) {
-            throw std::runtime_error("failed to allocate buffer for DSV4 compressor state");
-        }
-
-        ggml_backend_buffer_clear(buf, 0);
+        ggml_backend_buffer_t buf = ctxs_bufs.emplace(std::move(ctx), buft, "DSV4 compressor state");
 
         LLAMA_LOG_INFO("%s: %10s DSV4 %s state buffer size = %8.2f MiB\n",
                 __func__, ggml_backend_buffer_name(buf), name, ggml_backend_buffer_get_size(buf)/1024.0/1024.0);
-
-        ctxs_bufs.emplace_back(std::move(ctx), buf);
     }
 
     LLAMA_LOG_INFO("%s: %s ratio = %u, state = %u x %u, streams = %u, rs_seq = %u, layers = %zu, size = %7.2f MiB\n",
@@ -1125,18 +1107,7 @@ uint32_t llama_dsv4_comp_state::get_n_rows() const {
 }
 
 std::map<ggml_backend_buffer_type_t, size_t> llama_dsv4_comp_state::memory_breakdown() const {
-    std::map<ggml_backend_buffer_type_t, size_t> ret;
-    for (const auto & [ctx, buf] : ctxs_bufs) {
-        ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get());
-
-        if (hparams.no_alloc) {
-            GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) == nullptr);
-            ret[buft] += ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft);
-        } else {
-            ret[buft] += ggml_backend_buffer_get_size(buf.get());
-        }
-    }
-    return ret;
+    return ctxs_bufs.memory_breakdown();
 }
 
 void llama_dsv4_comp_state::state_write(
@@ -1271,6 +1242,7 @@ size_t llama_dsv4_comp_state::total_size() const {
 
 llama_kv_cache_dsv4::llama_kv_cache_dsv4(
         const llama_model & model,
+                     bool   no_alloc,
                 ggml_type   type_k,
                 ggml_type   type_v,
                      bool   v_trans,
@@ -1315,7 +1287,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
     dsv4_make_k_only(hparams_raw);
 
     kv_raw = std::make_unique<llama_kv_cache_iswa>(
-            model, hparams_raw, type_k, type_v,
+            model, no_alloc, hparams_raw, type_k, type_v,
             v_trans, offload, swa_full, unified_raw, kv_size, n_seq_max, n_ubatch, n_pad,
             nullptr, filter_raw, reuse, nullptr);
 
@@ -1352,10 +1324,10 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
     // crosses devices. no mirrors without offload
     const uint32_t n_layer = model.hparams.n_layer();
     comp_read.assign(n_layer, -1);
-    comp_kv_src.resize(n_layer);
+    comp_il_src.resize(n_layer);
     for (uint32_t il = 0; il < n_layer; ++il) {
         const int32_t src = is_v41 ? model.hparams.dsv41_kv_src[il] : (int32_t) il;
-        comp_kv_src[il] = src >= 0 ? src : (int32_t) il;
+        comp_il_src[il] = src >= 0 ? src : (int32_t) il;
         if (model.hparams.dsv4_compress_ratios[il] == 0) {
             continue;
         }
@@ -1366,7 +1338,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
         // the first layer of this source on this device holds the mirror
         comp_read[il] = il;
         for (int32_t j = src + 1; j < (int32_t) il; ++j) {
-            if (comp_read[j] == j && comp_kv_src[j] == src && model.dev_layer(j) == model.dev_layer(il)) {
+            if (comp_read[j] == j && comp_il_src[j] == src && model.dev_layer(j) == model.dev_layer(il)) {
                 comp_read[il] = j;
                 break;
             }
@@ -1419,7 +1391,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
             __func__, dsv4_comp_size(kv_size, ratio_csa));
 
     kv_csa = std::make_unique<llama_kv_cache>(
-            model, hparams_csa, type_k, type_v,
+            model, no_alloc, hparams_csa, type_k, type_v,
             v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, ratio_csa), 256u), n_seq_max, n_pad,
             0, LLAMA_SWA_TYPE_NONE, nullptr, filter_csa, nullptr, nullptr);
 
@@ -1427,7 +1399,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
             __func__, dsv4_comp_size(kv_size, ratio_hca));
 
     kv_hca = std::make_unique<llama_kv_cache>(
-            model, hparams_hca, type_k, type_v,
+            model, no_alloc, hparams_hca, type_k, type_v,
             v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, ratio_hca), 256u), n_seq_max, n_pad,
             0, LLAMA_SWA_TYPE_NONE, nullptr, filter_hca, nullptr, nullptr);
 
@@ -1435,7 +1407,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
             __func__, dsv4_comp_size(kv_size, ratio_csa));
 
     kv_lid = std::make_unique<llama_kv_cache>(
-            model, hparams_lid, type_k, type_v,
+            model, no_alloc, hparams_lid, type_k, type_v,
             v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, ratio_csa), 256u), n_seq_max, n_pad,
             0, LLAMA_SWA_TYPE_NONE, nullptr, filter_lid, nullptr, nullptr);
 
@@ -1445,19 +1417,19 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
     const uint32_t csa_state_embd = is_v41 ? model.hparams.n_embd_head_k() : 2*model.hparams.n_embd_head_k();
 
     csa_state = std::make_unique<llama_dsv4_comp_state>(
-            model, offload, unified_compressed, n_seq_max, ratio_csa, csa_state_size,
+            model, no_alloc, offload, unified_compressed, n_seq_max, ratio_csa, csa_state_size,
             csa_state_embd, n_rs_seq, "csa", filter_csa_state);
 
     LLAMA_LOG_INFO("%s: creating DSV4 HCA compressor state\n", __func__);
 
     hca_state = std::make_unique<llama_dsv4_comp_state>(
-            model, offload, unified_compressed, n_seq_max, ratio_hca, ratio_hca,
+            model, no_alloc, offload, unified_compressed, n_seq_max, ratio_hca, ratio_hca,
             model.hparams.n_embd_head_k(), n_rs_seq, "hca", filter_hca_state);
 
     LLAMA_LOG_INFO("%s: creating DSV4 lightning-indexer compressor state\n", __func__);
 
     lid_state = std::make_unique<llama_dsv4_comp_state>(
-            model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO,
+            model, no_alloc, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO,
             2*model.hparams.indexer_head_size, n_rs_seq, "lid", filter_lid);
 
     // DSV4 attention reads compressed-K / compressor-state rows that the current
@@ -1811,9 +1783,9 @@ void llama_kv_cache_dsv4::state_write(llama_io_write_i & io, llama_seq_id seq_id
         const uint32_t n_rows_lid = seq_id >= 0 ?
             dsv4_state_n_used_k_rows(pos_max, lid_state->get_ratio(), kv_lid->get_size()) : kv_lid->get_size();
 
-        dsv4_state_write_k_cache(io, kv_csa.get(), seq_id, flags, n_rows_csa, comp_kv_src);
-        dsv4_state_write_k_cache(io, kv_hca.get(), seq_id, flags, n_rows_hca, comp_kv_src);
-        dsv4_state_write_k_cache(io, kv_lid.get(), seq_id, flags, n_rows_lid, comp_kv_src);
+        dsv4_state_write_k_cache(io, kv_csa.get(), seq_id, flags, n_rows_csa, comp_il_src);
+        dsv4_state_write_k_cache(io, kv_hca.get(), seq_id, flags, n_rows_hca, comp_il_src);
+        dsv4_state_write_k_cache(io, kv_lid.get(), seq_id, flags, n_rows_lid, comp_il_src);
     }
 
     csa_state->state_write(io, seq_id, flags, rs_idx);
@@ -1851,9 +1823,9 @@ void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id,
     if (!partial_only) {
         clear_compressed(seq_id, true);
 
-        dsv4_state_read_k_cache(io, kv_csa.get(), seq_id, flags, comp_kv_src);
-        dsv4_state_read_k_cache(io, kv_hca.get(), seq_id, flags, comp_kv_src);
-        dsv4_state_read_k_cache(io, kv_lid.get(), seq_id, flags, comp_kv_src);
+        dsv4_state_read_k_cache(io, kv_csa.get(), seq_id, flags, comp_il_src);
+        dsv4_state_read_k_cache(io, kv_hca.get(), seq_id, flags, comp_il_src);
+        dsv4_state_read_k_cache(io, kv_lid.get(), seq_id, flags, comp_il_src);
     }
 
     csa_state->state_read(io, seq_id, flags);

src/llama-kv-cache-dsv4.h

@@ -16,14 +16,15 @@ public:
 
     llama_dsv4_comp_state(
             const llama_model & model,
-            bool            offload,
-            bool            unified,
-            uint32_t        n_seq_max,
-            uint32_t        ratio,
-            uint32_t        state_size,
-            uint32_t        n_embd_state,
-            uint32_t        n_rs_seq,
-            const char    * name,
+                         bool   no_alloc,
+                         bool   offload,
+                         bool   unified,
+                     uint32_t   n_seq_max,
+                     uint32_t   ratio,
+                     uint32_t   state_size,
+                     uint32_t   n_embd_state,
+                     uint32_t   n_rs_seq,
+                   const char * name,
         const llama_memory_i::layer_filter_cb & filter);
 
     void clear(llama_seq_id seq_id, bool data);
@@ -61,15 +62,13 @@ private:
         std::vector<ggml_tensor *> score_stream;
     };
 
-    const llama_hparams & hparams;
-
     const uint32_t ratio;
     const uint32_t state_size;
     const uint32_t n_embd_state;
     const uint32_t n_stream;
     const uint32_t n_rs_seq;
 
-    std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
+    llama_memory_bufs ctxs_bufs;
 
     std::vector<layer> layers;
 
@@ -92,6 +91,7 @@ class llama_kv_cache_dsv4 : public llama_memory_i {
 public:
     llama_kv_cache_dsv4(
             const llama_model & model,
+                         bool   no_alloc,
                     ggml_type   type_k,
                     ggml_type   type_v,
                          bool   v_trans,
@@ -170,7 +170,7 @@ private:
     std::vector<uint32_t> rs_idx;
 
     // [n_layer] storage layer -> layer whose rows it holds (itself, or the source of a mirror)
-    std::vector<int32_t> comp_kv_src;
+    std::vector<int32_t> comp_il_src;
 
     // [n_layer] layer -> layer whose compressed-KV storage it reads (itself, its V4.1 kv source or a mirror), -1 = none
     std::vector<int32_t> comp_read;

src/llama-kv-cache-iswa.cpp

@@ -13,6 +13,7 @@
 
 llama_kv_cache_iswa::llama_kv_cache_iswa(
         const llama_model & model,
+                     bool   no_alloc,
                 ggml_type   type_k,
                 ggml_type   type_v,
                      bool   v_trans,
@@ -27,13 +28,14 @@ llama_kv_cache_iswa::llama_kv_cache_iswa(
     const layer_filter_cb & filter,
     const  layer_reuse_cb & reuse,
     const  layer_share_cb & share) :
-    llama_kv_cache_iswa(model, model.hparams, type_k, type_v, v_trans, offload, swa_full, unified,
+    llama_kv_cache_iswa(model, no_alloc, model.hparams, type_k, type_v, v_trans, offload, swa_full, unified,
             kv_size, n_seq_max, n_ubatch, n_pad, mem_other, filter, reuse, share) {
 }
 
 llama_kv_cache_iswa::llama_kv_cache_iswa(
         const llama_model & model,
-        const llama_hparams & hparams,
+                     bool   no_alloc,
+      const llama_hparams & hparams,
                 ggml_type   type_k,
                 ggml_type   type_v,
                      bool   v_trans,
@@ -93,14 +95,14 @@ llama_kv_cache_iswa::llama_kv_cache_iswa(
     }
 
     kv_base = std::make_unique<llama_kv_cache>(
-            model, hparams, type_k, type_v,
+            model, no_alloc, hparams, type_k, type_v,
             v_trans, offload, unified, size_base, n_seq_max, n_pad,
             0, LLAMA_SWA_TYPE_NONE, mem_other_base, filter_base, reuse, share);
 
     LLAMA_LOG_INFO("%s: creating     SWA KV cache, size = %u cells\n", __func__, size_swa);
 
     kv_swa = std::make_unique<llama_kv_cache>(
-            model, hparams, type_k, type_v,
+            model, no_alloc, hparams, type_k, type_v,
             v_trans, offload, unified, size_swa, n_seq_max, n_pad,
             hparams.n_swa, hparams.swa_type, mem_other_swa, filter_swa, reuse, share);
 }

src/llama-kv-cache-iswa.h

@@ -15,6 +15,7 @@ class llama_kv_cache_iswa : public llama_memory_i {
 public:
     llama_kv_cache_iswa(
             const llama_model & model,
+                         bool   no_alloc,
                     ggml_type   type_k,
                     ggml_type   type_v,
                          bool   v_trans,
@@ -32,7 +33,8 @@ public:
 
     llama_kv_cache_iswa(
             const llama_model & model,
-            const llama_hparams & hparams,
+                         bool   no_alloc,
+          const llama_hparams & hparams,
                     ggml_type   type_k,
                     ggml_type   type_v,
                          bool   v_trans,

src/llama-kv-cache-msa.cpp

@@ -12,6 +12,7 @@
 
 llama_kv_cache_msa::llama_kv_cache_msa(
         const llama_model & model,
+                     bool   no_alloc,
                 ggml_type   type_k,
                 ggml_type   type_v,
                      bool   v_trans,
@@ -32,7 +33,7 @@ llama_kv_cache_msa::llama_kv_cache_msa(
     LLAMA_LOG_INFO("%s: creating main KV cache, size = %u cells\n", __func__, kv_size);
 
     kv_base = std::make_unique<llama_kv_cache>(
-            model, model.hparams, type_k, type_v,
+            model, no_alloc, model.hparams, type_k, type_v,
             v_trans, offload, unified, kv_size, n_seq_max, n_pad,
             n_swa, swa_type, nullptr, filter, reuse, nullptr);
 
@@ -44,7 +45,7 @@ llama_kv_cache_msa::llama_kv_cache_msa(
     LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);
 
     kv_idx = std::make_unique<llama_kv_cache>(
-            model, hparams_idx, type_k, type_v,
+            model, no_alloc, hparams_idx, type_k, type_v,
             v_trans, offload, unified, kv_size, n_seq_max, n_pad,
             n_swa, swa_type, nullptr, filter_idx, reuse, nullptr);
 }

src/llama-kv-cache-msa.h

@@ -15,6 +15,7 @@ class llama_kv_cache_msa : public llama_memory_i {
 public:
     llama_kv_cache_msa(
             const llama_model & model,
+                         bool   no_alloc,
                     ggml_type   type_k,
                     ggml_type   type_v,
                          bool   v_trans,

src/llama-kv-cache.cpp

@@ -64,7 +64,8 @@ static void ggml_gen_hadamard(ggml_tensor * tensor) {
 
 llama_kv_cache::llama_kv_cache(
         const llama_model & model,
-        const llama_hparams & hparams,
+                     bool   no_alloc,
+      const llama_hparams & hparams,
                 ggml_type   type_k,
                 ggml_type   type_v,
                      bool   v_trans,
@@ -79,9 +80,10 @@ llama_kv_cache::llama_kv_cache(
     const layer_filter_cb & filter,
     const  layer_reuse_cb & reuse,
     const  layer_share_cb & share,
-             const char *   name_tag) :
+               const char * name_tag) :
     model(model), hparams(hparams), v_trans(v_trans),
     n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type),
+    ctxs_bufs(no_alloc),
     other(static_cast<llama_kv_cache *>(mem_other)),
     v_cells_impl(other ? other->v_cells_impl : std::make_shared<llama_kv_cells_vec>()),
     v_cells(*v_cells_impl) {
@@ -273,25 +275,11 @@ llama_kv_cache::llama_kv_cache(
         }
     }
 
-    // allocate tensors and initialize the buffers to avoid NaNs in the padding
+    // allocate the tensors and initialize the buffers to avoid NaNs in the padding
     for (auto & [buft, ctx] : ctx_map) {
-        ggml_backend_buffer_t buf;
-        if (hparams.no_alloc) {
-            buf = ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0); // dummy buffer
-            for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) {
-                t->buffer = buf; // set dummy buffer for KV cache so that the backend scheduler won't try to allocate it
-            }
-        } else {
-            buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer
-        }
-        if (!buf) {
-            throw std::runtime_error("failed to allocate buffer for kv cache");
-        }
+        ggml_backend_buffer_t buf = ctxs_bufs.emplace(std::move(ctx), buft, "kv cache");
 
         LLAMA_LOG_INFO("%s: %10s KV buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0);
-
-        ggml_backend_buffer_clear(buf, 0);
-        ctxs_bufs.emplace_back(std::move(ctx), buf);
     }
 
     {
@@ -705,20 +693,7 @@ llama_pos llama_kv_cache::seq_pos_max(llama_seq_id seq_id) const {
 }
 
 std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache::memory_breakdown() const {
-    std::map<ggml_backend_buffer_type_t, size_t> ret;
-    for (const auto & [ctx, buf] : ctxs_bufs) {
-        ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get());
-
-        if (hparams.no_alloc) {
-            GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) == nullptr);
-            ret[buft] += ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft);
-        } else {
-            // GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) != nullptr); // multi_buffer does not have a defined base
-            ret[buft] += ggml_backend_buffer_get_size(buf.get());
-        }
-    }
-
-    return ret;
+    return ctxs_bufs.memory_breakdown();
 }
 
 llama_memory_context_ptr llama_kv_cache::init_batch(

src/llama-kv-cache.h

@@ -98,6 +98,7 @@ public:
     //       likely through `struct llama_memory_params`
     llama_kv_cache(
             const llama_model & model,
+                         bool   no_alloc,
           const llama_hparams & hparams,
                     ggml_type   type_k,
                     ggml_type   type_v,
@@ -296,7 +297,7 @@ private:
     const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;
 
     // ggml contexts for the KV cache along with the allocated backend buffers:
-    std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
+    llama_memory_bufs ctxs_bufs;
 
     // the current index from where we start searching for a free slot in the ring buffer of KV cells (see find_slot())
     // note: this is not part of the KV state and it's only used to speed-up the find_slot() method

src/llama-memory-hybrid-idx.cpp

@@ -1,20 +1,16 @@
 #include "llama-memory-hybrid-idx.h"
 
-#include <algorithm>
-#include <cmath>
-#include <type_traits>
-
 #include "llama-impl.h"
 #include "llama-batch.h"
 #include "llama-io.h"
 #include "llama-model.h"
 
-
 #include <algorithm>
 #include <cassert>
 #include <cmath>
 #include <iterator>
 #include <stdexcept>
+#include <type_traits>
 
 //
 // llama_memory_hybrid_idx
@@ -22,6 +18,7 @@
 
 llama_memory_hybrid_idx::llama_memory_hybrid_idx(
         const llama_model & model,
+                     bool   no_alloc,
                             /* attn */
                 ggml_type   type_k,
                 ggml_type   type_v,
@@ -45,6 +42,7 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx(
     const layer_filter_cb & filter_idx) :
     llama_memory_hybrid(
         model,
+        no_alloc,
         type_k, type_v, v_trans, kv_size, n_pad, n_swa, swa_type,
         type_r, type_s, rs_size,
         n_seq_max, n_rs_seq, offload, unified,
@@ -68,7 +66,7 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx(
         LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);
 
         return new llama_kv_cache(
-            model, hparams_idx, type_k, type_v, v_trans, offload, unified,
+            model, no_alloc, hparams_idx, type_k, type_v, v_trans, offload, unified,
             kv_size, n_seq_max, n_pad, n_swa, swa_type,
             nullptr, filter_idx, nullptr, nullptr, "idx_");
     }()) {}

src/llama-memory-hybrid-idx.h

@@ -18,6 +18,7 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid {
 public:
     llama_memory_hybrid_idx(
         const llama_model & model,
+                     bool   no_alloc,
                             /* attn */
                 ggml_type   type_k,
                 ggml_type   type_v,

src/llama-memory-hybrid-iswa.cpp

@@ -10,6 +10,7 @@
 
 llama_memory_hybrid_iswa::llama_memory_hybrid_iswa(
         const llama_model & model,
+                      bool   no_alloc,
                             /* attn */
                 ggml_type   type_k,
                 ggml_type   type_v,
@@ -33,6 +34,7 @@ llama_memory_hybrid_iswa::llama_memory_hybrid_iswa(
     hparams(model.hparams),
     mem_attn(new llama_kv_cache_iswa(
         model,
+        no_alloc,
         type_k,
         type_v,
         v_trans,
@@ -52,6 +54,7 @@ llama_memory_hybrid_iswa::llama_memory_hybrid_iswa(
     )),
     mem_recr(new llama_memory_recurrent(
         model,
+        no_alloc,
         type_r,
         type_s,
         offload,

src/llama-memory-hybrid-iswa.h

@@ -20,6 +20,7 @@ class llama_memory_hybrid_iswa : public llama_memory_i {
 public:
     llama_memory_hybrid_iswa(
         const llama_model & model,
+                      bool   no_alloc,
                             /* attn */
                 ggml_type   type_k,
                 ggml_type   type_v,

src/llama-memory-hybrid.cpp

@@ -10,6 +10,7 @@
 
 llama_memory_hybrid::llama_memory_hybrid(
         const llama_model & model,
+                     bool   no_alloc,
                             /* attn */
                 ggml_type   type_k,
                 ggml_type   type_v,
@@ -33,6 +34,7 @@ llama_memory_hybrid::llama_memory_hybrid(
     hparams(model.hparams),
     mem_attn(new llama_kv_cache(
         model,
+        no_alloc,
         model.hparams,
         type_k,
         type_v,
@@ -53,6 +55,7 @@ llama_memory_hybrid::llama_memory_hybrid(
     )),
     mem_recr(new llama_memory_recurrent(
         model,
+        no_alloc,
         type_r,
         type_s,
         offload,

src/llama-memory-hybrid.h

@@ -20,6 +20,7 @@ class llama_memory_hybrid : public llama_memory_i {
 public:
     llama_memory_hybrid(
         const llama_model & model,
+                     bool   no_alloc,
                             /* attn */
                 ggml_type   type_k,
                 ggml_type   type_v,

src/llama-memory-recurrent.cpp

@@ -19,13 +19,14 @@
 
 llama_memory_recurrent::llama_memory_recurrent(
         const llama_model & model,
+                 bool   no_alloc,
                 ggml_type   type_r,
                 ggml_type   type_s,
                      bool   offload,
                  uint32_t   mem_size,
                  uint32_t   n_seq_max,
                  uint32_t   n_rs_seq,
-    const layer_filter_cb & filter) : hparams(model.hparams), n_seq_max(n_seq_max) {
+    const layer_filter_cb & filter) : hparams(model.hparams), n_seq_max(n_seq_max), ctxs_bufs(no_alloc) {
     const int32_t n_layer = hparams.n_layer();
 
     head = 0;
@@ -114,23 +115,11 @@ llama_memory_recurrent::llama_memory_recurrent(
         }
     }
 
-    // allocate tensors and initialize the buffers to avoid NaNs in the padding
+    // allocate the tensors and initialize the buffers to avoid NaNs in the padding
     for (auto & [buft, ctx] : ctx_map) {
-        ggml_backend_buffer_t buf;
-        if (hparams.no_alloc) {
-            buf = ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0); // dummy buffer
-            for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) {
-                t->buffer = buf; // set dummy buffer for RS cache so that the backend scheduler won't try to allocate it
-            }
-        } else {
-            buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer
-        }
-        if (!buf) {
-            throw std::runtime_error("failed to allocate buffer for rs cache");
-        }
-        ggml_backend_buffer_clear(buf, 0);
+        ggml_backend_buffer_t buf = ctxs_bufs.emplace(std::move(ctx), buft, "rs cache");
+
         LLAMA_LOG_INFO("%s: %10s RS buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0);
-        ctxs_bufs.emplace_back(std::move(ctx), buf);
     }
 
     if (is_empty()) {
@@ -442,18 +431,7 @@ void llama_memory_recurrent::set_rs_idx(llama_seq_id seq_id, uint32_t idx) {
 }
 
 std::map<ggml_backend_buffer_type_t, size_t> llama_memory_recurrent::memory_breakdown() const {
-    std::map<ggml_backend_buffer_type_t, size_t> ret;
-    for (const auto & [ctx, buf] : ctxs_bufs) {
-        ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get());
-
-        if (hparams.no_alloc) {
-            GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) == nullptr);
-            ret[buft] += ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft);
-        } else {
-            ret[buft] += ggml_backend_buffer_get_size(buf.get());
-        }
-    }
-    return ret;
+    return ctxs_bufs.memory_breakdown();
 }
 
 llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) {

src/llama-memory-recurrent.h

@@ -18,6 +18,7 @@ class llama_memory_recurrent : public llama_memory_i {
 public:
     llama_memory_recurrent(
             const llama_model & model,
+                         bool   no_alloc,
                     ggml_type   type_r,
                     ggml_type   type_s,
                          bool   offload,
@@ -120,8 +121,8 @@ private:
 
     const uint32_t n_seq_max = 1;
 
-    // ggml contexts for the KV cache along with the allocated backend buffers:
-    std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
+    // ggml contexts for the RS cache along with the allocated backend buffers:
+    llama_memory_bufs ctxs_bufs;
 
     // true if no layers - can happen if the layer filter removes all layers
     bool is_empty() const;

src/llama-memory.cpp

@@ -4,6 +4,8 @@
 #include "llama-io.h"
 
 #include <set>
+#include <stdexcept>
+#include <string>
 #include <vector>
 
 llama_memory_status llama_memory_status_combine(llama_memory_status s0, llama_memory_status s1) {
@@ -64,6 +66,52 @@ bool llama_memory_status_is_fail(llama_memory_status status) {
     return false;
 }
 
+//
+// llama_memory_bufs
+//
+
+ggml_backend_buffer_ptr llama_memory_bufs::alloc_buf(ggml_context * ctx, ggml_backend_buffer_type_t buft, bool no_alloc, const char * name) {
+    ggml_backend_buffer_ptr buf;
+    if (no_alloc) {
+        buf.reset(ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0)); // dummy buffer
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+            t->buffer = buf.get(); // set dummy buffer so that the backend scheduler won't try to allocate the tensors
+        }
+    } else {
+        buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx, buft)); // real buffer
+    }
+    if (!buf) {
+        throw std::runtime_error(std::string("failed to allocate buffer for ") + name);
+    }
+    return buf;
+}
+
+ggml_backend_buffer_t llama_memory_bufs::emplace(ggml_context_ptr ctx, ggml_backend_buffer_type_t buft, const char * name) {
+    ggml_backend_buffer_ptr buf = alloc_buf(ctx.get(), buft, no_alloc, name);
+
+    // initialize the buffer to avoid NaNs in the padding
+    ggml_backend_buffer_clear(buf.get(), 0);
+
+    bufs.emplace_back(std::move(ctx), std::move(buf));
+
+    return bufs.back().second.get();
+}
+
+std::map<ggml_backend_buffer_type_t, size_t> llama_memory_bufs::memory_breakdown() const {
+    std::map<ggml_backend_buffer_type_t, size_t> ret;
+    for (const auto & [ctx, buf] : bufs) {
+        ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get());
+
+        if (no_alloc) {
+            GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) == nullptr);
+            ret[buft] += ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft);
+        } else {
+            ret[buft] += ggml_backend_buffer_get_size(buf.get());
+        }
+    }
+    return ret;
+}
+
 // collects the buffers of the tensors that a state write visits, without keeping the data
 struct llama_memory_buffers_writer : public llama_io_write_i {
     std::set<ggml_backend_buffer_t> buffers;
@@ -84,6 +132,10 @@ struct llama_memory_buffers_writer : public llama_io_write_i {
 };
 
 std::vector<ggml_backend_buffer_t> llama_memory_get_buffers(const llama_memory_t mem) {
+    if (mem == nullptr) {
+        return {};
+    }
+
     llama_memory_buffers_writer writer;
 
     mem->state_write(writer);

src/llama-memory.h

@@ -6,6 +6,7 @@
 #include <map>
 #include <memory>
 #include <functional>
+#include <vector>
 
 struct llama_ubatch;
 
@@ -25,6 +26,36 @@ struct llama_memory_params {
     llama_context_type ctx_type;
 
     llama_memory_t mem_other;
+
+    // only simulate the memory allocations
+    bool no_alloc;
+};
+
+// the tensor contexts and backend buffers of a memory instance
+// with no_alloc, dummy buffers are attached instead of allocating the tensor data, so the memory use can still be measured
+struct llama_memory_bufs {
+    explicit llama_memory_bufs(bool no_alloc) : no_alloc(no_alloc) {}
+
+    // allocate the tensors of ctx on buft, or attach a dummy buffer if no_alloc is set
+    // returns the buffer, for logging
+    ggml_backend_buffer_t emplace(ggml_context_ptr ctx, ggml_backend_buffer_type_t buft, const char * name);
+
+    // the memory used by the tensors: the buffer sizes, or the simulated sizes if no_alloc is set
+    std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const;
+
+    // allocate a single ctx buffer - for the memory instances that keep their own buffer layout
+    static ggml_backend_buffer_ptr alloc_buf(ggml_context * ctx, ggml_backend_buffer_type_t buft, bool no_alloc, const char * name);
+
+    bool empty() const { return bufs.empty(); }
+
+    auto begin()       { return bufs.begin(); }
+    auto end()         { return bufs.end();   }
+    auto begin() const { return bufs.begin(); }
+    auto end()   const { return bufs.end();   }
+
+    std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> bufs;
+
+    const bool no_alloc;
 };
 
 enum llama_memory_status {

src/llama-model-saver.cpp

@@ -4,44 +4,23 @@
 #include "gguf.h"
 
 #include "llama-arch.h"
+#include "llama-ext.h"
 #include "llama.h"
 #include "llama-hparams.h"
 #include "llama-model.h"
 #include "llama-vocab.h"
-#include "models/models.h"
 
 #include <cstdint>
 #include <string>
 #include <unordered_set>
 
-bool llama_model_saver_supports_arch(llm_arch arch) {
-    switch (arch) {
-        case LLM_ARCH_SENSENOVA_U1:
-        case LLM_ARCH_GEMMA3N:
-        case LLM_ARCH_BITNET:
-        case LLM_ARCH_T5:
-        case LLM_ARCH_APERTUS:
-        case LLM_ARCH_STEP35:
-        case LLM_ARCH_CLEF: // the head tensors are not saved
-        case LLM_ARCH_VLLM_DECISION2:
-            return false;
-        default:
-            return true;
-    }
-}
-
 llama_model_saver::llama_model_saver(const struct llama_model * model) :
-        gguf_ctx(gguf_init_empty()), gguf_ctx_owned(true), model(model), llm_kv(model->arch) {
-    GGML_ASSERT(llama_model_saver_supports_arch(model->arch));
+        gguf_ctx(gguf_init_empty()), model(model), llm_kv(model->arch) {
+    GGML_ASSERT(llama_arch_supports_save(model->arch));
 }
 
-llama_model_saver::llama_model_saver(enum llm_arch arch, struct gguf_context * gguf_ctx) :
-        gguf_ctx(gguf_ctx == nullptr ? gguf_init_empty() : gguf_ctx), gguf_ctx_owned(gguf_ctx == nullptr), model(nullptr), llm_kv(arch) {}
-
 llama_model_saver::~llama_model_saver() {
-    if (gguf_ctx_owned) {
-        gguf_free(gguf_ctx);
-    }
+    gguf_free(gguf_ctx);
 }
 
 void llama_model_saver::add_kv(const enum llm_kv key, const uint32_t value) {
@@ -119,11 +98,6 @@ void llama_model_saver::add_kv(const enum llm_kv key, const Container & value, c
         GGML_ABORT("fatal error");
     }
 }
-// instantiate for external usage:
-template void llama_model_saver::add_kv<std::vector<uint32_t>>(const enum llm_kv, const std::vector<uint32_t> &, const bool);
-template void llama_model_saver::add_kv<std::vector<float>>(const enum llm_kv, const std::vector<float> &, const bool);
-template void llama_model_saver::add_kv<std::vector<uint64_t>>(const enum llm_kv, const std::vector<uint64_t> &, const bool);
-template void llama_model_saver::add_kv<std::vector<int32_t>>(const enum llm_kv, const std::vector<int32_t> &, const bool);
 
 void llama_model_saver::add_kv(const enum llm_kv key, const std::vector<std::string> & value) {
     std::vector<const char *> tmp(value.size());
@@ -342,8 +316,6 @@ void llama_model_saver::add_kv_from_model() {
     add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK,             hparams.hc_low_rank);
 
     if (model->arch == LLM_ARCH_DEEPSEEK41) {
-        const auto * m41 = dynamic_cast<const llama_model_deepseek41 *>(model);
-        GGML_ASSERT(m41 != nullptr);
         std::vector<uint32_t> kv_src;
         std::vector<uint32_t> idx_src;
         for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
@@ -362,8 +334,8 @@ void llama_model_saver::add_kv_from_model() {
             add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE,   hparams.dsv41_cand_block_size);
         }
         if (hparams.ple_n_heads > 0) {
-            add_kv(LLM_KV_ENGRAM_PAD_TOKEN_ID, m41->engram_pad_id);
-            add_kv(LLM_KV_ENGRAM_TOKEN_MAP,    m41->engram_token_map);
+            add_kv(LLM_KV_ENGRAM_PAD_TOKEN_ID, hparams.engram_pad_id);
+            add_kv(LLM_KV_ENGRAM_TOKEN_MAP,    model->engram_token_map);
         }
     }
 
@@ -551,7 +523,3 @@ void llama_model_saver::add_tensors_from_model() {
 void llama_model_saver::save(const std::string & path_model) {
     gguf_write_to_file(gguf_ctx, path_model.c_str(), false);
 }
-
-void llama_model_saver::save(FILE * file) {
-    gguf_write_to_file_ptr(gguf_ctx, file, false);
-}

src/llama-model-saver.h

@@ -6,17 +6,12 @@
 
 #include <vector>
 
-// FIXME temporary function for better error messages
-bool llama_model_saver_supports_arch(llm_arch arch);
-
 struct llama_model_saver {
     struct gguf_context * gguf_ctx = nullptr;
-    const bool gguf_ctx_owned;
     const struct llama_model * model;
     const struct LLM_KV llm_kv;
 
     llama_model_saver(const struct llama_model * model);
-    llama_model_saver(enum llm_arch arch, struct gguf_context * gguf_ctx);
     ~llama_model_saver();
 
     void add_kv(enum llm_kv key, uint32_t     value);
@@ -41,5 +36,4 @@ struct llama_model_saver {
     void add_tensors_from_model();
 
     void save(const std::string & path_model);
-    void save(FILE * file);
 };

src/llama-model.cpp

@@ -2448,10 +2448,14 @@ llama_split_mode llama_model::split_mode() const {
     return params.split_mode;
 }
 
+bool llama_model::no_alloc() const {
+    return params.no_alloc;
+}
+
 std::map<ggml_backend_buffer_type_t, size_t> llama_model::memory_breakdown() const {
     std::map<ggml_backend_buffer_type_t, size_t> ret;
     for (const auto & [ctx, bufs] : pimpl->ctxs_bufs) {
-        if (hparams.no_alloc) {
+        if (no_alloc()) {
             GGML_ASSERT(bufs.size() == 1);
             ggml_backend_buffer_t buf = bufs[0].get();
             GGML_ASSERT(ggml_backend_buffer_get_base(buf) == nullptr);
@@ -2506,7 +2510,7 @@ void llama_model::print_info() const {
     // hparams
     LLAMA_LOG_INFO("%s: arch                  = %s\n",     __func__, arch_name().c_str());
     LLAMA_LOG_INFO("%s: vocab_only            = %d\n",     __func__, hparams.vocab_only);
-    LLAMA_LOG_INFO("%s: no_alloc              = %d\n",     __func__, hparams.no_alloc);
+    LLAMA_LOG_INFO("%s: no_alloc              = %d\n",     __func__, no_alloc());
 
     if (!hparams.vocab_only) {
         LLAMA_LOG_INFO("%s: n_ctx_train           = %u\n",     __func__, hparams.n_ctx_train);
@@ -2830,6 +2834,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
 
                 res = new llama_kv_cache_msa(
                         *this,
+                        params.no_alloc,
                         params.type_k,
                         params.type_v,
                         !cparams.flash_attn,
@@ -2856,6 +2861,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
 
                     res = new llama_kv_cache(
                             *this,
+                            params.no_alloc,
                             hparams,
                             params.type_k,
                             params.type_v,
@@ -2882,6 +2888,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
 
                     res = new llama_kv_cache_dsa(
                             *this,
+                            params.no_alloc,
                             params.type_k,
                             params.type_v,
                             !cparams.flash_attn,
@@ -2923,6 +2930,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
 
                 res = new llama_memory_hybrid_idx(
                     /* model             */ *this,
+                    /* no_alloc          */ params.no_alloc,
                     /* attn_type_k       */ params.type_k,
                     /* attn_type_v       */ params.type_v,
                     /* attn_v_trans      */ !cparams.flash_attn,
@@ -2947,6 +2955,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                     // full-attention checkpoint: no indexer, so no indexer key cache
                     res = new llama_kv_cache(
                             *this,
+                            params.no_alloc,
                             hparams,
                             params.type_k,
                             params.type_v,
@@ -2968,6 +2977,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
 
                     res = new llama_kv_cache_dsa(
                             *this,
+                            params.no_alloc,
                             params.type_k,
                             params.type_v,
                             !cparams.flash_attn,
@@ -2994,6 +3004,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
 
                     res = new llama_kv_cache(
                             *this,
+                            params.no_alloc,
                             hparams,
                             params.type_k,
                             params.type_v,
@@ -3019,6 +3030,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
 
                     res = new llama_kv_cache_dsa_iswa(
                             *this,
+                            params.no_alloc,
                             params.type_k,
                             params.type_v,
                             !cparams.flash_attn,
@@ -3046,6 +3058,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
 
                     res = new llama_kv_cache_iswa(
                             *this,
+                            params.no_alloc,
                             params.type_k,
                             params.type_v,
                             !cparams.flash_attn,
@@ -3063,6 +3076,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                 } else {
                     res = new llama_kv_cache_dsv4(
                             *this,
+                            params.no_alloc,
                             params.type_k,
                             params.type_v,
                             !cparams.flash_attn,
@@ -3086,6 +3100,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
 
                     res = new llama_kv_cache_iswa(
                             *this,
+                            params.no_alloc,
                             params.type_k,
                             params.type_v,
                             !cparams.flash_attn,
@@ -3120,6 +3135,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                 if (llm_arch_is_recurrent(arch)) {
                     res = new llama_memory_recurrent(
                             *this,
+                            params.no_alloc,
                             GGML_TYPE_F32,
                             GGML_TYPE_F32,
                             cparams.offload_kqv,
@@ -3175,6 +3191,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                         // Use hybrid-iswa for hybrid models with SWA
                         res = new llama_memory_hybrid_iswa(
                             /* model             */ *this,
+                            /* no_alloc          */ params.no_alloc,
                             /* attn_type_k       */ params.type_k,
                             /* attn_type_v       */ params.type_v,
                             /* attn_v_trans      */ !cparams.flash_attn,
@@ -3195,6 +3212,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                         // sparse attention over a per-token indexer cache, in its own memory type
                         res = new llama_memory_hybrid_idx(
                             /* model             */ *this,
+                            /* no_alloc          */ params.no_alloc,
                             /* attn_type_k       */ params.type_k,
                             /* attn_type_v       */ params.type_v,
                             /* attn_v_trans      */ !cparams.flash_attn,
@@ -3215,6 +3233,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                     } else {
                         res = new llama_memory_hybrid(
                             /* model             */ *this,
+                            /* no_alloc          */ params.no_alloc,
                             /* attn_type_k       */ params.type_k,
                             /* attn_type_v       */ params.type_v,
                             /* attn_v_trans      */ !cparams.flash_attn,
@@ -3281,6 +3300,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
 
                             res = new llama_kv_cache_iswa(
                                     *this,
+                                    params.no_alloc,
                                     params.type_k,
                                     params.type_v,
                                     !cparams.flash_attn,
@@ -3298,6 +3318,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                         } else {
                             res = new llama_kv_cache_iswa(
                                     *this,
+                                    params.no_alloc,
                                     params.type_k,
                                     params.type_v,
                                     !cparams.flash_attn,
@@ -3318,6 +3339,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
 
                         res = new llama_kv_cache(
                                 *this,
+                                params.no_alloc,
                                 hparams,
                                 params.type_k,
                                 params.type_v,

src/llama-model.h

@@ -623,6 +623,10 @@ struct llama_layer {
     struct ggml_tensor * hc_ffn_down    = nullptr;
     struct ggml_tensor * hc_ffn_up      = nullptr;
     struct ggml_tensor * hc_ffn_inject  = nullptr;
+    struct ggml_tensor * hc_attn_down_inject = nullptr; // [hc_attn_down | hc_attn_inject] in one weight
+    struct ggml_tensor * hc_ffn_down_inject  = nullptr; // [hc_ffn_down | hc_ffn_inject] in one weight
+    struct ggml_tensor * wqkv_z              = nullptr; // [wqkv | wqkv_gate] in one weight
+    struct ggml_tensor * ssm_beta_alpha_cat  = nullptr; // [ssm_beta | ssm_alpha] in one weight, unlike the grouped ssm_beta_alpha
 
     struct ggml_tensor * ple_key        = nullptr;
     struct ggml_tensor * ple_value      = nullptr;
@@ -730,6 +734,9 @@ struct llama_model {
     std::vector<std::string> classifier_labels;
     std::vector<std::string> decision_labels;
 
+    // deepseek41 Engram: token id -> id of its normalized form, hashed into the n-grams
+    std::vector<int32_t> engram_token_map;
+
     struct ggml_tensor * tok_embd   = nullptr;
     struct ggml_tensor * type_embd  = nullptr;
     struct ggml_tensor * pos_embd   = nullptr;
@@ -879,6 +886,7 @@ struct llama_model {
 
     uint32_t n_gpu_layers() const;
     llama_split_mode split_mode() const;
+    bool no_alloc() const;
 
     std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const;
 

src/llama-moe-cache.cpp

@@ -1,6 +1,7 @@
 #include "llama-moe-cache.h"
 
 #include "llama-impl.h"
+#include "llama-memory.h"
 #include "llama-model.h"
 
 #include "ggml-cpp.h"
@@ -252,7 +253,7 @@ struct llama_moe_cache::impl {
     // views used by copy_experts
     ggml_context_ptr ctx_views;
 
-    impl(const llama_model & model, const std::vector<ggml_backend_t> & backends, const std::vector<ggml_backend_buffer_type_t> & bufts, size_t size) :
+    impl(const llama_model & model, bool no_alloc, const std::vector<ggml_backend_t> & backends, const std::vector<ggml_backend_buffer_type_t> & bufts, size_t size) :
             n_expert_used(model.hparams.n_expert_used_max()), layers(model.layers.size()) {
         for (size_t i = 0; i < backends.size(); ++i) {
             const auto dev_type = ggml_backend_dev_type(ggml_backend_get_device(backends[i]));
@@ -418,37 +419,19 @@ struct llama_moe_cache::impl {
             devices[g.id].buf_size += alloc_size(g, g.n_slots);
         }
 
-        if (model.hparams.no_alloc) {
-            // only used to measure the memory use, see llama_context::memory_breakdown
-            for (device & d : devices) {
-                if (!d.ctx) {
-                    continue;
-                }
-                d.buf.reset(ggml_backend_buft_alloc_buffer(d.buft, 0));
-                for (ggml_tensor * t = ggml_get_first_tensor(d.ctx.get()); t != nullptr; t = ggml_get_next_tensor(d.ctx.get(), t)) {
-                    t->buffer = d.buf.get();
-                }
+        // with no_alloc, dummy buffers are attached so the memory use can still be measured, see llama_context::memory_breakdown
+        for (device & d : devices) {
+            if (!d.ctx) {
+                continue;
             }
-            buf_host.reset(ggml_backend_buft_alloc_buffer(buft_host, 0));
-            for (ggml_tensor * t = ggml_get_first_tensor(ctx_host.get()); t != nullptr; t = ggml_get_next_tensor(ctx_host.get(), t)) {
-                t->buffer = buf_host.get();
-            }
-        } else {
-            for (device & d : devices) {
-                if (!d.ctx) {
-                    continue;
-                }
-                d.buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(d.ctx.get(), d.buft));
-                if (!d.buf) {
-                    throw std::runtime_error("failed to allocate the MoE cache buffers");
-                }
+            d.buf = llama_memory_bufs::alloc_buf(d.ctx.get(), d.buft, no_alloc, "MoE cache");
+            if (!no_alloc) {
                 ggml_backend_buffer_clear(d.buf.get(), 0);
                 d.buf_size = ggml_backend_buffer_get_size(d.buf.get());
             }
-            buf_host.reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx_host.get(), buft_host));
-            if (!buf_host) {
-                throw std::runtime_error("failed to allocate the MoE cache buffers");
-            }
+        }
+        buf_host = llama_memory_bufs::alloc_buf(ctx_host.get(), buft_host, no_alloc, "MoE cache slot maps");
+        if (!no_alloc) {
             ggml_backend_buffer_clear(buf_host.get(), 0xff); // all slots are -1
             buf_host_size = ggml_backend_buffer_get_size(buf_host.get());
 
@@ -633,8 +616,8 @@ struct llama_moe_cache::impl {
     }
 };
 
-llama_moe_cache::llama_moe_cache(const llama_model & model, const std::vector<ggml_backend_t> & backends, const std::vector<ggml_backend_buffer_type_t> & bufts, size_t size) :
-    pimpl(new impl(model, backends, bufts, size)) {
+llama_moe_cache::llama_moe_cache(const llama_model & model, bool no_alloc, const std::vector<ggml_backend_t> & backends, const std::vector<ggml_backend_buffer_type_t> & bufts, size_t size) :
+    pimpl(new impl(model, no_alloc, backends, bufts, size)) {
 }
 
 llama_moe_cache::~llama_moe_cache() = default;

src/llama-moe-cache.h

@@ -13,7 +13,7 @@ struct llama_model;
 class llama_moe_cache {
 public:
     // backends are all the backends of the context, each GPU gets its own cache of the given size for the layers assigned to it
-    llama_moe_cache(const llama_model & model, const std::vector<ggml_backend_t> & backends, const std::vector<ggml_backend_buffer_type_t> & bufts, size_t size);
+    llama_moe_cache(const llama_model & model, bool no_alloc, const std::vector<ggml_backend_t> & backends, const std::vector<ggml_backend_buffer_type_t> & bufts, size_t size);
     ~llama_moe_cache();
 
     // the device that caches layer il

src/llama-quant.cpp

@@ -125,7 +125,8 @@ static tensor_category tensor_get_category(const std::string & tensor_name) {
     if (tensor_name_match_token_embd(tensor_name.c_str())) {
         return tensor_category::TOKEN_EMBD;
     }
-    if (tensor_name.find("attn_qkv.weight") != std::string::npos) {
+    // attn_qkv_gate is [attn_qkv | attn_gate] in one weight, the qkv rows decide its type
+    if (tensor_name.find("attn_qkv.weight") != std::string::npos || tensor_name.find("attn_qkv_gate.weight") != std::string::npos) {
         return tensor_category::ATTENTION_QKV;
     }
     if (tensor_name.find("attn_kv_b.weight") != std::string::npos) {
@@ -161,6 +162,12 @@ static tensor_category tensor_get_category(const std::string & tensor_name) {
     return tensor_category::OTHER;
 }
 
+// check if tensor is a MoE expert tensor (used by GGML_OP_MUL_MAT_ID)
+static bool tensor_is_expert(const std::string & tensor_name) {
+    return tensor_name.find("_exps.")   != std::string::npos ||
+           tensor_name.find("_chexps.") != std::string::npos;
+}
+
 // check if category is for attention-v-like tensors (more sensitive to quantization)
 static bool category_is_attn_v(tensor_category cat) {
     return cat == tensor_category::ATTENTION_V     ||
@@ -296,7 +303,7 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param
     // trivial checks first -- no string ops needed
     if (params->only_copy)       return false;
 
-    // quantize only 2D and 3D tensors (experts)
+    // quantize only 2D+ tensors (matrices, experts, conv kernels)
     if (ggml_n_dims(tensor) < 2) return false;
 
     const std::string name = ggml_get_name(tensor);
@@ -541,12 +548,7 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type
     } else if (ftype == LLAMA_FTYPE_MOSTLY_MXFP4_MOE) {
         // MoE   tensors -> MXFP4
         // other tensors -> Q8_0
-        // MLA projection tensors are also 3D, so match expert tensor roles explicitly.
-        const bool is_bailingmoe3_expert = arch == LLM_ARCH_BAILINGMOE3 &&
-            (category == tensor_category::FFN_UP ||
-             category == tensor_category::FFN_GATE ||
-             category == tensor_category::FFN_DOWN);
-        if (tensor->ne[2] > 1 && (arch != LLM_ARCH_BAILINGMOE3 || is_bailingmoe3_expert)) {
+        if (tensor_is_expert(name)) {
             new_type = GGML_TYPE_MXFP4;
         } else {
             new_type = GGML_TYPE_Q8_0;
@@ -1307,11 +1309,11 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
                     if (it == imatrix_data->end()) {
                         LLAMA_LOG_INFO("\n====== %s: did not find weights for %s\n", __func__, tensor->name);
                     } else {
-                        if (it->second.size() == (size_t)tensor->ne[0]*tensor->ne[2]) {
+                        if (it->second.size() == (size_t)tensor->ne[0]*tensor->ne[2]*tensor->ne[3]) {
                             imatrix = it->second.data();
                         } else {
                             LLAMA_LOG_INFO("\n====== %s: imatrix size %d is different from tensor size %d for %s\n", __func__,
-                                    int(it->second.size()), int(tensor->ne[0]*tensor->ne[2]), tensor->name);
+                                    int(it->second.size()), int(tensor->ne[0]*tensor->ne[2]*tensor->ne[3]), tensor->name);
 
                             // this can happen when quantizing an old mixtral model with split tensors with a new incompatible imatrix
                             // this is a significant error and it may be good idea to abort the process if this happens,
@@ -1319,7 +1321,7 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
                             // tok_embd should be ignored in this case, since it always causes this warning
                             if (!tensor_name_match_token_embd(tensor->name)) {
                                 throw std::runtime_error(format("imatrix size %d is different from tensor size %d for %s",
-                                        int(it->second.size()), int(tensor->ne[0]*tensor->ne[2]), tensor->name));
+                                        int(it->second.size()), int(tensor->ne[0]*tensor->ne[2]*tensor->ne[3]), tensor->name));
                             }
                         }
                     }
@@ -1341,7 +1343,7 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
 
                 const int64_t n_per_row = tensor->ne[0];
                 const int64_t nrows_per_expert = tensor->ne[1];
-                const int64_t nrows_total = tensor->ne[1] * tensor->ne[2];
+                const int64_t nrows_total = ggml_nrows(tensor);
 
                 const size_t row_size_src = ggml_row_size(tensor->type, n_per_row);
                 const size_t row_size_dst = ggml_row_size(new_type,     n_per_row);

src/llama-sampler.cpp

@@ -18,6 +18,7 @@
 #include <numeric>
 #include <random>
 #include <unordered_map>
+#include <unordered_set>
 #include <stdexcept>
 
 // the ring buffer works similarly to std::deque, but with a fixed capacity
@@ -3333,9 +3334,13 @@ struct llama_sampler_dry {
 };
 
 // Ported from Koboldcpp, original PR: https://github.com/LostRuins/koboldcpp/pull/982 (Original author: pi6am)
-static void get_overlapping_token_sequences(const llama_vocab & vocab, const std::string& str, std::unordered_multimap<llama_token, std::vector<llama_token>>& token_sequences, int max_tail_len = -1) {
-    for (llama_token token_id = 0; token_id < (llama_token) vocab.n_tokens(); token_id++) {
-        std::string word = vocab.detokenize({token_id}, true);
+static void get_overlapping_token_sequences(const llama_vocab & vocab, const std::vector<std::string> & vocab_words, const std::string& str, std::unordered_multimap<llama_token, std::vector<llama_token>>& token_sequences, int max_tail_len = -1) {
+    // the tail tokenization only depends on how many characters of str were matched, so compute each one at most once
+    std::vector<std::vector<llama_token>> tail_cache(str.size() + 1);
+    std::vector<bool> tail_cached(str.size() + 1, false);
+
+    for (llama_token token_id = 0; token_id < (llama_token) vocab_words.size(); token_id++) {
+        const std::string & word = vocab_words[token_id];
         if (word.find(str) != std::string::npos) {
             token_sequences.emplace(token_id, std::vector<llama_token>());
         } else {
@@ -3352,10 +3357,14 @@ static void get_overlapping_token_sequences(const llama_vocab & vocab, const std
                     }
                 }
                 if (match) {
-                    std::vector<llama_token> tokenization = vocab.tokenize(str.substr(i), false, false);
-                    if (max_tail_len >= 0 && tokenization.size() > (size_t)max_tail_len) {
-                        tokenization.resize(max_tail_len);
+                    if (!tail_cached[i]) {
+                        tail_cache[i] = vocab.tokenize(str.substr(i), false, false);
+                        if (max_tail_len >= 0 && tail_cache[i].size() > (size_t)max_tail_len) {
+                            tail_cache[i].resize(max_tail_len);
+                        }
+                        tail_cached[i] = true;
                     }
+                    const std::vector<llama_token> & tokenization = tail_cache[i];
 
                     // Ensure we don't already have a duplicate matching tokenization
                     auto its = token_sequences.equal_range(token_id);
@@ -3656,6 +3665,16 @@ struct llama_sampler * llama_sampler_init_dry(const struct llama_vocab * vocab,
 
     if (dry_enabled && seq_breakers != nullptr && num_breakers > 0) {
         // Process sequence breakers
+        // each distinct breaker costs a full vocabulary scan: drop duplicates and stop at the cap
+        std::unordered_set<std::string> seen_breakers;
+
+        // detokenize the vocabulary once instead of rescanning it for every breaker
+        std::vector<std::string> vocab_words;
+        vocab_words.reserve(vocab->n_tokens());
+        for (llama_token token_id = 0; token_id < (llama_token) vocab->n_tokens(); ++token_id) {
+            vocab_words.push_back(vocab->detokenize({token_id}, true));
+        }
+
         for (size_t i = 0; i < num_breakers; ++i) {
             if (seq_breakers[i] == nullptr || std::strlen(seq_breakers[i]) == 0) {
                 LLAMA_LOG_WARN("skipping null or empty DRY sequence breaker at index %zu\n", i);
@@ -3663,17 +3682,22 @@ struct llama_sampler * llama_sampler_init_dry(const struct llama_vocab * vocab,
             }
 
             std::string sequence_break(seq_breakers[i]);
-            if (sequence_break.empty()) {
-                LLAMA_LOG_WARN("skipping empty DRY sequence breaker\n");
-                continue;
-            }
-
             if (sequence_break.size() > MAX_CHAR_LEN) {
                 LLAMA_LOG_WARN("truncating DRY sequence breaker to %d characters\n", MAX_CHAR_LEN);
                 sequence_break.resize(MAX_CHAR_LEN);
             }
 
-            get_overlapping_token_sequences(*vocab, sequence_break, processed_breakers, MAX_SEQ_LEN);
+            if (!seen_breakers.insert(sequence_break).second) {
+                continue;
+            }
+
+            if (seen_breakers.size() > LLAMA_DRY_MAX_SEQ_BREAKERS) {
+                LLAMA_LOG_WARN("ignoring DRY sequence breaker at index %zu: more than %d distinct breakers\n",
+                        i, LLAMA_DRY_MAX_SEQ_BREAKERS);
+                break;
+            }
+
+            get_overlapping_token_sequences(*vocab, vocab_words, sequence_break, processed_breakers, MAX_SEQ_LEN);
         }
     }
 

src/llama-vocab.cpp

@@ -29,31 +29,25 @@
 struct naive_trie {
     naive_trie() : has_value(false), value(0) {
     }
+    // no recursion: a token from the GGUF can be long
     void insert(const char * key, size_t len, int32_t value = 0) {
-        if (len == 0) {
-            this->has_value = true;
-            this->value = value;
-            return;
-        }
-        char c = key[0];
-        auto res = children.find(c);
-        if (res != children.end()) {
-            res->second.insert(key + 1, len - 1, value);
-        } else {
-            auto res = children.insert(std::make_pair(c, naive_trie()));
-            res.first->second.insert(key + 1, len - 1, value);
+        naive_trie * node = this;
+        for (size_t i = 0; i < len; i++) {
+            node = &node->children[key[i]];
         }
+        node->has_value = true;
+        node->value = value;
     }
     std::pair<const char *, size_t> get_longest_prefix(const char * key, size_t len, size_t offset = 0) const {
-        if (len == 0 || offset == len) {
-            return std::make_pair(key, offset);
+        const naive_trie * node = this;
+        while (offset < len) {
+            auto res = node->children.find(key[offset]);
+            if (res == node->children.end()) {
+                break;
+            }
+            node = &res->second;
+            offset++;
         }
-        char c = key[offset];
-        auto res = children.find(c);
-        if (res != children.end()) {
-            return res->second.get_longest_prefix(key, len, offset + 1);
-        }
-
         return std::make_pair(key, offset);
     }
     const struct naive_trie * traverse(const char c) const {
@@ -174,30 +168,37 @@ struct llm_tokenizer_spm_session {
     }
 
 private:
+    // no recursion: the merges of a long token can be deep
     void resegment(llm_symbol & symbol, std::vector<llama_token> & output) {
-        auto text = std::string(symbol.text, symbol.n);
-        auto token = vocab.text_to_token(text);
+        std::vector<const llm_symbol *> todo = { &symbol };
+        while (!todo.empty()) {
+            const llm_symbol & cur = *todo.back();
+            todo.pop_back();
 
-        // Do we need to support is_unused?
-        if (token != LLAMA_TOKEN_NULL) {
-            output.push_back(token);
-            return;
-        }
+            auto text = std::string(cur.text, cur.n);
+            auto token = vocab.text_to_token(text);
 
-        const auto p = rev_merge.find(text);
-
-        if (p == rev_merge.end()) {
-            // output any symbols that did not form tokens as bytes.
-            output.reserve(output.size() + symbol.n);
-            for (int j = 0; j < (int)symbol.n; ++j) {
-                llama_token id = vocab.byte_to_token(symbol.text[j]);
-                output.push_back(id);
+            // Do we need to support is_unused?
+            if (token != LLAMA_TOKEN_NULL) {
+                output.push_back(token);
+                continue;
             }
-            return;
-        }
 
-        resegment(symbols[p->second.first], output);
-        resegment(symbols[p->second.second], output);
+            const auto p = rev_merge.find(text);
+
+            if (p == rev_merge.end()) {
+                // output any symbols that did not form tokens as bytes.
+                output.reserve(output.size() + cur.n);
+                for (int j = 0; j < (int)cur.n; ++j) {
+                    llama_token id = vocab.byte_to_token(cur.text[j]);
+                    output.push_back(id);
+                }
+                continue;
+            }
+
+            todo.push_back(&symbols[p->second.second]);
+            todo.push_back(&symbols[p->second.first]);
+        }
     }
 
     void try_add_bigram(int left, int right) {

src/llama.cpp

@@ -342,7 +342,6 @@ static std::pair<int, llama_model *> llama_model_load(struct gguf_context * meta
         model->t_start_us = tm.t_start_us;
 
         model->hparams.vocab_only = params.vocab_only;
-        model->hparams.no_alloc   = params.no_alloc;
 
         try {
             model->load_hparams(ml);

src/models/deepseek4.cpp

@@ -342,25 +342,23 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_sinkhorn(
     return comb;
 }
 
-ggml_tensor * llama_model_deepseek4::graph::build_hc_pre(
+llama_model_deepseek4::graph::hc_pre
+llama_model_deepseek4::graph::build_hc_pre(
         ggml_tensor * x,
         ggml_tensor * hc_fn,
         ggml_tensor * hc_scale,
         ggml_tensor * hc_base,
-        ggml_tensor ** post,
-        ggml_tensor ** comb,
         int il) const {
-    return build_hc_pre(x, build_hc_mixes(x, hc_fn, hc_scale, hc_base, post, comb, false, il), il);
+    hc_mixes mixes = build_hc_mixes(x, hc_fn, hc_scale, hc_base, il);
+    return { build_hc_pre(x, mixes.pre, il), mixes.post, mixes.comb };
 }
 
-ggml_tensor * llama_model_deepseek4::graph::build_hc_mixes(
+llama_model_deepseek4::graph::hc_mixes
+llama_model_deepseek4::graph::build_hc_mixes(
         ggml_tensor * x,
         ggml_tensor * hc_fn,
         ggml_tensor * hc_scale,
         ggml_tensor * hc_base,
-        ggml_tensor ** post,
-        ggml_tensor ** comb,
-        bool f32_src,
         int il) const {
     const int64_t hc         = hparams.dsv4_hc_mult;
     const int64_t hc_dim     = hc*n_embd;
@@ -373,10 +371,6 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_mixes(
     ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt);
     ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, norm_rms_eps);
     ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm);
-    if (f32_src && !ggml_is_quantized(hc_fn->type)) {
-        // F32 weights (model.py hc_pre runs in fp32): request F32 inputs, the mixes are near-tie sensitive
-        ggml_prec_set_src(mixes, GGML_PREC_F32, 1);
-    }
     cb(mixes, "hc_mixes", il);
 
     ggml_tensor * scale_pre  = dsv4_view_1d(ctx0, hc_scale, 1, 0);
@@ -391,28 +385,29 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_mixes(
     pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps);
     cb(pre, "hc_pre", il);
 
-    *post = dsv4_view_2d(ctx0, mixes, hc, nt, hc);
-    *post = dsv4_hc_affine(ctx0, *post, scale_post, base_post);
-    *post = ggml_sigmoid(ctx0, *post);
-    *post = ggml_scale(ctx0, *post, 2.0f);
-    cb(*post, "hc_post", il);
+    ggml_tensor * post = dsv4_view_2d(ctx0, mixes, hc, nt, hc);
+    post = dsv4_hc_affine(ctx0, post, scale_post, base_post);
+    post = ggml_sigmoid(ctx0, post);
+    post = ggml_scale(ctx0, post, 2.0f);
+    cb(post, "hc_post", il);
 
+    ggml_tensor * comb = nullptr;
     if (cparams.fused_dsv4_hc_comb) {
-        *comb = ggml_dsv4_hc_comb(ctx0, mixes, hc_scale, hc_base, hparams.dsv4_hc_eps,
+        comb = ggml_dsv4_hc_comb(ctx0, mixes, hc_scale, hc_base, hparams.dsv4_hc_eps,
                 (int32_t) hparams.dsv4_hc_sinkhorn_iters);
-        res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, *comb, il});
+        res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, comb, il});
     } else {
         ggml_tensor * scale_comb = dsv4_view_1d(ctx0, hc_scale, 1, 2);
         ggml_tensor * base_comb  = dsv4_view_1d(ctx0, hc_base, hc*hc, 2*hc);
 
-        *comb = dsv4_view_2d(ctx0, mixes, hc*hc, nt, 2*hc);
-        *comb = dsv4_hc_affine(ctx0, *comb, scale_comb, base_comb);
-        *comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt);
-        *comb = build_hc_sinkhorn(*comb, il);
+        comb = dsv4_view_2d(ctx0, mixes, hc*hc, nt, 2*hc);
+        comb = dsv4_hc_affine(ctx0, comb, scale_comb, base_comb);
+        comb = ggml_reshape_3d(ctx0, comb, hc, hc, nt);
+        comb = build_hc_sinkhorn(comb, il);
     }
-    cb(*comb, "hc_comb", il);
+    cb(comb, "hc_comb", il);
 
-    return pre;
+    return { pre, post, comb };
 }
 
 ggml_tensor * llama_model_deepseek4::graph::build_hc_post(
@@ -1257,14 +1252,12 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p
         }
 
         ggml_tensor * residual = inpL;
-        ggml_tensor * post = nullptr;
-        ggml_tensor * comb = nullptr;
-
-        cur = build_hc_pre(inpL,
+        const hc_pre attn_hc = build_hc_pre(inpL,
                 model.layers[il].hc_attn_fn,
                 model.layers[il].hc_attn_scale,
                 model.layers[il].hc_attn_base,
-                &post, &comb, il);
+                il);
+        cur = attn_hc.cur;
         cb(cur, "hc_attn_pre", il);
 
         cur = build_norm(cur, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);
@@ -1272,20 +1265,21 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p
 
         cur = build_attention(model, inp_dsv4, cur, inp_pos, il);
 
-        inpL = build_hc_post(cur, residual, post, comb, il);
+        inpL = build_hc_post(cur, residual, attn_hc.post, attn_hc.comb, il);
         cb(inpL, "hc_attn_post", il);
 
         residual = inpL;
-        cur = build_hc_pre(inpL,
+        const hc_pre ffn_hc = build_hc_pre(inpL,
                 model.layers[il].hc_ffn_fn,
                 model.layers[il].hc_ffn_scale,
                 model.layers[il].hc_ffn_base,
-                &post, &comb, il);
+                il);
+        cur = ffn_hc.cur;
         cb(cur, "hc_ffn_pre", il);
 
         ggml_build_forward_expand(gf, residual);
-        ggml_build_forward_expand(gf, post);
-        ggml_build_forward_expand(gf, comb);
+        ggml_build_forward_expand(gf, ffn_hc.post);
+        ggml_build_forward_expand(gf, ffn_hc.comb);
 
         cur = build_norm(cur, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il);
         cb(cur, "ffn_norm", il);
@@ -1335,7 +1329,7 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p
         cur = ggml_add(ctx0, moe_out, ffn_shexp);
         cb(cur, "ffn_out", il);
 
-        inpL = build_hc_post(cur, residual, post, comb, il);
+        inpL = build_hc_post(cur, residual, ffn_hc.post, ffn_hc.comb, il);
         inpL = build_cvec(inpL, il);
         cb(inpL, "l_last", il);
     }
@@ -1435,14 +1429,12 @@ llama_model_deepseek4::graph_mtp::graph_mtp(const llama_model & model, const llm
     cb(inpL, "mtp_eh_proj", il);
 
     ggml_tensor * residual = inpL;
-    ggml_tensor * post = nullptr;
-    ggml_tensor * comb = nullptr;
-
-    ggml_tensor * cur = build_hc_pre(inpL,
+    hc_pre attn_hc = build_hc_pre(inpL,
             layer.hc_attn_fn,
             layer.hc_attn_scale,
             layer.hc_attn_base,
-            &post, &comb, il);
+            il);
+    ggml_tensor * cur = attn_hc.cur;
     cb(cur, "mtp_hc_attn_pre", il);
 
     cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
@@ -1450,15 +1442,16 @@ llama_model_deepseek4::graph_mtp::graph_mtp(const llama_model & model, const llm
 
     cur = build_attention(model, inp_attn, cur, inp_pos, il);
 
-    inpL = build_hc_post(cur, residual, post, comb, il);
+    inpL = build_hc_post(cur, residual, attn_hc.post, attn_hc.comb, il);
     cb(inpL, "mtp_hc_attn_post", il);
 
     residual = inpL;
-    cur = build_hc_pre(inpL,
+    hc_pre ffn_hc = build_hc_pre(inpL,
             layer.hc_ffn_fn,
             layer.hc_ffn_scale,
             layer.hc_ffn_base,
-            &post, &comb, il);
+            il);
+    cur = ffn_hc.cur;
     cb(cur, "mtp_hc_ffn_pre", il);
 
     cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
@@ -1488,7 +1481,7 @@ llama_model_deepseek4::graph_mtp::graph_mtp(const llama_model & model, const llm
     cur = ggml_add(ctx0, moe_out, ffn_shexp);
     cb(cur, "mtp_ffn_out", il);
 
-    inpL = build_hc_post(cur, residual, post, comb, il);
+    inpL = build_hc_post(cur, residual, ffn_hc.post, ffn_hc.comb, il);
     inpL = build_cvec(inpL, il);
     cb(inpL, "mtp_l_out", il);
 

src/models/deepseek41.cpp

@@ -11,9 +11,7 @@
 #include <string>
 
 // ref: deepseek-ai/DeepSeek-V4.1-Flash-1017 inference/model.py, engram.py
-// compressed layers attend over [window | all compressed rows] with top-k as a mask. with flash attention the default
-// is sparse: fused indexer, FA bounded to window + top-k, decode-sized ubatches gather only the selected rows (same
-// keys and mask, different key order). LLAMA_DSV41_SPARSE=0 restores the reference path
+// LLAMA_DSV41_SPARSE=0 disables the sparse attention path (see build_attention41)
 
 static constexpr int64_t DSV41_ENGRAM_DEAD = -1;
 
@@ -96,7 +94,7 @@ void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) {
         ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS,      mult);
         ml.get_arr(LLM_KV_PLE_HEAD_OFFSETS,           offsets);
         ml.get_arr(LLM_KV_PLE_HEAD_VOCAB_SIZES,       primes);
-        ml.get_key(LLM_KV_ENGRAM_PAD_TOKEN_ID,        engram_pad_id);
+        ml.get_key(LLM_KV_ENGRAM_PAD_TOKEN_ID,        hparams.engram_pad_id);
         ml.get_arr(LLM_KV_ENGRAM_TOKEN_MAP,           engram_token_map);
 
         const size_t n_e    = e_layers.size();
@@ -105,7 +103,7 @@ void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) {
         if (n_cols == 0 || hparams.ple_head_dim == 0 ||
                 mult.size() != n_e*n_gram || primes.size() != n_e*n_cols || offsets.size() != n_e*n_cols ||
                 mult.size() > LLAMA_MAX_PLE_NGRAM || primes.size() > LLAMA_MAX_PLE_HEADS ||
-                engram_pad_id >= engram_token_map.size()) {
+                hparams.engram_pad_id >= engram_token_map.size()) {
             throw std::runtime_error("DeepSeek-V4.1 engram metadata is inconsistent");
         }
         for (size_t e = 0; e < n_e; ++e) {
@@ -247,7 +245,7 @@ void llama_model_deepseek41::engram_hash(int e, const llama_ubatch & ubatch, con
 
     const uint64_t * mult   = hparams.ple_layer_multipliers.data() + (size_t) e*n_gram;
     const uint32_t * primes = hparams.ple_head_vocab_sizes.data()  + (size_t) e*n_cols;
-    const int64_t    pad    = engram_token_map[engram_pad_id];
+    const int64_t    pad    = engram_token_map[hparams.engram_pad_id];
     const uint32_t * offs   = hparams.ple_head_offsets.data()      + (size_t) e*n_cols;
 
     const auto cid = [&](llama_token t) -> int64_t {
@@ -279,8 +277,7 @@ void llama_model_deepseek41::engram_hash(int e, const llama_ubatch & ubatch, con
     }
 }
 
-// cand_pin: +1e30 (the reference uses +inf) on the block holding each query's newest compressed row, so a partly
-// filled block always survives; a finite pin keeps fully masked blocks at -inf instead of NaN
+// cand_pin: +1e30 on the block with the newest compressed row of each query, so a partly filled block is always kept; the reference uses +inf, but a finite pin keeps fully masked blocks at -inf instead of NaN
 class llm_graph_input_dsv41 : public llm_graph_input_i {
 public:
     llm_graph_input_dsv41(const llama_model_deepseek41 & model, const llama_kv_cache_dsv4_context * mctx) :
@@ -535,8 +532,7 @@ ggml_tensor * llama_model_deepseek41::graph::build_attention41(
         if (hparams.dsv41_topk_src[il] == il) {
             ggml_tensor * ik = ggml_view_4d(ctx0, comp_all, n_idx_embd, 1, n_comp, n_stream,
                     comp_all->nb[1], comp_all->nb[2], comp_all->nb[3], ggml_row_size(comp_all->type, n_embd_head));
-            // the fused indexer reads the strided view in place, the unfused path needs it contiguous; quantized
-            // caches are dequantized straight from the strided view (CUDA has no same-type copy of quantized views)
+            // the fused indexer reads the strided view in place, the unfused path needs it contiguous; CUDA has no same-type copy of quantized views, so dequantize them
             const bool fused_idx = sparse && cparams.fused_lid && ci.kq_mask->type == GGML_TYPE_F16;
             if (ggml_is_quantized(ik->type)) {
                 ik = ggml_cast(ctx0, ik, GGML_TYPE_F32);
@@ -551,8 +547,7 @@ ggml_tensor * llama_model_deepseek41::graph::build_attention41(
             ggml_tensor * iw = ggml_scale(ctx0, build_lora_mm(layer.indexer_proj, cur), 1.0f/sqrtf(float(n_idx_embd*n_idx_head)));
             iw = ggml_view_4d(ctx0, iw, iw->ne[0], nt/n_stream, 1, n_stream, iw->nb[1], iw->nb[1]*(nt/n_stream), iw->nb[1]*(nt/n_stream), 0);
 
-            // chunk queries so the [n_ih, tokens, n_comp] scores stay near 2^26 elements (at most 64 chunks); the fused
-            // indexer has no such intermediate and runs as one chunk
+            // chunk queries so the [n_ih, tokens, n_comp] scores stay near 2^26 elements (at most 64 chunks); the fused indexer runs as one chunk
             const int64_t n_tok = nt/n_stream;
             GGML_ASSERT(ci.kq_mask->ne[1] == n_tok);
             const int64_t chunk = fused_idx ? n_tok :
@@ -582,9 +577,6 @@ ggml_tensor * llama_model_deepseek41::graph::build_attention41(
                     ggml_tensor * iqc = ggml_view_4d(ctx0, iq, iq->ne[0], iq->ne[1], nc, n_stream, iq->nb[1], iq->nb[2], iq->nb[3], t0*iq->nb[2]);
 
                     ggml_tensor * kq = ggml_mul_mat(ctx0, ik_p, ggml_permute(ctx0, iqc, 0, 2, 1, 3));
-                    if (!ggml_is_quantized(ik->type)) {
-                        ggml_prec_set_src(kq, GGML_PREC_F32, 1);
-                    }
                     kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 2, 1, 0, 3));
                     score = ggml_sum_rows(ctx0, ggml_mul(ctx0, ggml_relu(ctx0, kq), rows(iw)));
                     ggml_tensor * mask  = rows(ci.kq_mask);
@@ -656,14 +648,12 @@ ggml_tensor * llama_model_deepseek41::graph::build_attention41(
         const int64_t k_sel    = top_k->ne[0];
         const int64_t n_raw    = raw_k->ne[2];
 
-        // few queries per stream: gather each query's top-k rows and attend over [window | selected], one FA stream
-        // per query, instead of a full n_comp mask per layer. same keys and mask as dense, different key order
+        // few queries per stream: attend over [window | top-k rows of the query], one FA stream per query; same keys and mask as dense, different key order
         const bool gather = sparse && cparams.flash_attn && raw_mask->type == GGML_TYPE_F16 && ci.kq_mask->type == GGML_TYPE_F16 &&
                 ((n_stream == 1 && nt <= 32) || nt == n_stream);
 
         if (gather) {
-            // ggml_top_k returns the rows in no particular order (CUB DeviceTopK can vary between runs) and FA sums the
-            // gathered keys in that order: sort the rows once per top-k source so the result is deterministic
+            // ggml_top_k returns the rows in no fixed order and FA sums the keys in that order: sort them so the result is deterministic
             ggml_tensor *& sel = gather_sel_of[topk_src];
             if (!sel) {
                 GGML_ASSERT(n_comp <= (INT64_C(1) << 24)); // row ids are exact in F32
@@ -714,8 +704,7 @@ ggml_tensor * llama_model_deepseek41::graph::build_attention41(
             }
             ggml_tensor * k_all = ggml_concat(ctx0, raw_k, comp_k, 2);
 
-            // mask: window row, then comp mask at the selected rows (-inf where top-k picked an invisible row);
-            // one per top-k source, shared by its layers
+            // window mask, then the comp mask at the selected rows (-inf where top-k picked an invisible row); one per top-k source
             ggml_tensor *& kq_mask = gather_mask_of[topk_src];
             if (!kq_mask) {
                 ggml_tensor * sel_mask = ggml_get_rows(ctx0, ggml_reshape_3d(ctx0, ci.kq_mask, 1, n_comp, nt), sel);
@@ -730,8 +719,7 @@ ggml_tensor * llama_model_deepseek41::graph::build_attention41(
             ggml_tensor * comp_k = ggml_view_4d(ctx0, comp_all, n_embd_head, 1, n_comp, n_stream,
                     comp_all->nb[1], comp_all->nb[2], comp_all->nb[3], 0);
             if (k_rot) {
-                // compressed rows are stored unrotated: rotate in F32, then cast to the cache type (dequantized from
-                // the strided view: CUDA has no same-type quantized copy)
+                // compressed rows are stored unrotated: rotate in F32, then cast to the cache type
                 comp_k = comp_k->type != GGML_TYPE_F32 ? ggml_cast(ctx0, comp_k, GGML_TYPE_F32) : ggml_cont(ctx0, comp_k);
                 comp_k = llama_mul_mat_hadamard(ctx0, comp_k, k_rot);
                 if (comp_k->type != raw_k->type) {
@@ -859,10 +847,7 @@ llama_model_deepseek41::graph::graph(const llama_model & model_base, const llm_g
         }
 
         ggml_tensor * residual = inpL;
-        ggml_tensor * post = nullptr;
-        ggml_tensor * comb = nullptr;
-
-        ggml_tensor * attn_pre = build_hc_mixes(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, &post, &comb, true, il);
+        const hc_mixes attn_mix = build_hc_mixes(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, il);
 
         // the first layer's attention collapses with the identity pre-mix: copy 0
         cur = pre_prev ? build_hc_pre(inpL, pre_prev, il) : ggml_cont(ctx0, ggml_view_2d(ctx0, inpL, n_embd, n_tokens, inpL->nb[2], 0));
@@ -873,28 +858,26 @@ llama_model_deepseek41::graph::graph(const llama_model & model_base, const llm_g
 
         cur = build_attention41(model, inp_dsv4, cur, inp_pos, cand_pin, il);
 
-        inpL = build_hc_post(cur, residual, post, comb, il);
+        inpL = build_hc_post(cur, residual, attn_mix.post, attn_mix.comb, il);
         cb(inpL, "hc_attn_post", il);
 
         residual = inpL;
-        pre_prev = build_hc_mixes(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, &post, &comb, true, il);
+        const hc_mixes ffn_mix = build_hc_mixes(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, il);
+        pre_prev = ffn_mix.pre;
 
-        cur = build_hc_pre(inpL, attn_pre, il);
+        cur = build_hc_pre(inpL, attn_mix.pre, il);
         cb(cur, "hc_ffn_pre", il);
 
         ggml_build_forward_expand(gf, residual);
-        ggml_build_forward_expand(gf, post);
-        ggml_build_forward_expand(gf, comb);
+        ggml_build_forward_expand(gf, ffn_mix.post);
+        ggml_build_forward_expand(gf, ffn_mix.comb);
 
         cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
         cb(cur, "ffn_norm", il);
 
-        // F32 router inputs: expert selection is near-tie sensitive
+        // F32 router accumulation: expert selection is near-tie sensitive
         ggml_tensor * logits = build_lora_mm(layer.ffn_gate_inp, cur);
         ggml_prec_set_acc(logits, GGML_PREC_F32);
-        if (!ggml_is_quantized(layer.ffn_gate_inp->type)) {
-            ggml_prec_set_src(logits, GGML_PREC_F32, 1);
-        }
         cb(logits, "ffn_moe_logits", il);
 
         // image (embedding) tokens use the vision router bias, as DeepSeek-V4
@@ -927,7 +910,7 @@ llama_model_deepseek41::graph::graph(const llama_model & model_base, const llm_g
         cur = ggml_add(ctx0, moe_out, ffn_shexp);
         cb(cur, "ffn_out", il);
 
-        inpL = build_hc_post(cur, residual, post, comb, il);
+        inpL = build_hc_post(cur, residual, ffn_mix.post, ffn_mix.comb, il);
         inpL = build_cvec(inpL, il);
         cb(inpL, "l_out", il);
     }

src/models/dflash.cpp

@@ -1204,14 +1204,12 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
         const auto & layer = model.layers[il];
 
         ggml_tensor * residual = inpL;
-        ggml_tensor * post = nullptr;
-        ggml_tensor * comb = nullptr;
-
-        ggml_tensor * cur = build_hc_pre(inpL,
+        hc_pre attn_hc = build_hc_pre(inpL,
                 layer.hc_attn_fn,
                 layer.hc_attn_scale,
                 layer.hc_attn_base,
-                &post, &comb, il);
+                il);
+        ggml_tensor * cur = attn_hc.cur;
         cb(cur, "hc_attn_pre", il);
 
         cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
@@ -1219,15 +1217,16 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
 
         cur = build_attention(model, inp_attn, cur, inp_pos, il);
 
-        inpL = build_hc_post(cur, residual, post, comb, il);
+        inpL = build_hc_post(cur, residual, attn_hc.post, attn_hc.comb, il);
         cb(inpL, "hc_attn_post", il);
 
         residual = inpL;
-        cur = build_hc_pre(inpL,
+        hc_pre ffn_hc = build_hc_pre(inpL,
                 layer.hc_ffn_fn,
                 layer.hc_ffn_scale,
                 layer.hc_ffn_base,
-                &post, &comb, il);
+                il);
+        cur = ffn_hc.cur;
         cb(cur, "hc_ffn_pre", il);
 
         cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
@@ -1256,7 +1255,7 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
         cur = ggml_add(ctx0, moe_out, ffn_shexp);
         cb(cur, "ffn_out", il);
 
-        inpL = build_hc_post(cur, residual, post, comb, il);
+        inpL = build_hc_post(cur, residual, ffn_hc.post, ffn_hc.comb, il);
         cb(inpL, "l_out", il);
     }
 
@@ -1403,10 +1402,7 @@ llama_model_dflash::graph_dsv41::graph_dsv41(const llama_model & model, const ll
         const auto & layer = model.layers[il];
 
         ggml_tensor * residual = inpL;
-        ggml_tensor * post = nullptr;
-        ggml_tensor * comb = nullptr;
-
-        ggml_tensor * attn_pre = build_hc_mixes(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, &post, &comb, false, il);
+        const hc_mixes attn_mix = build_hc_mixes(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, il);
 
         ggml_tensor * cur;
         if (pre_prev == nullptr) {
@@ -1461,13 +1457,14 @@ llama_model_dflash::graph_dsv41::graph_dsv41(const llama_model & model, const ll
             cb(cur, "attn_out", il);
         }
 
-        inpL = build_hc_post(cur, residual, post, comb, il);
+        inpL = build_hc_post(cur, residual, attn_mix.post, attn_mix.comb, il);
         cb(inpL, "hc_attn_post", il);
 
         residual = inpL;
-        ggml_tensor * ffn_pre = build_hc_mixes(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, &post, &comb, false, il);
+        const hc_mixes ffn_mix = build_hc_mixes(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, il);
+        pre_prev = ffn_mix.pre;
 
-        cur = build_hc_pre(inpL, attn_pre, il);
+        cur = build_hc_pre(inpL, attn_mix.pre, il);
         cb(cur, "hc_ffn_pre", il);
 
         cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
@@ -1496,13 +1493,11 @@ llama_model_dflash::graph_dsv41::graph_dsv41(const llama_model & model, const ll
         cur = ggml_add(ctx0, moe_out, ffn_shexp);
         cb(cur, "ffn_out", il);
 
-        inpL = build_hc_post(cur, residual, post, comb, il);
+        inpL = build_hc_post(cur, residual, ffn_mix.post, ffn_mix.comb, il);
 
         ggml_tensor * l_out = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
         cb(l_out, "l_out", il);
         ggml_build_forward_expand(gf, l_out);
-
-        pre_prev = ffn_pre;
     }
 
     // no hc_head: collapse with the last FFN's pre-mix

src/models/glm5-next.cpp

@@ -442,13 +442,12 @@ ggml_tensor * llama_model_glm5_next::graph::build_hc_sinkhorn(
     return comb;
 }
 
-ggml_tensor * llama_model_glm5_next::graph::build_hc_pre(
+llama_model_glm5_next::graph::hc_pre
+llama_model_glm5_next::graph::build_hc_pre(
         ggml_tensor * x,
         ggml_tensor * hc_fn,
         ggml_tensor * hc_scale,
         ggml_tensor * hc_base,
-        ggml_tensor ** post,
-        ggml_tensor ** comb,
         int il) const {
     const int64_t hc         = hparams.dsv4_hc_mult;
     const int64_t hc_dim     = hc*n_embd;
@@ -475,29 +474,29 @@ ggml_tensor * llama_model_glm5_next::graph::build_hc_pre(
     pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps);
     cb(pre, "hc_pre", il);
 
-    *post = dsv4_view_2d(ctx0, mixes, hc, nt, hc);
-    *post = dsv4_hc_affine(ctx0, *post, scale_post, base_post);
-    *post = ggml_sigmoid(ctx0, *post);
-    *post = ggml_scale(ctx0, *post, 2.0f);
-    cb(*post, "hc_post", il);
+    ggml_tensor * post = dsv4_view_2d(ctx0, mixes, hc, nt, hc);
+    post = dsv4_hc_affine(ctx0, post, scale_post, base_post);
+    post = ggml_sigmoid(ctx0, post);
+    post = ggml_scale(ctx0, post, 2.0f);
+    cb(post, "hc_post", il);
 
+    ggml_tensor * comb = nullptr;
     if (cparams.fused_dsv4_hc_comb) {
-        *comb = ggml_dsv4_hc_comb(ctx0, mixes, hc_scale, hc_base, hparams.dsv4_hc_eps,
+        comb = ggml_dsv4_hc_comb(ctx0, mixes, hc_scale, hc_base, hparams.dsv4_hc_eps,
                 (int32_t) hparams.dsv4_hc_sinkhorn_iters);
-        res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, *comb, il});
+        res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, comb, il});
     } else {
         ggml_tensor * scale_comb = dsv4_view_1d(ctx0, hc_scale, 1, 2);
         ggml_tensor * base_comb  = dsv4_view_1d(ctx0, hc_base, hc*hc, 2*hc);
 
-        *comb = dsv4_view_2d(ctx0, mixes, hc*hc, nt, 2*hc);
-        *comb = dsv4_hc_affine(ctx0, *comb, scale_comb, base_comb);
-        *comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt);
-        *comb = build_hc_sinkhorn(*comb, il);
+        comb = dsv4_view_2d(ctx0, mixes, hc*hc, nt, 2*hc);
+        comb = dsv4_hc_affine(ctx0, comb, scale_comb, base_comb);
+        comb = ggml_reshape_3d(ctx0, comb, hc, hc, nt);
+        comb = build_hc_sinkhorn(comb, il);
     }
-    cb(*comb, "hc_comb", il);
+    cb(comb, "hc_comb", il);
 
-    ggml_tensor * result = build_hc_pre(x, pre, il);
-    return result;
+    return { build_hc_pre(x, pre, il), post, comb };
 }
 
 ggml_tensor * llama_model_glm5_next::graph::build_hc_post(
@@ -609,8 +608,9 @@ llama_model_glm5_next::graph_mtp::graph_mtp(const llama_model & model, const llm
     cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
     cb(cur, "mtp_attn_norm", il);
 
-    ggml_tensor * prev_sel = nullptr;
-    cur = build_dsa_layer(cur, layer, mctx_hyb, inp_attn, inp_kpool, &prev_sel, il);
+    // the NextN block has a full indexer, nothing to carry in
+    dsa_out dsa = build_dsa_layer(cur, layer, mctx_hyb, inp_attn, inp_kpool, nullptr, il);
+    cur = dsa.cur;
     cb(cur, "mtp_attn_out", il);
 
     // narrow to the output tokens before the position-wise FFN; unmasked nextn embeddings need all rows
@@ -711,10 +711,8 @@ llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_p
         const auto & layer = model.layers[il];
 
         ggml_tensor * residual = inpL;
-        ggml_tensor * post = nullptr;
-        ggml_tensor * comb = nullptr;
-
-        cur = build_hc_pre(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, &post, &comb, il);
+        hc_pre attn_hc = build_hc_pre(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, il);
+        cur = attn_hc.cur;
         cb(cur, "hc_attn_pre", il);
 
         cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
@@ -725,14 +723,17 @@ llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_p
             cur = build_kda_layer(cur, layer, inp_rs, d_conv, head_dim, n_head_kda,
                                   d_inner, n_seq_tokens, n_seqs, il);
         } else {
-            cur = build_dsa_layer(cur, layer, mctx_hyb, inp_attn, inp_kpool, &prev_sel, il);
+            dsa_out dsa = build_dsa_layer(cur, layer, mctx_hyb, inp_attn, inp_kpool, prev_sel, il);
+            cur = dsa.cur;
+            prev_sel = dsa.sel;
         }
 
-        inpL = build_hc_post(cur, residual, post, comb, il);
+        inpL = build_hc_post(cur, residual, attn_hc.post, attn_hc.comb, il);
         cb(inpL, "hc_attn_post", il);
 
         residual = inpL;
-        cur = build_hc_pre(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, &post, &comb, il);
+        hc_pre ffn_hc = build_hc_pre(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, il);
+        cur = ffn_hc.cur;
         cb(cur, "hc_ffn_pre", il);
 
         cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
@@ -770,7 +771,7 @@ llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_p
             cb(cur, "ffn_out", il);
         }
 
-        inpL = build_hc_post(cur, residual, post, comb, il);
+        inpL = build_hc_post(cur, residual, ffn_hc.post, ffn_hc.comb, il);
         inpL = build_cvec(inpL, il);
         cb(inpL, "l_out", il);
     }
@@ -1039,10 +1040,11 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
 
 // Nope MLA layer with sparse attention over the indexer selection
 
-ggml_tensor * llama_model_glm5_next::graph::build_dsa_layer(
+llama_model_glm5_next::graph::dsa_out
+llama_model_glm5_next::graph::build_dsa_layer(
         ggml_tensor * cur, const llama_layer & layer,
         const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_attn_k * inp_attn,
-        llm_graph_input_kpool * inp_kpool, ggml_tensor ** prev_sel, int il) {
+        llm_graph_input_kpool * inp_kpool, ggml_tensor * prev_sel, int il) {
 
     const auto * mctx_mla = mctx_hyb->get_attn();
 
@@ -1073,13 +1075,11 @@ ggml_tensor * llama_model_glm5_next::graph::build_dsa_layer(
 
     ggml_tensor * kq_mask = inp_attn->get_kq_mask();
 
-    ggml_tensor * sel = nullptr;
+    ggml_tensor * sel = prev_sel;
     if (il >= (int) hparams.n_layer() || hparams.is_indexer_full(il)) { // the NextN block always has a full indexer
         sel = build_kpool_select(cur, qr, kq_mask, layer, mctx_hyb, inp_kpool, il);
-        *prev_sel = sel;
     } else {
-        GGML_ASSERT(*prev_sel != nullptr && "shared indexer layer must follow a full indexer layer");
-        sel = *prev_sel;
+        GGML_ASSERT(sel != nullptr && "shared indexer layer must follow a full indexer layer");
     }
 
     ggml_build_forward_expand(gf, q_absorbed);
@@ -1099,5 +1099,5 @@ ggml_tensor * llama_model_glm5_next::graph::build_dsa_layer(
     out = ggml_mul_mat(ctx0, layer.wo, out);
     cb(out, "attn_out", il);
 
-    return out;
+    return { out, sel };
 }

src/models/hy-v4.cpp

@@ -172,12 +172,12 @@ static ggml_tensor * hy_v4_hc_reduce(ggml_context * ctx0, ggml_tensor * x, ggml_
     return ggml_cast(ctx0, result, out_type);
 }
 
-ggml_tensor * llama_model_hy_v4::graph::build_hc_pre(
+llama_model_hy_v4::graph::hc_pre
+llama_model_hy_v4::graph::build_hc_pre(
         ggml_tensor * x,
         ggml_tensor * hc_fn,
         ggml_tensor * hc_scale,
         ggml_tensor * hc_base,
-        ggml_tensor ** post,
         int il) const {
     const int64_t hc  = hparams.dsv4_hc_mult;
     const int64_t nt  = x->ne[2];
@@ -202,16 +202,15 @@ ggml_tensor * llama_model_hy_v4::graph::build_hc_pre(
     cb(pre, "hc_pre", il);
 
     // post = magnitude*sigmoid(mixes[hc:2hc]*scale_post + base_post) + eps
-    ggml_tensor * po = hy_v4_view_2d(ctx0, mixes, hc, nt, hc);
-    po = ggml_mul(ctx0, po, scale_post);
-    po = ggml_add(ctx0, po, base_post);
-    po = ggml_sigmoid(ctx0, po);
-    po = ggml_scale(ctx0, po, hparams.hc_magnitude);
-    po = ggml_scale_bias(ctx0, po, 1.0f, hparams.dsv4_hc_eps);
-    *post = po;
-    cb(po, "hc_post_gate", il);
+    ggml_tensor * post = hy_v4_view_2d(ctx0, mixes, hc, nt, hc);
+    post = ggml_mul(ctx0, post, scale_post);
+    post = ggml_add(ctx0, post, base_post);
+    post = ggml_sigmoid(ctx0, post);
+    post = ggml_scale(ctx0, post, hparams.hc_magnitude);
+    post = ggml_scale_bias(ctx0, post, 1.0f, hparams.dsv4_hc_eps);
+    cb(post, "hc_post_gate", il);
 
-    return hy_v4_hc_reduce(ctx0, x, pre, hc, n_embd, nt, x->type);
+    return { hy_v4_hc_reduce(ctx0, x, pre, hc, n_embd, nt, x->type), post };
 }
 
 ggml_tensor * llama_model_hy_v4::graph::build_hc_post(
@@ -418,12 +417,13 @@ ggml_tensor * llama_model_hy_v4::graph::build_indexer_top_k(
     return ggml_cont(ctx0, ggml_top_k(ctx0, score, n_top_k));
 }
 
-ggml_tensor * llama_model_hy_v4::graph::build_attention_dsa(
+llama_model_hy_v4::graph::attn_dsa
+llama_model_hy_v4::graph::build_attention_dsa(
         const llama_model & model,
         llm_graph_input_attn_k_dsa * inp_attn_dsa,
         ggml_tensor * cur,
         ggml_tensor * inp_pos,
-        ggml_tensor ** last_top_k,
+        ggml_tensor * last_top_k,
         float kq_scale,
         int il) const {
     const auto & layer = model.layers[il];
@@ -436,11 +436,12 @@ ggml_tensor * llama_model_hy_v4::graph::build_attention_dsa(
     ggml_tensor * qr = ggml_mul_mat(ctx0, layer.wq_a, cur);
     qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
 
+    ggml_tensor * top_k = last_top_k;
     if (hparams.is_indexer_full(il)) {
-        *last_top_k = build_indexer_top_k(model, inp_attn_dsa, cur, qr, inp_pos, il);
-        cb(*last_top_k, "top_k", il);
+        top_k = build_indexer_top_k(model, inp_attn_dsa, cur, qr, inp_pos, il);
+        cb(top_k, "top_k", il);
     }
-    GGML_ASSERT(*last_top_k != nullptr);
+    GGML_ASSERT(top_k != nullptr);
 
     ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_b, qr);
 
@@ -480,7 +481,7 @@ ggml_tensor * llama_model_hy_v4::graph::build_attention_dsa(
 
     ggml_tensor * attn = build_attn(inp_attn_dsa,
             nullptr, nullptr, nullptr,
-            Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, layer.wv_b, *last_top_k, kq_scale, il);
+            Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, layer.wv_b, top_k, kq_scale, il);
     cb(attn, "attn_kqv", il);
 
     ggml_tensor * gate = ggml_mul_mat(ctx0, layer.wqkv_gate, cur);
@@ -491,7 +492,7 @@ ggml_tensor * llama_model_hy_v4::graph::build_attention_dsa(
     ggml_tensor * out = build_lora_mm(layer.wo, attn);
     cb(out, "attn_out", il);
 
-    return out;
+    return { out, top_k };
 }
 
 llama_model_hy_v4::graph::graph(const llama_model & model, const llm_graph_params & params) :
@@ -520,23 +521,27 @@ llama_model_hy_v4::graph::graph(const llama_model & model, const llm_graph_param
 
     for (int il = 0; il < n_layer; ++il) {
         ggml_tensor * residual = inpL;
-        ggml_tensor * post = nullptr;
-
-        cur = build_hc_pre(inpL, model.layers[il].hc_attn_fn, model.layers[il].hc_attn_scale,
-                model.layers[il].hc_attn_base, &post, il);
+        hc_pre attn_hc = build_hc_pre(inpL, model.layers[il].hc_attn_fn, model.layers[il].hc_attn_scale,
+                model.layers[il].hc_attn_base, il);
+        cur = attn_hc.cur;
         cur = build_norm(cur, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);
         cb(cur, "attn_norm", il);
 
-        cur = is_dsa
-            ? build_attention_dsa(model, inp_attn_dsa, cur, inp_pos, &last_top_k, kq_scale, il)
-            : build_attention(model, inp_attn, cur, inp_pos, kq_scale, il);
+        if (is_dsa) {
+            attn_dsa attn = build_attention_dsa(model, inp_attn_dsa, cur, inp_pos, last_top_k, kq_scale, il);
+            cur = attn.cur;
+            last_top_k = attn.top_k;
+        } else {
+            cur = build_attention(model, inp_attn, cur, inp_pos, kq_scale, il);
+        }
 
-        inpL = build_hc_post(cur, residual, post, il);
+        inpL = build_hc_post(cur, residual, attn_hc.post, il);
         cb(inpL, "hc_attn_out", il);
 
         residual = inpL;
-        cur = build_hc_pre(inpL, model.layers[il].hc_ffn_fn, model.layers[il].hc_ffn_scale,
-                model.layers[il].hc_ffn_base, &post, il);
+        hc_pre ffn_hc = build_hc_pre(inpL, model.layers[il].hc_ffn_fn, model.layers[il].hc_ffn_scale,
+                model.layers[il].hc_ffn_base, il);
+        cur = ffn_hc.cur;
         cur = build_norm(cur, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il);
         cb(cur, "ffn_norm", il);
 
@@ -575,7 +580,7 @@ llama_model_hy_v4::graph::graph(const llama_model & model, const llm_graph_param
             cb(cur, "ffn_out", il);
         }
 
-        inpL = build_hc_post(cur, residual, post, il);
+        inpL = build_hc_post(cur, residual, ffn_hc.post, il);
         cb(inpL, "l_out", il);
     }
 

src/models/models.h

@@ -1358,25 +1358,33 @@ struct llama_model_deepseek4 : public llama_model_base {
         // mean over the hyper-connection streams: [n_embd, hc, n_tokens] -> [n_embd, n_tokens]
         ggml_tensor * build_hc_mean(ggml_tensor * x) const;
 
-        // returns the collapsed input and fills the post / comb weights
-        ggml_tensor * build_hc_pre(
+        // all outputs of the hc mix builders:
+        //   build_hc_pre   - cur is the collapsed input, post / comb feed build_hc_post
+        //   build_hc_mixes - pre is the pre weights, post / comb feed build_hc_post
+        struct hc_pre {
+            ggml_tensor * cur;
+            ggml_tensor * post;
+            ggml_tensor * comb;
+        };
+
+        struct hc_mixes {
+            ggml_tensor * pre;
+            ggml_tensor * post;
+            ggml_tensor * comb;
+        };
+
+        hc_pre build_hc_pre(
                 ggml_tensor * x,
                 ggml_tensor * hc_fn,
                 ggml_tensor * hc_scale,
                 ggml_tensor * hc_base,
-                ggml_tensor ** post,
-                ggml_tensor ** comb,
                 int il) const;
 
-        // returns the pre weights and fills post / comb
-        ggml_tensor * build_hc_mixes(
+        hc_mixes build_hc_mixes(
                 ggml_tensor * x,
                 ggml_tensor * hc_fn,
                 ggml_tensor * hc_scale,
                 ggml_tensor * hc_base,
-                ggml_tensor ** post,
-                ggml_tensor ** comb,
-                bool f32_src,
                 int il) const;
 
         ggml_tensor * build_hc_post(
@@ -1498,11 +1506,7 @@ struct llama_model_deepseek41 : public llama_model_deepseek4 {
     void load_arch_hparams(llama_model_loader & ml) override;
     void load_arch_tensors(llama_model_loader & ml) override;
 
-    // Engram: the hash constants are hparams.ple_* (one module per Engram layer), hashed on the host (ggml has no
-    // int64 mul/xor/mod); the token map and its pad id are V4.1 only
-    uint32_t engram_pad_id = 0;
-    std::vector<int32_t> engram_token_map;
-
+    // Engram: the hash constants are hparams.ple_* (one module per Engram layer), hashed on the host (ggml has no int64 mul/xor/mod)
     int engram_index(int il) const;
 
     void engram_hash(int e, const llama_ubatch & ubatch, const std::vector<llama_token> & prev, std::vector<int32_t> & out) const;
@@ -2268,12 +2272,16 @@ struct llama_model_hy_v4 : public llama_model_base {
         // iHC (independent Hyper-Connections): pre reduces the hc streams to one and returns the
         // per-stream post gates, post writes the sublayer output back into the streams, head
         // collapses the streams before the final norm.
-        ggml_tensor * build_hc_pre(
+        struct hc_pre {
+            ggml_tensor * cur;   // the collapsed stream
+            ggml_tensor * post;  // per-stream post gates, feeds build_hc_post
+        };
+
+        hc_pre build_hc_pre(
                 ggml_tensor * x,
                 ggml_tensor * hc_fn,
                 ggml_tensor * hc_scale,
                 ggml_tensor * hc_base,
-                ggml_tensor ** post,
                 int il) const;
 
         ggml_tensor * build_hc_post(
@@ -2306,12 +2314,18 @@ struct llama_model_hy_v4 : public llama_model_base {
                 ggml_tensor * inp_pos,
                 int il) const;
 
-        ggml_tensor * build_attention_dsa(
+        // the attention output and the top-k to carry to the following "shared" layers
+        struct attn_dsa {
+            ggml_tensor * cur;
+            ggml_tensor * top_k;
+        };
+
+        attn_dsa build_attention_dsa(
                 const llama_model & model,
                 llm_graph_input_attn_k_dsa * inp_attn_dsa,
                 ggml_tensor * cur,
                 ggml_tensor * inp_pos,
-                ggml_tensor ** last_top_k,
+                ggml_tensor * last_top_k,
                 float kq_scale,
                 int il) const;
     };
@@ -2735,13 +2749,18 @@ struct llama_model_qwen4exp : public llama_model_base {
             llm_build_delta_net_base(params), model(model) {}
 
         // HC replaces every layer norm: residual is [n_embd, hc, n_tokens]
-        ggml_tensor * build_hc_mix(
+        struct hc_mix {
+            ggml_tensor * cur;     // the mixed [n_embd, T] stream
+            ggml_tensor * inject;  // [hc, T] scatter weights, null when the inject weights are null
+        };
+
+        hc_mix build_hc_mix(
                     ggml_tensor * x,
                     ggml_tensor * w_norm,
                     ggml_tensor * w_down,
                     ggml_tensor * w_up,
                     ggml_tensor * w_inject,
-                    ggml_tensor ** inject,
+                    ggml_tensor * w_down_inject, // [w_down | w_inject] in one weight, can be null
                             int   il);
 
         ggml_tensor * build_hc_combine(
@@ -2973,14 +2992,18 @@ struct llama_model_glm5_next : public llama_model_base {
         // mean over the hyper-connection streams: [n_embd, hc, n_tokens] -> [n_embd, n_tokens]
         ggml_tensor * build_hc_mean(ggml_tensor * x) const;
 
-        // returns the collapsed input and fills the post / comb weights
-        ggml_tensor * build_hc_pre(
+        // all outputs of the hc mix builder: cur is the collapsed input, post / comb feed build_hc_post
+        struct hc_pre {
+            ggml_tensor * cur;
+            ggml_tensor * post;
+            ggml_tensor * comb;
+        };
+
+        hc_pre build_hc_pre(
                 ggml_tensor * x,
                 ggml_tensor * hc_fn,
                 ggml_tensor * hc_scale,
                 ggml_tensor * hc_base,
-                ggml_tensor ** post,
-                ggml_tensor ** comb,
                 int il) const;
 
         ggml_tensor * build_hc_post(
@@ -3006,9 +3029,15 @@ struct llama_model_glm5_next : public llama_model_base {
         ggml_tensor * build_kpool_select(ggml_tensor * cur, ggml_tensor * qr, ggml_tensor * kq_mask, const llama_layer & layer,
                                          const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_kpool * inp_kpool, int il);
 
-        ggml_tensor * build_dsa_layer(ggml_tensor * cur, const llama_layer & layer,
-                                      const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_attn_k * inp_attn,
-                                      llm_graph_input_kpool * inp_kpool, ggml_tensor ** prev_sel, int il);
+        // attention output (cur) and the selection to carry to the next shared layer (sel)
+        struct dsa_out {
+            ggml_tensor * cur;
+            ggml_tensor * sel;
+        };
+
+        dsa_out build_dsa_layer(ggml_tensor * cur, const llama_layer & layer,
+                                const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_attn_k * inp_attn,
+                                llm_graph_input_kpool * inp_kpool, ggml_tensor * prev_sel, int il);
 
     };
 

src/models/qwen4exp.cpp

@@ -231,15 +231,32 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) {
         const int64_t value_dim  = head_v_dim * n_v_heads;
         const int64_t conv_dim   = key_dim * 2 + value_dim;
 
+        // matmuls that read the same input: try the merged [a | b] weight the converter writes, else the separate ones
+        auto create_merged = [&](llm_tensor t_merged, llm_tensor t_a, llm_tensor t_b, int64_t n_in, int64_t n_a, int64_t n_b,
+                ggml_tensor *& merged, ggml_tensor *& a, ggml_tensor *& b) {
+            if (flags & TENSOR_SKIP) {
+                const int skip = TENSOR_NOT_REQUIRED | TENSOR_SKIP;
+                create_tensor(tn(t_merged, "weight", il), { n_in, n_a + n_b }, skip | TENSOR_SKIP_IF_VIRTUAL);
+                create_tensor(tn(t_a,      "weight", il), { n_in, n_a },       skip);
+                create_tensor(tn(t_b,      "weight", il), { n_in, n_b },       skip);
+                return;
+            }
+            merged = create_tensor(tn(t_merged, "weight", il), { n_in, n_a + n_b }, flags | TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL);
+            if (!merged) {
+                a = create_tensor(tn(t_a, "weight", il), { n_in, n_a }, flags);
+                b = create_tensor(tn(t_b, "weight", il), { n_in, n_b }, flags);
+            }
+        };
+
         // two HC modules per layer: before the token mixer, before the MoE
         layer.hc_attn_norm   = create_tensor(tn(LLM_TENSOR_HC_ATTN_NORM,   "weight", il), { n_embd, hc }, flags | TENSOR_ALLOW_RESHAPE);
-        layer.hc_attn_down   = create_tensor(tn(LLM_TENSOR_HC_ATTN_DOWN,   "weight", il), { hc_dim, hc_lr }, flags);
         layer.hc_attn_up     = create_tensor(tn(LLM_TENSOR_HC_ATTN_UP,     "weight", il), { hc_lr, hc_dim }, flags);
-        layer.hc_attn_inject = create_tensor(tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", il), { hc_dim, hc }, flags);
+        create_merged(LLM_TENSOR_HC_ATTN_DOWN_INJECT, LLM_TENSOR_HC_ATTN_DOWN, LLM_TENSOR_HC_ATTN_INJECT, hc_dim, hc_lr, hc,
+                layer.hc_attn_down_inject, layer.hc_attn_down, layer.hc_attn_inject);
         layer.hc_ffn_norm    = create_tensor(tn(LLM_TENSOR_HC_FFN_NORM,    "weight", il), { n_embd, hc }, flags | TENSOR_ALLOW_RESHAPE);
-        layer.hc_ffn_down    = create_tensor(tn(LLM_TENSOR_HC_FFN_DOWN,    "weight", il), { hc_dim, hc_lr }, flags);
         layer.hc_ffn_up      = create_tensor(tn(LLM_TENSOR_HC_FFN_UP,      "weight", il), { hc_lr, hc_dim }, flags);
-        layer.hc_ffn_inject  = create_tensor(tn(LLM_TENSOR_HC_FFN_INJECT,  "weight", il), { hc_dim, hc }, flags);
+        create_merged(LLM_TENSOR_HC_FFN_DOWN_INJECT, LLM_TENSOR_HC_FFN_DOWN, LLM_TENSOR_HC_FFN_INJECT, hc_dim, hc_lr, hc,
+                layer.hc_ffn_down_inject, layer.hc_ffn_down, layer.hc_ffn_inject);
 
         if (!hparams.is_recr(il)) {
             // full attention: wq holds [q|gate] interleaved per head
@@ -255,13 +272,13 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) {
             layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", il), { idx_dim }, flags);
             layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), { idx_dim }, flags);
         } else {
-            layer.wqkv       = create_tensor(tn(LLM_TENSOR_ATTN_QKV,   "weight", il), { n_embd, key_dim * 2 + value_dim }, flags);
-            layer.wqkv_gate  = create_tensor(tn(LLM_TENSOR_ATTN_GATE,  "weight", il), { n_embd, value_dim }, flags);
+            create_merged(LLM_TENSOR_ATTN_QKV_GATE, LLM_TENSOR_ATTN_QKV, LLM_TENSOR_ATTN_GATE, n_embd, key_dim * 2 + value_dim, value_dim,
+                    layer.wqkv_z, layer.wqkv, layer.wqkv_gate);
             layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, flags);
             layer.ssm_dt     = create_tensor(tn(LLM_TENSOR_SSM_DT,     "bias",   il), { hparams.ssm_dt_rank }, flags);
             layer.ssm_a      = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN,         il), { hparams.ssm_dt_rank }, flags);
-            layer.ssm_beta   = create_tensor(tn(LLM_TENSOR_SSM_BETA,   "weight", il), { n_embd, n_v_heads }, flags);
-            layer.ssm_alpha  = create_tensor(tn(LLM_TENSOR_SSM_ALPHA,  "weight", il), { n_embd, n_v_heads }, flags);
+            create_merged(LLM_TENSOR_SSM_BETA_ALPHA_CAT, LLM_TENSOR_SSM_BETA, LLM_TENSOR_SSM_ALPHA, n_embd, n_v_heads, n_v_heads,
+                    layer.ssm_beta_alpha_cat, layer.ssm_beta, layer.ssm_alpha);
             layer.ssm_norm   = create_tensor(tn(LLM_TENSOR_SSM_NORM,   "weight", il), { head_v_dim }, flags);
             layer.ssm_out    = create_tensor(tn(LLM_TENSOR_SSM_OUT,    "weight", il), { value_dim, n_embd }, flags);
         }
@@ -326,14 +343,15 @@ static ggml_tensor * hc_stream_mean(ggml_context * ctx, ggml_tensor * x) {
 }
 
 // Hyper-connections keep hc parallel residual streams [n_embd, hc, T] in place of layer norms.
-// Returns the mixed [n_embd, T] stream; `inject` gets the [hc, T] scatter weights.
-ggml_tensor * llama_model_qwen4exp::graph::build_hc_mix(
+// build_hc_mix returns the mixed [n_embd, T] stream and the [hc, T] scatter weights.
+llama_model_qwen4exp::graph::hc_mix
+llama_model_qwen4exp::graph::build_hc_mix(
         ggml_tensor *  x,
         ggml_tensor *  w_norm,
         ggml_tensor *  w_down,
         ggml_tensor *  w_up,
         ggml_tensor *  w_inject,
-        ggml_tensor ** inject,
+        ggml_tensor *  w_down_inject,
         int            il) {
     const int64_t hc     = hparams.dsv4_hc_mult;
     const int64_t hc_dim = hc * n_embd;
@@ -345,8 +363,21 @@ ggml_tensor * llama_model_qwen4exp::graph::build_hc_mix(
     xn = ggml_reshape_2d(ctx0, xn, hc_dim, nt);
     cb(xn, "hc_norm", il);
 
-    ggml_tensor * lo = build_lora_mm(w_down, xn);
-    lo = ggml_silu(ctx0, ggml_scale(ctx0, lo, 1.0f / (float) hc));
+    ggml_tensor * lo = nullptr;
+    ggml_tensor * inject = nullptr;
+    if (w_down_inject) {
+        // one matmul for both: the hc_down rows, then the hc_inject rows
+        const int64_t n_lo = w_down_inject->ne[1] - hc;
+        ggml_tensor * lo_inject = build_lora_mm(w_down_inject, xn);
+        // scale the combined rows before the split; both lo and inject use 1/hc
+        lo_inject = ggml_scale(ctx0, lo_inject, 1.0f / (float) hc);
+        lo     = ggml_view_2d(ctx0, lo_inject, n_lo, nt, lo_inject->nb[1], 0);
+        inject = ggml_view_2d(ctx0, lo_inject, hc, nt, lo_inject->nb[1], ggml_row_size(lo_inject->type, n_lo));
+        cb(inject, "hc_inject", il);
+    } else {
+        lo = ggml_scale(ctx0, build_lora_mm(w_down, xn), 1.0f / (float) hc);
+    }
+    lo = ggml_silu(ctx0, lo);
     ggml_tensor * gate = build_lora_mm(w_up, lo);
     cb(gate, "hc_gate", il);
 
@@ -375,12 +406,12 @@ ggml_tensor * llama_model_qwen4exp::graph::build_hc_mix(
     }
     cb(mixed, "hc_mixed", il);
 
-    if (inject) {
-        *inject = build_lora_mm(w_inject, xn);
-        cb(*inject, "hc_inject", il);
+    if (!w_down_inject && w_inject) {
+        inject = ggml_scale(ctx0, build_lora_mm(w_inject, xn), 1.0f / (float) hc);
+        cb(inject, "hc_inject", il);
     }
 
-    return mixed;
+    return { mixed, inject };
 }
 
 ggml_tensor * llama_model_qwen4exp::graph::build_hc_combine(
@@ -392,7 +423,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_hc_combine(
     const int64_t nt = residual->ne[2];
 
     // 2*sigmoid centres the scatter weights on 1, so a zero injection is a plain residual add
-    ggml_tensor * w = ggml_sigmoid(ctx0, ggml_scale(ctx0, inject, 1.0f / (float) hc));
+    ggml_tensor * w = ggml_sigmoid(ctx0, inject);
     w = ggml_scale(ctx0, w, 2.0f);
 
     ggml_tensor * cur = nullptr;
@@ -474,13 +505,14 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa
             res_hc = build_ple(inp->get_recr(), ple_emb, res_hc, il);
         }
 
-        ggml_tensor * inject = nullptr;
-        ggml_tensor * cur = build_hc_mix(res_hc,
+        hc_mix attn_mix = build_hc_mix(res_hc,
                 model.layers[il].hc_attn_norm,
                 model.layers[il].hc_attn_down,
                 model.layers[il].hc_attn_up,
                 model.layers[il].hc_attn_inject,
-                &inject, il);
+                model.layers[il].hc_attn_down_inject,
+                il);
+        ggml_tensor * cur = attn_mix.cur;
 
         ggml_build_forward_expand(gf, cur);
 
@@ -492,27 +524,29 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa
 
         if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
             // everything below is per token, so drop the rows that produce no output
-            cur    = ggml_get_rows(ctx0, cur,    inp_out_ids);
-            inject = ggml_get_rows(ctx0, inject, inp_out_ids);
+            cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+            attn_mix.inject = ggml_get_rows(ctx0, attn_mix.inject, inp_out_ids);
 
             res_hc = ggml_reshape_2d(ctx0, res_hc, n_embd*hc, res_hc->ne[2]);
             res_hc = ggml_get_rows(ctx0, res_hc, inp_out_ids);
             res_hc = ggml_reshape_3d(ctx0, res_hc, n_embd, hc, res_hc->ne[1]);
         }
 
-        res_hc = build_hc_combine(res_hc, cur, inject, il);
+        res_hc = build_hc_combine(res_hc, cur, attn_mix.inject, il);
 
-        cur = build_hc_mix(res_hc,
+        hc_mix ffn_mix = build_hc_mix(res_hc,
                 model.layers[il].hc_ffn_norm,
                 model.layers[il].hc_ffn_down,
                 model.layers[il].hc_ffn_up,
                 model.layers[il].hc_ffn_inject,
-                &inject, il);
+                model.layers[il].hc_ffn_down_inject,
+                il);
+        cur = ffn_mix.cur;
 
         cur = build_layer_ffn(cur, il);
         cb(cur, "ffn_out", il);
 
-        res_hc = build_hc_combine(res_hc, cur, inject, il);
+        res_hc = build_hc_combine(res_hc, cur, ffn_mix.inject, il);
 
         // "l_last" is the layer output name that build_cvec and imatrix look for
         cb(res_hc, "l_last", il);
@@ -530,9 +564,10 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa
     }
 
     // the final mixer is the output norm: there is no separate one
-    ggml_tensor * cur = build_hc_mix(res_hc,
+    hc_mix head_mix = build_hc_mix(res_hc,
             model.hc_head_norm, model.hc_head_down, model.hc_head_up,
             nullptr, nullptr, -1);
+    ggml_tensor * cur = head_mix.cur;
 
     cb(cur, "result_norm", -1);
     res->t_embd = cur;
@@ -609,14 +644,15 @@ llama_model_qwen4exp::graph_mtp::graph_mtp(const llama_model & model, const llm_
     ggml_tensor * res_hc = build_lora_mm(layer.nextn.eh_proj, ggml_concat(ctx0, e_norm, h_norm, 0)); // [n_embd, hc, n_tokens]
     cb(res_hc, "mtp_eh_proj", il);
 
-    ggml_tensor * inject = nullptr;
-    ggml_tensor * cur = build_hc_mix(res_hc, layer.hc_attn_norm, layer.hc_attn_down, layer.hc_attn_up, layer.hc_attn_inject, &inject, il);
+    hc_mix attn_mix = build_hc_mix(res_hc, layer.hc_attn_norm, layer.hc_attn_down, layer.hc_attn_up, layer.hc_attn_inject, layer.hc_attn_down_inject, il);
+    ggml_tensor * cur    = attn_mix.cur;
     cur    = build_layer_attn(inp_hyb->get_attn(), mctx_hyb, inp_kpool, cur, inp_pos, sections, il);
-    res_hc = build_hc_combine(res_hc, cur, inject, il);
+    res_hc = build_hc_combine(res_hc, cur, attn_mix.inject, il);
 
-    cur    = build_hc_mix(res_hc, layer.hc_ffn_norm, layer.hc_ffn_down, layer.hc_ffn_up, layer.hc_ffn_inject, &inject, il);
+    hc_mix ffn_mix = build_hc_mix(res_hc, layer.hc_ffn_norm, layer.hc_ffn_down, layer.hc_ffn_up, layer.hc_ffn_inject, layer.hc_ffn_down_inject, il);
+    cur    = ffn_mix.cur;
     cur    = build_layer_ffn(cur, il);
-    res_hc = build_hc_combine(res_hc, cur, inject, il);
+    res_hc = build_hc_combine(res_hc, cur, ffn_mix.inject, il);
 
     // the next draft step reads this residual as its h
     ggml_tensor * flat     = ggml_reshape_2d(ctx0, res_hc, n_embd*hc, n_tokens);
@@ -625,9 +661,10 @@ llama_model_qwen4exp::graph_mtp::graph_mtp(const llama_model & model, const llm_
     cb(res->t_h_nextn, "h_nextn", il);
     ggml_build_forward_expand(gf, res->t_h_nextn);
 
-    cur = build_hc_mix(ggml_reshape_3d(ctx0, flat_out, n_embd, hc, flat_out->ne[1]),
+    hc_mix head_mix = build_hc_mix(ggml_reshape_3d(ctx0, flat_out, n_embd, hc, flat_out->ne[1]),
             layer.nextn.hc_head_norm, layer.nextn.hc_head_down, layer.nextn.hc_head_up,
             nullptr, nullptr, il);
+    cur = head_mix.cur;
     cb(cur, "result_norm", -1);
     res->t_embd = cur;
 
@@ -642,6 +679,23 @@ std::pair<ggml_tensor *, ggml_tensor *> llama_model_qwen4exp::graph::build_qkvz(
     const int64_t n_seqs       = ubatch.n_seqs;
     const int64_t n_seq_tokens = ubatch.n_seq_tokens;
 
+    if (model.layers[il].wqkv_z) {
+        // qkv and z read the same input: one matmul, then the qkv rows and the z rows
+        ggml_tensor * qkvz = build_lora_mm(model.layers[il].wqkv_z, input);
+        cb(qkvz, "linear_attn_qkvz", il);
+
+        const int64_t n_z   = hparams.ssm_d_state * hparams.ssm_dt_rank;
+        const int64_t n_qkv = qkvz->ne[0] - n_z;
+
+        ggml_tensor * qkv_mixed = ggml_view_3d(ctx0, qkvz, n_qkv, n_seq_tokens, n_seqs, qkvz->nb[1], qkvz->nb[1] * n_seq_tokens, 0);
+        cb(qkv_mixed, "linear_attn_qkv_mixed", il);
+
+        ggml_tensor * z = ggml_view_2d(ctx0, qkvz, n_z, n_seq_tokens * n_seqs, qkvz->nb[1], ggml_row_size(qkvz->type, n_qkv));
+        cb(z, "z", il);
+
+        return { qkv_mixed, z };
+    }
+
     ggml_tensor * qkv_mixed = build_lora_mm(model.layers[il].wqkv, input, model.layers[il].wqkv_s);
     qkv_mixed = ggml_reshape_3d(ctx0, qkv_mixed, qkv_mixed->ne[0], n_seq_tokens, n_seqs);
     cb(qkv_mixed, "linear_attn_qkv_mixed", il);
@@ -1072,17 +1126,30 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear(
     ggml_tensor * qkv_mixed = qkvz.first;
     ggml_tensor * z         = qkvz.second;
 
-    ggml_tensor * beta = build_lora_mm(model.layers[il].ssm_beta, cur, model.layers[il].ssm_beta_s);
-    beta = ggml_reshape_4d(ctx0, beta, 1, num_v_heads, n_seq_tokens, n_seqs);
+    ggml_tensor * beta  = nullptr;
+    ggml_tensor * alpha = nullptr;
+    if (model.layers[il].ssm_beta_alpha_cat) {
+        // beta and alpha read the same input: one matmul, then the beta rows and the alpha rows
+        ggml_tensor * ba = build_lora_mm(model.layers[il].ssm_beta_alpha_cat, cur);
+        cb(ba, "beta_alpha", il);
+
+        beta  = ggml_view_4d(ctx0, ba, 1, num_v_heads, n_seq_tokens, n_seqs,
+                ggml_element_size(ba), ba->nb[1], ba->nb[1] * n_seq_tokens, 0);
+        alpha = ggml_view_3d(ctx0, ba, num_v_heads, n_seq_tokens, n_seqs,
+                ba->nb[1], ba->nb[1] * n_seq_tokens, ggml_row_size(ba->type, num_v_heads));
+    } else {
+        beta = build_lora_mm(model.layers[il].ssm_beta, cur, model.layers[il].ssm_beta_s);
+        beta = ggml_reshape_4d(ctx0, beta, 1, num_v_heads, n_seq_tokens, n_seqs);
+
+        alpha = build_lora_mm(model.layers[il].ssm_alpha, cur, model.layers[il].ssm_alpha_s);
+        alpha = ggml_reshape_3d(ctx0, alpha, num_v_heads, n_seq_tokens, n_seqs);
+    }
     cb(beta, "beta", il);
+    cb(alpha, "alpha", il);
 
     beta = ggml_sigmoid(ctx0, beta);
     cb(beta, "beta_sigmoid", il);
 
-    ggml_tensor * alpha = build_lora_mm(model.layers[il].ssm_alpha, cur, model.layers[il].ssm_alpha_s);
-    alpha = ggml_reshape_3d(ctx0, alpha, num_v_heads, n_seq_tokens, n_seqs);
-    cb(alpha, "alpha", il);
-
     ggml_tensor * alpha_biased   = ggml_add(ctx0, alpha, model.layers[il].ssm_dt);
     ggml_tensor * alpha_softplus = ggml_softplus(ctx0, alpha_biased);
     cb(alpha_softplus, "a_softplus", il);
@@ -1160,7 +1227,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear(
 
     ggml_tensor * output = build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il);
 
-    ggml_tensor * z_2d = ggml_reshape_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs);
+    ggml_tensor * z_2d = ggml_view_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs,
+            ggml_row_size(z->type, head_v_dim), z->nb[1], z->nb[1] * n_seq_tokens, 0);
 
     // 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);

tests/CMakeLists.txt

@@ -223,58 +223,57 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
 
     # llama_build_and_test(test-double-float.cpp) # SLOW
 
-    # TODO: after `test-llama-archs` is refactored to not link internal symbols, we can move all related tests out of
-    #       this branch and run them on all platforms
-    llama_build(test-llama-archs.cpp)
-
-    set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/")
-    file(MAKE_DIRECTORY "${MODEL_DIR}")
-
-    llama_test(
-        test-llama-archs
-        NAME test-generate-models
-        LABEL "main behavior backend"
-        ARGS -o "${MODEL_DIR}"
-    )
-    set_tests_properties(test-generate-models PROPERTIES
-        FIXTURES_SETUP generate-models
-    )
-
-    # Target model for test-speculative-rejection-sampling
-    set(SPEC_TARGET_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-speculative-models/target")
-    file(MAKE_DIRECTORY "${SPEC_TARGET_DIR}")
-    llama_test(
-        test-llama-archs
-        NAME test-generate-speculative-target
-        LABEL main
-        ARGS -a llama -s 999 -o "${SPEC_TARGET_DIR}"
-    )
-    set_tests_properties(
-        test-generate-speculative-target
-        PROPERTIES FIXTURES_SETUP speculative-model
-    )
-
-    # To avoid recreating the model use the fixture exclude any option:
-    # -FA "^speculative-model$"
-    llama_build_and_test(test-speculative-rejection-sampling.cpp ARGS
-        -m "${SPEC_TARGET_DIR}/llama-dense.gguf")
-    set_tests_properties(test-speculative-rejection-sampling PROPERTIES FIXTURES_REQUIRED speculative-model)
-
-    # Test context state handling across all architectures, using the generated dummy models
-    llama_test(
-        test-llama-context
-        LABEL "main behavior backend"
-        ARGS --models "${MODEL_DIR}"
-    )
-    # abort on graph reallocations that the initial reserve should have covered
-    set_tests_properties(test-llama-context PROPERTIES
-        FIXTURES_REQUIRED generate-models
-        ENVIRONMENT "GGML_SCHED_DEBUG_REALLOC=1"
-    )
-
-    llama_build(test-fusion.cpp)
 endif()
 
+llama_build(test-llama-archs.cpp)
+
+set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/")
+file(MAKE_DIRECTORY "${MODEL_DIR}")
+
+llama_test(
+    test-llama-archs
+    NAME test-generate-models
+    LABEL "main behavior backend"
+    ARGS -o "${MODEL_DIR}"
+)
+set_tests_properties(test-generate-models PROPERTIES
+    FIXTURES_SETUP generate-models
+)
+
+# Target model for test-speculative-rejection-sampling
+set(SPEC_TARGET_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-speculative-models/target")
+file(MAKE_DIRECTORY "${SPEC_TARGET_DIR}")
+llama_test(
+    test-llama-archs
+    NAME test-generate-speculative-target
+    LABEL main
+    ARGS -a llama -s 999 -o "${SPEC_TARGET_DIR}"
+)
+set_tests_properties(
+    test-generate-speculative-target
+    PROPERTIES FIXTURES_SETUP speculative-model
+)
+
+# To avoid recreating the model use the fixture exclude any option:
+# -FA "^speculative-model$"
+llama_build_and_test(test-speculative-rejection-sampling.cpp ARGS
+    -m "${SPEC_TARGET_DIR}/llama-dense.gguf")
+set_tests_properties(test-speculative-rejection-sampling PROPERTIES FIXTURES_REQUIRED speculative-model)
+
+# Test context state handling across all architectures, using the generated dummy models
+llama_test(
+    test-llama-context
+    LABEL "main behavior backend"
+    ARGS --models "${MODEL_DIR}"
+)
+# abort on graph reallocations that the initial reserve should have covered
+set_tests_properties(test-llama-context PROPERTIES
+    FIXTURES_REQUIRED generate-models
+    ENVIRONMENT "GGML_SCHED_DEBUG_REALLOC=1"
+)
+
+llama_build(test-fusion.cpp)
+
 llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp  LABEL behavior)
 llama_build_and_test(test-jinja.cpp  LABEL behavior)
 llama_test(test-jinja  NAME test-jinja-py  ARGS -py  LABEL "python behavior")

tests/llama-context/cases-kv.cpp

@@ -718,6 +718,54 @@ struct case_state_restore_fragmented : ctx_test {
     }
 };
 
+//
+// no_alloc
+//
+
+// - load a second instance of the model with no_alloc, so the tensor data stays on disk
+// - create a context from it, which reserves the compute graphs over the dummy buffers
+// - the memory use must still be reported, and nothing real may be allocated for the memory
+struct case_noalloc : ctx_test {
+    const char * name() const override { return "noalloc"; }
+
+    const char * description() const override {
+        return "with no_alloc the memory use is simulated: non-zero breakdown, null buffer bases";
+    }
+
+    case_result run(model_run & mr) const override {
+        const std::string & path = mr.params.model.path;
+        CTX_REQUIRE(!path.empty(), "the model path is not known");
+
+        llama_model_params mparams = llama_model_default_params();
+        mparams.no_alloc  = true;
+        mparams.load_mode = LLAMA_LOAD_MODE_NONE;
+
+        llama_model_ptr model(llama_model_load_from_file(path.c_str(), mparams));
+        CTX_REQUIRE(model, "failed to load the model with no_alloc");
+
+        auto ctx = init_ctx(model.get(), common_context_params_to_llama(mr.params), 0);
+        CTX_REQUIRE(ctx, "failed to create the context with no_alloc");
+
+        size_t n_ctx_mem = 0;
+        for (const auto & [_, data] : llama_get_memory_breakdown(ctx.get())) {
+            n_ctx_mem += data.context;
+        }
+
+        // models without a memory module (BERT, diffusion) have a null memory
+        if (llama_get_memory(ctx.get()) != nullptr) {
+            CTX_REQUIRE(n_ctx_mem > 0, "the simulated memory breakdown is zero");
+        } else {
+            CTX_REQUIRE(n_ctx_mem == 0, "a model without memory reported context memory");
+        }
+
+        for (ggml_backend_buffer_t buf : llama_memory_get_buffers(llama_get_memory(ctx.get()))) {
+            CTX_REQUIRE(ggml_backend_buffer_get_base(buf) == nullptr, "a memory buffer was really allocated");
+        }
+
+        return case_result::pass();
+    }
+};
+
 std::vector<ctx_test_ptr> make_kv_cases() {
     std::vector<ctx_test_ptr> cases;
 
@@ -733,6 +781,7 @@ std::vector<ctx_test_ptr> make_kv_cases() {
     cases.push_back(std::make_unique<case_state_restore_failure>());
     cases.push_back(std::make_unique<case_state_rotation>());
     cases.push_back(std::make_unique<case_state_restore_fragmented>());
+    cases.push_back(std::make_unique<case_noalloc>());
 
     return cases;
 }

tests/peg-parser/test-json-parser.cpp

@@ -130,4 +130,23 @@ void test_json_parser(testing &t) {
             t.assert_true("fail", result.fail());
         });
     });
+
+    // json-value is entered once per level
+    t.test("deep JSON", [](testing &t) {
+        auto json = build_peg_parser([](common_peg_parser_builder & p) { return p.json(); });
+
+        const auto parse = [&](size_t depth) {
+            std::string input = std::string(depth, '[') + std::string(depth, ']');
+            common_peg_parse_context ctx(input);
+            try {
+                return json.parse(ctx).success() ? std::string() : std::string("fail");
+            } catch (const common_peg_recursion_error & e) {
+                return std::string(e.what());
+            }
+        };
+
+        t.assert_equal("64 levels", std::string(), parse(64));
+        t.assert_true("65 levels", parse(65).find("nested more than 64 levels deep") != std::string::npos);
+        t.assert_true("100000 levels", parse(100000).find("nested more than 64 levels deep") != std::string::npos);
+    });
 }

tests/snapshots/deepseek-v3.1.schema

@@ -2927,430 +2927,552 @@ blk.60.attn_k_b.weight q4_0
 [MXFP4_MOE] mxfp4
 output.weight q8_0
 token_embd.weight q8_0
+blk.0.attn_k_b.weight q8_0
 blk.0.attn_kv_a_mqa.weight q8_0
 blk.0.attn_output.weight q8_0
 blk.0.attn_q_a.weight q8_0
 blk.0.attn_q_b.weight q8_0
+blk.0.attn_v_b.weight q8_0
 blk.0.ffn_down.weight q8_0
 blk.0.ffn_gate.weight q8_0
 blk.0.ffn_up.weight q8_0
+blk.1.attn_k_b.weight q8_0
 blk.1.attn_kv_a_mqa.weight q8_0
 blk.1.attn_output.weight q8_0
 blk.1.attn_q_a.weight q8_0
 blk.1.attn_q_b.weight q8_0
+blk.1.attn_v_b.weight q8_0
 blk.1.ffn_down.weight q8_0
 blk.1.ffn_gate.weight q8_0
 blk.1.ffn_up.weight q8_0
+blk.2.attn_k_b.weight q8_0
 blk.2.attn_kv_a_mqa.weight q8_0
 blk.2.attn_output.weight q8_0
 blk.2.attn_q_a.weight q8_0
 blk.2.attn_q_b.weight q8_0
+blk.2.attn_v_b.weight q8_0
 blk.2.ffn_down.weight q8_0
 blk.2.ffn_gate.weight q8_0
 blk.2.ffn_up.weight q8_0
+blk.3.attn_k_b.weight q8_0
 blk.3.attn_kv_a_mqa.weight q8_0
 blk.3.attn_output.weight q8_0
 blk.3.attn_q_a.weight q8_0
 blk.3.attn_q_b.weight q8_0
+blk.3.attn_v_b.weight q8_0
 blk.3.ffn_down_shexp.weight q8_0
 blk.3.ffn_gate_shexp.weight q8_0
 blk.3.ffn_up_shexp.weight q8_0
+blk.4.attn_k_b.weight q8_0
 blk.4.attn_kv_a_mqa.weight q8_0
 blk.4.attn_output.weight q8_0
 blk.4.attn_q_a.weight q8_0
 blk.4.attn_q_b.weight q8_0
+blk.4.attn_v_b.weight q8_0
 blk.4.ffn_down_shexp.weight q8_0
 blk.4.ffn_gate_shexp.weight q8_0
 blk.4.ffn_up_shexp.weight q8_0
+blk.5.attn_k_b.weight q8_0
 blk.5.attn_kv_a_mqa.weight q8_0
 blk.5.attn_output.weight q8_0
 blk.5.attn_q_a.weight q8_0
 blk.5.attn_q_b.weight q8_0
+blk.5.attn_v_b.weight q8_0
 blk.5.ffn_down_shexp.weight q8_0
 blk.5.ffn_gate_shexp.weight q8_0
 blk.5.ffn_up_shexp.weight q8_0
+blk.6.attn_k_b.weight q8_0
 blk.6.attn_kv_a_mqa.weight q8_0
 blk.6.attn_output.weight q8_0
 blk.6.attn_q_a.weight q8_0
 blk.6.attn_q_b.weight q8_0
+blk.6.attn_v_b.weight q8_0
 blk.6.ffn_down_shexp.weight q8_0
 blk.6.ffn_gate_shexp.weight q8_0
 blk.6.ffn_up_shexp.weight q8_0
+blk.7.attn_k_b.weight q8_0
 blk.7.attn_kv_a_mqa.weight q8_0
 blk.7.attn_output.weight q8_0
 blk.7.attn_q_a.weight q8_0
 blk.7.attn_q_b.weight q8_0
+blk.7.attn_v_b.weight q8_0
 blk.7.ffn_down_shexp.weight q8_0
 blk.7.ffn_gate_shexp.weight q8_0
 blk.7.ffn_up_shexp.weight q8_0
+blk.8.attn_k_b.weight q8_0
 blk.8.attn_kv_a_mqa.weight q8_0
 blk.8.attn_output.weight q8_0
 blk.8.attn_q_a.weight q8_0
 blk.8.attn_q_b.weight q8_0
+blk.8.attn_v_b.weight q8_0
 blk.8.ffn_down_shexp.weight q8_0
 blk.8.ffn_gate_shexp.weight q8_0
 blk.8.ffn_up_shexp.weight q8_0
+blk.9.attn_k_b.weight q8_0
 blk.9.attn_kv_a_mqa.weight q8_0
 blk.9.attn_output.weight q8_0
 blk.9.attn_q_a.weight q8_0
 blk.9.attn_q_b.weight q8_0
+blk.9.attn_v_b.weight q8_0
 blk.9.ffn_down_shexp.weight q8_0
 blk.9.ffn_gate_shexp.weight q8_0
 blk.9.ffn_up_shexp.weight q8_0
+blk.10.attn_k_b.weight q8_0
 blk.10.attn_kv_a_mqa.weight q8_0
 blk.10.attn_output.weight q8_0
 blk.10.attn_q_a.weight q8_0
 blk.10.attn_q_b.weight q8_0
+blk.10.attn_v_b.weight q8_0
 blk.10.ffn_down_shexp.weight q8_0
 blk.10.ffn_gate_shexp.weight q8_0
 blk.10.ffn_up_shexp.weight q8_0
+blk.11.attn_k_b.weight q8_0
 blk.11.attn_kv_a_mqa.weight q8_0
 blk.11.attn_output.weight q8_0
 blk.11.attn_q_a.weight q8_0
 blk.11.attn_q_b.weight q8_0
+blk.11.attn_v_b.weight q8_0
 blk.11.ffn_down_shexp.weight q8_0
 blk.11.ffn_gate_shexp.weight q8_0
 blk.11.ffn_up_shexp.weight q8_0
+blk.12.attn_k_b.weight q8_0
 blk.12.attn_kv_a_mqa.weight q8_0
 blk.12.attn_output.weight q8_0
 blk.12.attn_q_a.weight q8_0
 blk.12.attn_q_b.weight q8_0
+blk.12.attn_v_b.weight q8_0
 blk.12.ffn_down_shexp.weight q8_0
 blk.12.ffn_gate_shexp.weight q8_0
 blk.12.ffn_up_shexp.weight q8_0
+blk.13.attn_k_b.weight q8_0
 blk.13.attn_kv_a_mqa.weight q8_0
 blk.13.attn_output.weight q8_0
 blk.13.attn_q_a.weight q8_0
 blk.13.attn_q_b.weight q8_0
+blk.13.attn_v_b.weight q8_0
 blk.13.ffn_down_shexp.weight q8_0
 blk.13.ffn_gate_shexp.weight q8_0
 blk.13.ffn_up_shexp.weight q8_0
+blk.14.attn_k_b.weight q8_0
 blk.14.attn_kv_a_mqa.weight q8_0
 blk.14.attn_output.weight q8_0
 blk.14.attn_q_a.weight q8_0
 blk.14.attn_q_b.weight q8_0
+blk.14.attn_v_b.weight q8_0
 blk.14.ffn_down_shexp.weight q8_0
 blk.14.ffn_gate_shexp.weight q8_0
 blk.14.ffn_up_shexp.weight q8_0
+blk.15.attn_k_b.weight q8_0
 blk.15.attn_kv_a_mqa.weight q8_0
 blk.15.attn_output.weight q8_0
 blk.15.attn_q_a.weight q8_0
 blk.15.attn_q_b.weight q8_0
+blk.15.attn_v_b.weight q8_0
 blk.15.ffn_down_shexp.weight q8_0
 blk.15.ffn_gate_shexp.weight q8_0
 blk.15.ffn_up_shexp.weight q8_0
+blk.16.attn_k_b.weight q8_0
 blk.16.attn_kv_a_mqa.weight q8_0
 blk.16.attn_output.weight q8_0
 blk.16.attn_q_a.weight q8_0
 blk.16.attn_q_b.weight q8_0
+blk.16.attn_v_b.weight q8_0
 blk.16.ffn_down_shexp.weight q8_0
 blk.16.ffn_gate_shexp.weight q8_0
 blk.16.ffn_up_shexp.weight q8_0
+blk.17.attn_k_b.weight q8_0
 blk.17.attn_kv_a_mqa.weight q8_0
 blk.17.attn_output.weight q8_0
 blk.17.attn_q_a.weight q8_0
 blk.17.attn_q_b.weight q8_0
+blk.17.attn_v_b.weight q8_0
 blk.17.ffn_down_shexp.weight q8_0
 blk.17.ffn_gate_shexp.weight q8_0
 blk.17.ffn_up_shexp.weight q8_0
+blk.18.attn_k_b.weight q8_0
 blk.18.attn_kv_a_mqa.weight q8_0
 blk.18.attn_output.weight q8_0
 blk.18.attn_q_a.weight q8_0
 blk.18.attn_q_b.weight q8_0
+blk.18.attn_v_b.weight q8_0
 blk.18.ffn_down_shexp.weight q8_0
 blk.18.ffn_gate_shexp.weight q8_0
 blk.18.ffn_up_shexp.weight q8_0
+blk.19.attn_k_b.weight q8_0
 blk.19.attn_kv_a_mqa.weight q8_0
 blk.19.attn_output.weight q8_0
 blk.19.attn_q_a.weight q8_0
 blk.19.attn_q_b.weight q8_0
+blk.19.attn_v_b.weight q8_0
 blk.19.ffn_down_shexp.weight q8_0
 blk.19.ffn_gate_shexp.weight q8_0
 blk.19.ffn_up_shexp.weight q8_0
+blk.20.attn_k_b.weight q8_0
 blk.20.attn_kv_a_mqa.weight q8_0
 blk.20.attn_output.weight q8_0
 blk.20.attn_q_a.weight q8_0
 blk.20.attn_q_b.weight q8_0
+blk.20.attn_v_b.weight q8_0
 blk.20.ffn_down_shexp.weight q8_0
 blk.20.ffn_gate_shexp.weight q8_0
 blk.20.ffn_up_shexp.weight q8_0
+blk.21.attn_k_b.weight q8_0
 blk.21.attn_kv_a_mqa.weight q8_0
 blk.21.attn_output.weight q8_0
 blk.21.attn_q_a.weight q8_0
 blk.21.attn_q_b.weight q8_0
+blk.21.attn_v_b.weight q8_0
 blk.21.ffn_down_shexp.weight q8_0
 blk.21.ffn_gate_shexp.weight q8_0
 blk.21.ffn_up_shexp.weight q8_0
+blk.22.attn_k_b.weight q8_0
 blk.22.attn_kv_a_mqa.weight q8_0
 blk.22.attn_output.weight q8_0
 blk.22.attn_q_a.weight q8_0
 blk.22.attn_q_b.weight q8_0
+blk.22.attn_v_b.weight q8_0
 blk.22.ffn_down_shexp.weight q8_0
 blk.22.ffn_gate_shexp.weight q8_0
 blk.22.ffn_up_shexp.weight q8_0
+blk.23.attn_k_b.weight q8_0
 blk.23.attn_kv_a_mqa.weight q8_0
 blk.23.attn_output.weight q8_0
 blk.23.attn_q_a.weight q8_0
 blk.23.attn_q_b.weight q8_0
+blk.23.attn_v_b.weight q8_0
 blk.23.ffn_down_shexp.weight q8_0
 blk.23.ffn_gate_shexp.weight q8_0
 blk.23.ffn_up_shexp.weight q8_0
+blk.24.attn_k_b.weight q8_0
 blk.24.attn_kv_a_mqa.weight q8_0
 blk.24.attn_output.weight q8_0
 blk.24.attn_q_a.weight q8_0
 blk.24.attn_q_b.weight q8_0
+blk.24.attn_v_b.weight q8_0
 blk.24.ffn_down_shexp.weight q8_0
 blk.24.ffn_gate_shexp.weight q8_0
 blk.24.ffn_up_shexp.weight q8_0
+blk.25.attn_k_b.weight q8_0
 blk.25.attn_kv_a_mqa.weight q8_0
 blk.25.attn_output.weight q8_0
 blk.25.attn_q_a.weight q8_0
 blk.25.attn_q_b.weight q8_0
+blk.25.attn_v_b.weight q8_0
 blk.25.ffn_down_shexp.weight q8_0
 blk.25.ffn_gate_shexp.weight q8_0
 blk.25.ffn_up_shexp.weight q8_0
+blk.26.attn_k_b.weight q8_0
 blk.26.attn_kv_a_mqa.weight q8_0
 blk.26.attn_output.weight q8_0
 blk.26.attn_q_a.weight q8_0
 blk.26.attn_q_b.weight q8_0
+blk.26.attn_v_b.weight q8_0
 blk.26.ffn_down_shexp.weight q8_0
 blk.26.ffn_gate_shexp.weight q8_0
 blk.26.ffn_up_shexp.weight q8_0
+blk.27.attn_k_b.weight q8_0
 blk.27.attn_kv_a_mqa.weight q8_0
 blk.27.attn_output.weight q8_0
 blk.27.attn_q_a.weight q8_0
 blk.27.attn_q_b.weight q8_0
+blk.27.attn_v_b.weight q8_0
 blk.27.ffn_down_shexp.weight q8_0
 blk.27.ffn_gate_shexp.weight q8_0
 blk.27.ffn_up_shexp.weight q8_0
+blk.28.attn_k_b.weight q8_0
 blk.28.attn_kv_a_mqa.weight q8_0
 blk.28.attn_output.weight q8_0
 blk.28.attn_q_a.weight q8_0
 blk.28.attn_q_b.weight q8_0
+blk.28.attn_v_b.weight q8_0
 blk.28.ffn_down_shexp.weight q8_0
 blk.28.ffn_gate_shexp.weight q8_0
 blk.28.ffn_up_shexp.weight q8_0
+blk.29.attn_k_b.weight q8_0
 blk.29.attn_kv_a_mqa.weight q8_0
 blk.29.attn_output.weight q8_0
 blk.29.attn_q_a.weight q8_0
 blk.29.attn_q_b.weight q8_0
+blk.29.attn_v_b.weight q8_0
 blk.29.ffn_down_shexp.weight q8_0
 blk.29.ffn_gate_shexp.weight q8_0
 blk.29.ffn_up_shexp.weight q8_0
+blk.30.attn_k_b.weight q8_0
 blk.30.attn_kv_a_mqa.weight q8_0
 blk.30.attn_output.weight q8_0
 blk.30.attn_q_a.weight q8_0
 blk.30.attn_q_b.weight q8_0
+blk.30.attn_v_b.weight q8_0
 blk.30.ffn_down_shexp.weight q8_0
 blk.30.ffn_gate_shexp.weight q8_0
 blk.30.ffn_up_shexp.weight q8_0
+blk.31.attn_k_b.weight q8_0
 blk.31.attn_kv_a_mqa.weight q8_0
 blk.31.attn_output.weight q8_0
 blk.31.attn_q_a.weight q8_0
 blk.31.attn_q_b.weight q8_0
+blk.31.attn_v_b.weight q8_0
 blk.31.ffn_down_shexp.weight q8_0
 blk.31.ffn_gate_shexp.weight q8_0
 blk.31.ffn_up_shexp.weight q8_0
+blk.32.attn_k_b.weight q8_0
 blk.32.attn_kv_a_mqa.weight q8_0
 blk.32.attn_output.weight q8_0
 blk.32.attn_q_a.weight q8_0
 blk.32.attn_q_b.weight q8_0
+blk.32.attn_v_b.weight q8_0
 blk.32.ffn_down_shexp.weight q8_0
 blk.32.ffn_gate_shexp.weight q8_0
 blk.32.ffn_up_shexp.weight q8_0
+blk.33.attn_k_b.weight q8_0
 blk.33.attn_kv_a_mqa.weight q8_0
 blk.33.attn_output.weight q8_0
 blk.33.attn_q_a.weight q8_0
 blk.33.attn_q_b.weight q8_0
+blk.33.attn_v_b.weight q8_0
 blk.33.ffn_down_shexp.weight q8_0
 blk.33.ffn_gate_shexp.weight q8_0
 blk.33.ffn_up_shexp.weight q8_0
+blk.34.attn_k_b.weight q8_0
 blk.34.attn_kv_a_mqa.weight q8_0
 blk.34.attn_output.weight q8_0
 blk.34.attn_q_a.weight q8_0
 blk.34.attn_q_b.weight q8_0
+blk.34.attn_v_b.weight q8_0
 blk.34.ffn_down_shexp.weight q8_0
 blk.34.ffn_gate_shexp.weight q8_0
 blk.34.ffn_up_shexp.weight q8_0
+blk.35.attn_k_b.weight q8_0
 blk.35.attn_kv_a_mqa.weight q8_0
 blk.35.attn_output.weight q8_0
 blk.35.attn_q_a.weight q8_0
 blk.35.attn_q_b.weight q8_0
+blk.35.attn_v_b.weight q8_0
 blk.35.ffn_down_shexp.weight q8_0
 blk.35.ffn_gate_shexp.weight q8_0
 blk.35.ffn_up_shexp.weight q8_0
+blk.36.attn_k_b.weight q8_0
 blk.36.attn_kv_a_mqa.weight q8_0
 blk.36.attn_output.weight q8_0
 blk.36.attn_q_a.weight q8_0
 blk.36.attn_q_b.weight q8_0
+blk.36.attn_v_b.weight q8_0
 blk.36.ffn_down_shexp.weight q8_0
 blk.36.ffn_gate_shexp.weight q8_0
 blk.36.ffn_up_shexp.weight q8_0
+blk.37.attn_k_b.weight q8_0
 blk.37.attn_kv_a_mqa.weight q8_0
 blk.37.attn_output.weight q8_0
 blk.37.attn_q_a.weight q8_0
 blk.37.attn_q_b.weight q8_0
+blk.37.attn_v_b.weight q8_0
 blk.37.ffn_down_shexp.weight q8_0
 blk.37.ffn_gate_shexp.weight q8_0
 blk.37.ffn_up_shexp.weight q8_0
+blk.38.attn_k_b.weight q8_0
 blk.38.attn_kv_a_mqa.weight q8_0
 blk.38.attn_output.weight q8_0
 blk.38.attn_q_a.weight q8_0
 blk.38.attn_q_b.weight q8_0
+blk.38.attn_v_b.weight q8_0
 blk.38.ffn_down_shexp.weight q8_0
 blk.38.ffn_gate_shexp.weight q8_0
 blk.38.ffn_up_shexp.weight q8_0
+blk.39.attn_k_b.weight q8_0
 blk.39.attn_kv_a_mqa.weight q8_0
 blk.39.attn_output.weight q8_0
 blk.39.attn_q_a.weight q8_0
 blk.39.attn_q_b.weight q8_0
+blk.39.attn_v_b.weight q8_0
 blk.39.ffn_down_shexp.weight q8_0
 blk.39.ffn_gate_shexp.weight q8_0
 blk.39.ffn_up_shexp.weight q8_0
+blk.40.attn_k_b.weight q8_0
 blk.40.attn_kv_a_mqa.weight q8_0
 blk.40.attn_output.weight q8_0
 blk.40.attn_q_a.weight q8_0
 blk.40.attn_q_b.weight q8_0
+blk.40.attn_v_b.weight q8_0
 blk.40.ffn_down_shexp.weight q8_0
 blk.40.ffn_gate_shexp.weight q8_0
 blk.40.ffn_up_shexp.weight q8_0
+blk.41.attn_k_b.weight q8_0
 blk.41.attn_kv_a_mqa.weight q8_0
 blk.41.attn_output.weight q8_0
 blk.41.attn_q_a.weight q8_0
 blk.41.attn_q_b.weight q8_0
+blk.41.attn_v_b.weight q8_0
 blk.41.ffn_down_shexp.weight q8_0
 blk.41.ffn_gate_shexp.weight q8_0
 blk.41.ffn_up_shexp.weight q8_0
+blk.42.attn_k_b.weight q8_0
 blk.42.attn_kv_a_mqa.weight q8_0
 blk.42.attn_output.weight q8_0
 blk.42.attn_q_a.weight q8_0
 blk.42.attn_q_b.weight q8_0
+blk.42.attn_v_b.weight q8_0
 blk.42.ffn_down_shexp.weight q8_0
 blk.42.ffn_gate_shexp.weight q8_0
 blk.42.ffn_up_shexp.weight q8_0
+blk.43.attn_k_b.weight q8_0
 blk.43.attn_kv_a_mqa.weight q8_0
 blk.43.attn_output.weight q8_0
 blk.43.attn_q_a.weight q8_0
 blk.43.attn_q_b.weight q8_0
+blk.43.attn_v_b.weight q8_0
 blk.43.ffn_down_shexp.weight q8_0
 blk.43.ffn_gate_shexp.weight q8_0
 blk.43.ffn_up_shexp.weight q8_0
+blk.44.attn_k_b.weight q8_0
 blk.44.attn_kv_a_mqa.weight q8_0
 blk.44.attn_output.weight q8_0
 blk.44.attn_q_a.weight q8_0
 blk.44.attn_q_b.weight q8_0
+blk.44.attn_v_b.weight q8_0
 blk.44.ffn_down_shexp.weight q8_0
 blk.44.ffn_gate_shexp.weight q8_0
 blk.44.ffn_up_shexp.weight q8_0
+blk.45.attn_k_b.weight q8_0
 blk.45.attn_kv_a_mqa.weight q8_0
 blk.45.attn_output.weight q8_0
 blk.45.attn_q_a.weight q8_0
 blk.45.attn_q_b.weight q8_0
+blk.45.attn_v_b.weight q8_0
 blk.45.ffn_down_shexp.weight q8_0
 blk.45.ffn_gate_shexp.weight q8_0
 blk.45.ffn_up_shexp.weight q8_0
+blk.46.attn_k_b.weight q8_0
 blk.46.attn_kv_a_mqa.weight q8_0
 blk.46.attn_output.weight q8_0
 blk.46.attn_q_a.weight q8_0
 blk.46.attn_q_b.weight q8_0
+blk.46.attn_v_b.weight q8_0
 blk.46.ffn_down_shexp.weight q8_0
 blk.46.ffn_gate_shexp.weight q8_0
 blk.46.ffn_up_shexp.weight q8_0
+blk.47.attn_k_b.weight q8_0
 blk.47.attn_kv_a_mqa.weight q8_0
 blk.47.attn_output.weight q8_0
 blk.47.attn_q_a.weight q8_0
 blk.47.attn_q_b.weight q8_0
+blk.47.attn_v_b.weight q8_0
 blk.47.ffn_down_shexp.weight q8_0
 blk.47.ffn_gate_shexp.weight q8_0
 blk.47.ffn_up_shexp.weight q8_0
+blk.48.attn_k_b.weight q8_0
 blk.48.attn_kv_a_mqa.weight q8_0
 blk.48.attn_output.weight q8_0
 blk.48.attn_q_a.weight q8_0
 blk.48.attn_q_b.weight q8_0
+blk.48.attn_v_b.weight q8_0
 blk.48.ffn_down_shexp.weight q8_0
 blk.48.ffn_gate_shexp.weight q8_0
 blk.48.ffn_up_shexp.weight q8_0
+blk.49.attn_k_b.weight q8_0
 blk.49.attn_kv_a_mqa.weight q8_0
 blk.49.attn_output.weight q8_0
 blk.49.attn_q_a.weight q8_0
 blk.49.attn_q_b.weight q8_0
+blk.49.attn_v_b.weight q8_0
 blk.49.ffn_down_shexp.weight q8_0
 blk.49.ffn_gate_shexp.weight q8_0
 blk.49.ffn_up_shexp.weight q8_0
+blk.50.attn_k_b.weight q8_0
 blk.50.attn_kv_a_mqa.weight q8_0
 blk.50.attn_output.weight q8_0
 blk.50.attn_q_a.weight q8_0
 blk.50.attn_q_b.weight q8_0
+blk.50.attn_v_b.weight q8_0
 blk.50.ffn_down_shexp.weight q8_0
 blk.50.ffn_gate_shexp.weight q8_0
 blk.50.ffn_up_shexp.weight q8_0
+blk.51.attn_k_b.weight q8_0
 blk.51.attn_kv_a_mqa.weight q8_0
 blk.51.attn_output.weight q8_0
 blk.51.attn_q_a.weight q8_0
 blk.51.attn_q_b.weight q8_0
+blk.51.attn_v_b.weight q8_0
 blk.51.ffn_down_shexp.weight q8_0
 blk.51.ffn_gate_shexp.weight q8_0
 blk.51.ffn_up_shexp.weight q8_0
+blk.52.attn_k_b.weight q8_0
 blk.52.attn_kv_a_mqa.weight q8_0
 blk.52.attn_output.weight q8_0
 blk.52.attn_q_a.weight q8_0
 blk.52.attn_q_b.weight q8_0
+blk.52.attn_v_b.weight q8_0
 blk.52.ffn_down_shexp.weight q8_0
 blk.52.ffn_gate_shexp.weight q8_0
 blk.52.ffn_up_shexp.weight q8_0
+blk.53.attn_k_b.weight q8_0
 blk.53.attn_kv_a_mqa.weight q8_0
 blk.53.attn_output.weight q8_0
 blk.53.attn_q_a.weight q8_0
 blk.53.attn_q_b.weight q8_0
+blk.53.attn_v_b.weight q8_0
 blk.53.ffn_down_shexp.weight q8_0
 blk.53.ffn_gate_shexp.weight q8_0
 blk.53.ffn_up_shexp.weight q8_0
+blk.54.attn_k_b.weight q8_0
 blk.54.attn_kv_a_mqa.weight q8_0
 blk.54.attn_output.weight q8_0
 blk.54.attn_q_a.weight q8_0
 blk.54.attn_q_b.weight q8_0
+blk.54.attn_v_b.weight q8_0
 blk.54.ffn_down_shexp.weight q8_0
 blk.54.ffn_gate_shexp.weight q8_0
 blk.54.ffn_up_shexp.weight q8_0
+blk.55.attn_k_b.weight q8_0
 blk.55.attn_kv_a_mqa.weight q8_0
 blk.55.attn_output.weight q8_0
 blk.55.attn_q_a.weight q8_0
 blk.55.attn_q_b.weight q8_0
+blk.55.attn_v_b.weight q8_0
 blk.55.ffn_down_shexp.weight q8_0
 blk.55.ffn_gate_shexp.weight q8_0
 blk.55.ffn_up_shexp.weight q8_0
+blk.56.attn_k_b.weight q8_0
 blk.56.attn_kv_a_mqa.weight q8_0
 blk.56.attn_output.weight q8_0
 blk.56.attn_q_a.weight q8_0
 blk.56.attn_q_b.weight q8_0
+blk.56.attn_v_b.weight q8_0
 blk.56.ffn_down_shexp.weight q8_0
 blk.56.ffn_gate_shexp.weight q8_0
 blk.56.ffn_up_shexp.weight q8_0
+blk.57.attn_k_b.weight q8_0
 blk.57.attn_kv_a_mqa.weight q8_0
 blk.57.attn_output.weight q8_0
 blk.57.attn_q_a.weight q8_0
 blk.57.attn_q_b.weight q8_0
+blk.57.attn_v_b.weight q8_0
 blk.57.ffn_down_shexp.weight q8_0
 blk.57.ffn_gate_shexp.weight q8_0
 blk.57.ffn_up_shexp.weight q8_0
+blk.58.attn_k_b.weight q8_0
 blk.58.attn_kv_a_mqa.weight q8_0
 blk.58.attn_output.weight q8_0
 blk.58.attn_q_a.weight q8_0
 blk.58.attn_q_b.weight q8_0
+blk.58.attn_v_b.weight q8_0
 blk.58.ffn_down_shexp.weight q8_0
 blk.58.ffn_gate_shexp.weight q8_0
 blk.58.ffn_up_shexp.weight q8_0
+blk.59.attn_k_b.weight q8_0
 blk.59.attn_kv_a_mqa.weight q8_0
 blk.59.attn_output.weight q8_0
 blk.59.attn_q_a.weight q8_0
 blk.59.attn_q_b.weight q8_0
+blk.59.attn_v_b.weight q8_0
 blk.59.ffn_down_shexp.weight q8_0
 blk.59.ffn_gate_shexp.weight q8_0
 blk.59.ffn_up_shexp.weight q8_0
+blk.60.attn_k_b.weight q8_0
 blk.60.attn_kv_a_mqa.weight q8_0
 blk.60.attn_output.weight q8_0
 blk.60.attn_q_a.weight q8_0
 blk.60.attn_q_b.weight q8_0
+blk.60.attn_v_b.weight q8_0
 blk.60.ffn_down_shexp.weight q8_0
 blk.60.ffn_gate_shexp.weight q8_0
 blk.60.ffn_up_shexp.weight q8_0

tests/test-backend-ops.cpp

@@ -7126,7 +7126,7 @@ struct test_concat : public test_case {
     const std::array<int64_t, 4> ne_a;
     const int64_t ne_b_d;
     const int dim;
-    const int v; // view (1 << 0: non-cont a (first 3 dim), 1 << 1: non-cont b (first 3 dim), 1 << 2: non-cont a (last 2 dim), 1 << 3: non-cont b (last 2 dim))
+    const int v; // view (1 << 0: non-cont a (first 3 dim), 1 << 1: non-cont b (first 3 dim), 1 << 2: non-cont a (last 2 dim), 1 << 3: non-cont b (last 2 dim), 1 << 4: transposed b)
 
     std::string vars() override {
         return VARS_TO_STR5(type, ne_a, ne_b_d, dim, v);
@@ -7175,6 +7175,13 @@ struct test_concat : public test_case {
 
             b = ggml_view_4d(ctx, b, ne_b[0], ne_b[1], ne_b[2], ne_b[3], b->nb[1], b->nb[2], b->nb[3], 0);
             ggml_set_name(b, "view_of_b");
+        } else if (v & 16) {
+            auto ne = ne_b; std::swap(ne[0], ne[1]);
+            b = ggml_new_tensor(ctx, type, 4, ne.data());
+            ggml_set_name(b, "b");
+
+            b = ggml_transpose(ctx, b);
+            ggml_set_name(b, "transpose_of_b");
         } else {
             b = ggml_new_tensor(ctx, type, 4, ne_b.data());
             ggml_set_name(b, "b");
@@ -9996,6 +10003,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     std::default_random_engine rng(0);
 
     // unary ops
+    // TODO: add GGML_UNARY_OP_SOFTPLUS test that is sensitive to using `log1p(exp(x))` vs `log(exp(x) + 1)`
     for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32, GGML_TYPE_BF16}) {
         for (int v : {0, 1}) {
             for (int op = 0; op < GGML_UNARY_OP_COUNT; op++) {
@@ -10892,9 +10900,20 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_l2_norm     (GGML_TYPE_F32, {4, 1, 65536, 1}, 1e-12f, false, false));
     test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, {4, 1, 65536, 1}, 1e-6f, false, false));
 
+    for (uint32_t n : {1, 127, 129, 511, 513, 2816, 4096, 16385}) {
+        test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, {n, 3, 2, 1}));
+        test_cases.emplace_back(new test_rms_norm_mul_rope({n, 3, 2, 1}, 1e-6f, false, false, true, GGML_ROPE_TYPE_NORMAL, true, false));
+        test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, {n, 3, 2, 1}, 1e-6f, false, false, false, false, true));
+    }
+
     for (ggml_type set_rows_type : { GGML_TYPE_F32, GGML_TYPE_F16 }) {
         test_cases.emplace_back(new test_rms_norm_mul_rope({ 256, 1, 1, 1 }, 1e-6f, false, true, false, GGML_ROPE_TYPE_NORMAL, false, false, set_rows_type));
         test_cases.emplace_back(new test_rms_norm_mul_rope({ 128, 4, 3, 1 }, 1e-6f, false, true, false, GGML_ROPE_TYPE_NORMAL, false, false, set_rows_type));
+        for (int mode : {GGML_ROPE_TYPE_NORMAL, GGML_ROPE_TYPE_NEOX}) {
+            for (int64_t n : {64, 256, 512, 1024}) {
+                test_cases.emplace_back(new test_rms_norm_mul_rope({n, 4, 3, 1}, 1e-6f, false, true, true, mode, true, true, set_rows_type));
+            }
+        }
     }
 
     for (float eps : { 0.0f, 1e-6f, 1e-4f, 1e-1f, 1.0f }) {
@@ -11079,6 +11098,17 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_mul_mat_id_w4a8(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256));
     test_cases.emplace_back(new test_mul_mat_id_w4a4(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256));
 
+#if 0
+    // > 2^31-element dst (8.9 GB), about 45 GB peak memory. Too large to be enabled by default.
+    // kernel_mul_mm dst offset of the last batch is past INT32_MAX: rows 0-31 use full tiles, rows 32-39 use edge tiles.
+    test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 1048576, 40, 64, {1, 1}, {53, 1}));
+
+    // kernel_mul_mm_id: the index of tokens 32768-33791 does not fit in 16 bits (about 0.3 GB peak memory).
+    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 4, 2, false, 64, 33792, 64));
+    // kernel_mul_mm_id: > 2^31-element dst (8.7 GB), about 45 GB peak memory. Tokens 32514+ have a dst offset past INT32_MAX, tokens 32768+ also need a 32-bit token index.
+    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 16, 8, false, 8256, 32832, 64));
+#endif
+
 #if 0
     // > 4GB A matrix. Too slow to be enabled by default.
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16,  900000,  3, 2592, {1, 1}, {1, 1}));
@@ -11780,7 +11810,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_rope(GGML_TYPE_F32, {256,  8, 512, 1}, 256, GGML_ROPE_TYPE_NEOX,   512, 1.0f, 0.0f, 1.0f, false, 0, true)); // gemma4 E2B sliding
     test_cases.emplace_back(new test_rope(GGML_TYPE_F32, {512,  8, 512, 1}, 128, GGML_ROPE_TYPE_NEOX,   512, 1.0f, 0.0f, 1.0f, true,  0, true)); // gemma4 E4B global
 
-    for (int v : { 0, 1, 2, 3 }) {
+    for (int v : { 0, 1, 2, 3, 16 }) {
         for (int dim : { 0, 1, 2, 3, }) {
             test_cases.emplace_back(new test_concat(GGML_TYPE_F32, {11, 12, 13, 14}, 7, dim, v));
             test_cases.emplace_back(new test_concat(GGML_TYPE_F16, {11, 12, 13, 14}, 7, dim, v));
@@ -12508,6 +12538,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
 static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
     std::vector<std::unique_ptr<test_case>> test_cases;
 
+    // delta-net conv input: conv state followed by the transposed qkv of the ubatch
+    test_cases.emplace_back(new test_concat(GGML_TYPE_F32, {3, 8192, 1, 1}, 512, 0, 16));
+
     // SWIGLU at a 27B-class FFN width, fused [gate|up] vs split operands
     // note: same bytes either way, so a backend that indexes them differently shows it here
     for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32, GGML_TYPE_BF16}) {
@@ -12537,6 +12570,15 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
         }
     }
 
+    // RMS_NORM and fused variants at decode and prefill shapes
+    for (int64_t n_tokens : {1, 512}) {
+        test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, {4096, n_tokens, 1, 1}));
+        test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, {128, 32*n_tokens, 1, 1}));
+        test_cases.emplace_back(new test_rms_norm_mul_rope({4096, n_tokens, 1, 1}, 1e-6f, false, false, true, GGML_ROPE_TYPE_NORMAL, true, false));
+        test_cases.emplace_back(new test_rms_norm_mul_rope({128, 32, n_tokens, 1}, 1e-6f, false, false, true, GGML_ROPE_TYPE_NEOX));
+        test_cases.emplace_back(new test_rms_norm_mul_rope({128, 8, n_tokens, 1}, 1e-6f, false, true, true, GGML_ROPE_TYPE_NEOX));
+    }
+
     // 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 }) {

tests/test-grammar-integration.cpp

@@ -879,6 +879,18 @@ static void test_failure_left_recursion() {
         empty ::= "blah" | )""";
     assert(test_build_grammar_fails(hardest_str));
 
+    const auto chain = [](int n) {
+        std::string s = "root ::= r0\n";
+        for (int i = 0; i < n; i++) {
+            s += "r" + std::to_string(i) + " ::= r" + std::to_string(i + 1) + "\n";
+        }
+        return s + "r" + std::to_string(n) + " ::= \"a\"\n";
+    };
+    llama_grammar * grammar = build_grammar(chain(100));
+    assert(grammar != nullptr);
+    llama_grammar_free_impl(grammar);
+    assert(test_build_grammar_fails(chain(1000)));
+
     fprintf(stderr, "  ✅︎ Passed\n");
 }
 
@@ -1476,6 +1488,25 @@ static void test_json_schema() {
             // R"""({"productId": 1, "productName": "A green door", "price": 12.50, "tags": ["home", "green", "home"]})""",
         }
     );
+
+    // the deepest schema of json_schema_to_grammar must give a valid grammar
+    std::string deep_schema;
+    for (int i = 0; i < 127; i++) {
+        deep_schema += R"({"type": "array", "items": )";
+    }
+    deep_schema += R"({"type": "string"})" + std::string(127, '}');
+    test_schema(
+        "127 nested arrays",
+        deep_schema,
+        // Passing strings
+        {
+            std::string(127, '[') + R"("a")" + std::string(127, ']'),
+        },
+        // Failing strings
+        {
+            std::string(126, '[') + R"("a")" + std::string(126, ']'),
+        }
+    );
 }
 
 int main() {

tests/test-grammar-parser.cpp

@@ -165,6 +165,20 @@ int main()
         root ::= "a"{5000,6000}
     )""");
 
+    // Recursion depth limit exceeded (Issue #29495)
+    {
+        std::string deep_grammar = "root ::= ";
+        int n_depth = 300;
+        for (int i = 0; i < n_depth; ++i) {
+            deep_grammar += "(";
+        }
+        deep_grammar += "\"x\"";
+        for (int i = 0; i < n_depth; ++i) {
+            deep_grammar += ")";
+        }
+        verify_failure(deep_grammar.c_str());
+    }
+
     verify_parsing(R"""(
         root ::= "a"{0,5000}
     )""", {

tests/test-jinja.cpp

@@ -17,6 +17,7 @@
 using json = common_json;
 
 static void test_template(testing & t, const std::string & name, const std::string & tmpl, const json & vars, const std::string & expect);
+static void test_template_error(testing & t, const std::string & name, const std::string & tmpl, const std::string & expect);
 
 static void test_whitespace_control(testing & t);
 static void test_conditionals(testing & t);
@@ -606,6 +607,17 @@ static void test_expressions(testing & t) {
         {{"cond", true}},
         "yes"
     );
+
+    test_template(t, "nested parentheses",
+        "{{ " + std::string(30, '(') + "1" + std::string(30, ')') + " }}",
+        json::object(),
+        "1"
+    );
+
+    test_template_error(t, "too deeply nested parentheses",
+        "{{ " + std::string(1000, '(') + "1" + std::string(1000, ')') + " }}",
+        "template is nested more than 128 levels deep"
+    );
 }
 
 static void test_set_statement(testing & t) {
@@ -1029,6 +1041,11 @@ static void test_filters(testing & t) {
         {{"x", nullptr}},
         "None"
     );
+
+    test_template_error(t, "tojson of a list that contains itself",
+        "{% set arr = [] %}{% set _ = arr.append(arr) %}{{ arr|tojson }}",
+        "value is nested more than 128 levels deep"
+    );
 }
 
 static void test_literals(testing & t) {
@@ -1291,6 +1308,17 @@ static void test_macros(testing & t) {
         json::object(),
         "Dict {\n    root1: 1\n    root2: {\n        nest1: 1\n        nest2: {\n            nest3: 2\n        }\n    }\n}"
     );
+
+    test_template(t, "macro recursion",
+        "{% macro f(n) %}{% if n < 100 %}{{ f(n + 1) }}{% else %}{{ n }}{% endif %}{% endmacro %}{{ f(0) }}",
+        json::object(),
+        "100"
+    );
+
+    test_template_error(t, "endless macro recursion",
+        "{% macro f(n) %}{{ f(n + 1) }}{% endmacro %}{{ f(0) }}",
+        "macro calls are nested more than 128 levels deep"
+    );
 }
 
 static void test_namespace(testing & t) {
@@ -2554,6 +2582,22 @@ static void test_string_parts(testing & t) {
 
 }
 
+static void test_template_error(testing & t, const std::string & name, const std::string & tmpl, const std::string & expect) {
+    t.test(name, [&tmpl, &expect](testing & t) {
+        std::string error;
+        try {
+            jinja::lexer lexer;
+            jinja::program ast = jinja::parse_from_tokens(lexer.tokenize(tmpl));
+            jinja::context ctx(tmpl);
+            jinja::runtime runtime(ctx);
+            runtime.gather_string_parts(runtime.execute(ast));
+        } catch (const std::exception & e) {
+            error = e.what();
+        }
+        t.assert_true("expected error: " + expect + ", got: " + error, error.find(expect) != std::string::npos);
+    });
+}
+
 static void test_template_cpp(testing & t, const std::string & name, const std::string & tmpl, const json & vars, const std::string & expect) {
     t.test(name, [&tmpl, &vars, &expect](testing & t) {
         jinja::lexer lexer;

tests/test-json-schema-to-grammar.cpp

@@ -1513,6 +1513,23 @@ static void test_all(const std::string & title, std::function<void(const TestCas
         ""
     });
 
+    test({
+        FAILURE,
+        "allOf with a $ref cycle",
+        R"""({
+            "allOf": [{"$ref": "#/$defs/a"}],
+            "$defs": {"a": {"$ref": "#/$defs/a"}}
+        })""",
+        ""
+    });
+
+    test({
+        FAILURE,
+        "long property name with additionalProperties",
+        R"""({"type": "object", "properties": {")""" + std::string(1000, 'a') + R"""(": {"type": "string"}}, "additionalProperties": true})""",
+        ""
+    });
+
     // only the property with the bad pattern degrades
     test({
         SUCCESS,

tests/test-llama-archs.cpp

@@ -6,10 +6,7 @@
 #include "ggml-cpp.h"
 #include "llama.h"
 #include "llama-cpp.h"
-
-// TODO: replace with #include "llama-ext.h" in the future
-#include "../src/llama-arch.h"
-#include "../src/llama-model-saver.h"
+#include "../src/llama-ext.h"
 
 #include <cinttypes>
 #include <cmath>
@@ -17,6 +14,7 @@
 #include <cstdio>
 #include <cstring>
 #include <cstdint>
+#include <filesystem>
 #include <random>
 #include <regex>
 #include <stdexcept>
@@ -28,7 +26,7 @@ static bool arch_matches(const std::string & filter, llm_arch arch) {
     if (filter.empty()) {
         return true;
     }
-    return std::regex_search(llm_arch_name(arch), std::regex(filter));
+    return std::regex_search(llama_arch_name(arch), std::regex(filter));
 }
 
 // normalized mean squared error = mse(a, b) / mse(a, 0)
@@ -164,9 +162,50 @@ static bool host_experts_test(const llm_arch arch) {
     }
 }
 
+// writes arch-resolved GGUF keys into the fixture context, the test-side equivalent of llama_model_saver
+struct kv_writer {
+    struct gguf_context * gguf_ctx;
+    const llm_arch        arch;
+
+    kv_writer(enum llm_arch arch, struct gguf_context * gguf_ctx) : gguf_ctx(gguf_ctx), arch(arch) {}
+
+    void add_kv(enum llm_kv key, uint32_t     value) { gguf_set_val_u32 (gguf_ctx, llama_arch_kv_name(arch, key).c_str(), value); }
+    void add_kv(enum llm_kv key, int32_t      value) { gguf_set_val_i32 (gguf_ctx, llama_arch_kv_name(arch, key).c_str(), value); }
+    void add_kv(enum llm_kv key, uint64_t     value) { gguf_set_val_u64 (gguf_ctx, llama_arch_kv_name(arch, key).c_str(), value); }
+    void add_kv(enum llm_kv key, float        value) { gguf_set_val_f32 (gguf_ctx, llama_arch_kv_name(arch, key).c_str(), value); }
+    void add_kv(enum llm_kv key, bool         value) { gguf_set_val_bool(gguf_ctx, llama_arch_kv_name(arch, key).c_str(), value); }
+    void add_kv(enum llm_kv key, const char * value) { gguf_set_val_str (gguf_ctx, llama_arch_kv_name(arch, key).c_str(), value); }
+
+    template <typename T>
+    void add_kv(enum llm_kv key, const std::vector<T> & value) {
+        if (value.empty()) {
+            return;
+        }
+        if (std::is_same<T, uint32_t>::value) {
+            gguf_set_arr_data(gguf_ctx, llama_arch_kv_name(arch, key).c_str(), GGUF_TYPE_UINT32, value.data(), value.size());
+        } else if (std::is_same<T, int32_t>::value) {
+            gguf_set_arr_data(gguf_ctx, llama_arch_kv_name(arch, key).c_str(), GGUF_TYPE_INT32, value.data(), value.size());
+        } else if (std::is_same<T, uint64_t>::value) {
+            gguf_set_arr_data(gguf_ctx, llama_arch_kv_name(arch, key).c_str(), GGUF_TYPE_UINT64, value.data(), value.size());
+        } else if (std::is_same<T, float>::value) {
+            gguf_set_arr_data(gguf_ctx, llama_arch_kv_name(arch, key).c_str(), GGUF_TYPE_FLOAT32, value.data(), value.size());
+        } else {
+            GGML_ABORT("fatal error");
+        }
+    }
+
+    void add_kv(enum llm_kv key, const std::vector<std::string> & value) {
+        std::vector<const char *> tmp(value.size());
+        for (size_t i = 0; i < value.size(); ++i) {
+            tmp[i] = value[i].c_str();
+        }
+        gguf_set_arr_str(gguf_ctx, llama_arch_kv_name(arch, key).c_str(), tmp.data(), tmp.size());
+    }
+};
+
 static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
     gguf_context_ptr ret(gguf_init_empty());
-    llama_model_saver ms(arch, ret.get());
+    kv_writer ms(arch, ret.get());
     const uint32_t n_ctx = 256;
 
     uint32_t n_vocab = 128;
@@ -243,7 +282,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
     }
     const uint32_t n_embd_head = n_embd / n_head;
 
-    ms.add_kv(LLM_KV_GENERAL_ARCHITECTURE,      llm_arch_name(arch));
+    ms.add_kv(LLM_KV_GENERAL_ARCHITECTURE,      llama_arch_name(arch));
     ms.add_kv(LLM_KV_VOCAB_SIZE,                n_vocab);
     ms.add_kv(LLM_KV_CONTEXT_LENGTH,            n_ctx);
     ms.add_kv(LLM_KV_EMBEDDING_LENGTH,          n_embd);
@@ -638,6 +677,11 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/)
     return true;
 }
 
+static std::string get_tmp_gguf_path() {
+    static size_t id = 0;
+    return (std::filesystem::temp_directory_path() / ("test-llama-archs-" + std::to_string(++id) + ".gguf")).string();
+}
+
 static std::pair<llama_model_ptr, llama_context_ptr> get_model_and_ctx(
         struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const float stdev,
         const std::vector<ggml_backend_dev_t> & devs,
@@ -958,7 +1002,7 @@ static int save_models(const std::string & arch_filter, const size_t seed, const
         }
     }, &ud);
 
-    for (const llm_arch & arch : llm_arch_all()) {
+    for (const llm_arch & arch : llama_arch_all()) {
         if (arch == LLM_ARCH_UNKNOWN) {
             continue;
         }
@@ -978,14 +1022,14 @@ static int save_models(const std::string & arch_filter, const size_t seed, const
             if (!moe && moe_mandatory(arch)) {
                 continue;
             }
-            if (!llama_model_saver_supports_arch(arch) || !arch_supported(arch)) {
-                LOG_INF("%s: %s model (%s) is unsupported, skipping\n", __func__, llm_arch_name(arch), moe ? "MoE" : "dense");
+            if (!llama_arch_supports_save(arch) || !arch_supported(arch)) {
+                LOG_INF("%s: %s model (%s) is unsupported, skipping\n", __func__, llama_arch_name(arch), moe ? "MoE" : "dense");
                 continue;
             }
             gguf_context_ptr gguf_ctx = get_gguf_ctx(arch, moe);
             auto model_and_ctx = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, stdev, {});
-            const std::string path = dir + "/" + llm_arch_name(arch) + (moe ? "-moe.gguf" : "-dense.gguf");
-            LOG_INF("%s: Saving %s model (%s) to %s...\n", __func__, llm_arch_name(arch), moe ? "MoE" : "dense", path.c_str());
+            const std::string path = dir + "/" + llama_arch_name(arch) + (moe ? "-moe.gguf" : "-dense.gguf");
+            LOG_INF("%s: Saving %s model (%s) to %s...\n", __func__, llama_arch_name(arch), moe ? "MoE" : "dense", path.c_str());
             llama_model_save_to_file(model_and_ctx.first.get(), path.c_str());
         }
     }
@@ -1094,8 +1138,8 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con
     }
 
     size_t max_arch_name_length = 0;
-    for (const llm_arch & arch : llm_arch_all()) {
-        max_arch_name_length = std::max(max_arch_name_length, strlen(llm_arch_name(arch)));
+    for (const llm_arch & arch : llama_arch_all()) {
+        max_arch_name_length = std::max(max_arch_name_length, strlen(llama_arch_name(arch)));
     }
 
     const std::string template_header  = std::string("|%" + std::to_string(max_arch_name_length) + "s|%") + std::to_string(max_device_label_length) + "s|%6s|%15s|%9s|%15s|\n";
@@ -1116,7 +1160,7 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con
         LOG("-");
     }
     LOG("|------|---------------|---------|---------------|\n");
-    for (const llm_arch & arch : llm_arch_all()) {
+    for (const llm_arch & arch : llama_arch_all()) {
         if (arch == LLM_ARCH_UNKNOWN) {
             continue;
         }
@@ -1152,7 +1196,7 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con
                 const llama_model_tensor_buft_override * overrides = dc.host_experts ? host_experts_overrides : nullptr;
 
                 // print test config first; should anything fail during model loading or inference, at least we know which test case caused it
-                LOG(template_row_cfg.c_str(), llm_arch_name(arch), dc.label.c_str(), config_name.c_str());
+                LOG(template_row_cfg.c_str(), llama_arch_name(arch), dc.label.c_str(), config_name.c_str());
                 fflush(stdout);
 
                 std::pair<llama_model_ptr, llama_context_ptr> model_and_ctx_dev;
@@ -1171,7 +1215,7 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con
                         model_and_ctx_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, stdev, {}, LLAMA_SPLIT_MODE_LAYER, encode);
                         logits_cpu = get_logits(model_and_ctx_cpu.first.get(), model_and_ctx_cpu.second.get(), tokens, encode);
                     }
-                    if (dc.split_mode != LLAMA_SPLIT_MODE_TENSOR || llm_arch_supports_sm_tensor(arch)) {
+                    if (dc.split_mode != LLAMA_SPLIT_MODE_TENSOR || llama_arch_supports_sm_tensor(arch)) {
                         test_executed = true;
                         model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, stdev, dc.devs, dc.split_mode, encode, overrides, dc.moe_cache_size);
                         logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode);
@@ -1194,7 +1238,7 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con
                             }
                             std::vector<float> logits_mixed;
                             std::vector<float> logits_chunks;
-                            if (llm_arch_supports_mixed_batch(arch)) {
+                            if (llama_arch_supports_mixed_batch(arch)) {
                                 if (get_logits_mixed(model_and_ctx_dev.first.get(), lctx_dev, tokens, embd_mixed, false, logits_chunks) != 0 ||
                                     get_logits_mixed(model_and_ctx_dev.first.get(), lctx_dev, tokens, embd_mixed, true,  logits_mixed)  != 0) {
                                     throw std::runtime_error("failed to decode mixed batch");
@@ -1227,30 +1271,33 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con
                         }
                     }
 
-                    FILE * file = tmpfile(); // Can be null on Windows without administrator privileges.
-                    // FIXME: when adding a tensor to a gguf_context a copy is made, this changes the pointer which the meta backend
-                    //     in turn uses to map the tensors to their simple equivalents - this is fundamentally incompatible
-                    if (file != nullptr && llama_model_saver_supports_arch(arch) && dc.split_mode != LLAMA_SPLIT_MODE_TENSOR) {
+                    if (llama_arch_supports_save(arch) && dc.split_mode != LLAMA_SPLIT_MODE_TENSOR) {
                         test_executed = true;
                         GGML_ASSERT(model_and_ctx_dev.first && model_and_ctx_dev.second);
-                        llama_model_saver ms = llama_model_saver(model_and_ctx_dev.first.get());
-                        ms.add_kv_from_model();
-                        ms.add_tensors_from_model();
-                        ms.save(file);
-                        rewind(file);
+                        // FIXME: when adding a tensor to a gguf_context a copy is made, this changes the pointer which the meta backend
+                        //     in turn uses to map the tensors to their simple equivalents - this is fundamentally incompatible
+                        const std::string path = get_tmp_gguf_path();
+                        llama_model_save_to_file(model_and_ctx_dev.first.get(), path.c_str());
+                        FILE * file = fopen(path.c_str(), "rb");
+                        GGML_ASSERT(file != nullptr);
 
-                        auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, stdev, dc.devs, dc.split_mode, encode, overrides, dc.moe_cache_size);
-                        const std::vector<float> logits_roundtrip = get_logits(
-                            model_and_ctx_roundtrip.first.get(), model_and_ctx_roundtrip.second.get(), tokens, encode);
-                        status_roundtrip = "\033[1;32mOK\033[0m";
-                        GGML_ASSERT(logits_roundtrip.size() == logits_dev.size());
-                        for (size_t i = 0; i < logits_roundtrip.size(); i++) {
-                            if (logits_roundtrip[i] != logits_dev[i]) {
-                                test_ok = false;
-                                status_roundtrip = "\033[1;31mFAIL\033[0m";
-                                break;
+                        {
+                            auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, stdev, dc.devs, dc.split_mode, encode, overrides, dc.moe_cache_size);
+                            const std::vector<float> logits_roundtrip = get_logits(
+                                model_and_ctx_roundtrip.first.get(), model_and_ctx_roundtrip.second.get(), tokens, encode);
+                            status_roundtrip = "\033[1;32mOK\033[0m";
+                            GGML_ASSERT(logits_roundtrip.size() == logits_dev.size());
+                            for (size_t i = 0; i < logits_roundtrip.size(); i++) {
+                                if (logits_roundtrip[i] != logits_dev[i]) {
+                                    test_ok = false;
+                                    status_roundtrip = "\033[1;31mFAIL\033[0m";
+                                    break;
+                                }
                             }
                         }
+
+                        fclose(file);
+                        std::filesystem::remove(path);
                     }
                 }
 
@@ -1303,7 +1350,7 @@ int main(int argc, char ** argv) {
         } else if (strcmp(argv[i], "-a") == 0 || strcmp(argv[i], "--arch") == 0) {
             if (i + 1 < argc) {
                 const std::string arch_name = argv[++i];
-                if (llm_arch_from_string(arch_name) != LLM_ARCH_UNKNOWN) {
+                if (llama_arch_from_string(arch_name) != LLM_ARCH_UNKNOWN) {
                     // exact architecture name
                     arch_filter = "^" + arch_name + "$";
                 } else {

tests/test-llama-context.cpp

@@ -220,7 +220,7 @@ int main(int argc, char ** argv) {
     bool list_cases        = false;
     bool dump_logs_on_fail = true;
 
-    int parallel_workers = 1; // 0 = auto, same flag as in test-backend-ops
+    int parallel_workers = 0; // 0 = auto, same flag as in test-backend-ops
 
     auto parse_parallel_workers = [&](const char * flag, const char * value) {
         try {
@@ -361,12 +361,9 @@ int main(int argc, char ** argv) {
 
     llama_backend_init();
 
-    // the threads a context runs on, resolved by the common parser - in table mode every worker
-    //   builds a context of its own, so the run uses n_pool times this many
-    const int32_t n_threads = params.cpuparams.n_threads;
-
     // run the full suite for a single model
     if (models_dir.empty()) {
+        const int32_t n_threads = params.cpuparams.n_threads;
         LOG("threads: %d\n", n_threads);
 
         const model_result res = run_model(params.model.path, params, cases, logs_dir,
@@ -404,14 +401,14 @@ int main(int argc, char ** argv) {
         return 1;
     }
 
-    // the dummy models are tiny, but the same flag is meant to run sweeps of real checkpoints,
-    //   where every worker holds the weights it loaded - so do not scale with the cores
-    const unsigned n_hw   = std::thread::hardware_concurrency();
-    const size_t   n_auto = n_hw > 0 ? std::min<size_t>(n_hw, 4) : 1;
+    const unsigned n_hw_2 = std::thread::hardware_concurrency() / 2;
+    const size_t   n_auto = std::max<size_t>(1, std::min<size_t>(n_hw_2, 4));
     const size_t   n_pool = parallel_workers > 0
         ? std::min<size_t>((size_t) parallel_workers, models.size())
         : n_auto;
 
+    const int32_t n_threads = std::max<int32_t>(1, n_hw_2 / n_pool);
+
     // a column is as wide as its case name, but never narrower than the widest verdict
     const size_t verdict_width = std::max({ test_status_width(test_status::PASS),
                                             test_status_width(test_status::FAIL),

tests/test-server-component.cpp

@@ -2,6 +2,7 @@
 
 #include "../tools/server/server-common.h"
 #include "../tools/server/server-decision.h"
+#include "../tools/server/server-schema.h"
 
 #include <cstdlib>
 #include <regex>
@@ -490,6 +491,68 @@ static void test_decision_depth(testing & t) {
     });
 }
 
+// DRY sequence breakers
+//
+
+static void test_dry_sequence_breakers(testing & t) {
+    common_params params_base;
+
+    const auto eval = [&](const json & data) {
+        return server_schema::eval_llama_cmpl_schema(nullptr, params_base, {}, data);
+    };
+
+    {
+        json data;
+        data["dry_sequence_breakers"] = json::array({"ab", "ab", "cd", "ab"});
+        const auto params = eval(data);
+        const std::vector<std::string> expected = {"ab", "cd"};
+        t.assert_true("duplicates removed, order kept", params.sampling.dry_sequence_breakers == expected);
+    }
+
+    {
+        std::vector<std::string> many;
+        for (int i = 0; i <= LLAMA_DRY_MAX_SEQ_BREAKERS; ++i) {
+            many.push_back("b" + std::to_string(i));
+        }
+        json data;
+        data["dry_sequence_breakers"] = many;
+
+        bool thrown = false;
+        try {
+            eval(data);
+        } catch (const std::exception &) {
+            thrown = true;
+        }
+        t.assert_true("more than the cap is rejected", thrown);
+    }
+}
+
+//
+// glob_match
+//
+
+static void test_glob_match(testing & t) {
+    t.test("patterns", [](testing & t) {
+        t.assert_true("literal", glob_match("src/main.cpp", "src/main.cpp"));
+        t.assert_true("star", glob_match("src/*.cpp", "src/main.cpp"));
+        t.assert_true("star stops at /", !glob_match("*.cpp", "src/main.cpp"));
+        t.assert_true("double star", glob_match("**/*.cpp", "a/b/c/main.cpp"));
+        t.assert_true("double star at end", glob_match("src/**", "src/a/b"));
+        t.assert_true("question mark", glob_match("ma?n.cpp", "main.cpp"));
+        t.assert_true("question mark not /", !glob_match("a?b", "a/b"));
+        t.assert_true("class", glob_match("[a-c]x", "bx"));
+        t.assert_true("negated class", !glob_match("[!a-c]x", "bx"));
+        t.assert_true("unclosed class", glob_match("[ab", "[ab"));
+        t.assert_true("mismatch", !glob_match("main.cpp", "main.c"));
+    });
+    t.test("long input", [](testing & t) {
+        const std::string path(100000, 'a');
+        t.assert_true("long literal", glob_match(path, path));
+        t.assert_true("long path", glob_match("**", path));
+        t.assert_true("too many wildcards", !glob_match(std::string(1000, '*'), "a"));
+    });
+}
+
 int main(int argc, char ** argv) {
     testing t;
 
@@ -511,6 +574,8 @@ int main(int argc, char ** argv) {
     t.test("decision_escape", test_decision_escape);
     t.test("rerank_template", test_rerank_template);
     t.test("decision_depth",  test_decision_depth);
+    t.test("glob_match",      test_glob_match);
+    t.test("dry_sequence_breakers", test_dry_sequence_breakers);
 
     return t.summary();
 }

tools/mtmd/mtmd-cli.cpp

@@ -75,6 +75,7 @@ static void inject_test_response_marker() {
 
 struct mtmd_cli_context {
     mtmd::context_ptr ctx_vision;
+    mtmd_helper_cap   mtmd_cap = {};
     common_init_result_ptr llama_init;
 
     llama_model       * model;
@@ -120,7 +121,8 @@ struct mtmd_cli_context {
 
         init_vision_context(params);
 
-        if (!mtmd_helper_model_can_chat(lctx, ctx_vision.get())) {
+        mtmd_cap = mtmd_helper_get_capabilities(model, ctx_vision.get());
+        if (!mtmd_cap.chat) {
             LOG_ERR("Model does not support chat mode\n");
             LOG_ERR("Hint: for TTS models, please use llama-tts\n");
             exit(1);
@@ -565,7 +567,7 @@ int main(int argc, char ** argv) {
 
     if (g_is_interrupted) return 130;
 
-    if (mtmd_support_audio_stream(ctx.ctx_vision.get())) {
+    if (ctx.mtmd_cap.audio_stream) {
         if (params.image.size() != 1) {
             LOG_ERR("This model transcribes audio in realtime, pass exactly one --audio file\n");
             return 1;
@@ -624,7 +626,7 @@ int main(int argc, char ** argv) {
         if (mtmd_support_audio(ctx.ctx_vision.get())) {
             LOG("\n   /audio <path>    load an audio");
         }
-        if (mtmd_helper_support_video(ctx.ctx_vision.get())) {
+        if (ctx.mtmd_cap.video) {
             LOG("\n   /video <path>    load a video");
         }
         LOG("\n   /clear           clear the chat history");

tools/mtmd/mtmd-helper.cpp

@@ -501,15 +501,6 @@ mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(const mtmd_context
     return mtmd_helper_bitmap_init_from_buf(ctx, buf.data(), buf.size(), placeholder, opt);
 }
 
-bool mtmd_helper_support_video(const mtmd_context * ctx) {
-#ifdef MTMD_VIDEO
-    return mtmd_support_vision(ctx);
-#else
-    GGML_UNUSED(ctx);
-    return false;
-#endif
-}
-
 //
 // Video input helpers
 //
@@ -1033,17 +1024,40 @@ int32_t mtmd_helper_video_read_next(mtmd_helper_video * ctx,
 #endif
 }
 
-bool mtmd_helper_model_can_chat(const llama_context * lctx, const mtmd_context * mctx) {
+mtmd_helper_cap mtmd_helper_get_capabilities(const llama_model * model, const mtmd_context * mctx) {
+    mtmd_helper_cap cap = {};
+    cap.chat     = true;
+    cap.text_gen = true;
     if (!mctx) {
-        return true;
+        return cap;
     }
 
-    auto * model = llama_get_model(lctx);
-    auto * tmpl = llama_model_chat_template(model, nullptr);
-    auto info = mtmd_gen_audio_get_info(mctx);
+    cap.vision = mtmd_support_vision(mctx);
+    cap.audio  = mtmd_support_audio(mctx);
+    cap.audio_stream = mtmd_support_audio_stream(mctx);
+#ifdef MTMD_VIDEO
+    cap.video  = cap.vision;
+#endif
+    const auto gen_audio_type = mtmd_gen_audio_get_info(mctx).type;
+    cap.audio_gen = gen_audio_type != MTMD_GEN_AUDIO_TYPE_NONE;
+    cap.audio_gen_spk_ref = gen_audio_type == MTMD_GEN_AUDIO_TYPE_QWEN3TTS
+                         || gen_audio_type == MTMD_GEN_AUDIO_TYPE_POCKETTTS;
 
     // tts-only model cannot be used for chat (no chat template)
-    bool is_tts_only = info.type != MTMD_GEN_AUDIO_TYPE_NONE && tmpl == nullptr;
+    cap.chat = !cap.audio_gen || llama_model_chat_template(model, nullptr) != nullptr;
 
-    return !is_tts_only;
+    return cap;
+}
+
+bool mtmd_helper_model_can_chat(const llama_context * lctx, const mtmd_context * mctx) {
+    return mtmd_helper_get_capabilities(llama_get_model(lctx), mctx).chat;
+}
+
+bool mtmd_helper_support_video(const mtmd_context * ctx) {
+#ifdef MTMD_VIDEO
+    return mtmd_support_vision(ctx);
+#else
+    GGML_UNUSED(ctx);
+    return false;
+#endif
 }

tools/mtmd/mtmd-helper.h

@@ -45,9 +45,6 @@ MTMD_API struct mtmd_helper_init_opt mtmd_helper_init_opt_default(void);
 // Note: this also call mtmd_log_set() internally
 MTMD_API void mtmd_helper_log_set(ggml_log_callback log_callback, void * user_data);
 
-// Returns true if this build includes video support (MTMD_VIDEO was ON at compile time).
-MTMD_API bool mtmd_helper_support_video(const mtmd_context * ctx);
-
 struct mtmd_helper_bitmap_wrapper {
     mtmd_bitmap * bitmap;
     mtmd_helper_video * video_ctx;
@@ -186,8 +183,31 @@ MTMD_API int32_t mtmd_helper_video_read_next(mtmd_helper_video * ctx,
             mtmd_bitmap ** out_bitmap,
             char ** out_text);
 
-// return true if model can be used for chat
-MTMD_API bool mtmd_helper_model_can_chat(const struct llama_context * lctx, const struct mtmd_context * mctx);
+struct mtmd_helper_cap {
+    bool chat; // false for tts-only model (no chat template)
+
+    // input capabilities
+    bool vision;
+    bool audio;
+    bool video; // also depending on MTMD_VIDEO build option
+
+    // realtime capabilities
+    bool audio_stream; // realtime transcription, see mtmd_audio_stream_init()
+
+    // output capabilities
+    bool text_gen; // always true for now
+    bool audio_gen;
+    bool audio_gen_spk_ref; // accepts a speaker reference audio for voice cloning
+};
+typedef struct mtmd_helper_cap mtmd_helper_cap;
+
+// mctx can be NULL (text-only model)
+MTMD_API struct mtmd_helper_cap mtmd_helper_get_capabilities(const struct llama_model * model, const struct mtmd_context * mctx);
+
+DEPRECATED(MTMD_API bool mtmd_helper_model_can_chat(const struct llama_context * lctx, const struct mtmd_context * mctx),
+           "use mtmd_helper_get_capabilities() instead");
+DEPRECATED(MTMD_API bool mtmd_helper_support_video(const struct mtmd_context * ctx),
+           "use mtmd_helper_get_capabilities() instead");
 
 //
 // Audio generation helpers

tools/server/README.md

@@ -563,7 +563,7 @@ These words will not be included in the completion, so make sure to add them to
 
 `dry_penalty_last_n`: How many tokens to scan for repetitions. Default: `64`, where `0` is disabled.
 
-`dry_sequence_breakers`: Specify an array of sequence breakers for DRY sampling. Only a JSON array of strings is accepted. Default: `['\n', ':', '"', '*']`
+`dry_sequence_breakers`: Specify an array of sequence breakers for DRY sampling. Only a JSON array of strings is accepted. At most 64 distinct breakers, duplicates are ignored. Default: `['\n', ':', '"', '*']`
 
 `xtc_probability`: Set the chance for token removal via XTC sampler. Default: `0.0`, which is disabled.
 

tools/server/server-context.cpp

@@ -1012,16 +1012,17 @@ public:
     llama_model * model_tgt = nullptr;
 
     mtmd_context * mctx = nullptr;
+    mtmd_helper_cap mtmd_cap = {};
     // note: video_params.ffmpeg_bin_dir points into params_base, which outlives this struct
     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;
+        return mtmd_cap.audio_gen;
     }
 
     bool has_cap_chat() const {
-        return mctx == nullptr || mtmd_helper_model_can_chat(ctx_tgt, mctx);
+        return mtmd_cap.chat;
     }
 
     server_decision_context decision;
@@ -1397,6 +1398,8 @@ private:
             }
         }
 
+        mtmd_cap = mtmd_helper_get_capabilities(model_tgt, mctx);
+
         if (!llama_memory_can_shift(llama_get_memory(ctx_tgt))) {
             if (params_base.ctx_shift) {
                 params_base.ctx_shift = false;
@@ -1705,9 +1708,9 @@ private:
                 /* reasoning_format      */ params_base.reasoning_format,
                 /* chat_template_kwargs  */ params_base.default_template_kwargs,
                 /* tmpls                 */ std::move(chat_templates),
-                /* allow_image           */ mctx ? mtmd_support_vision(mctx) : false,
-                /* allow_audio           */ mctx ? mtmd_support_audio (mctx) : false,
-                /* allow_video           */ mctx ? mtmd_helper_support_video(mctx) : false,
+                /* allow_image           */ mtmd_cap.vision,
+                /* allow_audio           */ mtmd_cap.audio,
+                /* allow_video           */ mtmd_cap.video,
                 /* enable_thinking       */ enable_thinking,
                 /* reasoning_budget      */ params_base.sampling.reasoning_budget_tokens,
                 /* reasoning_budget_msg  */ params_base.sampling.reasoning_budget_message,
@@ -5127,7 +5130,7 @@ server_context_meta server_context::get_meta() const {
         /* 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),
+        /* has_audio_stream        */ impl->mtmd_cap.audio_stream,
         /* json_ui_settings        */ impl->json_ui_settings,
         /* slot_n_ctx              */ impl->n_ctx_slot(),
         /* pooling_type            */ llama_pooling_type(impl->ctx_tgt),

tools/server/server-models.cpp

@@ -533,7 +533,7 @@ void server_model_meta::update_caps(const common_params & base) {
     std::vector<std::string> output_modalities = {"text"};
     if (!params.model.path.empty()) {
         const common_gguf_info info = common_get_gguf_info(params.model.path);
-        // same rule as mtmd_helper_model_can_chat()
+        // same rule as mtmd_helper_get_capabilities()
         const bool can_chat = !gen_audio || info.has_chat_tmpl;
         output_modalities = server_model_output_modalities(info.decision_type, can_chat, gen_audio);
         n_ctx_train       = info.n_ctx_train;