From 26908739bc8ab9c5ccacbde5dc9c0b3609079e8f Mon Sep 17 00:00:00 2001 From: yarik <158273159+dazzywi@users.noreply.github.com> Date: Wed, 7 Oct 2026 08:58:22 +0000 Subject: [PATCH 01/40] imatrix : include clocale for std::setlocale (#30079) Co-authored-by: yarik --- tools/imatrix/imatrix.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/tools/imatrix/imatrix.cpp b/tools/imatrix/imatrix.cpp index c125e908b..5e9f911e0 100644 --- a/tools/imatrix/imatrix.cpp +++ b/tools/imatrix/imatrix.cpp @@ -10,6 +10,7 @@ #include #include +#include #include #include #include From b7dafa01e5f375c3010fb24b61a67329f957426a Mon Sep 17 00:00:00 2001 From: "Alessandro de Oliveira Faria (A.K.A.CABELO)" Date: Wed, 7 Oct 2026 06:18:33 -0300 Subject: [PATCH 02/40] vendor : update cpp-httplib to 0.60.0 (#30081) --- scripts/sync_vendor.py | 2 +- vendor/cpp-httplib/httplib.cpp | 181 +++++++++++++++++---------------- vendor/cpp-httplib/httplib.h | 6 +- 3 files changed, 99 insertions(+), 90 deletions(-) diff --git a/scripts/sync_vendor.py b/scripts/sync_vendor.py index bdfc8d62e..c667fcb0a 100755 --- a/scripts/sync_vendor.py +++ b/scripts/sync_vendor.py @@ -5,7 +5,7 @@ import os import sys import subprocess -HTTPLIB_VERSION = "refs/tags/v0.59.0" +HTTPLIB_VERSION = "refs/tags/v0.60.0" # used by examples/gguf-hash, these repos have no release tag, so we pin a commit XXHASH_COMMIT = "9f465f1ea932d6ad9a26cd77496311ffa544cd68" diff --git a/vendor/cpp-httplib/httplib.cpp b/vendor/cpp-httplib/httplib.cpp index 781544763..ac6f6cda7 100644 --- a/vendor/cpp-httplib/httplib.cpp +++ b/vendor/cpp-httplib/httplib.cpp @@ -169,7 +169,8 @@ void SSEClient::stop() { } bool SSEClient::parse_sse_line(const std::string &line, SSEMessage &msg, - int &retry_ms, bool &has_data) { + int &retry_ms, bool &has_data, + bool &has_id) { // Blank line signals end of event if (line.empty()) { return true; } @@ -199,6 +200,7 @@ bool SSEClient::parse_sse_line(const std::string &line, SSEMessage &msg, } else if (field == "id") { // Empty id is valid (clears the last event ID) msg.id = value; + has_id = true; } else if (field == "retry") { // Parse retry interval in milliseconds // Per the SSE spec, a value that is not all ASCII digits is ignored. @@ -269,7 +271,8 @@ void SSEClient::run_event_loop() { // Event receiving loop std::string buffer; SSEMessage current_msg; - bool has_data = false; + auto has_data = false; + auto has_id = false; while (running_.load() && result.next()) { buffer.append(result.data(), result.size()); @@ -288,13 +291,13 @@ void SSEClient::run_event_loop() { if (!line.empty() && line.back() == '\r') { line.pop_back(); } // Parse the line and check if event is complete - auto event_complete = - parse_sse_line(line, current_msg, reconnect_interval_ms_, has_data); + auto event_complete = parse_sse_line( + line, current_msg, reconnect_interval_ms_, has_data, has_id); if (event_complete) { // Update last_event_id for reconnection, even for an event that - // has no data - if (!current_msg.id.empty()) { last_event_id_ = current_msg.id; } + // has no data. An empty id clears it. + if (has_id) { last_event_id_ = current_msg.id; } // An event without a data field is not dispatched if (has_data) { dispatch_event(current_msg); } @@ -302,6 +305,7 @@ void SSEClient::run_event_loop() { // Reset the message for the next event either way current_msg.clear(); has_data = false; + has_id = false; } } @@ -3550,9 +3554,11 @@ struct WebSocketUpgradeResponse { std::string selected_subprotocol; }; -bool read_websocket_upgrade_response(Stream &strm, - const std::string &expected_accept, - WebSocketUpgradeResponse &upgrade) { +bool +read_websocket_upgrade_response(Stream &strm, + const std::string &expected_accept, + const std::string &offered_subprotocols, + WebSocketUpgradeResponse &upgrade) { // Read status line const auto bufsiz = 2048; char buf[bufsiz]; @@ -3609,6 +3615,22 @@ bool read_websocket_upgrade_response(Stream &strm, upgrade.selected_subprotocol = proto_it->second; } + // Verify the subprotocol is one the client offered (RFC 6455 4.1) + if (!upgrade.selected_subprotocol.empty()) { + auto was_offered = false; + split(offered_subprotocols.data(), + offered_subprotocols.data() + offered_subprotocols.size(), ',', + [&](const char *b, const char *e) { + if (std::string(b, e) == upgrade.selected_subprotocol) { + was_offered = true; + } + }); + if (!was_offered) { + upgrade.error = Error::WebSocketHandshake; + return false; + } + } + return true; } @@ -4254,6 +4276,9 @@ bool redirect(T &cli, Request &req, Response &res, new_req.method = "GET"; new_req.body.clear(); new_req.headers.clear(); + new_req.content_length_ = 0; + new_req.content_provider_ = nullptr; + new_req.is_chunked_content_provider_ = false; } Response new_res; @@ -5603,7 +5628,10 @@ bool perform_websocket_handshake(Stream &strm, Request &req, // Verify 101 response and Sec-WebSocket-Accept header auto expected_accept = websocket_accept_key(client_key); - return read_websocket_upgrade_response(strm, expected_accept, upgrade); + auto offered_subprotocols = + get_combined_header_value(req.headers, "Sec-WebSocket-Protocol"); + return read_websocket_upgrade_response(strm, expected_accept, + offered_subprotocols, upgrade); } bool is_ip_address(const std::string &host) { @@ -9946,6 +9974,12 @@ Server::process_request(Stream &strm, const std::string &remote_addr, protocols.emplace_back(b, e); }); selected_subprotocol = entry.sub_protocol_selector(protocols); + + // Ignore a selection the client did not offer (RFC 6455 4.2.2) + if (std::find(protocols.begin(), protocols.end(), + selected_subprotocol) == protocols.end()) { + selected_subprotocol.clear(); + } } } @@ -11016,7 +11050,7 @@ bool ClientImpl::redirect(Request &req, Response &res, Error &error) { if (next_host.empty()) { next_host = host_; } if (next_path.empty()) { next_path = "/"; } - auto path = decode_path_component(next_path) + next_query; + auto path = std::move(next_path) + next_query; // Same host redirect - use current client if (next_scheme == scheme && next_host == host_ && next_port == port_) { @@ -11271,9 +11305,9 @@ bool ClientImpl::write_request(Stream &strm, Request &req, // Write request line and headers if (detail::write_request_line(bstrm, req.method, path_with_query) < 0) { // A rejected method (not a token, e.g. carrying CR/LF) or target (e.g. - // CR/LF smuggled in via a decoded redirect Location under - // set_path_encode(false)) must fail the request cleanly instead of - // emitting a request-line-less, header-injecting request. + // CR/LF in a caller-supplied path under set_path_encode(false)) must + // fail the request cleanly instead of emitting a request-line-less, + // header-injecting request. error = Error::Write; rejected_locally = true; output_error_log(error, &req); @@ -16304,29 +16338,24 @@ bool verify_hostname(cert_t cert, const char *hostname) { auto ip_len = impl::parse_ip_address(host_str, ip_bytes); auto is_ip = ip_len > 0; - // Check Subject Alternative Names (SAN) - // In Mbed TLS 3.x, subject_alt_names contains raw values without ASN.1 tags - // - DNS names: raw string bytes - // - IP addresses: raw IP bytes (4 for IPv4, 16 for IPv6) + // Check Subject Alternative Names (SAN). Mbed TLS keeps the GeneralName type + // in buf.tag and the raw value in buf.p / buf.len. const mbedtls_x509_sequence *san = &mcert->subject_alt_names; while (san != nullptr && san->buf.p != nullptr && san->buf.len > 0) { const unsigned char *p = san->buf.p; size_t len = san->buf.len; + auto san_type = san->buf.tag & MBEDTLS_ASN1_TAG_VALUE_MASK; if (is_ip) { // For an IP host, only a matching iPAddress SAN of the same family // (4 bytes for IPv4, 16 bytes for IPv6) may authenticate it. - if (len == ip_len && memcmp(p, ip_bytes, ip_len) == 0) { return true; } - } else { - // Check if this SAN is a DNS name (printable ASCII string) - bool is_dns = len > 0; - for (size_t i = 0; i < len && is_dns; i++) { - if (p[i] < 32 || p[i] > 126) { is_dns = false; } - } - if (is_dns) { - std::string san_name(reinterpret_cast(p), len); - if (detail::match_hostname(san_name, host_str)) { return true; } + if (san_type == MBEDTLS_X509_SAN_IP_ADDRESS && len == ip_len && + memcmp(p, ip_bytes, ip_len) == 0) { + return true; } + } else if (san_type == MBEDTLS_X509_SAN_DNS_NAME) { + std::string san_name(reinterpret_cast(p), len); + if (detail::match_hostname(san_name, host_str)) { return true; } } san = san->next; } @@ -16405,65 +16434,45 @@ bool get_cert_sans(cert_t cert, std::vector &sans) { const mbedtls_x509_sequence *cur = &x509->subject_alt_names; while (cur != nullptr) { if (cur->buf.len > 0) { - // Mbed TLS stores SAN as ASN.1 sequences - // The tag byte indicates the type const unsigned char *p = cur->buf.p; - size_t len = cur->buf.len; + size_t value_len = cur->buf.len; - // First byte is the tag - unsigned char tag = *p; - p++; - len--; - - // Parse length (simple single-byte length assumed) - if (len > 0 && *p < 0x80) { - size_t value_len = *p; - p++; - len--; - - if (value_len <= len) { - SanEntry entry; - // ASN.1 context tags for GeneralName - switch (tag & 0x1F) { - case 2: // dNSName - entry.type = SanType::DNS; - entry.value = - std::string(reinterpret_cast(p), value_len); - break; - case 7: // iPAddress - entry.type = SanType::IP; - if (value_len == 4) { - // IPv4 - char buf[16]; - snprintf(buf, sizeof(buf), "%d.%d.%d.%d", p[0], p[1], p[2], p[3]); - entry.value = buf; - } else if (value_len == 16) { - // IPv6 - char buf[64]; - snprintf(buf, sizeof(buf), - "%02x%02x:%02x%02x:%02x%02x:%02x%02x:" - "%02x%02x:%02x%02x:%02x%02x:%02x%02x", - p[0], p[1], p[2], p[3], p[4], p[5], p[6], p[7], p[8], - p[9], p[10], p[11], p[12], p[13], p[14], p[15]); - entry.value = buf; - } - break; - case 1: // rfc822Name (email) - entry.type = SanType::EMAIL; - entry.value = - std::string(reinterpret_cast(p), value_len); - break; - case 6: // uniformResourceIdentifier - entry.type = SanType::URI; - entry.value = - std::string(reinterpret_cast(p), value_len); - break; - default: entry.type = SanType::OTHER; break; - } - - if (!entry.value.empty()) { sans.push_back(std::move(entry)); } + SanEntry entry; + switch (cur->buf.tag & MBEDTLS_ASN1_TAG_VALUE_MASK) { + case MBEDTLS_X509_SAN_DNS_NAME: + entry.type = SanType::DNS; + entry.value = std::string(reinterpret_cast(p), value_len); + break; + case MBEDTLS_X509_SAN_IP_ADDRESS: + entry.type = SanType::IP; + if (value_len == 4) { + // IPv4 + char buf[16]; + snprintf(buf, sizeof(buf), "%d.%d.%d.%d", p[0], p[1], p[2], p[3]); + entry.value = buf; + } else if (value_len == 16) { + // IPv6 + char buf[64]; + snprintf(buf, sizeof(buf), + "%02x%02x:%02x%02x:%02x%02x:%02x%02x:" + "%02x%02x:%02x%02x:%02x%02x:%02x%02x", + p[0], p[1], p[2], p[3], p[4], p[5], p[6], p[7], p[8], p[9], + p[10], p[11], p[12], p[13], p[14], p[15]); + entry.value = buf; } + break; + case MBEDTLS_X509_SAN_RFC822_NAME: + entry.type = SanType::EMAIL; + entry.value = std::string(reinterpret_cast(p), value_len); + break; + case MBEDTLS_X509_SAN_UNIFORM_RESOURCE_IDENTIFIER: + entry.type = SanType::URI; + entry.value = std::string(reinterpret_cast(p), value_len); + break; + default: entry.type = SanType::OTHER; break; } + + if (!entry.value.empty()) { sans.push_back(std::move(entry)); } } cur = cur->next; } @@ -18105,9 +18114,9 @@ void WebSocket::start_heartbeat() { if (ping_interval_sec_ == 0) { return; } ping_thread_ = std::thread([this]() { std::unique_lock lock(ping_mutex_); - while (!closed_) { - ping_cv_.wait_for(lock, std::chrono::seconds(ping_interval_sec_)); - if (closed_) { break; } + // The predicate keeps a spurious wakeup from sending a ping early + while (!ping_cv_.wait_for(lock, std::chrono::seconds(ping_interval_sec_), + [this]() { return closed_.load(); })) { // If the peer has failed to respond to the previous pings, give up. // RFC 6455 does not define a pong-timeout mechanism; this is an // opt-in liveness check controlled by max_missed_pongs_. diff --git a/vendor/cpp-httplib/httplib.h b/vendor/cpp-httplib/httplib.h index 283fa1387..553924015 100644 --- a/vendor/cpp-httplib/httplib.h +++ b/vendor/cpp-httplib/httplib.h @@ -8,8 +8,8 @@ #ifndef CPPHTTPLIB_HTTPLIB_H #define CPPHTTPLIB_HTTPLIB_H -#define CPPHTTPLIB_VERSION "0.59.0" -#define CPPHTTPLIB_VERSION_NUM "0x003b00" +#define CPPHTTPLIB_VERSION "0.60.0" +#define CPPHTTPLIB_VERSION_NUM "0x003c00" #ifdef _WIN32 #if defined(_WIN32_WINNT) && _WIN32_WINNT < 0x0A00 @@ -4377,7 +4377,7 @@ public: private: bool parse_sse_line(const std::string &line, SSEMessage &msg, int &retry_ms, - bool &has_data); + bool &has_data, bool &has_id); void run_event_loop(); void dispatch_event(const SSEMessage &msg); bool should_reconnect(int count) const; From ad2156533102a0d3c4e5fbdf422dc25fba4d03ba Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Wed, 7 Oct 2026 17:22:56 +0800 Subject: [PATCH 03/40] musa: use the tile lightning indexer kernel (#30080) Signed-off-by: Xiaodong Ye --- ggml/src/ggml-cuda/lightning-indexer.cu | 89 +++++++++++++++---------- 1 file changed, 53 insertions(+), 36 deletions(-) diff --git a/ggml/src/ggml-cuda/lightning-indexer.cu b/ggml/src/ggml-cuda/lightning-indexer.cu index 57cb1c1dc..429d2e200 100644 --- a/ggml/src/ggml-cuda/lightning-indexer.cu +++ b/ggml/src/ggml-cuda/lightning-indexer.cu @@ -239,6 +239,14 @@ static __global__ void lightning_indexer_kernel_wmma( // tokens scored per block by the tile kernel #define LIGHTNING_INDEXER_TILE_TOKENS 8 +// heads whose queries the tile kernel stages per pass, MUSA arch 21 caps static shared memory +// at 28 KB and the queries of four heads do not fit there next to the key tile +#if defined(GGML_USE_MUSA) && defined(__MUSA_ARCH__) && __MUSA_ARCH__ < 220 +#define LIGHTNING_INDEXER_TILE_HEADS_PER_PASS 2 +#else +#define LIGHTNING_INDEXER_TILE_HEADS_PER_PASS 4 +#endif + // TODO there is one ugly assumption used in this kernel - that WARP_SIZE is equal to 32 // thanks to that one warp operating on float4 processes whole indexer K/Q vectors // 32 * 4 = 128 (N_EMBD) @@ -406,9 +414,11 @@ static __global__ void lightning_indexer_kernel_tile( constexpr int KEY_LANES = THREADS_PER_BLOCK / TOKENS_PER_BLOCK; constexpr int KEYS_PER_THREAD = K_VECS_PER_BLOCK / KEY_LANES; constexpr int N_EMBD_H2 = N_EMBD / 2; + constexpr int HEADS_PER_PASS = N_HEAD < LIGHTNING_INDEXER_TILE_HEADS_PER_PASS ? N_HEAD : LIGHTNING_INDEXER_TILE_HEADS_PER_PASS; static_assert(THREADS_PER_BLOCK % TOKENS_PER_BLOCK == 0, "threads must cover the token tile"); static_assert(K_VECS_PER_BLOCK % KEY_LANES == 0, "key lanes must cover the key tile"); + static_assert(N_HEAD % HEADS_PER_PASS == 0, "head passes must cover the heads"); const int tid = threadIdx.y * WARP_SIZE + threadIdx.x; const int start_kv = blockIdx.x * K_VECS_PER_BLOCK; @@ -417,7 +427,7 @@ static __global__ void lightning_indexer_kernel_tile( // the row padding keeps the keys of consecutive threads in distinct banks __shared__ half2 k_shared[K_VECS_PER_BLOCK][N_EMBD_H2 + 1]; - __shared__ float2 q_shared[N_HEAD][TOKENS_PER_BLOCK][N_EMBD_H2]; + __shared__ float2 q_shared[HEADS_PER_PASS][TOKENS_PER_BLOCK][N_EMBD_H2]; __shared__ float w_shared[N_HEAD][TOKENS_PER_BLOCK]; // phase 1 - stage the key tile four elements at a time, rows past n_kv are zero @@ -451,22 +461,7 @@ static __global__ void lightning_indexer_kernel_tile( k_shared[r][2*c4 + 1] = hi; } - // phase 2 - stage the queries and weights of every head, tokens past n_batch are zero - -#pragma unroll - for (int i = tid; i < N_HEAD * TOKENS_PER_BLOCK * (N_EMBD / 4); i += THREADS_PER_BLOCK) { - const int h = i / (TOKENS_PER_BLOCK * (N_EMBD / 4)); - const int r = i / (N_EMBD / 4) % TOKENS_PER_BLOCK; - const int c4 = i % (N_EMBD / 4); - - float4 v = make_float4(0.0f, 0.0f, 0.0f, 0.0f); - if (start_batch + r < n_batch) { - v = *(const float4 *) ((const char *) Q + h*nbq1 + (start_batch + r)*nbq2 + i_stream*nbq3 + c4*sizeof(float4)); - } - - q_shared[h][r][2*c4 + 0] = make_float2(v.x, v.y); - q_shared[h][r][2*c4 + 1] = make_float2(v.z, v.w); - } + // phase 2 - stage the weights of every head, tokens past n_batch are zero if (tid < N_HEAD * TOKENS_PER_BLOCK) { const int h = tid / TOKENS_PER_BLOCK; @@ -475,33 +470,60 @@ static __global__ void lightning_indexer_kernel_tile( ((const float *) ((const char *) W + (start_batch + r)*nbw1 + i_stream*nbw3))[h] : 0.0f; } - __syncthreads(); - - // phase 3 - float products of the widened keys for every head, ReLU, weight - const int kl = tid % KEY_LANES; const int tl = tid / KEY_LANES; float qk[N_HEAD][KEYS_PER_THREAD] = { { 0.0f } }; -#pragma unroll 8 - for (int c = 0; c < N_EMBD_H2; ++c) { - float2 k_val[KEYS_PER_THREAD]; #pragma unroll - for (int j = 0; j < KEYS_PER_THREAD; ++j) { - k_val[j] = __half22float2(k_shared[kl + j*KEY_LANES][c]); + for (int h0 = 0; h0 < N_HEAD; h0 += HEADS_PER_PASS) { + // the previous pass is fully consumed before its queries are replaced + if (h0 > 0) { + __syncthreads(); } + + // phase 3 - stage the queries of the heads of this pass, tokens past n_batch are zero + #pragma unroll - for (int h = 0; h < N_HEAD; ++h) { - const float2 q_val = q_shared[h][tl][c]; + for (int i = tid; i < HEADS_PER_PASS * TOKENS_PER_BLOCK * (N_EMBD / 4); i += THREADS_PER_BLOCK) { + const int h = i / (TOKENS_PER_BLOCK * (N_EMBD / 4)); + const int r = i / (N_EMBD / 4) % TOKENS_PER_BLOCK; + const int c4 = i % (N_EMBD / 4); + + float4 v = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + if (start_batch + r < n_batch) { + v = *(const float4 *) ((const char *) Q + (h0 + h)*nbq1 + (start_batch + r)*nbq2 + i_stream*nbq3 + c4*sizeof(float4)); + } + + q_shared[h][r][2*c4 + 0] = make_float2(v.x, v.y); + q_shared[h][r][2*c4 + 1] = make_float2(v.z, v.w); + } + + __syncthreads(); + + // phase 4 - float products of the widened keys for the heads of this pass + +#pragma unroll 8 + for (int c = 0; c < N_EMBD_H2; ++c) { + float2 k_val[KEYS_PER_THREAD]; #pragma unroll for (int j = 0; j < KEYS_PER_THREAD; ++j) { - qk[h][j] = fmaf(k_val[j].x, q_val.x, qk[h][j]); - qk[h][j] = fmaf(k_val[j].y, q_val.y, qk[h][j]); + k_val[j] = __half22float2(k_shared[kl + j*KEY_LANES][c]); + } +#pragma unroll + for (int h = 0; h < HEADS_PER_PASS; ++h) { + const float2 q_val = q_shared[h][tl][c]; +#pragma unroll + for (int j = 0; j < KEYS_PER_THREAD; ++j) { + qk[h0 + h][j] = fmaf(k_val[j].x, q_val.x, qk[h0 + h][j]); + qk[h0 + h][j] = fmaf(k_val[j].y, q_val.y, qk[h0 + h][j]); + } } } } + // phase 5 - ReLU, weight, add the mask and write, consecutive threads write consecutive keys + float score[KEYS_PER_THREAD] = { 0.0f }; #pragma unroll @@ -512,8 +534,6 @@ static __global__ void lightning_indexer_kernel_tile( } } - // phase 4 - add the mask and write, consecutive threads write consecutive keys - const int i_batch = start_batch + tl; if (i_batch >= n_batch) { return; @@ -677,8 +697,6 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_F32) GGML_ABORT("fatal error"); } -#ifndef GGML_USE_MUSA - // MUSA archs 21 and 22 cap static shared memory at 28 KB, below what the tile kernel stages } else if (n_embd == 128 && n_head == 4 && n_batch >= LIGHTNING_INDEXER_TILE_TOKENS) { // too few heads for a wmma tile, the tile kernel shares the keys across the tokens constexpr int WARPS_PER_BLOCK = 8; @@ -698,9 +716,8 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_BF16) LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_F32) GGML_ABORT("fatal error"); -#endif // GGML_USE_MUSA } else if (n_embd == 128 && n_head == 4) { - // a batch smaller than a token tile, or MUSA, use vector kernel + // a batch smaller than a token tile, use vector kernel constexpr int K_VECS_PER_WARP = 8; constexpr int WARPS_PER_BLOCK = 8; constexpr int K_VECS_PER_BLOCK = K_VECS_PER_WARP * WARPS_PER_BLOCK; From 7481354a174714a3447ea0f8723c1dd1b3cbf8e2 Mon Sep 17 00:00:00 2001 From: Toki Nasin <141258697+tokinasin@users.noreply.github.com> Date: Wed, 7 Oct 2026 18:34:16 +0900 Subject: [PATCH 04/40] convert : fix token configuration for PLaMo-3 (#29843) * convert : Fix token configuration for PLaMo-3 The reasoning and tool calling tags in PLaMo-3 consist of three tokens each. For example, for reasoning: * reasoning start: `<|plamo:begin_`, `think`, `:plamo|>` * reasoning end: `<|plamo:end_`, `think`, `:plamo|>` Registering `<|plamo:begin_`, `<|plamo:end_`, and `:plamo|>` as `USER_DEFINED` so that they are parsed correctly. Also PLaMo-3 models use <|plamo:tag|> as EOT, while PLaMo-2 models use <|plamo:op|>. Take EOT token as a parameter and look it up so that PLaMo-2 and PLaMo-3 can use their appropriate ones. * use NORMAL instead of USER_DEFINED --- conversion/base.py | 18 ++++++++++++------ conversion/plamo.py | 7 +++++-- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/conversion/base.py b/conversion/base.py index 3b6fe04c2..306a9f0f2 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -2499,7 +2499,11 @@ class TextModel(ModelBase): if template is not None: self.gguf_writer.add_chat_template(template) - def _set_vocab_plamo(self): + def _set_vocab_plamo( + self, + eot_token: str, + normal_tokens: Iterable[str] = (), + ): # PLaMo models use a custom tokenizer with a .jsonl file tokenizer_jsonl_path = self.dir_model / "tokenizer.jsonl" tokenizer_config_path = self.dir_model / "tokenizer_config.json" @@ -2523,27 +2527,30 @@ class TextModel(ModelBase): tokens = [] scores = [] toktypes = [] + normal_tokens = set(normal_tokens) with open(tokenizer_jsonl_path, "r", encoding="utf-8") as f: for line_num, line in enumerate(f): if line.strip(): token_data = json.loads(line) # Format: [token, score, type, ?, ?, ?, ?] - token = token_data[0].encode("utf-8") + token_str = token_data[0] + token = token_str.encode("utf-8") score = float(token_data[1]) token_type_str = token_data[2] if len(token_data) > 2 else "NORMAL" tokens.append(token) scores.append(score) - if token_type_str == "UNKNOWN": + if token_str in normal_tokens: + toktypes.append(gguf.TokenType.NORMAL) + elif token_type_str == "UNKNOWN": toktypes.append(gguf.TokenType.UNKNOWN) elif token_type_str == "CONTROL": toktypes.append(gguf.TokenType.CONTROL) elif token_type_str == "BYTE": toktypes.append(gguf.TokenType.BYTE) else: - token_str = token_data[0] if token_str.startswith("<|plamo:") and token_str.endswith("|>"): toktypes.append(gguf.TokenType.CONTROL) else: @@ -2580,8 +2587,7 @@ class TextModel(ModelBase): token_id = tokens.index(tokenizer_config["unk_token"].encode("utf-8")) self.gguf_writer.add_unk_token_id(token_id) - # Add <|plamo:op|> as EOT to ensure appropriate end of generation - self.gguf_writer.add_eot_token_id(4) + self.gguf_writer.add_eot_token_id(tokens.index(eot_token.encode("utf-8"))) self.gguf_writer.add_add_space_prefix(False) diff --git a/conversion/plamo.py b/conversion/plamo.py index 52055e4a1..a843e059a 100644 --- a/conversion/plamo.py +++ b/conversion/plamo.py @@ -64,7 +64,7 @@ class Plamo2Model(TextModel): model_arch = gguf.MODEL_ARCH.PLAMO2 def set_vocab(self): - self._set_vocab_plamo() + self._set_vocab_plamo(eot_token="<|plamo:op|>") def set_gguf_parameters(self): hparams = self.hparams @@ -170,7 +170,10 @@ class Plamo3Model(TextModel): }) def set_vocab(self): - self._set_vocab_plamo() + self._set_vocab_plamo( + eot_token="<|plamo:tag|>", + normal_tokens=("<|plamo:begin_", "<|plamo:end_", ":plamo|>"), + ) tokenizer_config_path = self.dir_model / "tokenizer_config.json" tokenizer_config = {} From 42b021b4dc42be573f1e1463528532fc8294c650 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= Date: Wed, 7 Oct 2026 11:40:35 +0200 Subject: [PATCH 05/40] vocab : add plamo fim tokens (#30090) --- src/llama-vocab.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index de20c757f..dbadaf100 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2860,6 +2860,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { || t.first == "▁
"          // CodeLlama
                         || t.first == "<|code_prefix|>" // GLM-4.5
                         || t.first == "<|prefix|>"      // Falcon-H1-Tiny-Coder
+                        || t.first == "<|plamo:fim_prefix|>" // PLaMo-3
                         ) {
                     special_fim_pre_id = t.second;
                     if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
@@ -2881,6 +2882,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                         || t.first == "▁"         // CodeLlama
                         || t.first == "<|code_suffix|>" // GLM-4.5
                         || t.first == "<|suffix|>"      // Falcon-H1-Tiny-Coder
+                        || t.first == "<|plamo:fim_suffix|>" // PLaMo-3
                         ) {
                     special_fim_suf_id = t.second;
                     if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
@@ -2902,6 +2904,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                         || t.first == "▁"         // CodeLlama
                         || t.first == "<|code_middle|>" // GLM-4.5
                         || t.first == "<|middle|>"      // Falcon-H1-Tiny-Coder
+                        || t.first == "<|plamo:fim_middle|>" // PLaMo-3
                         ) {
                     special_fim_mid_id = t.second;
                     if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
@@ -2952,6 +2955,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
             if (special_fim_sep_id == LLAMA_TOKEN_NULL) {
                 if (false
                         || t.first == "<|file_sep|>" // Qwen
+                        || t.first == "<|plamo:file_separator|>" // PLaMo-2/3
                         ) {
                     special_fim_sep_id = t.second;
                     if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {

From d0b490f25edec147b4dcd69cce31c7d1099dff1d Mon Sep 17 00:00:00 2001
From: Hrishith Thadicherla <99313418+hthadicherla@users.noreply.github.com>
Date: Wed, 7 Oct 2026 03:34:18 -0700
Subject: [PATCH 06/40] sampling : use greedy selection for eligible
 temperature-zero chains (#29797)

* sampling : use greedy selection for eligible temperature-zero chains

Assisted-by: OpenAI Codex

* Apply suggestion from @ggerganov

Co-authored-by: Georgi Gerganov 

* sampling: simplify zero-temperature greedy eligibility

Allow the same greedy selection on CPU and grammar/reasoning-budget paths.
Keep distribution sampling for dynamic temperature and requested probabilities.
Cover the common sampler selection and probability behavior in the existing sampler tests.

Assisted-by: OpenAI Codex

* sampling : use greedy selection after final top-k with k=1

Assisted-by: OpenAI Codex

* cont : clean-up

---------

Co-authored-by: Georgi Gerganov 
---
 common/sampling.cpp            |   7 +-
 src/llama-sampler.cpp          |   6 ++
 tests/test-backend-sampler.cpp | 190 +++++++++++++++++++++++++--------
 3 files changed, 155 insertions(+), 48 deletions(-)

diff --git a/common/sampling.cpp b/common/sampling.cpp
index d9c508049..6cc8872f2 100644
--- a/common/sampling.cpp
+++ b/common/sampling.cpp
@@ -399,8 +399,11 @@ struct common_sampler * common_sampler_init(
             // only if user explicitly included adaptive-p sampler
             samplers.push_back(llama_sampler_init_adaptive_p(params.adaptive_target, params.adaptive_decay, params.seed));
         } else {
-            // default: sample from distribution
-            samplers.push_back(llama_sampler_init_dist(params.seed));
+            // Keep distribution sampling when callers request probabilities.
+            const bool greedy = params.n_probs == 0 && !params.samplers.empty() &&
+                ((params.samplers.back() == COMMON_SAMPLER_TYPE_TEMPERATURE && params.temp == 0.0f && params.dynatemp_range == 0.0f) ||
+                 (params.samplers.back() == COMMON_SAMPLER_TYPE_TOP_K && params.top_k == 1));
+            samplers.push_back(greedy ? llama_sampler_init_greedy() : llama_sampler_init_dist(params.seed));
         }
     } else if (params.mirostat == 1) {
         samplers.push_back(llama_sampler_init_temp(params.temp));
diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp
index 6c958a287..521c5718e 100644
--- a/src/llama-sampler.cpp
+++ b/src/llama-sampler.cpp
@@ -1086,6 +1086,12 @@ static void llama_sampler_greedy_backend_apply(
     struct ggml_tensor * curl = ggml_argmax(ctx, logits);
     ggml_set_name(curl, "greedy_argmax");
 
+    if (data->candidates != nullptr) {
+        struct ggml_tensor * candidates = ggml_reshape_2d(ctx, data->candidates, 1, ggml_nelements(data->candidates));
+        curl = ggml_get_rows(ctx, candidates, curl);
+        ggml_set_name(curl, "greedy_sampled_token");
+    }
+
     data->sampled = curl;
 }
 
diff --git a/tests/test-backend-sampler.cpp b/tests/test-backend-sampler.cpp
index 56736ac46..62c5ce40b 100644
--- a/tests/test-backend-sampler.cpp
+++ b/tests/test-backend-sampler.cpp
@@ -2,6 +2,7 @@
 #include "llama.h"
 #include "llama-cpp.h"
 #include "common.h"
+#include "sampling.h"
 
 #ifdef NDEBUG
 #undef NDEBUG
@@ -316,7 +317,7 @@ static llama_sampler * test_single_output_backend_sampler_init(
     return llama_sampler_init(&test_single_output_backend_sampler_i, ctx);
 }
 
-static void test_backend_greedy_sampling(const test_params & params) {
+static void test_greedy(const test_params & params) {
     const int seq_id = 0;
 
     struct llama_sampler_chain_params backend_sampler_params = llama_sampler_chain_default_params();
@@ -351,7 +352,95 @@ static void test_backend_greedy_sampling(const test_params & params) {
     }
 }
 
-static void test_backend_top_k_sampling(const test_params & params) {
+
+static void test_greedy_filtered_common(const test_params & params) {
+    auto check = [&](common_params_sampling sp, bool greedy) {
+        common_sampler_ptr sampler(common_sampler_init(params.model.get(), sp));
+        auto * chain = common_sampler_get(sampler.get());
+        auto * last = llama_sampler_chain_get(chain, llama_sampler_chain_n(chain) - 1);
+        GGML_ASSERT(strcmp(llama_sampler_name(last), greedy ? "greedy" : "dist") == 0);
+
+        llama_token_data tokens[] = {{0, -2.0f, 0.0f}, {1, -1.0f, 0.0f}, {2, 0.0f, 0.0f}, {3, 1.0f, 0.0f}};
+        llama_token_data_array candidates = {tokens, 4, -1, false};
+        llama_sampler_apply(chain, &candidates);
+        GGML_ASSERT(candidates.selected >= 0);
+        if (greedy || sp.n_probs > 0) {
+            GGML_ASSERT(candidates.data[candidates.selected].id == 3);
+        }
+        if (sp.n_probs > 0) {
+            GGML_ASSERT(candidates.data[candidates.selected].p == 1.0f);
+        }
+        if (sp.dynatemp_range > 0.0f) {
+            int positive = 0;
+            for (size_t i = 0; i < candidates.size; ++i) {
+                positive += candidates.data[i].p > 0.0f;
+            }
+            GGML_ASSERT(positive > 1);
+        }
+    };
+
+    common_params_sampling sp;
+    sp.temp = 0.0f;
+    sp.samplers = {COMMON_SAMPLER_TYPE_TOP_K, COMMON_SAMPLER_TYPE_TEMPERATURE};
+    for (bool backend : {false, true}) {
+        sp.backend_sampling = backend;
+        check(sp, true);
+        auto probabilities = sp;
+        probabilities.n_probs = 4;
+        check(probabilities, false);
+        auto dynamic = sp;
+        dynamic.dynatemp_range = 1.0f;
+        check(dynamic, false);
+    }
+    auto grammar = sp;
+    grammar.grammar = {COMMON_GRAMMAR_TYPE_USER, "root ::= [a-z]+"};
+    check(grammar, true);
+    auto budget = sp;
+    budget.reasoning_budget_start = {1};
+    budget.reasoning_budget_end = {{2}};
+    budget.reasoning_budget_tokens = 1;
+    check(budget, true);
+    sp.samplers = {COMMON_SAMPLER_TYPE_TEMPERATURE, COMMON_SAMPLER_TYPE_TOP_K};
+    check(sp, false);
+    sp.samplers.clear();
+    check(sp, false);
+
+    sp.temp = 0.8f;
+    sp.samplers = {COMMON_SAMPLER_TYPE_TEMPERATURE, COMMON_SAMPLER_TYPE_TOP_K};
+    for (bool backend : {false, true}) {
+        sp.backend_sampling = backend;
+        sp.top_k = 1;
+        check(sp, true);
+        auto probabilities = sp;
+        probabilities.n_probs = 4;
+        check(probabilities, false);
+        sp.top_k = 8;
+        check(sp, false);
+    }
+}
+
+static void test_greedy_filtered(const test_params & params) {
+    for (int k : {1, 8}) {
+        llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
+        llama_sampler_chain_add(chain.get(), llama_sampler_init_top_k(k));
+        llama_sampler_chain_add(chain.get(), llama_sampler_init_greedy());
+        std::vector configs = {{0, chain.get()}};
+        test_context test_ctx(params, configs);
+        GGML_ASSERT(test_ctx.decode({{0, "Write a Python function"}}));
+        for (int step = 0; step < 4; ++step) {
+            const int idx = test_ctx.idx_for_seq(0);
+            const auto * logits = llama_get_sampled_logits_ith(test_ctx.ctx.get(), idx);
+            const auto * ids = llama_get_sampled_candidates_ith(test_ctx.ctx.get(), idx);
+            GGML_ASSERT(llama_get_sampled_logits_count_ith(test_ctx.ctx.get(), idx) == (uint32_t) k);
+            GGML_ASSERT(llama_get_sampled_candidates_count_ith(test_ctx.ctx.get(), idx) == (uint32_t) k);
+            const auto expected = ids[std::max_element(logits, logits + k) - logits];
+            GGML_ASSERT(llama_get_sampled_token_ith(test_ctx.ctx.get(), idx) == expected);
+            GGML_ASSERT(test_ctx.decode_token(expected));
+        }
+    }
+}
+
+static void test_top_k(const test_params & params) {
     const int seq_id = 0;
     const int32_t k = 8;
     struct llama_sampler_chain_params backend_chain_params = llama_sampler_chain_default_params();
@@ -393,7 +482,7 @@ static void test_backend_top_k_sampling(const test_params & params) {
     printf("backend top-k hybrid sampling test PASSED\n");
 }
 
-static void test_backend_temp_sampling(const test_params & params) {
+static void test_temp(const test_params & params) {
     {
         const float temp_0 = 0.8f;
         struct llama_sampler_chain_params backend_chain_params_0 = llama_sampler_chain_default_params();
@@ -481,7 +570,7 @@ static void test_backend_temp_sampling(const test_params & params) {
     printf("backend temp sampling test PASSED\n");
 }
 
-static void test_backend_temp_ext_sampling(const test_params & params) {
+static void test_temp_ext(const test_params & params) {
     {
         int seq_id = 0;
         const float temp = 0.8f;
@@ -546,7 +635,7 @@ static void test_backend_temp_ext_sampling(const test_params & params) {
     printf("backend temp_ext sampling test PASSED\n");
 }
 
-static void test_backend_min_p_sampling(const test_params & params) {
+static void test_min_p(const test_params & params) {
     const int seq_id = 0;
     const float p = 0.1;
     struct llama_sampler_chain_params backend_chain_params = llama_sampler_chain_default_params();
@@ -598,7 +687,7 @@ static void test_backend_min_p_sampling(const test_params & params) {
     printf("min-p sampling test PASSED\n");
 }
 
-static void test_backend_top_p_sampling(const test_params & params) {
+static void test_top_p(const test_params & params) {
     const int seq_id = 0;
     const float p = 0.9;
     struct llama_sampler_chain_params backend_chain_params = llama_sampler_chain_default_params();
@@ -648,7 +737,7 @@ static void test_backend_top_p_sampling(const test_params & params) {
     printf("top-p sampling test PASSED\n");
 }
 
-static void test_backend_multi_sequence_sampling(const test_params & params) {
+static void test_multi_sequence(const test_params & params) {
     struct llama_sampler_chain_params chain_params_0 = llama_sampler_chain_default_params();
     llama_sampler_ptr sampler_chain_0(llama_sampler_chain_init(chain_params_0));
     llama_sampler_chain_add(sampler_chain_0.get(), llama_sampler_init_greedy());
@@ -714,7 +803,7 @@ static void test_backend_multi_sequence_sampling(const test_params & params) {
     printf("backend multi-sequence sampling test PASSED\n");
 }
 
-static void test_backend_dist_sampling(const test_params & params) {
+static void test_dist(const test_params & params) {
     const int seq_id = 0;
     const int32_t seed = 88;
 
@@ -742,7 +831,7 @@ static void test_backend_dist_sampling(const test_params & params) {
     printf("backend dist sampling test PASSED\n");
 }
 
-static void test_backend_dist_sampling_and_cpu(const test_params & params) {
+static void test_dist_and_cpu(const test_params & params) {
     const int seq_id = 0;
     const int32_t seed = 88;
 
@@ -772,7 +861,7 @@ static void test_backend_dist_sampling_and_cpu(const test_params & params) {
     printf("backend dist & cpu sampling test PASSED\n");
 }
 
-static void test_backend_logit_bias_sampling(const test_params & params) {
+static void test_logit_bias(const test_params & params) {
     const auto * model = params.model.get();
     const auto * vocab = llama_model_get_vocab(model);
 
@@ -1093,7 +1182,7 @@ static void compare_penalties_logits(
     GGML_ASSERT(stats.n_mismatch == 0);
 }
 
-static void test_penalty_parameter_values(const test_params & params) {
+static void check_penalty_parameter_values(const test_params & params) {
     struct penalty_test_case {
         const char * name;
         float repeat;
@@ -1292,7 +1381,7 @@ static void compare_masking_penalties_logits(
     GGML_ASSERT(stats.n_mismatch == 0);
 }
 
-static void test_backend_penalties_sampling(const test_params & params) {
+static void test_penalties(const test_params & params) {
     printf("Testing backend penalties (repeat + freq + presence)\n");
     compare_penalties_logits(params, 64, 1.1f, 0.5f, 0.25f, "Hello Hello world");
 
@@ -1371,14 +1460,14 @@ static void test_backend_penalties_sampling(const test_params & params) {
     }, 64, 1.0f, 0.0f, 0.25f, "Hello", penalties_position::after_filter, true);
 
     printf("Testing backend penalty parameter values\n");
-    test_penalty_parameter_values(params);
+    check_penalty_parameter_values(params);
 
     printf("backend penalties sampling test PASSED\n");
 }
 
 // This test verifies that it is possible to have two different backend samplers,
 // one that uses the backend dist sampler, and another that uses CPU dist sampler.
-static void test_backend_mixed_sampling(const test_params & params) {
+static void test_mixed(const test_params & params) {
     struct llama_sampler_chain_params chain_params_0 = llama_sampler_chain_default_params();
     llama_sampler_ptr sampler_chain_0(llama_sampler_chain_init(chain_params_0));
     llama_sampler_chain_add(sampler_chain_0.get(), llama_sampler_init_dist(88));
@@ -1428,7 +1517,7 @@ static void test_backend_mixed_sampling(const test_params & params) {
     printf("backend mixed sampling test PASSED\n");
 }
 
-static void test_backend_set_sampler(const test_params & params) {
+static void test_set_sampler(const test_params & params) {
     const int seq_id = 0;
     const int32_t seed = 88;
 
@@ -1493,7 +1582,7 @@ static void test_backend_set_sampler(const test_params & params) {
     printf("backend set sampler test PASSED\n");
 }
 
-static void test_backend_cpu_mixed_batch(const test_params & params) {
+static void test_cpu_mixed(const test_params & params) {
     // Sequence 0 uses backend sampling
     struct llama_sampler_chain_params chain_params_0 = llama_sampler_chain_default_params();
     llama_sampler_ptr sampler_chain_0(llama_sampler_chain_init(chain_params_0));
@@ -1581,7 +1670,7 @@ static void test_backend_cpu_mixed_batch(const test_params & params) {
     printf("backend-cpu mixed batch test PASSED\n");
 }
 
-static void test_backend_multi_output_limit(const test_params & params) {
+static void test_multi_output_limit(const test_params & params) {
     const llama_seq_id seq_id = 0;
 
     llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
@@ -1594,15 +1683,15 @@ static void test_backend_multi_output_limit(const test_params & params) {
         batch.add(llama_vocab_bos(test_ctx.vocab), i, seq_id, true);
     }
 
-    printf(">>> test_backend_multi_output_limit expected error start:\n");
+    printf(">>> test_multi_output_limit expected error start:\n");
     const int ret = llama_process(test_ctx.ctx.get(), LLAMA_PROCESS_TYPE_DECODE, batch.get());
     GGML_ASSERT(ret != 0 && "llama_decode should reject outputs above the per-sequence limit");
-    printf("<<< test_backend_multi_output_limit expected error end.\n");
+    printf("<<< test_multi_output_limit expected error end.\n");
 
     printf("backend multi-output limit test PASSED\n");
 }
 
-static void test_backend_multi_sequence_multi_output_dist(const test_params & params) {
+static void test_multi_output_multi_sequence_dist(const test_params & params) {
     const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
     const int32_t n_vocab = llama_vocab_n_tokens(vocab);
     const uint32_t seeds[] = { 88, 1337 };
@@ -1697,7 +1786,7 @@ static void test_backend_multi_sequence_multi_output_dist(const test_params & pa
     printf("backend multi-sequence multi-output dist test PASSED\n");
 }
 
-static void test_backend_multi_output_dist_transaction(const test_params & params) {
+static void test_multi_output_dist_transaction(const test_params & params) {
     const llama_seq_id seq_id = 0;
     const uint32_t seed = 95;
     const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
@@ -1762,7 +1851,7 @@ static void test_backend_multi_output_dist_transaction(const test_params & param
     printf("backend multi-output dist transaction test PASSED\n");
 }
 
-static void test_backend_multi_output_sampling_chain(const test_params & params) {
+static void test_multi_output_sampling_chain(const test_params & params) {
     const llama_seq_id seq_id = 0;
     const uint32_t seed = 88;
     const float p = 0.9f;
@@ -1913,7 +2002,7 @@ static void test_backend_multi_output_sampling_chain(const test_params & params)
     printf("backend multi-output sampling chain test PASSED\n");
 }
 
-static void test_backend_multi_output_cpu_suffix(const test_params & params) {
+static void test_multi_output_cpu_suffix(const test_params & params) {
     const llama_seq_id seq_id = 0;
     const int32_t k = 8;
     const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
@@ -1977,28 +2066,35 @@ struct backend_test_case {
     bool enabled_by_default;
 };
 
+// note: test names are "test_" and match the function implementing them
 static const backend_test_case BACKEND_TESTS[] = {
-    { "greedy",          test_backend_greedy_sampling,         true  },
-    { "logit_bias",      test_backend_logit_bias_sampling,     true  },
-    { "penalties",       test_backend_penalties_sampling,      true  },
-    { "temp",            test_backend_temp_sampling,           true  },
-    { "temp_ext",        test_backend_temp_ext_sampling,       true  },
-    { "top_k",           test_backend_top_k_sampling,          true  },
-    { "multi_sequence",  test_backend_multi_sequence_sampling, true  },
-    { "dist",            test_backend_dist_sampling,           true  },
-    { "dist_and_cpu",    test_backend_dist_sampling_and_cpu,   true  },
-    { "set_sampler",     test_backend_set_sampler,             true  },
-    { "multi_output_limit",    test_backend_multi_output_limit,      true },
-    { "multi_sequence_multi_output_dist", test_backend_multi_sequence_multi_output_dist, true },
-    { "multi_output_dist_transaction", test_backend_multi_output_dist_transaction, true },
-    { "multi_output_sampling_chain", test_backend_multi_output_sampling_chain, true },
-    { "multi_output_cpu",      test_backend_multi_output_cpu_suffix, true },
-    { "mixed",           test_backend_mixed_sampling,          true  },
-    { "min_p",           test_backend_min_p_sampling,          true  },
-    { "cpu_mixed",       test_backend_cpu_mixed_batch,         true  },
-    { "top_p",           test_backend_top_p_sampling,          true  },
+    // single sampler
+    { "test_greedy",                 test_greedy,                 true },
+    { "test_greedy_filtered",        test_greedy_filtered,        true },
+    { "test_greedy_filtered_common", test_greedy_filtered_common, true },
+    { "test_temp",                   test_temp,                   true },
+    { "test_temp_ext",               test_temp_ext,               true },
+    { "test_top_k",                  test_top_k,                  true },
+    { "test_top_p",                  test_top_p,                  true },
+    { "test_min_p",                  test_min_p,                  true },
+    { "test_logit_bias",             test_logit_bias,             true },
+    { "test_penalties",              test_penalties,              true },
+    // multiple sequences
+    { "test_multi_sequence",         test_multi_sequence,         true },
+    { "test_dist",                   test_dist,                   true },
+    { "test_dist_and_cpu",           test_dist_and_cpu,           true },
+    { "test_mixed",                  test_mixed,                  true },
+    { "test_cpu_mixed",              test_cpu_mixed,              true },
+    { "test_set_sampler",            test_set_sampler,            true },
+    // multiple outputs per sequence
+    { "test_multi_output_limit",                  test_multi_output_limit,               true },
+    { "test_multi_output_multi_sequence_dist",    test_multi_output_multi_sequence_dist, true },
+    { "test_multi_output_dist_transaction",       test_multi_output_dist_transaction,    true },
+    { "test_multi_output_sampling_chain",         test_multi_output_sampling_chain,      true },
+    { "test_multi_output_cpu_suffix",             test_multi_output_cpu_suffix,          true },
 };
 
+// TODO: add usage, examples
 static test_args parse_cli(int argc, char ** argv) {
     test_args out;
 
@@ -2082,10 +2178,12 @@ static std::vector collect_tests_to_run(const std::st
             }
 #ifdef GGML_USE_HIP
             // TODO: remove this when https://github.com/ggml-org/llama.cpp/pull/26592 is merged
-            if (test.name == "penalties" || test.name == "set_sampler" ||
-                test.name == "mixed"     || test.name == "top_p"       ||
-                test.name == "multi_output_sampling_chain" ||
-                test.name == "multi_output_cpu") {
+            if (test.name == "test_penalties"                   ||
+                test.name == "test_set_sampler"                 ||
+                test.name == "test_mixed"                       ||
+                test.name == "test_top_p"                       ||
+                test.name == "test_multi_output_sampling_chain" ||
+                test.name == "test_multi_output_cpu_suffix") {
                 fprintf(stderr, "Skipping test '%s' on HIP backend (no backend TOP_K support)\n", test.name.c_str());
                 continue;
             }

From 48499d2e1c86c1d5fe7a05f8bf083156e7069cc0 Mon Sep 17 00:00:00 2001
From: SIDDARTHA REDDY <75976672+SIDDARTHAREDDY8@users.noreply.github.com>
Date: Wed, 7 Oct 2026 06:57:48 -0400
Subject: [PATCH 07/40] qwen3tts : guard speaker_encoder_config patch for
 CustomVoice variant (#29179)

* Fix KeyError converting Qwen3-TTS CustomVoice variant without speaker_encoder_config (fixes #29088)

Signed-off-by: SIDDARTHA REDDY <75976672+SIDDARTHAREDDY8@users.noreply.github.com>

* Address review feedback: drop tests/test-convert-qwen3tts.py

Per reviewer feedback on PR #29179, remove the regression test file.
The fix itself is unchanged.

Signed-off-by: SIDDARTHA REDDY <75976672+SIDDARTHAREDDY8@users.noreply.github.com>

---------

Signed-off-by: SIDDARTHA REDDY <75976672+SIDDARTHAREDDY8@users.noreply.github.com>
---
 conversion/qwen3tts.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/conversion/qwen3tts.py b/conversion/qwen3tts.py
index 2c35799f7..cdb1e0c9b 100644
--- a/conversion/qwen3tts.py
+++ b/conversion/qwen3tts.py
@@ -218,8 +218,10 @@ class Qwen3TTSSpeakerEncoderModel(MmprojModel):
         if hparams is None:
             hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False)
         hparams["text_config"] = {"hidden_size": hparams["talker_config"]["hidden_size"]}
-        # ECAPA-TDNN has a fixed 4-stage backbone, but MmprojModel.__init__ needs a n_block_keys
-        hparams["speaker_encoder_config"]["n_layers"] = 4
+        # ECAPA-TDNN has a fixed 4-stage backbone, but MmprojModel.__init__ needs a n_block_keys.
+        # The CustomVoice variant ships no speaker encoder, so its config lacks this key entirely.
+        if "speaker_encoder_config" in hparams:
+            hparams["speaker_encoder_config"]["n_layers"] = 4
         super().__init__(dir_model, *args, hparams=hparams, **kwargs)
         self._wav_config_cache = None
 

From b9acf138a1e28ce1fc23b5a4fc4b12444b50f7ea Mon Sep 17 00:00:00 2001
From: "Piotr Wilkin (ilintar)" 
Date: Wed, 7 Oct 2026 13:16:53 +0200
Subject: [PATCH 08/40] feat: add GLM5Next MTP, optimize (#29928)

* llama : add GLM5-Next NextN (MTP) graph

Build the GLM5-Next multi-token-prediction head as graph_mtp: the NextN block
embeds enorm(tok)+hnorm(h) through eh_proj, runs one plain DSA layer and the
shared lm_head, reusing the trunk's builders through the no_build tag ctor.
llama_memory_recurrent also tolerates a partial seq_rm when the context holds
no recurrent layers, which is what the MTP draft context needs.

Assisted-by: Claude

* llama : glm5-next: skip dead compute in headless NextN forwards

A NextN forward with no output rows (the MTP catch-up and the draft-context
prefill) persists only through its cache writes, so the headless graph keeps
the MLA latent, indexer key|gate and pooled-key writes and drops the query
path, the indexer selection, the attention body, the FFN and the LM head. The
4-token catch-up falls from 6.9 ms to 0.33 ms of kernels; the greedy output
hashes and the draft acceptance are unchanged.

Assisted-by: Claude

* llama : glm5-next: fix NextN extraction contracts and shared-tail rollback

Three fixes from the architectural review. The headless graph prune now also
requires that no unmasked nextn extraction is live, because that mode reads
n_tokens hidden rows regardless of the logits flags. Masked extraction
publishes the hidden rows gathered by the output ids, so a batch whose output
flags are not a prefix exports the right rows. A partial recurrent rollback
whose tail cell is shared with another sequence is now rejected instead of
silently moving that sequence's tail.

Assisted-by: Claude

* llama : glm5-next: tidy comments in the MTP changes

Assisted-by: Claude

* llama : glm5-next: crop the MTP graph to the output rows instead of pruning it

Replace the headless NextN prune with the crop pattern the other MTP
graphs use: gather the attention output and the block input at the
output ids before the position-wise FFN and the shared head. A NextN
forward with no output rows (the MTP catch-up and the draft-context
prefill) then runs the FFN and the head over zero rows. The 4-token
catch-up falls from 6.9 ms to 2.9 ms of kernels; greedy output hashes
are unchanged.

Assisted-by: Claude

* glm5-next: use the nextn crop helpers in the MTP graph

Replace the local crop condition and the masked select of t_h_nextn
with crop_before_nextn and crop_after_nextn, so the MTP graph narrows
its rows the same way as the main graph and the other models.

Describe the shared cell and empty filter branches of the recurrent
partial rollback.

* glm5-next: load MTP-only and trunk-only GGUF files

Make the trunk tensors optional when the file only holds the NextN
layer, and the NextN tensors optional when the file only holds the
trunk, so the split MTP GGUF loads as a draft model.

Co-authored-by: Georgi Gerganov 

---------

Co-authored-by: Pascal 
Co-authored-by: Georgi Gerganov 
---
 src/llama-memory-recurrent.cpp |   4 +
 src/models/glm5-next.cpp       | 185 ++++++++++++++++++++++++++++-----
 src/models/models.h            |   9 ++
 3 files changed, 173 insertions(+), 25 deletions(-)

diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp
index 53130e868..2893d819e 100644
--- a/src/llama-memory-recurrent.cpp
+++ b/src/llama-memory-recurrent.cpp
@@ -199,6 +199,10 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos
 
             // partial rollback via per-token snapshot index (bounded by n_rs_seq)
             if (0 < p0 && p0 <= cell.pos && p1 > cell.pos) {
+                // a cell shared by several sequences cannot move back for only one of them
+                if (cell.seq_id.size() > 1) {
+                    return false;
+                }
                 // the filter kept no layer (e.g. an MTP draft context), so only the position moves back
                 if (is_empty()) {
                     cell.pos = p0 - 1;
diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp
index 955405c75..4343e6ac6 100644
--- a/src/models/glm5-next.cpp
+++ b/src/models/glm5-next.cpp
@@ -64,9 +64,13 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
     const int64_t hc         = hparams.dsv4_hc_mult;
     const int64_t hc_mix_dim = (2 + hc)*hc;
 
-    // the NextN block is loaded but only used by the MTP graph.
-    // Separated trunk_only/mtp_only handling TODO with DECODER_MTP graph in the MTP follow up
-    int mtp_flags = 0;
+    // TODO: consolidate the trunk_only/mtp_only detection and flags shared with glm4-moe and qwen35moe
+    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
+    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
+    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
+    const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
+    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+
     if (!ml.load_mtp) {
         mtp_flags |= TENSOR_SKIP;
     }
@@ -79,18 +83,19 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
     for (int i = 0; i < n_layer_all; ++i) {
         auto & layer = layers[i];
 
-        const int flags = (i >= n_layer) ? mtp_flags : 0;
+        const bool is_nextn = i >= n_layer;
+        const int flags = is_nextn ? mtp_flags : trunk_flags;
 
         layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags);
         layer.ffn_norm  = create_tensor(tn(LLM_TENSOR_FFN_NORM,  "weight", i), {n_embd}, flags);
 
         if (i < n_layer) {
-            layer.hc_attn_fn    = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN,    "weight", i), {hc*n_embd, hc_mix_dim}, 0);
-            layer.hc_attn_base  = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE,  "weight", i), {hc_mix_dim}, 0);
-            layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0);
-            layer.hc_ffn_fn     = create_tensor(tn(LLM_TENSOR_HC_FFN_FN,     "weight", i), {hc*n_embd, hc_mix_dim}, 0);
-            layer.hc_ffn_base   = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE,   "weight", i), {hc_mix_dim}, 0);
-            layer.hc_ffn_scale  = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE,  "weight", i), {3}, 0);
+            layer.hc_attn_fn    = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN,    "weight", i), {hc*n_embd, hc_mix_dim}, flags);
+            layer.hc_attn_base  = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE,  "weight", i), {hc_mix_dim}, flags);
+            layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, flags);
+            layer.hc_ffn_fn     = create_tensor(tn(LLM_TENSOR_HC_FFN_FN,     "weight", i), {hc*n_embd, hc_mix_dim}, flags);
+            layer.hc_ffn_base   = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE,   "weight", i), {hc_mix_dim}, flags);
+            layer.hc_ffn_scale  = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE,  "weight", i), {3}, flags);
         }
 
         const int64_t head_dim = hparams.n_embd_head_kda;
@@ -100,25 +105,25 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
         if (hparams.is_recr(i)) {
             auto conv = [&](llm_tensor tid) {
                 ggml_tensor * t = create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED);
-                return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, 0);
+                return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, flags);
             };
             layer.ssm_q_conv = conv(LLM_TENSOR_SSM_CONV1D_Q);
             layer.ssm_k_conv = conv(LLM_TENSOR_SSM_CONV1D_K);
             layer.ssm_v_conv = conv(LLM_TENSOR_SSM_CONV1D_V);
 
-            create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, 0);
+            create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, flags);
 
-            layer.ssm_f_a  = create_tensor(tn(LLM_TENSOR_SSM_F_A,  "weight", i), {n_embd, head_dim}, 0);
-            layer.ssm_f_b  = create_tensor(tn(LLM_TENSOR_SSM_F_B,  "weight", i), {head_dim, d_inner}, 0);
-            layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0);
+            layer.ssm_f_a  = create_tensor(tn(LLM_TENSOR_SSM_F_A,  "weight", i), {n_embd, head_dim}, flags);
+            layer.ssm_f_b  = create_tensor(tn(LLM_TENSOR_SSM_F_B,  "weight", i), {head_dim, d_inner}, flags);
+            layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, flags);
 
-            layer.ssm_a    = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), {n_head}, 0);
-            layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0);
+            layer.ssm_a    = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), {n_head}, flags);
+            layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, flags);
 
-            layer.ssm_g_a    = create_tensor(tn(LLM_TENSOR_SSM_G_A,  "weight", i), {n_embd, head_dim}, 0);
-            layer.ssm_g_b    = create_tensor(tn(LLM_TENSOR_SSM_G_B,  "weight", i), {head_dim, d_inner}, 0);
-            layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, 0);
-            layer.wo         = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, 0);
+            layer.ssm_g_a    = create_tensor(tn(LLM_TENSOR_SSM_G_A,  "weight", i), {n_embd, head_dim}, flags);
+            layer.ssm_g_b    = create_tensor(tn(LLM_TENSOR_SSM_G_B,  "weight", i), {head_dim, d_inner}, flags);
+            layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, flags);
+            layer.wo         = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, flags);
         } else {
             const int64_t q_lora_rank      = hparams.n_lora_q;
             const int64_t kv_lora_rank     = hparams.n_lora_kv;
@@ -155,9 +160,9 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
         }
 
         if (i < (int) hparams.n_layer_dense_lead) {
-            layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
-            layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
-            layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd, n_ff}, 0);
+            layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, flags);
+            layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, flags);
+            layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd, n_ff}, flags);
         } else {
             const int64_t n_ff_exp        = hparams.n_ff_exp(i);
             const int64_t n_expert_shared = hparams.n_expert_shared;
@@ -187,7 +192,7 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
 
 std::unique_ptr llama_model_glm5_next::build_arch_graph(const llm_graph_params & params) const {
     if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
-        throw std::runtime_error("GLM5-Next NextN graph not implemented yet");
+        return std::make_unique(*this, params);
     }
     return std::make_unique(*this, params);
 }
@@ -537,6 +542,136 @@ ggml_tensor * llama_model_glm5_next::graph::build_hc_post(
     return out;
 }
 
+// construct the graph helpers without building the trunk, so graph_mtp can reuse them
+llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_params & params, no_build) :
+    llm_build_delta_net_base(params), model(model) {
+}
+
+llama_model_glm5_next::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) :
+    graph(model, params, no_build{}) {
+
+    GGML_ASSERT(hparams.n_layer_nextn > 0 && "GLM5-Next MTP requires n_layer_nextn > 0");
+    GGML_ASSERT(hparams.n_layer_nextn == 1 && "GLM5-Next MTP currently supports a single NextN block");
+
+    const int il = hparams.n_layer() + cparams.nextn_layer_offset;
+    GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
+                cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
+                "nextn_layer_offset out of range [0, n_layer_nextn)");
+
+    const auto & layer = model.layers[il];
+
+    GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm &&
+                "GLM5-Next MTP block is missing - load the model with MTP enabled");
+
+    GGML_ASSERT(hparams.n_embd_out() == (uint32_t) n_embd && "GLM5-Next MTP hidden width mismatch");
+
+    const auto * mctx_hyb = static_cast(mctx);
+
+    auto * inp_hyb   = build_inp_mem_hybrid_k();
+    auto * inp_attn  = inp_hyb->get_attn();
+    auto * inp_kpool = build_inp_kpool(mctx_hyb);
+
+    ggml_build_forward_expand(gf, inp_hyb->get_recr()->s_copy);
+
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+    auto inp = std::make_unique(n_embd);
+
+    inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+    ggml_set_input(inp->tokens);
+
+    inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
+    ggml_set_input(inp->embd);
+
+    inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
+    ggml_set_input(inp->h);
+    ggml_set_name(inp->h, "mtp_h_input");
+
+    ggml_tensor * tok_embd = ggml_get_rows(ctx0,
+            layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd, inp->tokens);
+    cb(tok_embd, "mtp_tok_embd", il);
+
+    ggml_tensor * h = inp->h;
+
+    res->add_input(std::move(inp));
+
+    ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
+    cb(e_norm, "mtp_enorm", il);
+
+    ggml_tensor * h_norm = build_norm(h, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
+    cb(h_norm, "mtp_hnorm", il);
+
+    ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, ggml_concat(ctx0, e_norm, h_norm, 0), layer.nextn.eh_proj_s);
+    cb(cur, "mtp_eh_proj", il);
+
+    ggml_tensor * inpSA = cur;
+
+    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);
+    cb(cur, "mtp_attn_out", il);
+
+    // narrow to the output tokens before the position-wise FFN; unmasked nextn embeddings need all rows
+    if (crop_before_nextn(inp_out_ids)) {
+        cur   = ggml_get_rows(ctx0, cur,   inp_out_ids);
+        inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+    }
+
+    ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+    cb(ffn_inp, "mtp_ffn_inp", il);
+
+    cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+    cb(cur, "mtp_ffn_norm", il);
+
+    ggml_tensor * moe_out = build_moe_ffn(cur,
+            layer.ffn_gate_inp,
+            layer.ffn_up_exps,
+            layer.ffn_gate_exps,
+            layer.ffn_down_exps,
+            layer.ffn_exp_probs_b,
+            n_expert, n_expert_used,
+            LLM_FFN_SILU, hparams.expert_weights_norm,
+            hparams.expert_weights_scale,
+            (llama_expert_gating_func_type) hparams.expert_gating_func,
+            il);
+    cb(moe_out, "mtp_ffn_moe_out", il);
+
+    ggml_tensor * ffn_shexp = build_ffn(cur,
+            layer.ffn_up_shexp,   nullptr, nullptr,
+            layer.ffn_gate_shexp, nullptr, nullptr,
+            layer.ffn_down_shexp, nullptr, nullptr,
+            nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
+    cb(ffn_shexp, "mtp_ffn_shexp", il);
+
+    cur = ggml_add(ctx0, moe_out, ffn_shexp);
+    cb(cur, "mtp_ffn_out", il);
+
+    cur = ggml_add(ctx0, cur, ffn_inp);
+    cb(cur, "mtp_post_ffn", il);
+
+    ggml_tensor * head_norm = layer.nextn.shared_head_norm ? layer.nextn.shared_head_norm : model.output_norm;
+    GGML_ASSERT(head_norm && "GLM5-Next MTP: missing both nextn.shared_head_norm and output_norm");
+    cur = build_norm(cur, head_norm, nullptr, LLM_NORM_RMS, -1);
+    cb(cur, "h_nextn", -1);
+    res->t_h_nextn = cur;
+
+    if (crop_after_nextn(inp_out_ids)) {
+        cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+    }
+    cb(cur, "mtp_shared_head_norm", -1);
+
+    ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
+    ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
+    GGML_ASSERT(head_w && "GLM5-Next MTP: missing both nextn.shared_head_head and output");
+    cur = build_lora_mm(head_w, cur, head_s);
+    cb(cur, "result_output", -1);
+
+    res->t_logits = cur;
+    ggml_build_forward_expand(gf, cur);
+}
+
 llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_params & params) :
     llm_build_delta_net_base(params), model(model) {
 
diff --git a/src/models/models.h b/src/models/models.h
index 023ed3021..1ef0c5156 100644
--- a/src/models/models.h
+++ b/src/models/models.h
@@ -2741,6 +2741,10 @@ struct llama_model_glm5_next : public llama_model_base {
     struct graph : public llm_build_delta_net_base {
         graph(const llama_model & model, const llm_graph_params & params);
 
+        // build the helpers without the trunk, so graph_mtp can reuse them
+        struct no_build {};
+        graph(const llama_model & model, const llm_graph_params & params, no_build);
+
         // collapse the hc streams with per-stream weights
         ggml_tensor * build_hc_pre(
                 ggml_tensor * x,
@@ -2789,6 +2793,11 @@ struct llama_model_glm5_next : public llama_model_base {
 
     };
 
+    // the NextN / MTP draft head: one DSA block appended after the trunk
+    struct graph_mtp : public graph {
+        graph_mtp(const llama_model & model, const llm_graph_params & params);
+    };
+
     std::unique_ptr build_arch_graph(const llm_graph_params & params) const override;
 };
 

From 7e8324f5feba1d8be3f966e1791e84107e40cdd7 Mon Sep 17 00:00:00 2001
From: Niko Maroulis 
Date: Wed, 7 Oct 2026 09:41:50 -0400
Subject: [PATCH 09/40] metal : fix MUL_MAT+ADD fusion when the residual is
 itself a MUL_MAT (#30100)

* metal : fix MUL_MAT+ADD fusion when the residual is itself a MUL_MAT

ggml_metal_op_mul_mat_mma picks the residual of a fused MUL_MAT+ADD as
"the ADD operand whose op is not MUL_MAT". When both operands of the ADD
are mat-mul outputs (x = W1 @ u + W2 @ v), that test is true for both, so
the residual resolves to the fused mat-mul's own, never-written output and
the kernel adds whatever that buffer holds.

The fusion check (ggml_metal_mul_mat_add_operand) already selects the
operand by identity; make the encoder do the same.

Clef decision models hit this in their head (proj_option_context @ ctx +
proj_option_lexical @ lex, 9 option rows): on Metal, /v1/systemone
probabilities collapse toward uniform (billing 0.28 where the CPU backend
gives 0.977, Cloudflare_clef-flash Q8_0), deterministic per memory layout,
correct with GGML_METAL_FUSION_DISABLE=1. Not a quantization issue: the
same file is right on CPU.

Add a MUL_MAT_ADD mode to test-backend-ops where the residual is a second
mat-mul; on Metal it fails 27 of 28 cases before this change (the one pass
is f16 n=2, under the MMA row threshold, so nothing fuses).

* Update tests/test-backend-ops.cpp

Co-authored-by: Georgi Gerganov 

---------

Co-authored-by: Georgi Gerganov 
---
 ggml/src/ggml-metal/ggml-metal-ops.cpp |  3 ++-
 tests/test-backend-ops.cpp             | 10 +++++++++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp
index 2a14cd0c2..597d5300c 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp
@@ -2476,7 +2476,8 @@ static int ggml_metal_op_mul_mat_mma(ggml_metal_op_t ctx, int idx) {
 
     if (fuse_add) {
         dst = ctx->node(idx + n_fuse - 1);
-        res = dst->src[0]->op == GGML_OP_MUL_MAT ? dst->src[1] : dst->src[0];
+        // the residual is the other operand of the ADD, by identity: it can itself be a MUL_MAT output
+        res = dst->src[0] == op ? dst->src[1] : dst->src[0];
     }
 
     auto pipeline = ggml_metal_library_get_pipeline_mul_mv_mma_auto(lib, op, fuse_add);
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index b080ab496..efc1777ef 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -7233,6 +7233,7 @@ enum mul_mat_add_mode {
     MUL_MAT_ADD_ROW,        // mm + a one-row res, which a same-shape fusion must leave alone
     MUL_MAT_ADD_RES_INPLACE, // res += mm
     MUL_MAT_ADD_B_INPLACE,   // b += mm: the sum overwrites the mat-mul input (m == k)
+    MUL_MAT_ADD_MM_MM,       // mm2 + mm: the residual is itself a mat-mul output, so both operands are MUL_MAT
 };
 
 static std::string var_to_str(mul_mat_add_mode mode) {
@@ -7242,6 +7243,7 @@ static std::string var_to_str(mul_mat_add_mode mode) {
         case MUL_MAT_ADD_ROW:         return "mm+row";
         case MUL_MAT_ADD_RES_INPLACE: return "res+=mm";
         case MUL_MAT_ADD_B_INPLACE:   return "b+=mm";
+        case MUL_MAT_ADD_MM_MM:       return "mm2+mm";
     }
     return "unknown";
 }
@@ -7287,6 +7289,12 @@ struct test_mul_mat_add : public test_case {
             case MUL_MAT_ADD_RES_MM:      out = ggml_add(ctx, res, mm);         break;
             case MUL_MAT_ADD_RES_INPLACE: out = ggml_add_inplace(ctx, res, mm); break;
             case MUL_MAT_ADD_B_INPLACE:   out = ggml_add_inplace(ctx, b, mm);   break;
+            case MUL_MAT_ADD_MM_MM:
+                {
+                    ggml_tensor * a2  = ggml_new_tensor_2d(ctx, type_a, k, m);
+                    ggml_tensor * mm2 = ggml_mul_mat(ctx, a2, b);
+                    out = ggml_add(ctx, mm2, mm);
+                } break;
         }
         ggml_set_name(out, "out");
 
@@ -10409,7 +10417,7 @@ static std::vector> make_test_cases_eval() {
             }
         }
         for (ggml_type type_a : {GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q5_K, GGML_TYPE_F16}) {
-            for (mul_mat_add_mode mode : {MUL_MAT_ADD_MM_RES, MUL_MAT_ADD_RES_MM, MUL_MAT_ADD_ROW, MUL_MAT_ADD_RES_INPLACE}) {
+            for (mul_mat_add_mode mode : {MUL_MAT_ADD_MM_RES, MUL_MAT_ADD_RES_MM, MUL_MAT_ADD_ROW, MUL_MAT_ADD_RES_INPLACE, MUL_MAT_ADD_MM_MM}) {
                 test_cases.emplace_back(new test_mul_mat_add(type_a, 1000, n, 1024, mode));
             }
             test_cases.emplace_back(new test_mul_mat_add(type_a, 2048, n, 2048, MUL_MAT_ADD_B_INPLACE));

From 988190680d5a89fce97de3c20df2c2813731fd61 Mon Sep 17 00:00:00 2001
From: pratiknarola-t 
Date: Wed, 7 Oct 2026 19:19:37 +0530
Subject: [PATCH 10/40] metal : few-row MMA mat-mul for the remaining src0
 types (#30065)

The generic few-row MMA kernel works for any type with a 16-weight
dequantizer, so it now also takes BF16, Q1_0, Q2_0, MXFP4, Q2_K, Q3_K,
TQ2_0 and the IQ types. Each type starts at the row count where it beats
the current kernels on an M3 Ultra: 5 rows for TQ2_0, 4 for BF16, 3
for MXFP4, Q2_0, Q2_K and IQ4_NL, and 2 for the others.

test-backend-ops perf -o MUL_MAT, m=4096, k=14336, M3 Ultra, time of this
change over master (mean of two interleaved runs each): 0.23 to 0.98 from
the threshold to 8 rows, 0.24 to 0.33 at 9 to 16 rows, and 0.99 to 1.01 at
1 and 512 rows.
---
 ggml/src/ggml-metal/ggml-metal-common.cpp    | 24 ++++++++++++++++++++
 ggml/src/ggml-metal/kernels/mul_mv_mma.metal | 19 ++++++++++++++++
 tests/test-backend-ops.cpp                   |  9 +++++---
 3 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/ggml/src/ggml-metal/ggml-metal-common.cpp b/ggml/src/ggml-metal/ggml-metal-common.cpp
index 8b065f54a..a00f87c99 100644
--- a/ggml/src/ggml-metal/ggml-metal-common.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-common.cpp
@@ -80,6 +80,22 @@ static bool ggml_metal_mul_mv_mma_type_supported(enum ggml_type type) {
         case GGML_TYPE_Q4_K:
         case GGML_TYPE_Q5_K:
         case GGML_TYPE_Q6_K:
+        case GGML_TYPE_BF16:
+        case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q2_0:
+        case GGML_TYPE_MXFP4:
+        case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q3_K:
+        case GGML_TYPE_IQ2_XXS:
+        case GGML_TYPE_IQ2_XS:
+        case GGML_TYPE_IQ2_S:
+        case GGML_TYPE_IQ3_XXS:
+        case GGML_TYPE_IQ3_S:
+        case GGML_TYPE_IQ1_S:
+        case GGML_TYPE_IQ1_M:
+        case GGML_TYPE_IQ4_NL:
+        case GGML_TYPE_IQ4_XS:
+        case GGML_TYPE_TQ2_0:
             return true;
         default:
             return false;
@@ -105,10 +121,18 @@ static int64_t ggml_metal_mul_mv_mma_rows_min(enum ggml_type type) {
     switch (type) {
         case GGML_TYPE_F32:
             return 6;
+        case GGML_TYPE_TQ2_0:
+            return 5;
+        case GGML_TYPE_BF16:
+            return 4;
         case GGML_TYPE_F16:
+        case GGML_TYPE_Q2_0:
         case GGML_TYPE_Q4_K:
         case GGML_TYPE_Q5_0:
         case GGML_TYPE_Q5_1:
+        case GGML_TYPE_Q2_K:
+        case GGML_TYPE_IQ4_NL:
+        case GGML_TYPE_MXFP4:
             return 3;
         default:
             return 2;
diff --git a/ggml/src/ggml-metal/kernels/mul_mv_mma.metal b/ggml/src/ggml-metal/kernels/mul_mv_mma.metal
index fcf981a93..5cf245d52 100644
--- a/ggml/src/ggml-metal/kernels/mul_mv_mma.metal
+++ b/ggml/src/ggml-metal/kernels/mul_mv_mma.metal
@@ -605,4 +605,23 @@ MUL_MV_MMA_GEN("q5_1", block_q5_1, 2,     dequantize_q5_1)
 MUL_MV_MMA_GEN("q4_K", block_q4_K, QK_NL, dequantize_q4_K)
 MUL_MV_MMA_GEN("q6_K", block_q6_K, QK_NL, dequantize_q6_K)
 
+#if defined(GGML_METAL_HAS_BF16)
+MUL_MV_MMA_GEN("bf16",    bfloat4x4,     1,     dequantize_bf16)
+#endif
+MUL_MV_MMA_GEN("q1_0",    block_q1_0,    8,     dequantize_q1_0)
+MUL_MV_MMA_GEN("q2_0",    block_q2_0,    4,     dequantize_q2_0)
+MUL_MV_MMA_GEN("mxfp4",   block_mxfp4,   2,     dequantize_mxfp4)
+MUL_MV_MMA_GEN("q2_K",    block_q2_K,    QK_NL, dequantize_q2_K)
+MUL_MV_MMA_GEN("q3_K",    block_q3_K,    QK_NL, dequantize_q3_K)
+MUL_MV_MMA_GEN("iq2_xxs", block_iq2_xxs, QK_NL, dequantize_iq2_xxs)
+MUL_MV_MMA_GEN("iq2_xs",  block_iq2_xs,  QK_NL, dequantize_iq2_xs)
+MUL_MV_MMA_GEN("iq2_s",   block_iq2_s,   QK_NL, dequantize_iq2_s)
+MUL_MV_MMA_GEN("iq3_xxs", block_iq3_xxs, QK_NL, dequantize_iq3_xxs)
+MUL_MV_MMA_GEN("iq3_s",   block_iq3_s,   QK_NL, dequantize_iq3_s)
+MUL_MV_MMA_GEN("iq1_s",   block_iq1_s,   QK_NL, dequantize_iq1_s)
+MUL_MV_MMA_GEN("iq1_m",   block_iq1_m,   QK_NL, dequantize_iq1_m)
+MUL_MV_MMA_GEN("iq4_nl",  block_iq4_nl,  2,     dequantize_iq4_nl)
+MUL_MV_MMA_GEN("iq4_xs",  block_iq4_xs,  QK_NL, dequantize_iq4_xs)
+MUL_MV_MMA_GEN("tq2_0",   block_tq2_0,   QK_NL, dequantize_tq2_0)
+
 #undef MUL_MV_MMA_GEN
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index efc1777ef..89554d4d2 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -10404,8 +10404,11 @@ static std::vector> make_test_cases_eval() {
             test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, m, n, k, {1, 1}, {1, 1}));
         }
         test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 64, n, 256, {3, 2}, {2, 1}));
-        for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0,
-                                 GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K}) {
+        for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1,
+                                 GGML_TYPE_Q8_0, GGML_TYPE_Q1_0, GGML_TYPE_Q2_0, GGML_TYPE_MXFP4, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K,
+                                 GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_TQ2_0,
+                                 GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XS, GGML_TYPE_IQ2_S, GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ3_S,
+                                 GGML_TYPE_IQ1_S, GGML_TYPE_IQ1_M, GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_XS}) {
             test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 48,  n, 2560, {1, 1}, {1, 1}));
             test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 1000, n, 1024, {1, 1}, {1, 1}));
             test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 3000, n, 512,  {1, 1}, {1, 1}));
@@ -10428,7 +10431,7 @@ static std::vector> make_test_cases_eval() {
             test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 1000, n, 1024, {1, 1}, {1, 1}, {0, 1, 2, 3}, 1280));
         }
         // a src1 with a nonzero mean, for the zero points and mins of quantized src0 types
-        for (ggml_type type_a : {GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q8_0, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K}) {
+        for (ggml_type type_a : {GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q8_0, GGML_TYPE_Q2_K, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K}) {
             test_cases.emplace_back(new test_mul_mat_pos(type_a, GGML_TYPE_F32, 1000, n, 1024, {1, 1}, {1, 1}));
         }
     }

From 448147d42a71d0147c6dc844f99f9ffe7419183a Mon Sep 17 00:00:00 2001
From: Pascal 
Date: Wed, 7 Oct 2026 16:30:00 +0200
Subject: [PATCH 11/40] llama: share the nextn tensor flags between models
 (#30097)

* llama: share the nextn tensor flags between models

Follow-up of the TODO in glm5-next: move the trunk-only and MTP-only
detection that each model copied into a nextn_flags helper of
llama_model_base. It probes the first trunk layer and the first NextN
layer, and adds TENSOR_SKIP when MTP is not loaded. qwen4exp probes
hc_attn_norm since it has no attn_norm.

deepseek4, nemotron-h, qwen35, qwen35moe, qwen3next and qwen4exp now
also accept a trunk-only file, like the other models.

* llama: avoid capturing structured bindings in the nextn flags

Lambdas that capture structured bindings need C++20, and GCC 15
rejects them under -Werror, so the models read the trunk and MTP
flags into plain variables.
---
 src/llama-model.cpp        | 21 +++++++++++++++++++++
 src/llama-model.h          |  7 +++++++
 src/models/bailingmoe3.cpp | 12 +++---------
 src/models/cohere2moe.cpp  | 15 +++------------
 src/models/deepseek2.cpp   | 12 +++---------
 src/models/deepseek32.cpp  | 12 +++---------
 src/models/deepseek4.cpp   |  6 +++---
 src/models/glm-dsa.cpp     | 16 +++-------------
 src/models/glm4-moe.cpp    | 12 +++---------
 src/models/glm5-next.cpp   | 13 +++----------
 src/models/hy-v3.cpp       | 15 +++------------
 src/models/mimo2.cpp       | 12 +++---------
 src/models/nemotron-h.cpp  |  6 +++---
 src/models/qwen35.cpp      |  6 +++---
 src/models/qwen35moe.cpp   |  6 +++---
 src/models/qwen3next.cpp   |  6 +++---
 src/models/qwen4exp.cpp    |  7 +++----
 src/models/step35.cpp      | 15 +++------------
 18 files changed, 76 insertions(+), 123 deletions(-)

diff --git a/src/llama-model.cpp b/src/llama-model.cpp
index 1f3b80b08..9f05a66e4 100644
--- a/src/llama-model.cpp
+++ b/src/llama-model.cpp
@@ -3461,6 +3461,27 @@ void llama_model_base::create_tensor_qkv(llama_layer & layer, int bid,
     }
 }
 
+llama_model_base::nextn_flags_t llama_model_base::nextn_flags(llama_model_loader & ml, llm_tensor trunk_probe) const {
+    int trunk = 0;
+    int mtp   = 0;
+
+    // a file without the first trunk layer is MTP-only, a file without the first NextN layer is trunk-only
+    if (hparams.n_layer_nextn > 0) {
+        if (ml.get_weight(tn(trunk_probe, "weight", 0).str().c_str()) == nullptr) {
+            trunk = TENSOR_NOT_REQUIRED;
+        }
+        if (ml.get_weight(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", hparams.n_layer()).str().c_str()) == nullptr) {
+            mtp = TENSOR_NOT_REQUIRED;
+        }
+    }
+
+    if (!ml.load_mtp) {
+        mtp |= TENSOR_SKIP;
+    }
+
+    return { trunk, mtp };
+}
+
 void llama_model_base::load_swa_pattern(llama_model_loader & ml, uint32_t n_pattern, bool dense_first) {
     if (ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, false)) {
         return;
diff --git a/src/llama-model.h b/src/llama-model.h
index a86cc5693..82124471f 100644
--- a/src/llama-model.h
+++ b/src/llama-model.h
@@ -849,6 +849,13 @@ struct llama_model_base : public llama_model {
                 int64_t n_embd_, int64_t n_embd_q_, int64_t n_embd_k_, int64_t n_embd_v_,
                 int flags);
 
+    // helper: tensor flags for a file that holds the trunk, the NextN layers, or both
+    struct nextn_flags_t {
+        int trunk; // TENSOR_NOT_REQUIRED when the file holds only the NextN layers
+        int mtp;   // TENSOR_NOT_REQUIRED when the file holds only the trunk, TENSOR_SKIP when MTP is not loaded
+    };
+    nextn_flags_t nextn_flags(llama_model_loader & ml, llm_tensor trunk_probe = LLM_TENSOR_ATTN_NORM) const;
+
     // helper: read the SWA pattern as one flag per layer, or as a period expanded by set_swa_pattern
     void load_swa_pattern(llama_model_loader & ml, uint32_t n_pattern, bool dense_first = false);
 
diff --git a/src/models/bailingmoe3.cpp b/src/models/bailingmoe3.cpp
index b04b58449..2458b1c1a 100644
--- a/src/models/bailingmoe3.cpp
+++ b/src/models/bailingmoe3.cpp
@@ -64,15 +64,9 @@ void llama_model_bailingmoe3::load_arch_tensors(llama_model_loader & ml) {
     const int64_t qk_head_dim = hparams.n_embd_head_k_mla();
     const int64_t v_head_dim = hparams.n_embd_head_v_mla();
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
-    int       mtp_flags   = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     for (int il = 0; il < n_layer; ++il) {
         auto & layer = layers[il];
diff --git a/src/models/cohere2moe.cpp b/src/models/cohere2moe.cpp
index 3baf3950c..a379e2f60 100644
--- a/src/models/cohere2moe.cpp
+++ b/src/models/cohere2moe.cpp
@@ -40,18 +40,9 @@ void llama_model_cohere2moe::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_cohere2moe::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP
-    // tensors live in a separate file. Mark MTP tensors NOT_REQUIRED so the
-    // trunk loads cleanly.
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only  ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
 
diff --git a/src/models/deepseek2.cpp b/src/models/deepseek2.cpp
index 47fc2f1ac..f1067a516 100644
--- a/src/models/deepseek2.cpp
+++ b/src/models/deepseek2.cpp
@@ -56,15 +56,9 @@ void llama_model_deepseek2::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
     const int64_t n_expert_shared = hparams.n_expert_shared;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only  ? TENSOR_NOT_REQUIRED : 0;
-    int       mtp_flags   = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     const bool is_mla = hparams.is_mla();
 
diff --git a/src/models/deepseek32.cpp b/src/models/deepseek32.cpp
index ade9b9ed9..76f763c76 100644
--- a/src/models/deepseek32.cpp
+++ b/src/models/deepseek32.cpp
@@ -46,15 +46,9 @@ void llama_model_deepseek32::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_deepseek32::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only   ? TENSOR_NOT_REQUIRED : 0;
-    int       mtp_flags   = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     const bool is_mla = hparams.is_mla();
     if (!is_mla) {
diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp
index 4edf0c2b4..0974b63b3 100644
--- a/src/models/deepseek4.cpp
+++ b/src/models/deepseek4.cpp
@@ -93,9 +93,9 @@ void llama_model_deepseek4::load_arch_tensors(llama_model_loader & ml) {
     const int64_t hc_dim      = hc_mult * n_embd;
     const int64_t hc_mix_dim  = (2 + hc_mult) * hc_mult;
 
-    const bool mtp_only = (n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const int trunk_flags = mtp_only    ? TENSOR_NOT_REQUIRED : 0;
-    const int mtp_flags   = ml.load_mtp ? 0 : TENSOR_SKIP;
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 
diff --git a/src/models/glm-dsa.cpp b/src/models/glm-dsa.cpp
index 32581f840..3e28ef271 100644
--- a/src/models/glm-dsa.cpp
+++ b/src/models/glm-dsa.cpp
@@ -75,19 +75,9 @@ void llama_model_glm_dsa::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
     const int64_t n_expert_shared = hparams.n_expert_shared;
 
-    // MTP-only: the GGUF carries only the NextN/MTP block(s) (user split target/draft).
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP
-    // tensors live in a separate file (or were stripped at conversion). Mark
-    // MTP tensors NOT_REQUIRED so the trunk loads cleanly.
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only   ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     const bool is_mla = hparams.is_mla();
     if (!is_mla) {
diff --git a/src/models/glm4-moe.cpp b/src/models/glm4-moe.cpp
index d65cace1f..4b41b5958 100644
--- a/src/models/glm4-moe.cpp
+++ b/src/models/glm4-moe.cpp
@@ -29,15 +29,9 @@ void llama_model_glm4_moe::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
     const int64_t n_expert_shared = hparams.n_expert_shared;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only  ? TENSOR_NOT_REQUIRED : 0;
-    int       mtp_flags   = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     GGML_ASSERT(hparams.n_expert > 0 && "n_expert must be > 0 for GLM4_MOE MoE layers");
     GGML_ASSERT(hparams.n_expert_used() > 0 && "n_expert_used must be > 0 for GLM4_MOE MoE layers");
diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp
index 4343e6ac6..48b07f1e2 100644
--- a/src/models/glm5-next.cpp
+++ b/src/models/glm5-next.cpp
@@ -64,16 +64,9 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
     const int64_t hc         = hparams.dsv4_hc_mult;
     const int64_t hc_mix_dim = (2 + hc)*hc;
 
-    // TODO: consolidate the trunk_only/mtp_only detection and flags shared with glm4-moe and qwen35moe
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 
diff --git a/src/models/hy-v3.cpp b/src/models/hy-v3.cpp
index bb0e0a2b8..c4c55c3f9 100644
--- a/src/models/hy-v3.cpp
+++ b/src/models/hy-v3.cpp
@@ -22,18 +22,9 @@ void llama_model_hy_v3::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_hy_v3::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP
-    // tensors live in a separate file (e.g. user split target/draft). Mark
-    // MTP tensors NOT_REQUIRED so the trunk loads cleanly.
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only   ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 
diff --git a/src/models/mimo2.cpp b/src/models/mimo2.cpp
index 866959c83..a70350191 100644
--- a/src/models/mimo2.cpp
+++ b/src/models/mimo2.cpp
@@ -25,15 +25,9 @@ void llama_model_mimo2::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_mimo2::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 
diff --git a/src/models/nemotron-h.cpp b/src/models/nemotron-h.cpp
index cd90e24d2..cf78ebfbe 100644
--- a/src/models/nemotron-h.cpp
+++ b/src/models/nemotron-h.cpp
@@ -49,9 +49,9 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_nemotron_h::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only    = hparams.n_layer_nextn > 0 && ml.get_weight("blk.0.attn_norm.weight") == nullptr;
-    const int  trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
-    const int  mtp_flags   = !ml.load_mtp ? TENSOR_SKIP : 0;
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     // mamba2 Mixer SSM params
     // NOTE: int64_t for tensor dimensions
diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp
index e484850da..350cea808 100644
--- a/src/models/qwen35.cpp
+++ b/src/models/qwen35.cpp
@@ -33,9 +33,9 @@ void llama_model_qwen35::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
 
diff --git a/src/models/qwen35moe.cpp b/src/models/qwen35moe.cpp
index 6e28028a8..f9cc8a2c6 100644
--- a/src/models/qwen35moe.cpp
+++ b/src/models/qwen35moe.cpp
@@ -36,9 +36,9 @@ void llama_model_qwen35moe::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
 
diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp
index 2d56c38ca..620161ac7 100644
--- a/src/models/qwen3next.cpp
+++ b/src/models/qwen3next.cpp
@@ -35,9 +35,9 @@ void llama_model_qwen3next::load_arch_tensors(llama_model_loader & ml) {
         throw std::runtime_error(arch_name() + " model cannot have zero experts");
     }
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
 
diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp
index d05f72fc3..88e2e8891 100644
--- a/src/models/qwen4exp.cpp
+++ b/src/models/qwen4exp.cpp
@@ -175,10 +175,9 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) {
     const int64_t hc_dim = hc * n_embd;
     const int64_t hc_lr  = hparams.hc_low_rank;
 
-    // an MTP-only file carries the MTP block, the embeddings and the LM head, but no trunk
-    const bool mtp_only    = n_layer_nextn > 0 && ml.get_weight(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", 0).str().c_str()) == nullptr;
-    const int  trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
-    const int  mtp_flags   = ml.load_mtp ? 0 : TENSOR_SKIP;
+    const auto nf = nextn_flags(ml, LLM_TENSOR_HC_ATTN_NORM);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
 
diff --git a/src/models/step35.cpp b/src/models/step35.cpp
index f4660bf19..bfa80fab3 100644
--- a/src/models/step35.cpp
+++ b/src/models/step35.cpp
@@ -37,18 +37,9 @@ void llama_model_step35::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_step35::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP
-    // tensors live in a separate file (e.g. user split target/draft). Mark
-    // MTP tensors NOT_REQUIRED so the trunk loads cleanly.
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only  ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 

From 18b5f8b1862ebfe0f1c33d2355b81a54d2fec867 Mon Sep 17 00:00:00 2001
From: Aldehir Rojas 
Date: Wed, 7 Oct 2026 11:49:17 -0500
Subject: [PATCH 12/40] server : accumulate generated text and tokens as parse
 input (#29876)

* server : collect text and token input

* common : add a tokenize helper that aligns tokens with bytes

* common : simplify tokenize logic
---
 common/chat.cpp                 | 73 +++++++++++++++++++++++++++++----
 common/chat.h                   | 33 +++++++++++++--
 common/peg-parser.h             |  9 +++-
 tests/test-chat.cpp             |  4 +-
 tools/server/server-context.cpp | 40 +++++++++---------
 tools/server/server-schema.cpp  |  8 +++-
 tools/server/server-task.cpp    | 30 +++++++-------
 tools/server/server-task.h      | 10 ++---
 8 files changed, 149 insertions(+), 58 deletions(-)

diff --git a/common/chat.cpp b/common/chat.cpp
index 1c1533325..2c795c281 100644
--- a/common/chat.cpp
+++ b/common/chat.cpp
@@ -1452,14 +1452,70 @@ common_chat_params common_chat_templates_apply(const struct common_chat_template
                               common_chat_templates_apply_legacy(tmpls, inputs);
 }
 
-common_chat_msg common_chat_parse(const std::string &               input,
+void common_chat_input::append(const std::string & piece, llama_token token) {
+    if (piece.empty()) {
+        return;
+    }
+    tokens.push_back(token);
+    tokens.resize(tokens.size() + piece.size() - 1, LLAMA_TOKEN_NULL);
+    text += piece;
+}
+
+void common_chat_input::append(const common_chat_input & chunk) {
+    tokens.insert(tokens.end(), chunk.tokens.begin(), chunk.tokens.end());
+    text += chunk.text;
+}
+
+void common_chat_input::truncate(size_t pos) {
+    if (pos < text.size()) {
+        text.erase(pos);
+        tokens.resize(pos);
+    }
+}
+
+common_chat_input common_chat_input::substr(size_t pos, size_t n) const {
+    common_chat_input out;
+    out.text = text.substr(pos, n);
+    out.tokens.assign(tokens.begin() + pos, tokens.begin() + pos + out.size());
+    return out;
+}
+
+void common_chat_input::prepend(const std::string & prefix) {
+    tokens.insert(tokens.begin(), prefix.size(), LLAMA_TOKEN_NULL);
+    text = prefix + text;
+}
+
+void common_chat_input::prepend(const common_chat_input & prefix) {
+    tokens.insert(tokens.begin(), prefix.tokens.begin(), prefix.tokens.end());
+    text = prefix.text + text;
+}
+
+common_chat_input common_chat_input_tokenize(const llama_vocab * vocab, const std::string & text) {
+    common_chat_input input;
+    auto tokens = common_tokenize(vocab, text, false, true);
+    for (size_t i = 0; i < tokens.size(); i++) {
+        std::string piece = common_token_to_piece(vocab, tokens[i], true);
+        if (i == 0 && std::isspace(piece[0]) && !std::isspace(text[0])) {
+            // Some tokenizers will add a space before the first special token, need to exclude
+            continue;
+        }
+        input.append(piece, tokens[i]);
+    }
+    if (input.text != text) {
+        // the pieces do not give back the same text, keep the text without tokens
+        return common_chat_input(text);
+    }
+    return input;
+}
+
+common_chat_msg common_chat_parse(const common_chat_input &         input,
                                   bool                              is_partial,
                                   const common_chat_parser_params & params) {
     return common_chat_peg_parse(params.parser, input, is_partial, params);
 }
 
 common_chat_msg common_chat_peg_parse(const common_peg_arena &          src_parser,
-                                      const std::string &               input,
+                                      const common_chat_input &         input,
                                       bool                              is_partial,
                                       const common_chat_parser_params & params) {
     const common_peg_arena & parser = src_parser.empty() ?
@@ -1470,18 +1526,17 @@ common_chat_msg common_chat_peg_parse(const common_peg_arena &          src_pars
         LOG_DBG("No parser definition detected, assuming pure content parser.");
     }
 
-    const std::string effective_input = params.generation_prompt.empty()
-        ? input
-        : params.generation_prompt + input;
+    common_chat_input effective_input = input;
+    effective_input.prepend(params.generation_prompt);
 
-    //LOG_DBG("Parsing PEG input with format %s: %s\n", common_chat_format_name(params.format), effective_input.c_str());
+    //LOG_DBG("Parsing PEG input with format %s: %s\n", common_chat_format_name(params.format), effective_input.text.c_str());
 
     common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_LENIENT;
     if (params.debug) {
         flags |= COMMON_PEG_PARSE_FLAG_DEBUG;
     }
 
-    common_peg_parse_context ctx(effective_input, flags);
+    common_peg_parse_context ctx(std::move(effective_input.text), std::move(effective_input.tokens), flags);
     auto result = parser.parse(ctx);
 
     if (result.fail()) {
@@ -1507,8 +1562,8 @@ common_chat_msg common_chat_peg_parse(const common_peg_arena &          src_pars
             }
             return msg;
         }
-        LOG_WRN("%s: unparsed %s output: %s\n", __func__, common_chat_format_name(params.format), effective_input.substr(result.end).c_str());
-        LOG_DBG("%s: full %s output triggering error:\n=== BEGIN ===\n%s\n=== END ===\n", __func__, common_chat_format_name(params.format), effective_input.c_str());
+        LOG_WRN("%s: unparsed %s output: %s\n", __func__, common_chat_format_name(params.format), ctx.input.substr(result.end).c_str());
+        LOG_DBG("%s: full %s output triggering error:\n=== BEGIN ===\n%s\n=== END ===\n", __func__, common_chat_format_name(params.format), ctx.input.c_str());
         throw std::runtime_error(std::string("The model produced output that does not match the expected ") + common_chat_format_name(params.format) + " format");
     }
 
diff --git a/common/chat.h b/common/chat.h
index 0e1423a5a..4509a04cc 100644
--- a/common/chat.h
+++ b/common/chat.h
@@ -282,6 +282,31 @@ struct common_chat_params {
     common_chat_msg_delimiters          message_delimiters;
 };
 
+struct common_chat_input {
+    std::string              text;
+    std::vector tokens;
+
+    common_chat_input() = default;
+
+    // plain text, with no tokens
+    explicit common_chat_input(std::string text) : text(std::move(text)), tokens(this->text.size(), LLAMA_TOKEN_NULL) {}
+
+    size_t size() const { return text.size(); }
+    bool empty() const { return text.empty(); }
+
+    void append(const std::string & piece, llama_token token);
+    void append(const common_chat_input & chunk);
+
+    void prepend(const std::string & prefix);
+    void prepend(const common_chat_input & prefix);
+
+    void truncate(size_t pos);
+
+    common_chat_input substr(size_t pos, size_t n = std::string::npos) const;
+};
+
+common_chat_input common_chat_input_tokenize(const llama_vocab * vocab, const std::string & text);
+
 // per-message parsing syntax
 // should be derived from common_chat_params
 struct common_chat_parser_params {
@@ -289,7 +314,7 @@ struct common_chat_parser_params {
     common_reasoning_format reasoning_format     = COMMON_REASONING_FORMAT_NONE; // TODO: refactor this to "bool parse_reasoning"
     // Whether reasoning_content should be inlined in the content (e.g. for reasoning_format=deepseek in stream mode)
     bool                    reasoning_in_content = false;
-    std::string             generation_prompt;
+    common_chat_input       generation_prompt;
     bool                    parse_tool_calls     = true;
     bool                    is_continuation      = false;
     bool                    echo                 = false;  // Include assistant prefilled msg in output
@@ -298,7 +323,7 @@ struct common_chat_parser_params {
     common_chat_parser_params() = default;
     common_chat_parser_params(const common_chat_params & chat_params) {
         format  = chat_params.format;
-        generation_prompt = chat_params.generation_prompt;
+        generation_prompt = common_chat_input(chat_params.generation_prompt);
     }
 };
 
@@ -337,8 +362,8 @@ std::string common_chat_format_example(const struct common_chat_templates *
                                        const std::map & chat_template_kwargs);
 
 const char *    common_chat_format_name(common_chat_format format);
-common_chat_msg common_chat_parse(const std::string & input, bool is_partial, const common_chat_parser_params & params);
-common_chat_msg common_chat_peg_parse(const common_peg_arena & src_parser, const std::string & input, bool is_partial, const common_chat_parser_params & params);
+common_chat_msg common_chat_parse(const common_chat_input & input, bool is_partial, const common_chat_parser_params & params);
+common_chat_msg common_chat_peg_parse(const common_peg_arena & src_parser, const common_chat_input & input, bool is_partial, const common_chat_parser_params & params);
 
 // used by arg and server
 const char *            common_reasoning_format_name(common_reasoning_format format);
diff --git a/common/peg-parser.h b/common/peg-parser.h
index 888325945..816d007b0 100644
--- a/common/peg-parser.h
+++ b/common/peg-parser.h
@@ -2,6 +2,7 @@
 
 #include "json-schema.h"
 #include "json.h"
+#include "llama.h"
 
 #include 
 #include 
@@ -182,7 +183,8 @@ inline common_peg_parse_flags operator~(common_peg_parse_flags a) {
 }
 
 struct common_peg_parse_context {
-    std::string input;
+    std::string input;               // [h,   e,  l,  l,  o,  _,  w,  o,  r,  l,  d]
+    std::vector tokens; // [id, -1, -1, -1, -1, id, -1, -1, -1, -1, -1]
     common_peg_parse_flags flags;
     common_peg_ast_arena ast;
 
@@ -194,6 +196,11 @@ struct common_peg_parse_context {
     common_peg_parse_context(const std::string & input, common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE)
         : input(input), flags(flags), parse_depth(0) {}
 
+    common_peg_parse_context(std::string input, std::vector tokens, common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE)
+        : input(std::move(input)), tokens(std::move(tokens)), flags(flags), parse_depth(0) {
+        GGML_ASSERT(this->tokens.empty() || this->tokens.size() == this->input.size());
+    }
+
     bool is_lenient() const { return flags & COMMON_PEG_PARSE_FLAG_LENIENT; }
     bool is_debug() const { return flags & COMMON_PEG_PARSE_FLAG_DEBUG; }
 };
diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp
index 9f96e1084..e31c093b0 100644
--- a/tests/test-chat.cpp
+++ b/tests/test-chat.cpp
@@ -1105,7 +1105,7 @@ struct make_peg_parser {
     common_chat_msg parse(const std::string & msg, bool is_partial) const {
         common_chat_parser_params parser_params(params_);
         parser_params.debug = detailed_debug_;
-        return common_chat_peg_parse(arena_, msg, is_partial, parser_params);
+        return common_chat_peg_parse(arena_, common_chat_input(msg), is_partial, parser_params);
     }
 };
 
@@ -4643,7 +4643,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             bool got_out_of_range = false;
             std::string error_msg;
             try {
-                common_chat_peg_parse(arena, bad_input, /*is_partial=*/false, pp);
+                common_chat_peg_parse(arena, common_chat_input(bad_input), /*is_partial=*/false, pp);
             } catch (const std::out_of_range & e) {
                 got_out_of_range = true;
                 error_msg = e.what();
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index d1af984e6..23835e185 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -267,7 +267,7 @@ struct server_slot {
     bool spec_is_replay = false;
     std::mt19937 spec_synth_rng;
 
-    // TODO: move members that belong to the task (such as `generated_text`, `has_new_line`) to task_results_state
+    // TODO: move members that belong to the task (such as `generated`, `has_new_line`) to task_results_state
     //       see https://github.com/ggml-org/llama.cpp/pull/18283#issuecomment-3710175837
     std::unique_ptr task;
     std::unique_ptr task_prev; // used for debugging
@@ -285,7 +285,7 @@ struct server_slot {
 
     size_t last_nl_pos = 0;
 
-    std::string  generated_text;
+    common_chat_input generated;
     std::string  debug_generated_text;
     llama_tokens generated_tokens;
     size_t n_sent_text = 0; // number of sent text character (i.e. handle partial UTF-8 on streaming)
@@ -381,7 +381,7 @@ struct server_slot {
         spec_is_replay = false;
 
         last_nl_pos    = 0;
-        generated_text = "";
+        generated      = {};
         has_new_line   = false;
         truncated      = false;
         stop           = STOP_TYPE_NONE;
@@ -750,7 +750,7 @@ struct server_slot {
 
             if (!only_metrics) {
                 res["prompt"] = ptask->tokens.detokenize(ctx_tgt, true);
-                res["generated"] = generated_text.empty() ? debug_generated_text : generated_text;
+                res["generated"] = generated.empty() ? debug_generated_text : generated.text;
             }
         }
 
@@ -1968,28 +1968,26 @@ private:
         const std::string token_str = result.text_to_send;
         slot.sampled = result.tok;
 
-        slot.generated_text += token_str;
+        slot.generated.append(token_str, result.tok);
         if (slot.task->params.return_tokens) {
             slot.generated_tokens.push_back(result.tok);
         }
         slot.has_next_token = true;
 
         // check if there is incomplete UTF-8 character at the end
-        bool incomplete = validate_utf8(slot.generated_text) < slot.generated_text.size();
+        bool incomplete = validate_utf8(slot.generated.text) < slot.generated.text.size();
 
         // search stop word and delete it
         if (!incomplete) {
-            size_t pos = std::min(slot.n_sent_text, slot.generated_text.size());
+            size_t pos = std::min(slot.n_sent_text, slot.generated.text.size());
 
-            const std::string str_test = slot.generated_text.substr(pos);
+            const std::string str_test = slot.generated.text.substr(pos);
             bool send_text = true;
 
             size_t stop_pos = slot.find_stopping_strings(str_test, token_str.size(), true);
             if (stop_pos != std::string::npos) {
-                slot.generated_text.erase(
-                    slot.generated_text.begin() + pos + stop_pos,
-                    slot.generated_text.end());
-                pos = std::min(slot.n_sent_text, slot.generated_text.size());
+                slot.generated.truncate(pos + stop_pos);
+                pos = std::min(slot.n_sent_text, slot.generated.text.size());
             } else if (slot.has_next_token && !llama_vocab_is_eog(vocab, result.tok) ) {
                 stop_pos = slot.find_stopping_strings(str_test, token_str.size(), false);
                 send_text = stop_pos == std::string::npos;
@@ -1998,7 +1996,7 @@ private:
             // check if there is any token to predict
             if (send_text) {
                 // no send the stop word in the response
-                result.text_to_send = slot.generated_text.substr(pos, std::string::npos);
+                result.text_to_send = slot.generated.text.substr(pos, std::string::npos);
                 slot.n_sent_text += result.text_to_send.size();
                 // add the token to slot queue and cache
             } else {
@@ -2042,17 +2040,17 @@ private:
                     size_t pos = slot.last_nl_pos;
 
                     int n_indent = 0;
-                    while (pos < slot.generated_text.size() && (slot.generated_text[pos] == ' ' || slot.generated_text[pos] == '\t')) {
+                    while (pos < slot.generated.text.size() && (slot.generated.text[pos] == ' ' || slot.generated.text[pos] == '\t')) {
                         n_indent++;
                         pos++;
                     }
 
-                    if (pos < slot.generated_text.size() && n_indent < slot.task->params.n_indent) {
+                    if (pos < slot.generated.text.size() && n_indent < slot.task->params.n_indent) {
                         slot.stop           = STOP_TYPE_LIMIT;
                         slot.has_next_token = false;
 
                         // cut the last line
-                        slot.generated_text.erase(pos, std::string::npos);
+                        slot.generated.truncate(pos);
 
                         SLT_DBG(slot, "stopped by indentation limit, n_gen = %d, n_indent = %d\n", (int) slot.stats.n_gen, n_indent);
                     }
@@ -2060,7 +2058,7 @@ private:
 
                 // find the next new line
                 {
-                    const size_t pos = slot.generated_text.find('\n', slot.last_nl_pos);
+                    const size_t pos = slot.generated.text.find('\n', slot.last_nl_pos);
 
                     if (pos != std::string::npos) {
                         slot.last_nl_pos = pos + 1;
@@ -2192,7 +2190,7 @@ private:
         if (is_begin) {
             res->is_begin = true;
         } else {
-            res->content = tkn.text_to_send;
+            res->content = slot.generated.substr(slot.n_sent_text - tkn.text_to_send.size(), tkn.text_to_send.size());
             res->tokens.assign(1, tkn.tok);
         }
 
@@ -2229,15 +2227,15 @@ private:
 
         // keep copy of last generated text for debugging purposes
         if (slots_debug) {
-            slot.debug_generated_text = slot.generated_text;
+            slot.debug_generated_text = slot.generated.text;
         }
 
         // in stream mode, content and tokens are already in last partial chunk
         if (slot.task->params.stream) {
-            res->content     = "";
+            res->content     = {};
             res->tokens      = llama_tokens{};
         } else {
-            res->content     = std::move(slot.generated_text);
+            res->content     = std::move(slot.generated);
             res->tokens      = std::move(slot.generated_tokens);
         }
         res->stats           = slot.stats;
diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp
index 27ecafb7a..c7d0bc4de 100644
--- a/tools/server/server-schema.cpp
+++ b/tools/server/server-schema.cpp
@@ -311,8 +311,14 @@ std::vector> make_llama_cmpl_schema(const common_params &
         ->set_desc("Generation prompt appended to the chat template output")
         ->set_handler([&](field_eval_context & ctx, const json & data) {
             std::string s = data.at("generation_prompt").get();
-            ctx.params.chat_parser_params.generation_prompt = s;
             ctx.params.sampling.generation_prompt = s;
+
+            if (ctx.vocab == nullptr) {
+                ctx.params.chat_parser_params.generation_prompt = common_chat_input(s);
+                return;
+            }
+
+            ctx.params.chat_parser_params.generation_prompt = common_chat_input_tokenize(ctx.vocab, s);
         }));
 
     add((new field_bool("parse_tool_calls", params.chat_parser_params.parse_tool_calls))
diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp
index a5c33c056..21c1fe4e1 100644
--- a/tools/server/server-task.cpp
+++ b/tools/server/server-task.cpp
@@ -76,7 +76,7 @@ json task_params::to_json(bool only_metrics) const {
             {"chat_format",               common_chat_format_name(chat_parser_params.format)},
             {"reasoning_format",          common_reasoning_format_name(chat_parser_params.reasoning_format)},
             {"reasoning_in_content",      chat_parser_params.reasoning_in_content},
-            {"generation_prompt",         chat_parser_params.generation_prompt},
+            {"generation_prompt",         chat_parser_params.generation_prompt.text},
             {"samplers",                  samplers},
             {"speculative.types",         common_speculative_type_name_str(speculative.types)},
             {"timings_per_token",         timings_per_token},
@@ -135,7 +135,7 @@ json task_params::to_json(bool only_metrics) const {
         {"chat_format",               common_chat_format_name(chat_parser_params.format)},
         {"reasoning_format",          common_reasoning_format_name(chat_parser_params.reasoning_format)},
         {"reasoning_in_content",      chat_parser_params.reasoning_in_content},
-        {"generation_prompt",         chat_parser_params.generation_prompt},
+        {"generation_prompt",         chat_parser_params.generation_prompt.text},
         {"samplers",                  samplers},
         {"speculative.types",         common_speculative_type_name_str(speculative.types)},
         {"timings_per_token",         timings_per_token},
@@ -155,20 +155,20 @@ task_result_state::task_result_state(const common_chat_parser_params & chat_pars
     , oai_resp_message_id("msg_" + random_string()) {
     if (chat_parser_params.is_continuation && !chat_parser_params.echo) {
         // initialize chat_msg to avoid emitting a delta containing the assistant prefill
-        chat_msg = common_chat_parse("", true, chat_parser_params);
+        chat_msg = common_chat_parse(generated_input, true, chat_parser_params);
     }
 }
 
 common_chat_msg task_result_state::update_chat_msg(
-        const std::string & text_added,
+        const common_chat_input & added,
         bool is_partial,
         std::vector & diffs,
         bool filter_tool_calls) {
-    generated_text += text_added;
+    generated_input.append(added);
     auto msg_prv_copy = chat_msg;
-    //SRV_DBG("Parsing chat message: %s\n", generated_text.c_str());
+    //SRV_DBG("Parsing chat message: %s\n", generated_input.text.c_str());
     auto new_msg = common_chat_parse(
-        generated_text,
+        generated_input,
         is_partial,
         chat_parser_params);
     if (!new_msg.empty()) {
@@ -340,7 +340,7 @@ json server_task_result_cmpl_final::to_json() {
 json server_task_result_cmpl_final::to_json_non_oaicompat() {
     json res = json {
         {"index",               index},
-        {"content",             content},
+        {"content",             content.text},
         {"tokens",              tokens},
         {"id_slot",             id_slot},
         {"stop",                true},
@@ -386,7 +386,7 @@ json server_task_result_cmpl_final::to_json_oaicompat() {
     json res = json {
         {"choices",            json::array({
             json{
-                {"text",          content},
+                {"text",          content.text},
                 {"index",         index},
                 {"logprobs",      logprobs},
                 {"finish_reason", finish_reason},
@@ -418,7 +418,7 @@ json server_task_result_cmpl_final::to_json_oaicompat_chat() {
         msg = oaicompat_msg;
     } else {
         msg.role = "assistant";
-        msg.content = content;
+        msg.content = content.text;
     }
     if (stop == STOP_TYPE_WORD || stop == STOP_TYPE_EOS) {
         finish_reason = msg.tool_calls.empty() ? "stop" : "tool_calls";
@@ -531,7 +531,7 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp() {
         msg = oaicompat_msg;
     } else {
         msg.role = "assistant";
-        msg.content = content;
+        msg.content = content.text;
     }
 
     std::vector output;
@@ -741,7 +741,7 @@ json server_task_result_cmpl_final::to_json_anthropic() {
         msg = oaicompat_msg;
     } else {
         msg.role = "assistant";
-        msg.content = content;
+        msg.content = content.text;
     }
 
     // thinking block comes first (Anthropic extended thinking format)
@@ -1050,7 +1050,7 @@ json server_task_result_cmpl_partial::to_json_non_oaicompat() {
     // non-OAI-compat JSON
     json res = json {
         {"index",            index},
-        {"content",          content},
+        {"content",          content.text},
         {"tokens",           tokens},
         {"stop",             false},
         {"id_slot",          id_slot},
@@ -1081,7 +1081,7 @@ json server_task_result_cmpl_partial::to_json_oaicompat() {
     json res = json {
         {"choices",            json::array({
             json{
-                {"text",          content},
+                {"text",          content.text},
                 {"index",         index},
                 {"logprobs",      logprobs},
                 {"finish_reason", nullptr},
@@ -1316,7 +1316,7 @@ json server_task_result_cmpl_partial::to_json_oaicompat_resp() {
 json server_task_result_cmpl_partial::to_json_oaicompat_asr() {
     json event = json {
         {"type", "transcript.text.delta"},
-        {"delta", content},
+        {"delta", content.text},
     };
     return event;
 }
diff --git a/tools/server/server-task.h b/tools/server/server-task.h
index 752595b0f..552c03b07 100644
--- a/tools/server/server-task.h
+++ b/tools/server/server-task.h
@@ -108,7 +108,7 @@ struct task_result_state {
     std::vector diffs;
     common_chat_parser_params chat_parser_params;
     common_chat_msg chat_msg;
-    std::string generated_text; // append new chunks of generated text here
+    common_chat_input generated_input; // append new chunks of generated text here
     std::vector generated_tool_call_ids;
     std::unordered_set sent_tool_call_names;
 
@@ -128,7 +128,7 @@ struct task_result_state {
 
     // parse partial tool calls and update the internal state
     common_chat_msg update_chat_msg(
-        const std::string & text_added,
+        const common_chat_input & added,
         bool is_partial,
         std::vector & diffs,
         bool filter_tool_calls = false);
@@ -350,7 +350,7 @@ struct completion_token_output {
 };
 
 struct server_task_result_cmpl_final : server_task_result {
-    std::string content;
+    common_chat_input content;
     llama_tokens tokens;
 
     bool stream;
@@ -425,8 +425,8 @@ struct server_task_result_cmpl_final : server_task_result {
 };
 
 struct server_task_result_cmpl_partial : server_task_result {
-    std::string  content;
-    llama_tokens tokens;
+    common_chat_input content;
+    llama_tokens      tokens;
 
     int32_t n_decoded;
     int32_t n_prompt_tokens;

From 42c787e8c191d49c01c757f4b7029d69ec0bf2c8 Mon Sep 17 00:00:00 2001
From: Shobhit 
Date: Thu, 8 Oct 2026 01:53:43 +0800
Subject: [PATCH 13/40] cuda: update uncoalesced memory reads in pool2d
 (#29425)

* cuda: update uncoalesced memory reads in pool2d

* cuda: Added perf and __restrict__ in pool2d

* Update POOL2D_WARP_KERNEL_MIN_WINDOW macro

Co-authored-by: Georgi Gerganov 

* cuda: Fix compiler bugs

---------

Co-authored-by: Georgi Gerganov 
---
 ggml/src/ggml-cuda/pool2d.cu  | 66 +++++++++++++++++++++++++++++++++++
 ggml/src/ggml-cuda/pool2d.cuh |  1 +
 tests/test-backend-ops.cpp    | 15 ++++++++
 3 files changed, 82 insertions(+)

diff --git a/ggml/src/ggml-cuda/pool2d.cu b/ggml/src/ggml-cuda/pool2d.cu
index c6d51e4d6..83a49753c 100644
--- a/ggml/src/ggml-cuda/pool2d.cu
+++ b/ggml/src/ggml-cuda/pool2d.cu
@@ -50,6 +50,65 @@ static  __global__ void pool2d_nchw_kernel(
     o_ptr[cur_oh * ow + cur_ow] = res;
 }
 
+template 
+static __global__ void pool2d_nchw_kernel_warp(
+        const int ih, const int iw, const int oh, const int ow,
+        const int kh, const int kw, const int sh, const int sw,
+        const int ph, const int pw, const int parallel_elements,
+        const Ti * __restrict__ src, To * __restrict__ dst, const enum ggml_op_pool op) {
+    const int warp_id = (threadIdx.x + blockIdx.x * blockDim.x) / WARP_SIZE;
+    const int lane     = threadIdx.x % WARP_SIZE;
+    if (warp_id >= parallel_elements) {
+        return;
+    }
+
+    const int I_HW = ih * iw;
+    const int O_HW = oh * ow;
+    const int nc     = warp_id / O_HW;
+    const int cur_oh = warp_id % O_HW / ow;
+    const int cur_ow = warp_id % O_HW % ow;
+    const Ti* i_ptr = src + nc * I_HW;
+
+    const int start_h = cur_oh * sh - ph;
+    const int bh = max(0, start_h);
+    const int eh = min(ih, start_h + kh);
+    const int start_w = cur_ow * sw - pw;
+    const int bw = max(0, start_w);
+    const int ew = min(iw, start_w + kw);
+
+    const int win_w     = ew - bw;
+    const int win_elems = (eh - bh) * win_w;
+    const To scale = 1. / (kh * kw);
+
+    To res;
+    switch (op) {
+        case GGML_OP_POOL_AVG: res = 0; break;
+        case GGML_OP_POOL_MAX: res = -FLT_MAX; break;
+        default: res = 0; assert(false);
+    }
+
+    for (int t = lane; t < win_elems; t += WARP_SIZE) {
+        const int i = bh + t / win_w;
+        const int j = bw + t % win_w;
+        const Ti cur = i_ptr[i * iw + j];
+        switch (op) {
+            case GGML_OP_POOL_AVG: res += cur * scale; break;
+            case GGML_OP_POOL_MAX: res = max(res, (To)cur); break;
+            default: break;
+        }
+    }
+
+#pragma unroll
+    for (int offset = WARP_SIZE/2; offset > 0; offset >>= 1) {
+        const To other = __shfl_xor_sync(0xFFFFFFFF, res, offset, WARP_SIZE);
+        res = (op == GGML_OP_POOL_MAX) ? max(res, other) : res + other;
+    }
+
+    if (lane == 0) {
+        dst[nc * O_HW + cur_oh * ow + cur_ow] = res;
+    }
+}
+
 static void pool2d_nchw_kernel_f32_f32_cuda(
         const int ih, const int iw, const int oh, const int ow,
         const int kh, const int kw, const int sh, const int sw,
@@ -57,6 +116,13 @@ static void pool2d_nchw_kernel_f32_f32_cuda(
         const float * src, float * dst, const enum ggml_op_pool op,
         cudaStream_t stream) {
 
+    if (kh * kw >= CUDA_POOL2D_WARP_KERNEL_MIN_WINDOW) {
+        const int warps_per_block = CUDA_POOL2D_BLOCK_SIZE / WARP_SIZE;
+        const int num_blocks = (parallel_elements + warps_per_block - 1) / warps_per_block;
+        pool2d_nchw_kernel_warp<<>>(ih, iw, oh, ow, kh, kw, sh, sw, ph, pw, parallel_elements, src, dst, op);
+        return;
+    }
+
     const int num_blocks = (parallel_elements + CUDA_POOL2D_BLOCK_SIZE - 1) / CUDA_POOL2D_BLOCK_SIZE;
     dim3 block_nums(num_blocks);
     pool2d_nchw_kernel<<>>(ih, iw, oh, ow, kh, kw, sh, sw, ph, pw, parallel_elements, src, dst, op);
diff --git a/ggml/src/ggml-cuda/pool2d.cuh b/ggml/src/ggml-cuda/pool2d.cuh
index 7841292bc..e2c7fb441 100644
--- a/ggml/src/ggml-cuda/pool2d.cuh
+++ b/ggml/src/ggml-cuda/pool2d.cuh
@@ -1,5 +1,6 @@
 #include "common.cuh"
 
 #define CUDA_POOL2D_BLOCK_SIZE 256
+#define CUDA_POOL2D_WARP_KERNEL_MIN_WINDOW 32
 
 void ggml_cuda_op_pool2d(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 89554d4d2..cfe5b73e7 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -9558,6 +9558,12 @@ static std::vector> make_test_cases_eval() {
                     }
                 }
             }
+            // MSA indexer block-max shape: wide 1-D window along ne0, no padding, non-divisible width
+            for (int64_t iw : {4096, 4160}) {
+                for (int blk : {32, 64}) {
+                    test_cases.emplace_back(new test_pool2d(pool_type, type_input, {iw, 2, 1, 1}, blk, 1, blk, 1, 0, 0));
+                }
+            }
         }
     }
 
@@ -11652,6 +11658,15 @@ static std::vector> make_test_cases_perf() {
             GGML_TYPE_F32, {n_kv, 512, 64, 1}, false, {2, 1, 0, 3}));
     }
 
+    // POOL_2D max over MiniMax-M3 indexer block scores:
+    // sc is [n_ps, n_head=4, n_tokens] and gets ggml_pool_2d(.., MAX, blk=128, 1, blk=128, 1, 0, 0).
+    for (int64_t n_ps : { 8192, 32768 }) {
+        for (int64_t n_tokens : { 1, 512 }) {
+            test_cases.emplace_back(new test_pool2d(
+                GGML_OP_POOL_MAX, GGML_TYPE_F32, {n_ps, 4, n_tokens, 1}, 128, 1, 128, 1, 0, 0));
+        }
+    }
+
     // LEAKY_RELU at FFN activation width, for direct comparison with RELU
     for (int64_t n_tokens : {512, 2048}) {
         test_cases.emplace_back(new test_leaky_relu(GGML_TYPE_F32, { 17408, n_tokens, 1, 1 }, 0.1f));

From d6cf9acb25e6c657d6c6a4645ab7b69fce119995 Mon Sep 17 00:00:00 2001
From: Aman Gupta 
Date: Wed, 7 Oct 2026 23:37:45 +0530
Subject: [PATCH 14/40] llama : add a GPU cache for MoE experts kept in host
 memory (#29887)

* llama : add a GPU cache for MoE experts kept in host memory

Assisted-by: Claude

* use llama_moe_cache_ptr
---
 common/arg.cpp             |  10 +
 common/common.cpp          |   2 +
 common/common.h            |   2 +
 common/speculative.cpp     |   3 +
 include/llama.h            |   2 +
 src/CMakeLists.txt         |   1 +
 src/llama-context.cpp      |  60 +++-
 src/llama-context.h        |   5 +-
 src/llama-cparams.h        |   2 +
 src/llama-graph.cpp        |  60 +++-
 src/llama-graph.h          |  17 +-
 src/llama-moe-cache.cpp    | 550 +++++++++++++++++++++++++++++++++++++
 src/llama-moe-cache.h      |  39 +++
 tests/test-llama-archs.cpp |  26 +-
 14 files changed, 761 insertions(+), 18 deletions(-)
 create mode 100644 src/llama-moe-cache.cpp
 create mode 100644 src/llama-moe-cache.h

diff --git a/common/arg.cpp b/common/arg.cpp
index 14d82f695..182f76f45 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -2776,6 +2776,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             llm_add_n_cpu_ffn_overrides(value, LLM_FFN_EXPS_REGEX, params.tensor_buft_overrides);
         }
     ).set_env("LLAMA_ARG_N_CPU_MOE"));
+    add_opt(common_arg(
+        {"--moe-cache-mib"}, "N",
+        "GPU cache size in MiB for the MoE experts kept in the CPU (default: 0, disabled)",
+        [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("invalid value");
+            }
+            params.moe_cache_size = (size_t) value*1024*1024;
+        }
+    ).set_env("LLAMA_ARG_MOE_CACHE_MIB"));
     add_opt(common_arg(
         {"-ncffn", "--n-cpu-ffn"}, "N",
         "keep the dense FFN weights of the first N layers in the CPU\n"
diff --git a/common/common.cpp b/common/common.cpp
index e36f8ab50..768b2e9a5 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1722,6 +1722,8 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     cparams.type_k = params.cache_type_k;
     cparams.type_v = params.cache_type_v;
 
+    cparams.moe_cache_size = params.moe_cache_size;
+
     return cparams;
 }
 
diff --git a/common/common.h b/common/common.h
index 3e3eff379..0f912d901 100644
--- a/common/common.h
+++ b/common/common.h
@@ -593,6 +593,8 @@ struct common_params {
     ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K
     ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V
 
+    size_t moe_cache_size = 0; // GPU cache size in bytes for the MoE experts kept in the CPU
+
     common_conversation_mode conversation_mode = COMMON_CONVERSATION_MODE_AUTO;
 
     // multimodal models (see tools/mtmd)
diff --git a/common/speculative.cpp b/common/speculative.cpp
index a7f095a12..d9ddf44d2 100644
--- a/common/speculative.cpp
+++ b/common/speculative.cpp
@@ -2561,6 +2561,9 @@ common_params common_base_params_to_speculative(const common_params & params) {
     result.n_outputs_max = params.n_parallel;
     result.n_outputs_max_per_seq = 1;
 
+    // the MoE cache is only used by the target context
+    result.moe_cache_size = 0;
+
     // dflash/dspark decode the whole noise block in a single pass and sample every block position on the backend
     // TODO: refactor such properties to be announced by the speculative types
     //       something like `struct common_speculative_type_props common_speculative_type_get_props(...);`
diff --git a/include/llama.h b/include/llama.h
index 260247e82..77f527d5b 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -396,6 +396,8 @@ extern "C" {
         enum ggml_type type_k; // data type for K cache [EXPERIMENTAL]
         enum ggml_type type_v; // data type for V cache [EXPERIMENTAL]
 
+        size_t moe_cache_size; // device cache in bytes for the experts kept in host memory, 0 = disabled [EXPERIMENTAL]
+
         // Abort callback
         // if it returns true, execution of llama_decode() will be aborted
         // currently works only with CPU execution
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index afdaddc79..97250c2e8 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -28,6 +28,7 @@ set(LLAMA_CORE_SOURCES
     llama-kv-cache-msa.cpp
     llama-kv-cache-dsv4.cpp
     llama-memory.cpp
+    llama-moe-cache.cpp
     llama-memory-hybrid.cpp
     llama-memory-hybrid-iswa.cpp
     llama-memory-hybrid-idx.cpp
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index ff2ea461c..5f861f96c 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -9,6 +9,7 @@
 #include "llama-memory.h"
 #include "llama-mmap.h"
 #include "llama-model.h"
+#include "llama-moe-cache.h"
 #include "llama-ext.h"
 #include "llama-sampler.h"
 #include "llama.h"
@@ -272,8 +273,9 @@ llama_context::llama_context(
         }
     }
 
-    cparams.op_offload = params.op_offload;
-    cparams.kv_unified = params.kv_unified;
+    cparams.op_offload     = params.op_offload;
+    cparams.kv_unified     = params.kv_unified;
+    cparams.moe_cache_size = params.moe_cache_size;
 
     // initialized later
     cparams.pipeline_parallel = false;
@@ -462,6 +464,22 @@ llama_context::llama_context(
             LLAMA_LOG_INFO("%s: pipeline parallelism enabled\n", __func__);
         }
 
+        if (cparams.moe_cache_size > 0) {
+            if (cparams.pipeline_parallel || model.n_devices() > 1) {
+                throw std::runtime_error("MoE cache does not support multiple devices");
+            }
+            for (size_t i = 0; i < backend_ptrs.size(); ++i) {
+                const auto type = ggml_backend_dev_type(ggml_backend_get_device(backend_ptrs[i]));
+                if (type == GGML_BACKEND_DEVICE_TYPE_GPU || type == GGML_BACKEND_DEVICE_TYPE_IGPU) {
+                    moe_cache = std::make_unique(model, backend_ptrs[i], backend_buft[i], cparams.moe_cache_size);
+                    break;
+                }
+            }
+            if (!moe_cache) {
+                throw std::runtime_error("MoE cache requires a GPU backend");
+            }
+        }
+
         sched_reserve();
 
         if (!cparams.flash_attn) {
@@ -2605,6 +2623,7 @@ llm_graph_params llama_context::graph_params(
         /*.loras       =*/ loras.get(),
         /*.mctx        =*/ mctx,
         /*.cross       =*/ &cross,
+        /*.moe_cache   =*/ moe_cache.get(),
         /*.prec_policy =*/ &model.prec_policy,
         /*.samplers    =*/ sampling.samplers,
         /*.n_outputs   =*/ n_outputs,
@@ -2645,7 +2664,14 @@ ggml_status llama_context::graph_compute(
 }
 
 bool llama_context::sched_copy_experts(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph, void * user_data) {
-    auto & st = static_cast(user_data)->copy_experts;
+    auto * lctx = static_cast(user_data);
+
+    // the slot maps of the MoE cache
+    if (lctx->moe_cache && lctx->moe_cache->copy(backend, src, dst, graph)) {
+        return true;
+    }
+
+    auto & st = lctx->copy_experts;
 
     // the ids must be computed before the split starts, so only the first node of the split is considered
     if (ggml_graph_n_nodes(graph) == 0) {
@@ -2692,10 +2718,28 @@ bool llama_context::sched_copy_experts(ggml_backend_t backend, const ggml_tensor
             last++;
         }
 
+        // the experts in the MoE cache are copied from device memory, the others are uploaded
+        int64_t next = first;
+        for (int64_t e = first; e <= last && lctx->moe_cache; ) {
+            const int64_t n = lctx->moe_cache->copy_experts(backend, src, dst, e, last);
+            if (n == 0) {
+                e++;
+                continue;
+            }
+            if (next < e) {
+                ggml_backend_tensor_set_async(backend, dst, (const uint8_t *) src->data + next*expert_size, next*expert_size, (e - next)*expert_size);
+            }
+            e   += n;
+            next = e;
+        }
+
         // copy a bit extra to ensure there are no NaNs in the padding of the last expert, this is necessary for MMQ in the CUDA backend
-        const size_t offset  = first*expert_size;
+        const size_t offset  = next*expert_size;
         const size_t padding = last < n_expert - 1 ? std::min(expert_size, 512) : 0;
-        ggml_backend_tensor_set_async(backend, dst, (const uint8_t *) src->data + offset, offset, (last - first + 1)*expert_size + padding);
+        const size_t size    = (last + 1 - next)*expert_size + padding;
+        if (size > 0) {
+            ggml_backend_tensor_set_async(backend, dst, (const uint8_t *) src->data + offset, offset, size);
+        }
 
         first = last + 1;
     }
@@ -3562,6 +3606,11 @@ llama_memory_breakdown llama_context::memory_breakdown() const {
             ret[buft].context += size;
         }
     }
+    if (moe_cache) {
+        for (const auto & [buft, size] : moe_cache->memory_breakdown()) {
+            ret[buft].context += size;
+        }
+    }
     if (model.hparams.no_alloc) {
         for (size_t i = 0; i < backends.size(); ++i) {
             ggml_backend_t             backend = backends[i].get();
@@ -3851,6 +3900,7 @@ llama_context_params llama_context_default_params() {
         /*.cb_eval_user_data           =*/ nullptr,
         /*.type_k                      =*/ GGML_TYPE_F16,
         /*.type_v                      =*/ GGML_TYPE_F16,
+        /*.moe_cache_size              =*/ 0,
         /*.abort_callback              =*/ nullptr,
         /*.abort_callback_data         =*/ nullptr,
         /*.embeddings                  =*/ false,
diff --git a/src/llama-context.h b/src/llama-context.h
index 28d386ad7..69bc01c19 100644
--- a/src/llama-context.h
+++ b/src/llama-context.h
@@ -7,6 +7,7 @@
 #include "llama-adapter.h"
 #include "llama-impl.h"
 #include "llama-memory.h"
+#include "llama-moe-cache.h"
 
 #include "ggml-cpp.h"
 #include "ggml-opt.h"
@@ -17,6 +18,7 @@
 
 struct llama_model;
 class llama_batch_allocr;
+class llama_moe_cache;
 
 class llama_io_read_i;
 class llama_io_write_i;
@@ -271,7 +273,7 @@ private:
 
     llm_graph_cb graph_get_cb() const;
 
-    // ggml_backend_sched copy callback, copies only the experts used by MUL_MAT_ID
+    // ggml_backend_sched copy callback, copies only the experts used by MUL_MAT_ID and updates the MoE cache
     static bool sched_copy_experts(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph, void * user_data);
 
     // disable auto fused ops (Flash Attention, Gated Delta Net) whose op lands on a device
@@ -299,6 +301,7 @@ private:
     llama_cross cross; // TODO: tmp for handling cross-attention - need something better probably
 
     llama_memory_ptr memory;
+    llama_moe_cache_ptr moe_cache;
 
     // decode output (2-dimensional array: [n_outputs][n_vocab])
     buffer_view logits = {nullptr, 0};
diff --git a/src/llama-cparams.h b/src/llama-cparams.h
index 004fec5a6..f4eb181a7 100644
--- a/src/llama-cparams.h
+++ b/src/llama-cparams.h
@@ -55,6 +55,8 @@ struct llama_cparams {
     bool pipeline_parallel;
     bool training;           // set by llama_opt_init()
 
+    size_t moe_cache_size;
+
     std::vector embeddings_layer_inp; // [n_layer()] extract input embeddings for layer
 
     enum llama_context_type ctx_type;
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 1112ad885..e0b7a47a9 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -2,6 +2,7 @@
 
 #include "llama-impl.h"
 #include "llama-model.h"
+#include "llama-moe-cache.h"
 #include "llama-batch.h"
 #include "llama-cparams.h"
 #include "llama-sampler.h"
@@ -1523,6 +1524,7 @@ llm_graph_context::llm_graph_context(const llm_graph_params & params) :
     loras            (params.loras),
     mctx             (params.mctx),
     cross            (params.cross),
+    moe_cache        (params.moe_cache),
     prec_policy      (params.prec_policy),
     samplers         (params.samplers),
     cb_func          (params.cb),
@@ -1589,8 +1591,12 @@ ggml_tensor * llm_graph_context::build_lora_mm_id(
           ggml_tensor * w,   // ggml_tensor * as
           ggml_tensor * cur, // ggml_tensor * b
           ggml_tensor * ids,
-          ggml_tensor * w_s) const {
-    ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur, ids);
+          ggml_tensor * w_s,
+          ggml_tensor * slots) const {
+    // the experts in the MoE cache are selected by their slots
+    ggml_tensor * res = slots == nullptr ?
+        ggml_mul_mat_id(ctx0, w, cur, ids) :
+        ggml_mul_mat_id(ctx0, moe_cache->get_experts(w), cur, slots);
 
     if (prec_policy) {
         prec_policy->apply(res);
@@ -2205,6 +2211,9 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
     //call early so that topk-moe can be used
     ggml_build_forward_expand(gf, weights);
 
+    // the experts of host-resident layers may be read from the MoE cache
+    ggml_tensor * slots = build_moe_cache_slots(selected_experts, up_exps, gate_exps, down_exps, gate_up_exps, il);
+
     cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens);
 
     if (weight_before_ffn) {
@@ -2219,7 +2228,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
 
     if (gate_up_exps) {
         // merged gate_up path: one mul_mat_id, then split into gate and up views
-        ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts, up_exps_s); // [n_ff*2, n_expert_used, n_tokens]
+        ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts, up_exps_s, slots); // [n_ff*2, n_expert_used, n_tokens]
         cb(gate_up, "ffn_moe_gate_up", il);
 
         if (up_exps_s) {
@@ -2238,7 +2247,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
         cb(up, "ffn_moe_up", il);
     } else {
         // separate gate and up path
-        up = build_lora_mm_id(up_exps, cur, selected_experts, up_exps_s); // [n_ff, n_expert_used, n_tokens]
+        up = build_lora_mm_id(up_exps, cur, selected_experts, up_exps_s, slots); // [n_ff, n_expert_used, n_tokens]
         cb(up, "ffn_moe_up", il);
 
         if (up_exps_s) {
@@ -2251,7 +2260,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
         }
 
         if (gate_exps) {
-            cur = build_lora_mm_id(gate_exps, cur, selected_experts, gate_exps_s); // [n_ff, n_expert_used, n_tokens]
+            cur = build_lora_mm_id(gate_exps, cur, selected_experts, gate_exps_s, slots); // [n_ff, n_expert_used, n_tokens]
             cb(cur, "ffn_moe_gate", il);
         } else {
             cur = up;
@@ -2352,7 +2361,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
             GGML_ABORT("fatal error");
     }
 
-    experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s); // [n_embd, n_expert_used, n_tokens]
+    experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s, slots); // [n_embd, n_expert_used, n_tokens]
     if (arch == LLM_ARCH_MISTRAL4) {
         // src1 can exceed F16 range
         ggml_prec_set_src(experts, GGML_PREC_F32, 1);
@@ -2409,6 +2418,45 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
     return moe_out;
 }
 
+ggml_tensor * llm_graph_context::build_moe_cache_slots(
+         ggml_tensor * selected_experts,
+         ggml_tensor * up_exps,
+         ggml_tensor * gate_exps,
+         ggml_tensor * down_exps,
+         ggml_tensor * gate_up_exps,
+                 int   il) const {
+    if (moe_cache == nullptr) {
+        return nullptr;
+    }
+
+    ggml_tensor * slot_map = moe_cache->get_slot_map(il, selected_experts->ne[1], selected_experts->ne[0]);
+    if (slot_map == nullptr) {
+        return nullptr;
+    }
+    for (ggml_tensor * w : { up_exps, gate_exps, down_exps, gate_up_exps }) {
+        if (w != nullptr && moe_cache->get_experts(w) == nullptr) {
+            return nullptr;
+        }
+    }
+
+    ggml_tensor * ids = selected_experts;
+    if (!ggml_is_contiguous(ids)) {
+        ids = ggml_cont(ctx0, ids);
+    }
+    ids = ggml_reshape_1d(ctx0, ids, ggml_nelements(ids));
+
+    // the slot map is a host weight, so the scheduler starts a new split here and copies it with the copy callback
+    // the callback reads the selected experts, uploads the missing ones and updates the slot map
+    ggml_tensor * slots = ggml_get_rows(ctx0, slot_map, ids); // [1, n_expert_used*n_tokens]
+    if (!ggml_backend_supports_op(moe_cache->backend(), slots)) {
+        return nullptr;
+    }
+    ggml_backend_sched_set_tensor_backend(sched, slots, moe_cache->backend());
+    cb(slots, "ffn_moe_slots", il);
+
+    return ggml_reshape_2d(ctx0, slots, selected_experts->ne[0], selected_experts->ne[1]); // [n_expert_used, n_tokens]
+}
+
 // input embeddings with optional lora
 ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd, float tok_scale) const {
     const int64_t n_embd_inp = hparams.n_embd_inp();
diff --git a/src/llama-graph.h b/src/llama-graph.h
index 838544576..79e8409ac 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -21,6 +21,8 @@ struct llama_cparams;
 struct llama_layer;
 struct llama_prec_policy;
 
+class llama_moe_cache;
+
 struct llama_memory_context_i;
 
 class llama_kv_cache_context;
@@ -793,6 +795,7 @@ struct llm_graph_params {
     const llama_adapter_loras    * loras;
     const llama_memory_context_i * mctx;
     const llama_cross            * cross;
+    const llama_moe_cache        * moe_cache;
 
     const llama_prec_policy * prec_policy = nullptr;
 
@@ -1036,6 +1039,7 @@ struct llm_graph_context {
     const llama_adapter_loras    * loras;
     const llama_memory_context_i * mctx;
     const llama_cross            * cross;
+    const llama_moe_cache        * moe_cache;
 
     const llama_prec_policy * prec_policy;
 
@@ -1078,11 +1082,13 @@ struct llm_graph_context {
               ggml_tensor * w_s = nullptr) const;
 
     // do mat_mul_id, while optionally apply lora and per-expert scale
+    // if slots is set, the experts are read from the MoE cache at these slots (see build_moe_cache_slots)
     ggml_tensor * build_lora_mm_id(
               ggml_tensor * w,   // ggml_tensor * as
               ggml_tensor * cur, // ggml_tensor * b
               ggml_tensor * ids,
-              ggml_tensor * w_s = nullptr) const;
+              ggml_tensor * w_s   = nullptr,
+              ggml_tensor * slots = nullptr) const;
 
     ggml_tensor * build_norm(
              ggml_tensor * cur,
@@ -1179,6 +1185,15 @@ struct llm_graph_context {
              ggml_tensor * down_exps_s = nullptr,
              ggml_tensor * selected_experts_in = nullptr) const;
 
+    // the slots of the selected experts in the MoE cache, nullptr if the experts of the layer are not read from the cache
+    ggml_tensor * build_moe_cache_slots(
+             ggml_tensor * selected_experts,
+             ggml_tensor * up_exps,
+             ggml_tensor * gate_exps,
+             ggml_tensor * down_exps,
+             ggml_tensor * gate_up_exps,
+                     int   il) const;
+
     //
     // inputs
     //
diff --git a/src/llama-moe-cache.cpp b/src/llama-moe-cache.cpp
new file mode 100644
index 000000000..7292a1e57
--- /dev/null
+++ b/src/llama-moe-cache.cpp
@@ -0,0 +1,550 @@
+#include "llama-moe-cache.h"
+
+#include "llama-impl.h"
+#include "llama-model.h"
+
+#include "ggml-cpp.h"
+
+#include 
+#include 
+#include 
+#include 
+
+namespace {
+
+// LRU of the experts of a group of layers, the slot of each expert is kept in the slot map of its layer
+struct moe_cache_lru {
+    int32_t n_expert = 0;
+    int32_t n_slots  = 0;
+
+    std::vector slot_map; // [n_layer] data of the slot maps, -1 if the expert is not cached
+    std::vector   key_of;   // [n_slots] il*n_expert + expert, -1 if empty
+
+    // doubly linked list of the slots, head is the least recently used
+    std::vector prev;
+    std::vector next;
+    int32_t head = -1;
+    int32_t tail = -1;
+
+    std::vector seen; // [n_expert]
+    uint32_t seen_gen = 0;
+    std::vector uniq;
+
+    void init(int32_t n_layer, int32_t n_expert, int32_t n_slots) {
+        this->n_expert = n_expert;
+        this->n_slots  = n_slots;
+        slot_map.assign(n_layer, nullptr);
+        key_of.assign(n_slots, -1);
+        prev.resize(n_slots);
+        next.resize(n_slots);
+        for (int32_t s = 0; s < n_slots; ++s) {
+            prev[s] = s - 1;
+            next[s] = s + 1 < n_slots ? s + 1 : -1;
+        }
+        head = 0;
+        tail = n_slots - 1;
+        seen.assign(n_expert, 0);
+    }
+
+    // move slot s to the tail (most recently used)
+    void touch(int32_t s) {
+        if (s == tail) {
+            return;
+        }
+        if (prev[s] >= 0) {
+            next[prev[s]] = next[s];
+        } else {
+            head = next[s];
+        }
+        prev[next[s]] = prev[s];
+
+        prev[s] = tail;
+        next[s] = -1;
+        next[tail] = s;
+        tail = s;
+    }
+
+    struct fill {
+        int32_t expert;
+        int32_t slot;
+    };
+
+    // give a slot to each expert selected by ids in layer il, the misses evict the least recently used experts
+    // returns false if the ids select more distinct experts than there are slots
+    bool plan(int32_t il, const int32_t * ids, size_t n_ids, std::vector & fills, size_t & n_hit) {
+        fills.clear();
+        n_hit = 0;
+
+        if (++seen_gen == 0) {
+            std::fill(seen.begin(), seen.end(), 0);
+            seen_gen = 1;
+        }
+        uniq.clear();
+        for (size_t i = 0; i < n_ids; ++i) {
+            GGML_ASSERT(ids[i] >= 0 && ids[i] < n_expert);
+            if (seen[ids[i]] != seen_gen) {
+                seen[ids[i]] = seen_gen;
+                uniq.push_back(ids[i]);
+            }
+        }
+        if (uniq.size() > (size_t) n_slots) {
+            return false;
+        }
+
+        int32_t * slots = slot_map[il];
+
+        // hits go to the tail first, so the head can be evicted below
+        for (int32_t e : uniq) {
+            if (slots[e] >= 0) {
+                touch(slots[e]);
+                n_hit++;
+            }
+        }
+        // sorted misses usually get consecutive slots, so the uploads can be merged
+        std::sort(uniq.begin(), uniq.end());
+        for (int32_t e : uniq) {
+            if (slots[e] >= 0) {
+                continue;
+            }
+            const int32_t s = head;
+            if (key_of[s] >= 0) {
+                slot_map[key_of[s] / n_expert][key_of[s] % n_expert] = -1;
+            }
+            key_of[s] = il*n_expert + e;
+            slots[e] = s;
+            touch(s);
+            fills.push_back({ e, s });
+        }
+        return true;
+    }
+};
+
+// gate, up, down or gate_up, down
+static std::vector llama_moe_cache_layer_experts(const llama_layer & layer) {
+    std::vector res;
+    for (ggml_tensor * t : { layer.ffn_gate_up_exps, layer.ffn_gate_exps, layer.ffn_up_exps, layer.ffn_down_exps }) {
+        if (t != nullptr) {
+            res.push_back(t);
+        }
+    }
+    return res;
+}
+
+static bool llama_moe_cache_same_layout(const std::vector & a, const std::vector & b) {
+    if (a.size() != b.size()) {
+        return false;
+    }
+    for (size_t i = 0; i < a.size(); ++i) {
+        if (a[i]->type != b[i]->type || !ggml_are_same_shape(a[i], b[i]) || a[i]->nb[2] != b[i]->nb[2]) {
+            return false;
+        }
+    }
+    return true;
+}
+
+static bool llama_moe_cache_is_host_weight(const ggml_tensor * t) {
+    return t->buffer != nullptr &&
+        ggml_backend_buffer_get_usage(t->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS &&
+        ggml_backend_buffer_is_host(t->buffer);
+}
+
+}
+
+struct llama_moe_cache::impl {
+    // layers with the same expert tensor layout share the banks and the LRU of a group
+    struct group {
+        std::vector ref;   // expert tensors of the first layer
+        std::vector layers;
+        std::vector banks; // device storage of all slots, one per expert tensor
+        size_t host_bytes = 0;
+        int32_t n_slots = 0;
+        moe_cache_lru lru;
+    };
+
+    struct layer {
+        int32_t ig = -1;                    // -1 if the layer is not cached
+        ggml_tensor * slot_map = nullptr;   // I32 [1, n_expert] in host memory
+        std::vector experts; // host expert tensors, in the order of the banks
+    };
+
+    struct binding {
+        int32_t il;
+        int32_t ip;           // index of the bank
+        ggml_tensor * cached; // view of the bank used in place of the host experts
+    };
+
+    struct stats {
+        size_t hits   = 0;
+        size_t misses = 0;
+        size_t bytes  = 0;
+    };
+
+    static constexpr int64_t max_batch = 32;
+
+    ggml_backend_t backend;
+    int32_t n_expert_used;
+
+    stats stats_small; // up to 8 tokens per ubatch
+    stats stats_large;
+    stats stats_copy;  // experts copied from the cache for large batches
+
+    std::vector groups;
+    std::vector layers;
+    std::unordered_map bindings; // host experts -> cached experts
+    std::unordered_map layer_of; // slot map -> layer
+
+    std::vector ids;
+    std::vector fills;
+
+    // banks and their views on the device
+    ggml_context_ptr ctx;
+    ggml_backend_buffer_ptr buf;
+    size_t buf_size = 0;
+
+    // slot maps in host memory
+    ggml_context_ptr ctx_host;
+    ggml_backend_buffer_ptr buf_host;
+    size_t buf_host_size = 0;
+
+    // views used by copy_experts
+    ggml_context_ptr ctx_views;
+
+    impl(const llama_model & model, ggml_backend_t backend, ggml_backend_buffer_type_t buft, size_t size) :
+            backend(backend), n_expert_used(model.hparams.n_expert_used_max()), layers(model.layers.size()) {
+        ggml_backend_dev_t dev = ggml_backend_get_device(backend);
+        const auto dev_type = ggml_backend_dev_type(dev);
+        if (dev_type != GGML_BACKEND_DEVICE_TYPE_GPU && dev_type != GGML_BACKEND_DEVICE_TYPE_IGPU) {
+            throw std::runtime_error("MoE cache requires a GPU backend");
+        }
+        if (model.split_mode() == LLAMA_SPLIT_MODE_TENSOR) {
+            throw std::runtime_error("MoE cache does not support tensor parallelism");
+        }
+        if (model.hparams.n_expert == 0 || n_expert_used == 0) {
+            throw std::runtime_error("MoE cache requires a MoE model");
+        }
+
+        // only cache layers that keep all of their experts in host memory
+        size_t host_bytes = 0;
+        for (size_t il = 0; il < model.layers.size(); ++il) {
+            auto experts = llama_moe_cache_layer_experts(model.layers[il]);
+            if (experts.empty() || model.dev_layer(il) != dev ||
+                !std::all_of(experts.begin(), experts.end(), llama_moe_cache_is_host_weight)) {
+                continue;
+            }
+            auto it = std::find_if(groups.begin(), groups.end(), [&](const group & g) { return llama_moe_cache_same_layout(g.ref, experts); });
+            if (it == groups.end()) {
+                groups.emplace_back();
+                it = groups.end() - 1;
+                it->ref = experts;
+            }
+            it->layers.push_back(il);
+            for (const ggml_tensor * t : experts) {
+                it->host_bytes += ggml_nbytes(t);
+                host_bytes     += ggml_nbytes(t);
+            }
+        }
+        if (groups.empty()) {
+            LLAMA_LOG_WARN("%s: no layer has all of its experts in host memory, MoE cache is disabled\n", __func__);
+            return;
+        }
+
+        // one extra slot at the end, CUDA MMQ can read past the last expert
+        const size_t alignment = ggml_backend_buft_get_alignment(buft);
+        auto alloc_size = [&](const group & g, int32_t n_slots) {
+            size_t res = 0;
+            for (const ggml_tensor * t : g.ref) {
+                res += GGML_PAD(t->nb[2]*(n_slots + 1), alignment);
+            }
+            return res;
+        };
+
+        // split the budget by the size of the experts, so each group caches the same fraction of its experts
+        size_t n_tensors      = 0;
+        size_t n_tensors_host = 0;
+        for (group & g : groups) {
+            const int32_t n_expert  = g.ref[0]->ne[2];
+            const size_t  budget    = (size_t) ((double) size*g.host_bytes/host_bytes);
+            const int32_t max_slots = g.layers.size()*n_expert;
+            while (g.n_slots < max_slots && alloc_size(g, g.n_slots + 1) <= budget) {
+                g.n_slots++;
+            }
+            if (g.n_slots < n_expert_used) {
+                LLAMA_LOG_WARN("%s: MoE cache budget is too small for %zu layers, they are not cached\n", __func__, g.layers.size());
+                g.n_slots = 0;
+                continue;
+            }
+            g.lru.init(model.layers.size(), n_expert, g.n_slots);
+            n_tensors      += g.ref.size()*(1 + g.layers.size());
+            n_tensors_host += g.layers.size();
+        }
+        if (n_tensors == 0) {
+            throw std::runtime_error("MoE cache is too small to hold the experts of one token");
+        }
+
+        auto init_ctx = [](size_t n_tensors) {
+            ggml_init_params params = {
+                /*.mem_size   =*/ n_tensors*ggml_tensor_overhead(),
+                /*.mem_buffer =*/ nullptr,
+                /*.no_alloc   =*/ true,
+            };
+            ggml_context_ptr res(ggml_init(params));
+            if (!res) {
+                throw std::runtime_error("failed to create the MoE cache context");
+            }
+            return res;
+        };
+        ctx       = init_ctx(n_tensors);
+        ctx_host  = init_ctx(n_tensors_host);
+        ctx_views = init_ctx(2);
+
+        ggml_backend_buffer_type_t buft_host = ggml_backend_cpu_buffer_type();
+        const size_t alignment_host = ggml_backend_buft_get_alignment(buft_host);
+
+        for (size_t ig = 0; ig < groups.size(); ++ig) {
+            group & g = groups[ig];
+            if (g.n_slots == 0) {
+                continue;
+            }
+            for (const ggml_tensor * t : g.ref) {
+                ggml_tensor * bank = ggml_new_tensor_3d(ctx.get(), t->type, t->ne[0], t->ne[1], g.n_slots + 1);
+                GGML_ASSERT(bank->nb[2] == t->nb[2]);
+                ggml_format_name(bank, "moe_cache.%zu.%s", ig, t->name);
+                g.banks.push_back(bank);
+            }
+            for (int32_t il : g.layers) {
+                layer & l = layers[il];
+                l.ig      = (int32_t) ig;
+                l.experts = llama_moe_cache_layer_experts(model.layers[il]);
+                for (size_t ip = 0; ip < l.experts.size(); ++ip) {
+                    ggml_tensor * bank   = g.banks[ip];
+                    ggml_tensor * cached = ggml_view_3d(ctx.get(), bank, bank->ne[0], bank->ne[1], g.n_slots, bank->nb[1], bank->nb[2], 0);
+                    ggml_format_name(cached, "moe_cache.%s", l.experts[ip]->name);
+                    bindings[l.experts[ip]] = { il, (int32_t) ip, cached };
+                }
+                l.slot_map = ggml_new_tensor_2d(ctx_host.get(), GGML_TYPE_I32, 1, g.ref[0]->ne[2]);
+                ggml_format_name(l.slot_map, "moe_cache.slot_map-%d", il);
+                layer_of[l.slot_map] = il;
+                buf_host_size += GGML_PAD(ggml_nbytes(l.slot_map), alignment_host);
+            }
+            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
+            buf.reset(ggml_backend_buft_alloc_buffer(buft, 0));
+            buf_host.reset(ggml_backend_buft_alloc_buffer(buft_host, 0));
+            for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) {
+                t->buffer = buf.get();
+            }
+            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 {
+            buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft));
+            buf_host.reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx_host.get(), buft_host));
+            if (!buf || !buf_host) {
+                throw std::runtime_error("failed to allocate the MoE cache buffers");
+            }
+            ggml_backend_buffer_clear(buf.get(), 0);
+            ggml_backend_buffer_clear(buf_host.get(), 0xff); // all slots are -1
+            buf_size      = ggml_backend_buffer_get_size(buf.get());
+            buf_host_size = ggml_backend_buffer_get_size(buf_host.get());
+
+            for (group & g : groups) {
+                for (int32_t il : g.layers) {
+                    if (layers[il].slot_map != nullptr) {
+                        g.lru.slot_map[il] = (int32_t *) layers[il].slot_map->data;
+                    }
+                }
+            }
+        }
+
+        // as weights, the ops that read the banks run on the device and the slot maps are copied with the copy callback
+        ggml_backend_buffer_set_usage(buf.get(),      GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+        ggml_backend_buffer_set_usage(buf_host.get(), GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+
+        LLAMA_LOG_INFO("%s: %10s MoE cache size = %8.2f MiB for %.2f MiB of host experts\n", __func__,
+            ggml_backend_buft_name(buft), buf_size/1024.0/1024.0, host_bytes/1024.0/1024.0);
+        for (const group & g : groups) {
+            LLAMA_LOG_INFO("%s: %2zu layers, %s: %5d slots (%.1f%%)\n", __func__,
+                g.layers.size(), ggml_type_name(g.ref.back()->type), g.n_slots, 100.0*g.n_slots/(g.layers.size()*g.ref[0]->ne[2]));
+        }
+    }
+
+    ~impl() {
+        log_stats();
+    }
+
+    ggml_tensor * get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used) const {
+        if (il < 0 || il >= (int32_t) layers.size() || layers[il].ig < 0) {
+            return nullptr;
+        }
+        const layer & l = layers[il];
+
+        // large batches use most experts of a layer, so they gain little from the cache and would evict the experts used in generation
+        if (n_tokens == 0 || n_tokens > max_batch || std::min(n_tokens*n_expert_used, l.slot_map->ne[1]) > groups[l.ig].n_slots) {
+            return nullptr;
+        }
+        return l.slot_map;
+    }
+
+    ggml_tensor * get_experts(const ggml_tensor * w) const {
+        const auto it = bindings.find(w);
+        return it != bindings.end() ? it->second.cached : nullptr;
+    }
+
+    int64_t copy_experts(ggml_backend_t backend, const ggml_tensor * w, ggml_tensor * dst, int64_t e, int64_t last) {
+        const auto it = bindings.find(w);
+        if (it == bindings.end() || backend != this->backend) {
+            return 0;
+        }
+        const binding & b = it->second;
+        const group & g = groups[layers[b.il].ig];
+
+        // large batches only read the cache, so the experts used in generation stay in it
+        const int32_t * slots = g.lru.slot_map[b.il];
+        if (slots == nullptr || slots[e] < 0) {
+            return 0;
+        }
+        int64_t n = 1;
+        while (e + n <= last && slots[e + n] == slots[e] + n) {
+            n++;
+        }
+
+        ggml_tensor * bank = g.banks[b.ip];
+        ggml_reset(ctx_views.get());
+        ggml_tensor * src_view = ggml_view_3d(ctx_views.get(), bank, bank->ne[0], bank->ne[1], n, bank->nb[1], bank->nb[2], slots[e]*bank->nb[2]);
+        ggml_tensor * dst_view = ggml_view_3d(ctx_views.get(), dst,  dst->ne[0],  dst->ne[1],  n, dst->nb[1],  dst->nb[2],  e*dst->nb[2]);
+        ggml_backend_view_init(src_view);
+        ggml_backend_view_init(dst_view);
+        ggml_backend_tensor_copy_async(backend, backend, src_view, dst_view);
+
+        stats_copy.hits  += n;
+        stats_copy.bytes += ggml_nbytes(src_view);
+
+        return n;
+    }
+
+    bool copy(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph) {
+        const auto it = layer_of.find(src);
+        if (it == layer_of.end()) {
+            return false;
+        }
+        const int32_t il = it->second;
+        const layer & l = layers[il];
+        group & g = groups[l.ig];
+
+        GGML_ASSERT(backend == this->backend);
+
+        // the get_rows that looks up the slots of the selected experts
+        const int n_nodes = ggml_graph_n_nodes(graph);
+        const ggml_tensor * lookup = nullptr;
+        for (int i = 0; i < n_nodes && lookup == nullptr; ++i) {
+            const ggml_tensor * node = ggml_graph_node(graph, i);
+            if (node->op == GGML_OP_GET_ROWS && node->src[0] == dst) {
+                lookup = node;
+            }
+        }
+        GGML_ASSERT(lookup != nullptr);
+
+        // the selected experts must be computed in an earlier split
+        // the scheduler starts a new split at the lookup because it reads a host weight, but only if the split already has inputs
+        const ggml_tensor * sel = lookup->src[1];
+        for (int i = 0; i < n_nodes; ++i) {
+            const ggml_tensor * node = ggml_graph_node(graph, i);
+            if (node == sel || node == sel->view_src) {
+                GGML_ABORT("the experts of layer %d are selected in the same split as their MoE cache lookup", il);
+            }
+        }
+        GGML_ASSERT(ggml_is_contiguous(sel));
+
+        ids.resize(ggml_nelements(sel));
+        ggml_backend_tensor_get_async(backend, sel, ids.data(), 0, ggml_nbytes(sel));
+        ggml_backend_synchronize(backend);
+
+        size_t n_hit = 0;
+        if (!g.lru.plan(il, ids.data(), ids.size(), fills, n_hit)) {
+            GGML_ABORT("the MoE cache is too small for the experts selected in layer %d", il);
+        }
+
+        // upload the missing experts, consecutive experts going to consecutive slots are uploaded together
+        size_t bytes = 0;
+        for (size_t ip = 0; ip < l.experts.size(); ++ip) {
+            const ggml_tensor * w    = l.experts[ip];
+            ggml_tensor       * bank = g.banks[ip];
+            const size_t expert_size = w->nb[2];
+            for (size_t i = 0; i < fills.size();) {
+                size_t n = 1;
+                while (i + n < fills.size() && fills[i + n].expert == fills[i].expert + (int32_t) n && fills[i + n].slot == fills[i].slot + (int32_t) n) {
+                    n++;
+                }
+                ggml_backend_tensor_set_async(backend, bank, (const uint8_t *) w->data + fills[i].expert*expert_size, fills[i].slot*expert_size, n*expert_size);
+                bytes += n*expert_size;
+                i += n;
+            }
+        }
+
+        stats & st = ids.size() <= (size_t) 8*n_expert_used ? stats_small : stats_large;
+        st.hits   += n_hit;
+        st.misses += fills.size();
+        st.bytes  += bytes;
+
+        // the next copy synchronizes the backend before it changes the slot map again
+        ggml_backend_tensor_set_async(backend, dst, src->data, 0, ggml_nbytes(src));
+
+        return true;
+    }
+
+    void log_stats() const {
+        auto log = [](const char * name, const stats & st) {
+            const size_t n = st.hits + st.misses;
+            if (n == 0) {
+                return;
+            }
+            LLAMA_LOG_INFO("llama_moe_cache: %s: hits = %zu, misses = %zu, hit rate = %.2f%%, uploaded = %.2f MiB\n",
+                name, st.hits, st.misses, 100.0*st.hits/n, st.bytes/1024.0/1024.0);
+        };
+        log("ubatch <= 8", stats_small);
+        log("ubatch  > 8", stats_large);
+        if (stats_copy.hits > 0) {
+            LLAMA_LOG_INFO("llama_moe_cache: large batches: %zu experts copied from the cache, %.2f MiB\n", stats_copy.hits, stats_copy.bytes/1024.0/1024.0);
+        }
+    }
+};
+
+llama_moe_cache::llama_moe_cache(const llama_model & model, ggml_backend_t backend, ggml_backend_buffer_type_t buft, size_t size) :
+    pimpl(new impl(model, backend, buft, size)) {
+}
+
+llama_moe_cache::~llama_moe_cache() = default;
+
+ggml_backend_t llama_moe_cache::backend() const {
+    return pimpl->backend;
+}
+
+ggml_tensor * llama_moe_cache::get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used) const {
+    return pimpl->get_slot_map(il, n_tokens, n_expert_used);
+}
+
+ggml_tensor * llama_moe_cache::get_experts(const ggml_tensor * w) const {
+    return pimpl->get_experts(w);
+}
+
+bool llama_moe_cache::copy(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph) {
+    return pimpl->copy(backend, src, dst, graph);
+}
+
+int64_t llama_moe_cache::copy_experts(ggml_backend_t backend, const ggml_tensor * w, ggml_tensor * dst, int64_t e, int64_t last) {
+    return pimpl->copy_experts(backend, w, dst, e, last);
+}
+
+std::map llama_moe_cache::memory_breakdown() const {
+    std::map res;
+    if (pimpl->buf) {
+        res[ggml_backend_buffer_get_type(pimpl->buf.get())] += pimpl->buf_size;
+    }
+    if (pimpl->buf_host) {
+        res[ggml_backend_buffer_get_type(pimpl->buf_host.get())] += pimpl->buf_host_size;
+    }
+    return res;
+}
diff --git a/src/llama-moe-cache.h b/src/llama-moe-cache.h
new file mode 100644
index 000000000..4872370d4
--- /dev/null
+++ b/src/llama-moe-cache.h
@@ -0,0 +1,39 @@
+#pragma once
+
+#include "ggml-backend.h"
+
+#include 
+#include 
+
+struct llama_model;
+
+// keeps the most recently used experts of host-resident MoE layers in a device buffer
+// each layer has a slot map in host memory: when the scheduler copies it to the device, the copy callback uploads the missing experts
+class llama_moe_cache {
+public:
+    llama_moe_cache(const llama_model & model, ggml_backend_t backend, ggml_backend_buffer_type_t buft, size_t size);
+    ~llama_moe_cache();
+
+    ggml_backend_t backend() const;
+
+    // the slot map of layer il, if its experts can be read from the cache for n_tokens tokens, nullptr otherwise
+    ggml_tensor * get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used) const;
+
+    // the experts of w in the cache, nullptr if w is not cached
+    ggml_tensor * get_experts(const ggml_tensor * w) const;
+
+    // ggml_backend_sched copy callback, returns false if src is not a slot map
+    bool copy(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph);
+
+    // for large batches: copy the experts of w that are in the cache, starting at expert e and up to expert last, to the copy dst of w
+    // returns the number of experts copied, 0 if expert e is not in the cache
+    int64_t copy_experts(ggml_backend_t backend, const ggml_tensor * w, ggml_tensor * dst, int64_t e, int64_t last);
+
+    std::map memory_breakdown() const;
+
+private:
+    struct impl;
+    std::unique_ptr pimpl;
+};
+
+using llama_moe_cache_ptr = std::unique_ptr;
diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp
index 54122629a..b54b3cc18 100644
--- a/tests/test-llama-archs.cpp
+++ b/tests/test-llama-archs.cpp
@@ -495,7 +495,7 @@ static std::pair get_model_and_ctx(
         struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const float stdev,
         const std::vector & devs,
         const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false,
-        const llama_model_tensor_buft_override * tensor_buft_overrides = nullptr) {
+        const llama_model_tensor_buft_override * tensor_buft_overrides = nullptr, const size_t moe_cache_size = 0) {
     GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr));
     llama_model_params model_params = llama_model_default_params();
     model_params.progress_callback = silent_model_load_progress;
@@ -512,6 +512,11 @@ static std::pair get_model_and_ctx(
     if (!encode) {
         ctx_params.n_ubatch = 64;
     }
+    if (moe_cache_size > 0) {
+        // the MoE cache is only used for small ubatches
+        ctx_params.moe_cache_size = moe_cache_size;
+        ctx_params.n_ubatch = 2;
+    }
 
     tensor_data_params tensor_params = { seed, stdev };
     llama_model_ptr model(gguf_ctx != nullptr ?
@@ -865,9 +870,10 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con
         std::string                     label;
         llama_split_mode                split_mode;
         bool                            host_experts; // keep the experts in host memory, see host_experts_test
+        size_t                          moe_cache_size;
 
-        device_config(std::vector devs, std::string name, llama_split_mode split_mode, bool host_experts = false)
-            : devs(std::move(devs)), label(std::move(name)), split_mode(split_mode), host_experts(host_experts) {}
+        device_config(std::vector devs, std::string name, llama_split_mode split_mode, bool host_experts = false, size_t moe_cache_size = 0)
+            : devs(std::move(devs)), label(std::move(name)), split_mode(split_mode), host_experts(host_experts), moe_cache_size(moe_cache_size) {}
     };
 
     const llama_model_tensor_buft_override host_experts_overrides[] = {
@@ -905,6 +911,16 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con
             dev_configs.emplace_back(devices_meta, "Host experts", LLAMA_SPLIT_MODE_LAYER, true);
             max_device_label_length = std::max(max_device_label_length, dev_configs.back().label.length());
         }
+
+        // the ops that use the host experts run on a GPU and read the experts from a cache
+        // the cache has only a few slots (4 for 288 KiB experts), so the experts are evicted and uploaded again
+        if (!devices_meta.empty()) {
+            const enum ggml_backend_dev_type type = ggml_backend_dev_type(devices_meta[0]);
+            if (type == GGML_BACKEND_DEVICE_TYPE_GPU || type == GGML_BACKEND_DEVICE_TYPE_IGPU) {
+                dev_configs.emplace_back(std::vector{devices_meta[0]}, "MoE cache", LLAMA_SPLIT_MODE_LAYER, true, 1536*1024);
+                max_device_label_length = std::max(max_device_label_length, dev_configs.back().label.length());
+            }
+        }
     }
 
     size_t max_arch_name_length = 0;
@@ -987,7 +1003,7 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con
                     }
                     if (dc.split_mode != LLAMA_SPLIT_MODE_TENSOR || llm_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);
+                        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);
                         const double nmse_val = nmse(logits_cpu, logits_dev);
                         snprintf(nmse_str, sizeof(nmse_str), "(%.2e)", nmse_val);
@@ -1053,7 +1069,7 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con
                         ms.save(file);
                         rewind(file);
 
-                        auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, stdev, dc.devs, dc.split_mode, encode, overrides);
+                        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 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";

From 50a6c5cf7c09ea5ea9c1937e4895ce88331c2869 Mon Sep 17 00:00:00 2001
From: Terrence Zhao <32208165+Terrencezzj@users.noreply.github.com>
Date: Wed, 7 Oct 2026 14:07:50 -0400
Subject: [PATCH 15/40] mtmd: add cohere2 vision support (#30062)

* cohere2 vision model

* address comments

* remove redundant mapping

* follow existing patterns

* fused linear_1
---
 conversion/__init__.py         |  1 +
 conversion/command_r.py        | 29 +++++++++++++--
 gguf-py/gguf/constants.py      |  1 +
 gguf-py/gguf/tensor_mapping.py |  1 +
 tools/mtmd/clip-impl.h         |  2 ++
 tools/mtmd/clip.cpp            | 20 +++++++++++
 tools/mtmd/models/siglip.cpp   | 10 ++++++
 tools/mtmd/mtmd-image.cpp      | 66 ++++++++++++++++++++++++++++++++++
 tools/mtmd/mtmd-image.h        |  8 +++++
 tools/mtmd/mtmd.cpp            | 10 ++++++
 10 files changed, 146 insertions(+), 2 deletions(-)

diff --git a/conversion/__init__.py b/conversion/__init__.py
index 96b018adf..9ee55d5a6 100644
--- a/conversion/__init__.py
+++ b/conversion/__init__.py
@@ -303,6 +303,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
     "AudioFlamingo3ForConditionalGeneration": "ultravox",
     "ClefModel": "clef",
     "CogVLMForCausalLM": "cogvlm",
+    "Cohere2VisionForConditionalGeneration": "command_r",
     "PplxDeciderModel": "pplx_decider",
     "DeepseekOCR2ForCausalLM": "deepseek",
     "DeepseekOCRForCausalLM": "deepseek",
diff --git a/conversion/command_r.py b/conversion/command_r.py
index 971f93ebd..0db6eb68a 100644
--- a/conversion/command_r.py
+++ b/conversion/command_r.py
@@ -1,14 +1,14 @@
 from __future__ import annotations
 
 import re
-from typing import Iterable, TYPE_CHECKING
+from typing import Callable, Iterable, TYPE_CHECKING
 
 import torch
 
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, gguf, logger
+from .base import MmprojModel, ModelBase, TextModel, gguf, logger
 
 
 @ModelBase.register("CohereForCausalLM")
@@ -180,3 +180,28 @@ class Cohere2MoeModel(TextModel):
         experts = [k for d in self._experts for k in d.keys()]
         if len(experts) > 0:
             raise ValueError(f"Unprocessed experts: {experts}")
+
+
+@ModelBase.register("Cohere2VisionForConditionalGeneration")
+# [TAG_HF_EXAMPLE_GATED] CohereLabs/command-a-vision-07-2025 is gated
+@ModelBase.example("CohereLabs/command-a-plus-05-2026-bf16")
+class Cohere2VisionModel(MmprojModel):
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.COHERE2V)
+        self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams["layer_norm_eps"])
+        self.gguf_writer.add_vision_projector_scale_factor(self.global_config["downsample_factor"])
+        self.gguf_writer.add_vision_preproc_max_tiles(self.preprocessor_config["max_patches"])
+        self.gguf_writer.add_vision_use_gelu(True)
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if ".embeddings." in name:
+            return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        if not name.startswith(("model.vision_tower.", "model.multi_modal_projector.")):
+            return None
+        return super().filter_tensors((name, gen))
diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index 9434e4120..842cd7891 100644
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -6130,6 +6130,7 @@ class VisionProjectorType:
     MIMO_AUDIO     = "mimo_audio"
     GRANITE4_VISION = "granite4_vision"
     MUSE_GLIMMER   = "muse-glimmer"
+    COHERE2V       = "cohere2v"
 
 
 # Items here are (block size, type size)
diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py
index 156ac3a50..a45e8e7d3 100644
--- a/gguf-py/gguf/tensor_mapping.py
+++ b/gguf-py/gguf/tensor_mapping.py
@@ -1622,6 +1622,7 @@ class TensorNameMap:
         MODEL_TENSOR.V_MMPROJ: (
             "aligner.w{bid}", # deepseek4v (w1 -> mm.1, w2 -> mm.2)
             "multi_modal_projector.linear_{bid}",
+            "model.multi_modal_projector.linear_{bid}", # cohere2v
             "mm_projector.proj.linear_{bid}", # Kimi-K2.5
             "visual.merger.mlp.{bid}", # qwen2vl
             "mlp_AR.linear_{bid}", # PaddleOCR-VL
diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h
index 9aec2a25b..e73b74aff 100644
--- a/tools/mtmd/clip-impl.h
+++ b/tools/mtmd/clip-impl.h
@@ -507,6 +507,7 @@ enum projector_type {
     PROJECTOR_TYPE_POCKETTTS_SPKENC,
     PROJECTOR_TYPE_POCKETTTS_GEN,
     PROJECTOR_TYPE_MUSE_GLIMMER,
+    PROJECTOR_TYPE_COHERE2V,
     PROJECTOR_TYPE_UNKNOWN,
 };
 
@@ -574,6 +575,7 @@ static std::map PROJECTOR_TYPE_NAMES = {
     { PROJECTOR_TYPE_POCKETTTS_SPKENC,  "pockettts_spkenc"},
     { PROJECTOR_TYPE_POCKETTTS_GEN,     "pockettts_gen"},
     { PROJECTOR_TYPE_MUSE_GLIMMER,      "muse-glimmer"},
+    { PROJECTOR_TYPE_COHERE2V,          "cohere2v"},
 };
 
 static projector_type clip_projector_type_from_string(const std::string & str) {
diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp
index f437040cf..06c562897 100644
--- a/tools/mtmd/clip.cpp
+++ b/tools/mtmd/clip.cpp
@@ -937,6 +937,7 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const
     switch (ctx->proj_type()) {
         case PROJECTOR_TYPE_GEMMA3:
         case PROJECTOR_TYPE_IDEFICS3:
+        case PROJECTOR_TYPE_COHERE2V:
         case PROJECTOR_TYPE_LFM2:
         case PROJECTOR_TYPE_JANUS_PRO:
         case PROJECTOR_TYPE_PHI4:
@@ -1514,6 +1515,15 @@ struct clip_model_loader {
                         get_u32(KEY_PREPROC_IMAGE_SIZE, hparams.image_longest_edge, false);
                         hparams.set_limit_image_tokens();
                     } break;
+                case PROJECTOR_TYPE_COHERE2V:
+                    {
+                        hparams.image_pad_rf = PAD_NONE;
+                        get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge);
+                        get_u32(KEY_PREPROC_MAX_TILES, hparams.preproc_max_tiles);
+                        if (hparams.preproc_max_tiles <= 0 || hparams.preproc_max_tiles > 256) {
+                            throw std::runtime_error(string_format("%s: preproc_max_tiles (%d) must be in range [1, 256]\n", __func__, hparams.preproc_max_tiles));
+                        }
+                    } break;
                 case PROJECTOR_TYPE_LFM2:
                     {
                         hparams.image_resize_algo    = RESIZE_ALGO_BILINEAR;
@@ -2770,6 +2780,13 @@ struct clip_model_loader {
                 {
                     model.mm_fc_w = get_tensor(string_format(TN_MM_PROJECTOR, "weight"));
                 } break;
+            case PROJECTOR_TYPE_COHERE2V:
+                {
+                    model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 1, "weight"));
+                    model.mm_1_b = get_tensor(string_format(TN_LLAVA_PROJ, 1, "bias"));
+                    model.mm_2_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight"));
+                    model.mm_2_b = get_tensor(string_format(TN_LLAVA_PROJ, 2, "bias"));
+                } break;
             case PROJECTOR_TYPE_LFM2:
                 {
                     model.mm_input_norm_w = get_tensor(TN_MM_INP_NORM, false);
@@ -4233,6 +4250,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
         case PROJECTOR_TYPE_GEMMA4V:
         case PROJECTOR_TYPE_GEMMA4UV:
         case PROJECTOR_TYPE_IDEFICS3:
+        case PROJECTOR_TYPE_COHERE2V:
         case PROJECTOR_TYPE_INTERNVL:
         case PROJECTOR_TYPE_NEMOTRON_V2_VL:
         case PROJECTOR_TYPE_LLAMA4:
@@ -5321,6 +5339,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
         case PROJECTOR_TYPE_GEMMA3:
         case PROJECTOR_TYPE_GEMMA3NV:
         case PROJECTOR_TYPE_IDEFICS3:
+        case PROJECTOR_TYPE_COHERE2V:
         case PROJECTOR_TYPE_INTERNVL:
         case PROJECTOR_TYPE_NEMOTRON_V2_VL:
         case PROJECTOR_TYPE_QWEN2A:
@@ -6055,6 +6074,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
         case PROJECTOR_TYPE_KIMIK25:
         case PROJECTOR_TYPE_YASA2:
         case PROJECTOR_TYPE_DEEPSEEK4V:
+        case PROJECTOR_TYPE_COHERE2V:
             return ctx->model.mm_2_w->ne[1];
         case PROJECTOR_TYPE_HUNYUANVL:
             return ctx->model.mm_model_proj->ne[1];
diff --git a/tools/mtmd/models/siglip.cpp b/tools/mtmd/models/siglip.cpp
index 7ef98eed0..96d21ceee 100644
--- a/tools/mtmd/models/siglip.cpp
+++ b/tools/mtmd/models/siglip.cpp
@@ -45,6 +45,16 @@ ggml_cgraph * clip_graph_siglip::build() {
         cur = build_patch_merge_permute(cur, scale_factor);
         cur = build_mm(model.mm_fc_w, cur);
 
+    } else if (proj_type == PROJECTOR_TYPE_COHERE2V) {
+        // tiles are square, so the pixel shuffle is the same as Idefics3
+        cur = build_patch_merge_permute(cur, model.hparams.n_merge);
+        cur = build_mm(model.mm_1_w, cur);
+        cur = ggml_add(ctx0, cur, model.mm_1_b);
+        // linear_1 output is [x, gate], HF computes silu(gate) * x
+        cur = ggml_swiglu_swapped(ctx0, cur);
+        cur = build_mm(model.mm_2_w, cur);
+        cur = ggml_add(ctx0, cur, model.mm_2_b);
+
     } else if (proj_type == PROJECTOR_TYPE_LFM2) {
         // pixel unshuffle block
         const int scale_factor = model.hparams.n_merge;
diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp
index ea67ec3b8..5ce0fa059 100644
--- a/tools/mtmd/mtmd-image.cpp
+++ b/tools/mtmd/mtmd-image.cpp
@@ -1142,6 +1142,72 @@ mtmd_image_preproc_out mtmd_image_preprocessor_idefics3::preprocess(const clip_i
     return output;
 }
 
+//
+// mtmd_image_preprocessor_cohere2v
+//
+
+mtmd_image_preproc_out mtmd_image_preprocessor_cohere2v::preprocess(const clip_image_u8 & img) const {
+    const auto inst = get_slice_instructions(img.get_size());
+    auto sliced = slice_image(img, inst);
+
+    mtmd_image_preproc_out output;
+    if (sliced.slices.empty()) {
+        output.append_overview(hparams, sliced.overview, true);
+        return output;
+    }
+    // slices first, then thumbnail
+    output.append(hparams, sliced.slices, true);
+    output.append_overview(hparams, sliced.overview, true);
+    output.grid_x = inst.grid_size.width;
+    output.grid_y = inst.grid_size.height;
+    return output;
+}
+
+mtmd_image_preprocessor_llava_uhd::slice_instructions mtmd_image_preprocessor_cohere2v::get_slice_instructions(const clip_image_size & original_size) const {
+    const int tile = hparams.image_size;
+
+    // pick the grid with the least upscale; if all grids need downscale, pick the one with the least downscale
+    // grids are visited by tile count, then by width, same order as HF for ties
+    double best_down = -1.0;
+    double best_up   = std::numeric_limits::max();
+    clip_image_size grid_down = { 1, 1 };
+    clip_image_size grid_up   = { 0, 0 };
+    for (int n = 1; n <= hparams.preproc_max_tiles; n++) {
+        for (int w = 1; w <= n; w++) {
+            if (n % w != 0) {
+                continue;
+            }
+            const clip_image_size g = { w, n / w };
+            const double scale = std::min(
+                (double) (g.width  * tile) / original_size.width,
+                (double) (g.height * tile) / original_size.height);
+            if (scale < 1.0) {
+                if (scale > best_down) {
+                    best_down = scale;
+                    grid_down = g;
+                }
+            } else if (scale < best_up) {
+                best_up = scale;
+                grid_up = g;
+            }
+        }
+    }
+    const clip_image_size grid = grid_up.width > 0 ? grid_up : grid_down;
+
+    slice_instructions inst;
+    inst.overview_size = { tile, tile };
+    inst.refined_size  = { tile * grid.width, tile * grid.height };
+    inst.grid_size     = grid;
+    if (grid.width * grid.height > 1) {
+        for (int y = 0; y < grid.height; y++) {
+            for (int x = 0; x < grid.width; x++) {
+                inst.slices.push_back({ x * tile, y * tile, { tile, tile } });
+            }
+        }
+    }
+    return inst;
+}
+
 //
 // mtmd_image_preprocessor_internvl
 //
diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h
index 4fa6207d0..2dff2b47d 100644
--- a/tools/mtmd/mtmd-image.h
+++ b/tools/mtmd/mtmd-image.h
@@ -191,6 +191,14 @@ struct mtmd_image_preprocessor_internvl : mtmd_image_preprocessor_llava_uhd {
     mtmd_image_preproc_out preprocess(const clip_image_u8 & img) const override;
 };
 
+// stretch the image to a grid of tiles, add a thumbnail if there is more than 1 tile
+// ref: https://github.com/huggingface/transformers/blob/main/src/transformers/models/cohere2_vision/image_processing_cohere2_vision.py
+struct mtmd_image_preprocessor_cohere2v : mtmd_image_preprocessor_llava_uhd {
+    mtmd_image_preprocessor_cohere2v(const clip_ctx * ctx) : mtmd_image_preprocessor_llava_uhd(ctx) {}
+    mtmd_image_preproc_out preprocess(const clip_image_u8 & img) const override;
+    slice_instructions get_slice_instructions(const clip_image_size & original_size) const override;
+};
+
 // DeepSeek-OCR (v1/v2) global view + optional local tile grid
 struct mtmd_image_preprocessor_deepseekocr : mtmd_image_preprocessor {
     mtmd_image_preprocessor_deepseekocr(const clip_ctx * ctx)
diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp
index 9e7b1cc80..b2319b901 100644
--- a/tools/mtmd/mtmd.cpp
+++ b/tools/mtmd/mtmd.cpp
@@ -798,6 +798,16 @@ struct mtmd_context {
                     image_preproc = std::make_unique(ctx_v);
                     ov_img_first = false;
                 } break;
+            case PROJECTOR_TYPE_COHERE2V:
+                {
+                    // <|START_OF_IMG|> (tile embeddings) <|IMG_LINE_BREAK|> ... <|END_OF_IMG|>
+                    img_beg = "<|START_OF_IMG|>";
+                    img_end = "<|END_OF_IMG|>";
+                    tok_sli_img_end = {lookup_token("<|IMG_LINE_BREAK|>")};
+                    tok_ov_img_end  = tok_sli_img_end;
+                    ov_img_first = false;
+                    image_preproc = std::make_unique(ctx_v);
+                } break;
             case PROJECTOR_TYPE_KIMIVL:
                 {
                     // <|media_start|> ... (image embeddings) ... <|media_end|>

From b86d2f07542b29ab099aed34fd6b6d1b2fd4b81c Mon Sep 17 00:00:00 2001
From: bri-prism <288398250+bri-prism@users.noreply.github.com>
Date: Wed, 7 Oct 2026 11:56:18 -0700
Subject: [PATCH 16/40] cuda: FWHT kernels for block widths above 512 (#29100)

The CUDA FWHT covers widths 64 to 512. It runs one row per warp and keeps N/32
values per lane, so wider blocks need more registers per lane than that layout
allows.

fwht_cuda_block runs one row per thread block with 256 threads, so each thread
keeps N/256 values. Stages below the warp width still shuffle, those up to the
block width go through shared memory, and the rest stay in registers. Same
butterfly and sign convention as the warp kernel.

Widths 64 to 512 keep the warp kernel. 1024 through 8192 use the new one, for
both F32 and F16 sources. ggml_cuda_op_mul_mat_use_fwht (the shared
supports_op/dispatch predicate added in #29096) does not check width, so it
needed no change here: any width it admits that ggml_cuda_op_fwht can't serve
already falls through correctly to the cuBLAS path.

Rebased onto current master with #29096's F16 commit underneath it, since this
depends on the same F16 template infrastructure; that commit applied cleanly,
the only conflict was in test-backend-ops.cpp where an unrelated intervening
commit's own test additions landed near this block.

test-backend-ops on an A10 (lambdalabs): MUL_MAT 1297/1297, including all
FWHT/Hadamard cases (18 existing, 4 new F32 wide, 4 new F16 wide, 2 new
many-rows, 1 too-big boundary moved to 16384).
---
 ggml/src/ggml-cuda/fwht.cu | 113 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 112 insertions(+), 1 deletion(-)

diff --git a/ggml/src/ggml-cuda/fwht.cu b/ggml/src/ggml-cuda/fwht.cu
index 67eea594d..d69e2ed81 100644
--- a/ggml/src/ggml-cuda/fwht.cu
+++ b/ggml/src/ggml-cuda/fwht.cu
@@ -2,6 +2,9 @@
 #include "convert.cuh"
 #include "fwht.cuh"
 
+// wide FWHT blocks use one row per thread block with this many threads
+#define GGML_CUDA_FWHT_BLOCK_NT 256
+
 template 
 __launch_bounds__(4*ggml_cuda_get_physical_warp_size(), 1)
 __global__ void fwht_cuda(const T * src, float * dst, const int64_t n_rows, const float scale) {
@@ -59,6 +62,87 @@ __global__ void fwht_cuda(const T * src, float * dst, const int64_t n_rows, cons
     }
 }
 
+// Wide blocks: one row per thread block instead of per warp, so each thread keeps N/NT
+// values rather than N/32. Stages below the warp width still shuffle, those up to the
+// block width go through shared memory, and the rest stay in registers.
+template 
+__launch_bounds__(NT, 1)
+__global__ void fwht_cuda_block(const T * src, float * dst, const int64_t n_rows, const float scale) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int NE        = N / NT;
+    static_assert(NE >= 1 && N % NT == 0 && NT % warp_size == 0, "bad FWHT block shape");
+
+    __shared__ float s[N];
+
+    const int64_t r = blockIdx.x;
+    if (r >= n_rows) {
+        return;
+    }
+
+    src += r * N;
+    dst += r * N;
+
+    const int tid  = threadIdx.x;
+    const int lane = tid % warp_size;
+
+    ggml_cuda_pdl_sync();
+
+    float reg[NE];
+#pragma unroll
+    for (int i = 0; i < NE; ++i) {
+        reg[i] = ggml_cuda_cast(src[i * NT + tid]) * scale;
+    }
+
+    // stages within a warp: partner differs in the lane bits
+#pragma unroll
+    for (int h = 1; h < warp_size; h *= 2) {
+#pragma unroll
+        for (int j = 0; j < NE; j++) {
+            const float val  = reg[j];
+            const float val2 = __shfl_xor_sync(0xFFFFFFFF, val, h, warp_size);
+            reg[j] = (lane & h) == 0 ? val + val2 : val2 - val;
+        }
+    }
+
+    // stages across warps: partner differs in the thread-index bits above the lane
+#pragma unroll
+    for (int h = warp_size; h < NT; h *= 2) {
+#pragma unroll
+        for (int j = 0; j < NE; j++) {
+            s[j * NT + tid] = reg[j];
+        }
+        __syncthreads();
+#pragma unroll
+        for (int j = 0; j < NE; j++) {
+            const float val  = reg[j];
+            const float val2 = s[j * NT + (tid ^ h)];
+            reg[j] = (tid & h) == 0 ? val + val2 : val2 - val;
+        }
+        __syncthreads();
+    }
+
+    // stages above the block width: partner is another register of the same thread
+#pragma unroll
+    for (int h = NT; h < N; h *= 2) {
+        const int step = h / NT;
+#pragma unroll
+        for (int j = 0; j < NE; j += 2 * step) {
+#pragma unroll
+            for (int k = 0; k < step; k++) {
+                const float x = reg[j + k];
+                const float y = reg[j + k + step];
+                reg[j + k]        = x + y;
+                reg[j + k + step] = x - y;
+            }
+        }
+    }
+
+#pragma unroll
+    for (int i = 0; i < NE; ++i) {
+        dst[i * NT + tid] = reg[i];
+    }
+}
+
 template 
 static bool ggml_cuda_op_fwht_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) {
     const int     n    = src->ne[0];
@@ -94,7 +178,34 @@ static bool ggml_cuda_op_fwht_impl(ggml_backend_cuda_context & ctx, const ggml_t
             ggml_cuda_kernel_launch(fwht_cuda<512, T>, launch_params, src_d, dst_d, rows, scale);
             return true;
         default:
-            return false;
+            break;
+    }
+
+    // wide blocks: one row per thread block
+    {
+        constexpr int nt = GGML_CUDA_FWHT_BLOCK_NT;
+
+        dim3 grid_dims_w(rows, 1, 1);
+        dim3 block_dims_w(nt, 1, 1);
+        const ggml_cuda_kernel_launch_params launch_params_w =
+            ggml_cuda_kernel_launch_params(grid_dims_w, block_dims_w, 0, stream);
+
+        switch (n) {
+            case 1024:
+                ggml_cuda_kernel_launch(fwht_cuda_block<1024, nt, T>, launch_params_w, src_d, dst_d, rows, scale);
+                return true;
+            case 2048:
+                ggml_cuda_kernel_launch(fwht_cuda_block<2048, nt, T>, launch_params_w, src_d, dst_d, rows, scale);
+                return true;
+            case 4096:
+                ggml_cuda_kernel_launch(fwht_cuda_block<4096, nt, T>, launch_params_w, src_d, dst_d, rows, scale);
+                return true;
+            case 8192:
+                ggml_cuda_kernel_launch(fwht_cuda_block<8192, nt, T>, launch_params_w, src_d, dst_d, rows, scale);
+                return true;
+            default:
+                return false;
+        }
     }
 }
 

From 88dcc460d628698bb8305b98c200c34f1edfdc04 Mon Sep 17 00:00:00 2001
From: Tarek Dakhran 
Date: Wed, 7 Oct 2026 22:04:39 +0200
Subject: [PATCH 17/40] model : add LiquidAI/d1-3B decision model (#30110)

* model : add LiquidAI/d1-3b decision model

mtmd : read LFM2 image resize algo from GGUF

Assisted-by: Claude Opus 5.5

* common : rename decision type d1 to lfm2-d1

Assisted-by: Claude Opus 5.5
---
 common/common.cpp                |   1 +
 common/common.h                  |   1 +
 conversion/__init__.py           |   1 +
 conversion/lfm2.py               |  73 +++++++++++++++++-
 gguf-py/gguf/constants.py        |   2 +
 gguf-py/gguf/gguf_writer.py      |   3 +
 tools/mtmd/clip-impl.h           |   1 +
 tools/mtmd/clip.cpp              |  17 ++++-
 tools/server/README.md           |   6 +-
 tools/server/server-common.cpp   |   1 +
 tools/server/server-context.cpp  |  10 +++
 tools/server/server-decision.cpp | 124 ++++++++++++++++++++++++++++++-
 tools/server/server-decision.h   |   4 +
 tools/server/server-task.h       |   5 +-
 14 files changed, 239 insertions(+), 10 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index 768b2e9a5..0e891cea6 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1170,6 +1170,7 @@ static const std::map COMMON_DECISION_TYPE_NA
     { COMMON_DECISION_TYPE_LAYA,           "laya"          },
     { COMMON_DECISION_TYPE_CLEF,           "clef"          },
     { COMMON_DECISION_TYPE_PPLX_DECIDER,   "pplx-decider"  },
+    { COMMON_DECISION_TYPE_LFM2_D1,        "lfm2-d1"       },
 };
 
 static common_decision_type common_decision_type_from_string(const std::string & str) {
diff --git a/common/common.h b/common/common.h
index 0f912d901..6f8acf31d 100644
--- a/common/common.h
+++ b/common/common.h
@@ -963,6 +963,7 @@ enum common_decision_type {
     COMMON_DECISION_TYPE_LAYA,    // score of one marker token per option, read from the embeddings output
     COMMON_DECISION_TYPE_CLEF,    // all questions in one prompt, score of option i read from the embeddings output at row i
     COMMON_DECISION_TYPE_PPLX_DECIDER, // same as openjev, label codes of 1 or 2 letters
+    COMMON_DECISION_TYPE_LFM2_D1, // same as openjev, the labels depend on the question type
     COMMON_DECISION_TYPE_UNKNOWN, // a decision model of a type that is not supported
 };
 
diff --git a/conversion/__init__.py b/conversion/__init__.py
index 9ee55d5a6..34c56bbee 100644
--- a/conversion/__init__.py
+++ b/conversion/__init__.py
@@ -161,6 +161,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "Lfm2BidirectionalForMaskedLM": "lfm2",
     "Lfm2BidirectionalModel": "lfm2",
     "Lfm2ForCausalLM": "lfm2",
+    "D1Model": "lfm2",
     "Lfm2Model": "lfm2",
     "Lfm2MoeForCausalLM": "lfm2",
     "Llama4ForCausalLM": "llama",
diff --git a/conversion/lfm2.py b/conversion/lfm2.py
index e50bbdd0f..d4343675e 100644
--- a/conversion/lfm2.py
+++ b/conversion/lfm2.py
@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+import json
+
+from pathlib import Path
 from typing import Any, Callable, Iterable, TYPE_CHECKING
 
 import torch
@@ -7,7 +10,7 @@ import torch
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import MmprojModel, ModelBase, TextModel, gguf
+from .base import MmprojModel, ModelBase, TextModel, gguf, jinja_str_or_json, logger
 
 from .gemma import ConformerAudioModel
 
@@ -65,6 +68,68 @@ class LFM2Model(TextModel):
         yield from super().modify_tensors(data_torch, name, bid)
 
 
+def _is_d1_checkpoint(dir_model: Path) -> bool:
+    if not (dir_model / "config.json").is_file():
+        return False
+    with open(dir_model / "config.json", encoding="utf-8") as f:
+        return json.load(f).get("auto_map", {}).get("AutoModel", "").endswith(".D1Model")
+
+
+@ModelBase.register_hparams_loader(_is_d1_checkpoint)
+def _load_d1_hparams(dir_model: Path) -> dict[str, Any]:
+    logger.info("gguf: detected d1 checkpoint")
+    hparams = ModelBase.load_hparams(dir_model, False, guess=False)
+    # the mmproj stays LFM2-VL
+    hparams["text_config"]["architectures"] = ["D1Model"]
+    return hparams
+
+
+@ModelBase.register("D1Model")
+@ModelBase.example("LiquidAI/d1-3b")
+class D1Model(LFM2Model):
+    model_arch = gguf.MODEL_ARCH.LFM2
+
+    def set_vocab(self):
+        super().set_vocab()
+        self.gguf_writer.add_chat_template([{"name": "systemone", "template": self._systemone_template()}])
+
+    @staticmethod
+    def _systemone_template() -> str:
+        # follows prompt.py of the model repo
+        description = jinja_str_or_json("o.description")
+        choice = (
+            "{{ '\\n\\nOptions:\\n' }}"
+            "{% for o in options %}{{ o.label }} {% if o.description %}" + description + "{% else %}{{ o.key | replace('_', ' ') }}{% endif %}"
+            "{% if not loop.last %}{{ '\\n' }}{% endif %}{% endfor %}"
+            "{{ '\\n\\nReply with the option code only.' }}"
+        )
+        # with criteria, a missing description is written as None
+        noul = (
+            "{% set ns = namespace(criteria=false) %}{% for o in options %}{% if o.description is not none %}{% set ns.criteria = true %}{% endif %}{% endfor %}"
+            "{% if ns.criteria %}"
+            "{% for o in options %}{{ '\\nYes: ' if o.key == 'true' else '\\nNo: ' }}"
+            "{% if o.description is none %}None{% else %}" + description + "{% endif %}{% endfor %}{% endif %}"
+            "{{ '\\n\\nReply with yes or no only.' }}"
+        )
+        score = (
+            "{{ '\\n\\n' }}{% for o in options %}{{ o.key }} " + description + "{{ '\\n' }}{% endfor %}"
+            "{{ '\\nReply with a single digit 0-' }}{{ options | length - 1 }}{{ ' only.' }}"
+        )
+        return (
+            "<|startoftext|><|im_start|>user\n"
+            "{% for image in images %}{{ image }}{% endfor %}"
+            "{% if state is not none %}{% if state is string %}{{ state }}{% else %}{{ state | tojson(indent=2) }}{% endif %}"
+            "{{ '\\n\\n\\nQUESTION:\\n' }}{% endif %}"
+            + jinja_str_or_json("instructions")
+            + "{% if type == 'choice' %}" + choice + "{% elif type == 'noul' %}" + noul + "{% else %}" + score + "{% endif %}"
+            "{{ '<|im_end|>\\n<|im_start|>assistant\\n' }}"
+        )
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_decision_type(gguf.DecisionType.LFM2_D1)
+
+
 @ModelBase.register("Lfm2Model", "Lfm2BidirectionalModel", "Lfm2BidirectionalForMaskedLM")
 @ModelBase.example("LiquidAI/LFM2.5-ColBERT-350M", "LiquidAI/LFM2.5-Embedding-350M", "LiquidAI/LFM2.5-Encoder-350M", "LiquidAI/LFM2.5-Encoder-230M")
 class LFM2ColBertModel(LFM2Model):
@@ -188,6 +253,12 @@ class LFM2VLModel(MmprojModel):
         # python notation, e.g. for vision_feature_layer == -1, we pick last layer -> vision_feature_layers_to_drop = 0
         vision_feature_layers_to_drop = -(self.global_config.get("vision_feature_layer", -1) + 1)
         self.gguf_writer.add_vision_block_count(self.find_vparam(self.n_block_keys) - vision_feature_layers_to_drop)
+        # PIL resample enum
+        if (resample := self.preprocessor_config.get("resample")) is not None:
+            resize_algo = {1: "lanczos", 2: "bilinear", 3: "bicubic"}.get(resample)
+            if resize_algo is None:
+                raise ValueError(f"unsupported resample: {resample}")
+            self.gguf_writer.add_vision_image_resize_algo(resize_algo)
 
     @classmethod
     def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index 842cd7891..c181cb44a 100644
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -408,6 +408,7 @@ class Keys:
         BLOCK_COUNT           = "clip.vision.block_count"
         IMAGE_MEAN            = "clip.vision.image_mean"
         IMAGE_STD             = "clip.vision.image_std"
+        IMAGE_RESIZE_ALGO     = "clip.vision.image_resize_algo"
         SPATIAL_MERGE_SIZE    = "clip.vision.spatial_merge_size"
         SWIGLU_CLAMP          = "clip.vision.swiglu_clamp"
         EXPERT_COUNT_PER_LAYER = "clip.vision.expert_count_per_layer" # dots3note pyramid MoE, 0 = dense layer
@@ -6070,6 +6071,7 @@ class DecisionType:
     NIMBLE  = "nimble"   # same as openjev, the prompt lists all the questions of the request
     CLEF    = "clef"     # joint head over all questions, one score per option
     PPLX_DECIDER = "pplx-decider"  # same as openjev, label codes of 1 or 2 letters
+    LFM2_D1 = "lfm2-d1"  # same as openjev, the labels depend on the question type
 
 
 class VisionProjectorType:
diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py
index e814691d7..f33a8a551 100644
--- a/gguf-py/gguf/gguf_writer.py
+++ b/gguf-py/gguf/gguf_writer.py
@@ -1432,6 +1432,9 @@ class GGUFWriter:
     def add_vision_image_mean(self, values: Sequence[float]) -> None:
         self.add_array(Keys.ClipVision.IMAGE_MEAN, values)
 
+    def add_vision_image_resize_algo(self, value: str) -> None:
+        self.add_string(Keys.ClipVision.IMAGE_RESIZE_ALGO, value)
+
     def add_vision_image_std(self, values: Sequence[float]) -> None:
         self.add_array(Keys.ClipVision.IMAGE_STD, values)
 
diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h
index e73b74aff..abf75d998 100644
--- a/tools/mtmd/clip-impl.h
+++ b/tools/mtmd/clip-impl.h
@@ -58,6 +58,7 @@
 #define KEY_PATCH_SIZE              "clip.vision.patch_size"
 #define KEY_IMAGE_MEAN              "clip.vision.image_mean"
 #define KEY_IMAGE_STD               "clip.vision.image_std"
+#define KEY_IMAGE_RESIZE_ALGO       "clip.vision.image_resize_algo"
 #define KEY_PROJ_SCALE_FACTOR       "clip.vision.projector.scale_factor"
 #define KEY_PROJ_SAMPLE_QUERY_SIDE  "clip.vision.projector.query_side"
 #define KEY_PROJ_SAMPLE_WINDOW_SIDE "clip.vision.projector.window_side"
diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp
index 06c562897..b7302e4bf 100644
--- a/tools/mtmd/clip.cpp
+++ b/tools/mtmd/clip.cpp
@@ -1526,9 +1526,20 @@ struct clip_model_loader {
                     } break;
                 case PROJECTOR_TYPE_LFM2:
                     {
-                        hparams.image_resize_algo    = RESIZE_ALGO_BILINEAR;
-                        hparams.image_resize_algo_rf = RESIZE_ALGO_BILINEAR;
-                        hparams.image_resize_algo_ov = RESIZE_ALGO_BILINEAR;
+                        // default for older GGUFs
+                        std::string resize_algo = "bilinear";
+                        get_string(KEY_IMAGE_RESIZE_ALGO, resize_algo, false);
+                        if (resize_algo == "bilinear") {
+                            hparams.image_resize_algo = RESIZE_ALGO_BILINEAR;
+                        } else if (resize_algo == "bicubic") {
+                            hparams.image_resize_algo = RESIZE_ALGO_BICUBIC;
+                        } else if (resize_algo == "lanczos") {
+                            hparams.image_resize_algo = RESIZE_ALGO_LANCZOS;
+                        } else {
+                            throw std::runtime_error(string_format("%s: unsupported image resize algo: %s\n", __func__, resize_algo.c_str()));
+                        }
+                        hparams.image_resize_algo_rf = hparams.image_resize_algo;
+                        hparams.image_resize_algo_ov = hparams.image_resize_algo;
                         get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge, false);
                         // ref: https://huggingface.co/LiquidAI/LFM2.5-VL-1.6B/blob/main/processor_config.json
                         hparams.set_limit_image_tokens(64, 256);
diff --git a/tools/server/README.md b/tools/server/README.md
index 79def4079..8b271395b 100644
--- a/tools/server/README.md
+++ b/tools/server/README.md
@@ -1703,7 +1703,7 @@ Follows the [TypeSafe API](https://docs.typesafe.ai/api), streaming is not suppo
 
 *Options:*
 
-`state`: The content to evaluate. Can be a string, an object or an array. A value that is not a string is given to the model as JSON text.
+`state`: The content to evaluate. Can be a string, an object or an array. A value that is not a string is given to the model as JSON text. For lfm2-d1, it can be `null`, for example to ask about images only.
 
 `images`: Optional. An array of images, the maximum number may be limited depending on the model. Each one is a data URL (`data:image/...;base64,...`). See the image input section below.
 
@@ -1718,13 +1718,13 @@ Follows the [TypeSafe API](https://docs.typesafe.ai/api), streaming is not suppo
 
 The questions of a request are answered independently, an answer does not depend on the other questions. The exception is clef: it reads all the questions in one prompt and decides them jointly.
 
-The number of options of a `choice` question is limited by the model, for example: 52 for openjev, 255 for laya, clef and pplx-decider. For laya, long questions and options are truncated to the token budget the model was trained with.
+The number of options of a `choice` question is limited by the model, for example: 52 for openjev, 255 for laya, clef, pplx-decider and lfm2-d1. For laya, long questions and options are truncated to the token budget the model was trained with.
 
 For laya and clef, the whole prompt is evaluated in one batch: it must fit in `--ubatch-size`. A server that runs clef only serves this endpoint, text generation is not available.
 
 *Image input:*
 
-Image input needs a model that supports it (for example: openjev, clef, pplx-decider) and its multimodal projector, see `--mmproj`.
+Image input needs a model that supports it (for example: openjev, clef, pplx-decider, lfm2-d1) and its multimodal projector, see `--mmproj`.
 
 Images can be given in two ways, and both can be used in the same request:
 
diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp
index 77dbec5e1..f56595231 100644
--- a/tools/server/server-common.cpp
+++ b/tools/server/server-common.cpp
@@ -156,6 +156,7 @@ std::vector server_model_output_modalities(common_decision_type dec
         case COMMON_DECISION_TYPE_LAYA:
         case COMMON_DECISION_TYPE_CLEF:
         case COMMON_DECISION_TYPE_PPLX_DECIDER:
+        case COMMON_DECISION_TYPE_LFM2_D1:
             return {"decisions"};
         default:
             // fallback when there is no decision type or the metadata is bad
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index 23835e185..6f3519215 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -2344,6 +2344,16 @@ private:
                 GGML_ASSERT(label >= 0 && label < n_vocab);
                 res->scores.push_back(logits[label]);
             }
+            if (!decision.label_groups.empty()) {
+                std::vector scores;
+                size_t i = 0;
+                for (const int32_t n : decision.label_groups) {
+                    GGML_ASSERT(n > 0 && i + n <= res->scores.size());
+                    scores.push_back(*std::max_element(res->scores.begin() + i, res->scores.begin() + i + n));
+                    i += n;
+                }
+                res->scores = std::move(scores);
+            }
         } else {
             // the outputs of this slot in this batch are the last tokens of the prompt
             std::vector idx;
diff --git a/tools/server/server-decision.cpp b/tools/server/server-decision.cpp
index 9ed68009f..71657a359 100644
--- a/tools/server/server-decision.cpp
+++ b/tools/server/server-decision.cpp
@@ -3,6 +3,7 @@
 #include "../../src/llama-ext.h" // staging API: llama_decision_order
 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -122,6 +123,9 @@ void server_decision_context::init(const llama_model * model) {
         n_options_max   = 255;
         noul_true_first = true;
         choice_sorted   = true;
+    } else if (model_type == COMMON_DECISION_TYPE_LFM2_D1) {
+        n_options_max   = 255;
+        noul_true_first = true;
     } else {
         throw std::runtime_error("unsupported decision model type: " + type_name);
     }
@@ -135,7 +139,8 @@ void server_decision_context::init(const llama_model * model) {
 //
 
 std::vector server_decision_context::parse_questions(const json & body) const {
-    if (!body.contains("state") || body.at("state").is_null()) {
+    // d1 accepts a null state (images only)
+    if (!body.contains("state") || (body.at("state").is_null() && type != COMMON_DECISION_TYPE_LFM2_D1)) {
         throw std::invalid_argument("\"state\" must be provided");
     }
     if (!body.contains("questions") || !body.at("questions").is_object() || body.at("questions").empty()) {
@@ -374,9 +379,112 @@ size_t server_decision_context::n_outputs(const server_decision_question & quest
     return question.options.size();
 }
 
+// label codes follow prompt.py of the model repo
+void server_decision_context::d1_labels(const server_decision_question & question, std::vector & texts, std::vector & groups) const {
+    const size_t n_options = question.options.size();
+
+    auto get_single_tokens = [&](const std::vector & forms) {
+        llama_tokens out;
+        for (const auto & form : forms) {
+            const auto toks = common_tokenize(vocab, form, false, false);
+            if (toks.size() == 1 && std::find(out.begin(), out.end(), toks[0]) == out.end()) {
+                out.push_back(toks[0]);
+            }
+        }
+        return out;
+    };
+
+    if (question.type != SERVER_DECISION_QUESTION_CHOICE) {
+        for (const auto & opt : question.options) {
+            llama_tokens group;
+            if (question.type == SERVER_DECISION_QUESTION_SCORE) {
+                group = get_single_tokens({opt.key});
+            } else if (opt.key == "true") {
+                group = get_single_tokens({"yes", "Yes", "YES"});
+            } else {
+                group = get_single_tokens({"no", "No", "NO"});
+            }
+            if (group.empty()) {
+                throw std::runtime_error("decision label is not a single token: " + opt.key);
+            }
+            texts.push_back(opt.key);
+            groups.push_back(group);
+        }
+        return;
+    }
+
+    bool is_letters = true;
+    for (const auto & opt : question.options) {
+        is_letters = is_letters && opt.key.size() == 1 && std::isalpha((unsigned char) opt.key[0]);
+    }
+
+    std::vector codes;
+    for (size_t i = 0; i < n_options; i++) {
+        if (is_letters) {
+            codes.push_back(question.options[i].key);
+        } else if (n_options <= 26) {
+            codes.push_back(std::string(1, 'A' + i));
+        } else {
+            codes.push_back(string_format("%02zu", i));
+        }
+    }
+
+    std::vector pool;
+    for (char c = 'A'; c <= 'Z'; c++) {
+        pool.push_back(std::string(1, c));
+    }
+    for (int i = 0; i < 100; i++) {
+        pool.push_back(string_format("%02d", i));
+    }
+    for (char c = 'a'; c <= 'z'; c++) {
+        pool.push_back(std::string(1, c));
+    }
+    for (int i = 0; i < 200; i++) {
+        pool.push_back(string_format("#%d", i));
+    }
+    for (char a = 'A'; a <= 'Z'; a++) {
+        for (char b = 'A'; b <= 'Z'; b++) {
+            pool.push_back(std::string{a, b});
+        }
+    }
+
+    llama_tokens used;
+    auto take = [&](const std::string & code) {
+        const auto toks = common_tokenize(vocab, code, false, false);
+        if (toks.size() != 1 || std::find(used.begin(), used.end(), toks[0]) != used.end()) {
+            return false;
+        }
+        used.push_back(toks[0]);
+        llama_tokens group = {toks[0]};
+        for (const llama_token tok : get_single_tokens({" " + code})) {
+            if (tok != toks[0]) {
+                group.push_back(tok);
+            }
+        }
+        texts.push_back(code);
+        groups.push_back(group);
+        return true;
+    };
+    for (const auto & code : codes) {
+        bool is_taken = take(code);
+        for (size_t i = 0; !is_taken && i < pool.size(); i++) {
+            is_taken = take(pool[i]);
+        }
+        if (!is_taken) {
+            throw std::invalid_argument(string_format("no single-token label left for %zu options", n_options));
+        }
+    }
+}
+
 json server_decision_context::render_options(const server_decision_question & question, size_t variant) const {
     const size_t n_options = question.options.size();
 
+    std::vector  d1_texts;
+    std::vector d1_groups;
+    if (type == COMMON_DECISION_TYPE_LFM2_D1) {
+        d1_labels(question, d1_texts, d1_groups);
+    }
+
     // the second variant shows the options in the reverse order
     json options = json::array();
     for (size_t i = 0; i < n_options; i++) {
@@ -394,6 +502,9 @@ json server_decision_context::render_options(const server_decision_question & qu
         if (!label_texts.empty()) {
             option["label"] = label_texts[i];
         }
+        if (!d1_texts.empty()) {
+            option["label"] = d1_texts[i];
+        }
         options.push_back(option);
     }
     return options;
@@ -474,6 +585,17 @@ void server_decision_context::fill_task(
     if (type == COMMON_DECISION_TYPE_OPENJEV || type == COMMON_DECISION_TYPE_LEV || type == COMMON_DECISION_TYPE_NIMBLE || type == COMMON_DECISION_TYPE_PPLX_DECIDER) {
         // lev reads the ratings of a noul question at its first labels, not at the digits
         task.decision.labels.assign(labels.begin(), labels.begin() + n_outputs(question));
+    }
+    if (type == COMMON_DECISION_TYPE_LFM2_D1) {
+        std::vector  texts;
+        std::vector groups;
+        d1_labels(question, texts, groups);
+        for (const auto & group : groups) {
+            task.decision.labels.insert(task.decision.labels.end(), group.begin(), group.end());
+            task.decision.label_groups.push_back(group.size());
+        }
+    }
+    if (type == COMMON_DECISION_TYPE_OPENJEV || type == COMMON_DECISION_TYPE_LEV || type == COMMON_DECISION_TYPE_NIMBLE || type == COMMON_DECISION_TYPE_PPLX_DECIDER || type == COMMON_DECISION_TYPE_LFM2_D1) {
         if (!files.empty()) {
             task.tokens = process_mtmd_prompt(mctx, prompt, files, init_opt);
             return;
diff --git a/tools/server/server-decision.h b/tools/server/server-decision.h
index f29211941..8f357f374 100644
--- a/tools/server/server-decision.h
+++ b/tools/server/server-decision.h
@@ -43,6 +43,7 @@ struct server_decision_context {
             case COMMON_DECISION_TYPE_KEV:
             case COMMON_DECISION_TYPE_NIMBLE:
             case COMMON_DECISION_TYPE_PPLX_DECIDER:
+            case COMMON_DECISION_TYPE_LFM2_D1:
                 return true;
             default:
                 return false;
@@ -60,6 +61,7 @@ struct server_decision_context {
             case COMMON_DECISION_TYPE_OPENJEV:
             case COMMON_DECISION_TYPE_CLEF:
             case COMMON_DECISION_TYPE_PPLX_DECIDER:
+            case COMMON_DECISION_TYPE_LFM2_D1:
                 return true;
             default:
                 return false;
@@ -129,6 +131,8 @@ private:
             size_t n_images) const;
     json render_options(const server_decision_question & question, size_t variant) const;
     size_t n_outputs(const server_decision_question & question) const;
+    // LFM2_D1: label text and tokens of each option
+    void d1_labels(const server_decision_question & question, std::vector & texts, std::vector & groups) const;
     void fill_task_laya(llama_tokens & tokens, const server_decision_question & question, server_task & task) const;
 
     float get_temperature(const server_decision_question & question) const;
diff --git a/tools/server/server-task.h b/tools/server/server-task.h
index 552c03b07..0cbc5f353 100644
--- a/tools/server/server-task.h
+++ b/tools/server/server-task.h
@@ -178,8 +178,9 @@ struct server_task {
     // used by SERVER_TASK_TYPE_DECISION
     // where to read the model output of each option, exactly one of the two lists is used
     struct decision {
-        std::vector labels;  // logits of these tokens, at the last prompt token
-        std::vector     markers; // embeddings[column] at these prompt positions
+        std::vector labels;       // logits of these tokens, at the last prompt token
+        std::vector     label_groups; // if set, number of labels per output, the output is their max
+        std::vector     markers;      // embeddings[column] at these prompt positions
         int32_t                  column = 0;
         // if set, embeddings is [q | k], and the output is instead the scaled dot product of q[pointer] and k[marker]
         int32_t                  pointer = -1;

From 5de733437bc5f6d6b2714b39c4023fd3d3b16698 Mon Sep 17 00:00:00 2001
From: Frost-54 <91168148+Frost-54@users.noreply.github.com>
Date: Thu, 8 Oct 2026 06:37:51 +1000
Subject: [PATCH 18/40] chat : name tool and argument parser rules by index
 (#30088)

* bugfix: infinite recursion caused by a tool named 'call'(#29967)

* chat : use index for schema and argument rules

* tests : remove tests

* tests : add expect_rules to peg test parser

---------

Co-authored-by: Alde Rojas 
---
 common/chat-auto-parser-generator.cpp | 16 ++++++------
 common/chat-peg-parser.cpp            | 34 ++++++++++++++-----------
 common/parsers/deepseek.cpp           | 10 ++++----
 common/parsers/functionary-v3-2.cpp   |  6 ++---
 common/parsers/gemma4.cpp             |  4 +--
 common/parsers/gigachat-v3.cpp        |  7 +++---
 common/parsers/gpt-oss.cpp            |  6 ++---
 common/parsers/k2-horizon.cpp         | 10 ++++----
 common/parsers/kimi-k2.cpp            |  6 ++---
 common/parsers/kimi-k3.cpp            |  7 +++---
 common/parsers/ling3.cpp              | 10 ++++----
 common/parsers/llm-jp-harmony.cpp     |  6 ++---
 common/parsers/minicpm5.cpp           |  8 +++---
 common/parsers/minimax-m3.cpp         | 11 ++++----
 common/parsers/ministral3.cpp         |  6 ++---
 common/parsers/muse-glimmer.cpp       |  8 +++---
 common/parsers/parsers.cpp            | 13 +++++-----
 common/parsers/parsers.h              |  8 +++---
 common/parsers/qwen3-coder.cpp        | 12 ++++-----
 docs/development/parsing.md           |  7 +++---
 tests/test-chat.cpp                   | 36 ++++++++++++++++++++++++++-
 21 files changed, 138 insertions(+), 93 deletions(-)

diff --git a/common/chat-auto-parser-generator.cpp b/common/chat-auto-parser-generator.cpp
index b78789d8c..9e3b6d640 100644
--- a/common/chat-auto-parser-generator.cpp
+++ b/common/chat-auto-parser-generator.cpp
@@ -291,7 +291,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_json(parser_build_context
 
     common_peg_parser tool_choice = p.choice();
 
-    foreach_function(inputs.tools, [&](const json & tool) {
+    foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
         const auto & func   = tool.at("function");
         std::string  name   = func.at("name");
         const auto   schema = common_chat_tool_parameters(func);
@@ -308,7 +308,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_json(parser_build_context
             }
             have_call_id = true;
         }
-        auto args_parser = p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema));
+        auto args_parser = p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", schema));
         if (!arguments.start.empty()) {
             args_parser = p.literal(arguments.start) + args_parser;
         }
@@ -318,7 +318,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_json(parser_build_context
 
         auto atomic_peek = !arguments.start.empty() ? std::optional(p.peek(p.literal(arguments.start))) : std::nullopt;
         auto func_parser = build_func_parser(p, name, call_id_section, have_call_id, args_parser, atomic_peek);
-        tool_choice |= p.rule("tool-" + name, func_parser);
+        tool_choice |= p.rule("tool-" + std::to_string(tool_index), func_parser);
     });
 
     auto require_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED;
@@ -364,14 +364,14 @@ common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_conte
 
     common_peg_parser tool_choice = p.choice();
 
-    foreach_function(inputs.tools, [&](const json & tool) {
+    foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
         const auto & func = tool.at("function");
         std::string  name = func.at("name");
 
         // Build parser for each argument, separating required and optional
         std::vector required_parsers;
         std::vector optional_parsers;
-        foreach_parameter(func, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
+        foreach_parameter(func, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
             auto arg =
                 p.tool_arg(p.tool_arg_open(arguments.name_prefix + p.tool_arg_name(p.literal(param.name)) +
                                            arguments.name_suffix) +
@@ -380,10 +380,10 @@ common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_conte
                                 p.ac(p.tool_arg_string_value(until_suffix) +
                                     p.tool_arg_close(p.literal(arguments.value_suffix)), arguments.value_suffix) :
                                 (p.tool_arg_json_value(p.schema(
-                                    p.json(), "tool-" + name + "-arg-" + param.name + "-schema", doc, *param.schema)) +
+                                    p.json(), "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index) + "-schema", doc, *param.schema)) +
                                     p.tool_arg_close(p.literal(arguments.value_suffix)))));
 
-            auto named_arg = p.rule("tool-" + name + "-arg-" + param.name, arg);
+            auto named_arg = p.rule("tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index), arg);
             if (param.required) {
                 required_parsers.push_back(named_arg);
             } else {
@@ -434,7 +434,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_conte
         auto atomic_peek = (!arguments.name_prefix.empty() && !required_parsers.empty()) ?
             std::optional(p.peek(p.literal(arguments.name_prefix))) : std::nullopt;
         auto func_parser = build_func_parser(p, name, call_id_section, have_call_id, args_seq, atomic_peek);
-        tool_choice |= p.rule("tool-" + name, func_parser);
+        tool_choice |= p.rule("tool-" + std::to_string(tool_index), func_parser);
     });
 
     auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED;
diff --git a/common/chat-peg-parser.cpp b/common/chat-peg-parser.cpp
index ee9a05fe0..74b81dcdd 100644
--- a/common/chat-peg-parser.cpp
+++ b/common/chat-peg-parser.cpp
@@ -483,7 +483,9 @@ common_peg_parser common_chat_peg_builder::standard_constructed_tools(
     // Build tool choices for tagged format
     auto tool_choices = choice();
 
-    for (const auto & tool_def : tools) {
+    for (size_t i = 0; i < tools.size(); i++) {
+        const auto & tool_def = tools[i];
+
         if (!tool_def.contains("function")) {
             continue;
         }
@@ -513,7 +515,7 @@ common_peg_parser common_chat_peg_builder::standard_constructed_tools(
         auto tool_parser = tool(tool_open(literal(func_opener) + tool_name(literal(name)) + literal(func_name_suffix)) +
                                 space() + tool_args(args) + space() + tool_close(literal(func_closer)));
 
-        tool_choices |= rule("tool-" + name, tool_parser);
+        tool_choices |= rule("tool-" + std::to_string(i), tool_parser);
     }
 
     // Build the section with markers
@@ -560,7 +562,8 @@ common_peg_parser common_chat_peg_builder::python_style_tool_calls(
 
     auto tool_choices = choice();
 
-    for (const auto & tool_def : tools) {
+    for (size_t i = 0; i < tools.size(); i++) {
+        const auto & tool_def = tools[i];
         if (!tool_def.contains("function")) {
             continue;
         }
@@ -607,7 +610,7 @@ common_peg_parser common_chat_peg_builder::python_style_tool_calls(
             space() + tool_args(args) + space() + tool_close(literal(")"))
         );
 
-        tool_choices |= rule("tool-" + name, tool_parser);
+        tool_choices |= rule("tool-" + std::to_string(i), tool_parser);
     }
 
     if (parallel_tool_calls) {
@@ -635,7 +638,8 @@ common_peg_parser common_chat_peg_builder::build_json_tools_function_is_key(
 
     auto tool_choices = choice();
 
-    for (const auto & tool_def : tools) {
+    for (size_t i = 0; i < tools.size(); i++) {
+        const auto & tool_def = tools[i];
         if (!tool_def.contains("function")) {
             continue;
         }
@@ -668,10 +672,10 @@ common_peg_parser common_chat_peg_builder::build_json_tools_function_is_key(
         // Arguments — either wrapped in args_key or parsed directly
         common_peg_parser args_parser = eps();
         if (args_key.empty()) {
-            args_parser = tool_args(schema(json(), "tool-" + name + "-schema", params));
+            args_parser = tool_args(schema(json(), "tool-" + std::to_string(i) + "-schema", params));
         } else {
             args_parser = literal("\"" + effective_args_key + "\"") + space() + literal(":") + space() +
-                          tool_args(schema(json(), "tool-" + name + "-schema", params));
+                          tool_args(schema(json(), "tool-" + std::to_string(i) + "-schema", params));
         }
         inner_fields.push_back(args_parser);
 
@@ -698,7 +702,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_function_is_key(
             space() + tool_close(literal("}"))
         );
 
-        tool_choices |= rule("tool-" + name, tool_parser);
+        tool_choices |= rule("tool-" + std::to_string(i), tool_parser);
     }
 
     return tool_choices;
@@ -721,7 +725,8 @@ common_peg_parser common_chat_peg_builder::build_json_tools_nested_keys(
     std::string nested_name_field = !name_spec.first.empty() ? name_spec.second  : effective_name_key;
     std::string nested_args_field = !args_spec.first.empty() ? args_spec.second  : effective_args_key;
 
-    for (const auto & tool_def : tools) {
+    for (size_t i = 0; i < tools.size(); i++) {
+        const auto & tool_def = tools[i];
         if (!tool_def.contains("function")) {
             continue;
         }
@@ -732,7 +737,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_nested_keys(
         auto nested_name = literal("\"" + nested_name_field + "\"") + space() + literal(":") + space() +
                           atomic(literal("\"") + tool_name(literal(name)) + literal("\""));
         auto nested_args = literal("\"" + nested_args_field + "\"") + space() + literal(":") + space() +
-                          tool_args(schema(json(), "tool-" + name + "-schema", params));
+                          tool_args(schema(json(), "tool-" + std::to_string(i) + "-schema", params));
 
         auto nested_object = literal("{") + space() +
                             nested_name + space() + literal(",") + space() +
@@ -770,7 +775,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_nested_keys(
         auto nested_field = literal("\"" + nested_prefix + "\"") + space() + literal(":") + space() + nested_object;
         tool_parser_body = tool_parser_body + nested_field + space() + tool_close(literal("}"));
 
-        tool_choices |= rule("tool-" + name, tool(tool_parser_body));
+        tool_choices |= rule("tool-" + std::to_string(i), tool(tool_parser_body));
     }
 
     return tool_choices;
@@ -790,7 +795,8 @@ common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys(
     auto name_key_parser = literal("\"" + effective_name_key + "\"");
     auto args_key_parser = literal("\"" + effective_args_key + "\"");
 
-    for (const auto & tool_def : tools) {
+    for (size_t i = 0; i < tools.size(); i++) {
+        const auto & tool_def = tools[i];
         if (!tool_def.contains("function")) {
             continue;
         }
@@ -801,7 +807,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys(
         auto tool_name_ = name_key_parser + space() + literal(":") + space() +
                          atomic(literal("\"") + tool_name(literal(name)) + literal("\""));
         auto tool_args_ = args_key_parser + space() + literal(":") + space() +
-                         tool_args(schema(json(), "tool-" + name + "-schema", params));
+                         tool_args(schema(json(), "tool-" + std::to_string(i) + "-schema", params));
 
         // Build ID parsers if keys are provided
         common_peg_parser id_parser = eps();
@@ -861,7 +867,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys(
         }
         ordered_body = ordered_body + space() + tool_close(literal("}"));
 
-        tool_choices |= rule("tool-" + name, tool(ordered_body));
+        tool_choices |= rule("tool-" + std::to_string(i), tool(ordered_body));
     }
 
     return tool_choices;
diff --git a/common/parsers/deepseek.cpp b/common/parsers/deepseek.cpp
index 9ca4bb34c..88122eaab 100644
--- a/common/parsers/deepseek.cpp
+++ b/common/parsers/deepseek.cpp
@@ -152,13 +152,13 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
         // build tool call section first since we might need it in reasoning
         auto tool_choice = p.choice();
         if (has_tool_calls) {
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
 
                 std::vector required_parsers;
                 std::vector optional_parsers;
-                foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
+                foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
                     bool is_string = param.schema->may_be_string();
 
                     auto arg = p.tool_arg(
@@ -166,11 +166,11 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
                                         p.literal("\" string=\"" + std::string(is_string ? "true" : "false") + "\">")) +
                         (is_string ?
                              p.tool_arg_string_value(p.until(PARAM_END)) :
-                             p.tool_arg_json_value(p.schema(p.json(), "tool-" + name + "-arg-" + param.name + "-schema",
+                             p.tool_arg_json_value(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index) + "-schema",
                                                             doc, *param.schema))) +
                         p.tool_arg_close(p.literal(PARAM_END)));
 
-                    auto named_arg = p.rule("tool-" + name + "-arg-" + param.name, arg);
+                    auto named_arg = p.rule("tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index), arg);
                     if (param.required) {
                         required_parsers.push_back(named_arg);
                     } else {
@@ -199,7 +199,7 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
                                                                    p.tool_name(p.literal(name)) + p.literal("\">\n")) +
                                                        invoke_body + p.space() + p.tool_close(p.literal(INVOKE_END)));
 
-                tool_choice |= p.rule("tool-" + name, func_parser);
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), func_parser);
             });
         }
 
diff --git a/common/parsers/functionary-v3-2.cpp b/common/parsers/functionary-v3-2.cpp
index 9d47f0a33..56f7643f2 100644
--- a/common/parsers/functionary-v3-2.cpp
+++ b/common/parsers/functionary-v3-2.cpp
@@ -42,7 +42,7 @@ common_chat_params common_chat_params_init_functionary_v3_2(const common_chat_te
 
         // Build tool call parsers for each available function
         auto tool_choice = p.choice();
-        foreach_function(inputs.tools, [&](const json & tool) {
+        foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
             const auto & function = tool.at("function");
             std::string  name     = function.at("name");
             const auto   schema   = common_chat_tool_parameters(function);
@@ -50,10 +50,10 @@ common_chat_params common_chat_params_init_functionary_v3_2(const common_chat_te
             // Tool format: >>>function_name\n{json_args}
             auto tool_parser = p.tool(
                 p.tool_open(p.tool_name(p.literal(name)) + p.literal("\n")) +
-                p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema))
+                p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", schema))
             );
 
-            tool_choice |= p.rule("tool-" + name, tool_parser);
+            tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_parser);
         });
 
         auto content_only = content_until_end;
diff --git a/common/parsers/gemma4.cpp b/common/parsers/gemma4.cpp
index f43c5ae46..2223089da 100644
--- a/common/parsers/gemma4.cpp
+++ b/common/parsers/gemma4.cpp
@@ -254,13 +254,13 @@ common_chat_params common_chat_params_init_gemma4(const common_chat_template &
 
             auto tool_choice = p.choice();
 
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
                 // TODO @aldehir : need to extend json-schema-to-grammar to produce more than JSON rules
                 // const auto & params   = function.at("parameters");
 
-                tool_choice |= p.rule("tool-" + name, p.tool(p.sequence({
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), p.tool(p.sequence({
                     p.tool_open(p.tool_name(p.literal(name)) + p.peek(p.literal("{"))),
                     p.tool_args(p.ref("gemma4-dict")),
                 })));
diff --git a/common/parsers/gigachat-v3.cpp b/common/parsers/gigachat-v3.cpp
index 48abcb3f3..949942547 100644
--- a/common/parsers/gigachat-v3.cpp
+++ b/common/parsers/gigachat-v3.cpp
@@ -30,17 +30,18 @@ common_chat_params common_chat_params_init_gigachat_v3(
         if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
             // Build a choice of all available tools
             auto tool_choice = p.choice();
-            for (const auto & tool : inputs.tools) {
+            for (size_t i = 0; i < inputs.tools.size(); i++) {
+                const auto & tool = inputs.tools[i];
                 const auto & function = tool.at("function");
                 std::string name = function.at("name");
                 const auto  schema = common_chat_tool_parameters(function);
 
                 auto tool_name = p.json_member("name", "\"" + p.tool_name(p.literal(name)) + "\"");
-                auto tool_args = p.json_member("arguments", p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)));
+                auto tool_args = p.json_member("arguments", p.tool_args(p.schema(p.json(), "tool-" + std::to_string(i) + "-schema", schema)));
 
                 auto tool_open = p.tool_open(p.literal("{") << tool_name);
 
-                tool_choice |= p.rule("tool-" + name, tool_open << "," << tool_args << "}");
+                tool_choice |= p.rule("tool-" + std::to_string(i), tool_open << "," << tool_args << "}");
             }
 
             // Define the tool call structure
diff --git a/common/parsers/gpt-oss.cpp b/common/parsers/gpt-oss.cpp
index 00beb41a4..ed50524a9 100644
--- a/common/parsers/gpt-oss.cpp
+++ b/common/parsers/gpt-oss.cpp
@@ -106,14 +106,14 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
         if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
             auto tool_choice = p.choice();
 
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
                 const auto   params   = common_chat_tool_parameters(function);
 
                 auto func_name  = p.literal(" to=functions.") + p.tool_name(p.literal(name));
                 auto constraint = p.optional(p.space() + p.optional(p.literal("<|constrain|>")) + constrain_type);
-                auto args       = p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", params));
+                auto args       = p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", params));
 
                 // recipient in role header
                 //   <|start|>assistant to=functions.NAME<|channel|>(commentary|analysis)[constraint]<|message|>ARGS
@@ -123,7 +123,7 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
                 //   <|channel|>(commentary|analysis) to=functions.NAME[constraint]<|message|>ARGS
                 auto tool_in_channel = p.tool(p.tool_open(channel + func_name + constraint + p.literal("<|message|>")) + args);
 
-                tool_choice |= p.rule("tool-" + name, tool_in_role | tool_in_channel);
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_in_role | tool_in_channel);
             });
 
             auto tool_call  = p.trigger_rule("tool-call", tool_choice);
diff --git a/common/parsers/k2-horizon.cpp b/common/parsers/k2-horizon.cpp
index 687a893ce..bd64f6890 100644
--- a/common/parsers/k2-horizon.cpp
+++ b/common/parsers/k2-horizon.cpp
@@ -110,14 +110,14 @@ common_chat_params common_chat_params_init_k2_horizon(const common_chat_template
             // The models leave out  even when asked for xml_typed
             auto arg_type = call_format == "xml_typed" ? p.optional(ARG_TYPE + p.until(ARG_TYPE_END) + ARG_TYPE_END + p.space()) : p.eps();
 
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
 
                 std::vector required_args;
                 std::vector optional_args;
-                foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
-                    auto rule_name = "tool-" + name + "-arg-" + param.name;
+                foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
+                    auto rule_name = "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index);
                     auto types     = param.schema->value_types();
                     auto arg_value = arg_string;
                     if (!types.has(common_chat_schema::TYPE_STRING)) {
@@ -149,12 +149,12 @@ common_chat_params common_chat_params_init_k2_horizon(const common_chat_template
                     (param.required ? required_args : optional_args).push_back(p.rule(rule_name, arg));
                 });
 
-                auto args = p.permute("tool-" + name + "-args", required_args);
+                auto args = p.permute("tool-" + std::to_string(tool_index) + "-args", required_args);
                 if (!optional_args.empty()) {
                     args = args + p.zero_or_more(p.choice(optional_args));
                 }
 
-                tool_choice |= p.rule("tool-" + name, p.tool(
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), p.tool(
                     p.tool_open(CALL_START + p.tool_name(p.literal(name)) + "\n") + p.tool_args(args) << p.tool_close(p.literal(CALL_END))));
             });
         }
diff --git a/common/parsers/kimi-k2.cpp b/common/parsers/kimi-k2.cpp
index 5ee9121ab..9992968a5 100644
--- a/common/parsers/kimi-k2.cpp
+++ b/common/parsers/kimi-k2.cpp
@@ -79,7 +79,7 @@ common_chat_params common_chat_params_init_kimi_k2(const common_chat_template &
         // The ID format is: functions.:
         // We need to match: functions.:
         auto tool_choice = p.choice();
-        foreach_function(inputs.tools, [&](const json & tool) {
+        foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
             const auto & function = tool.at("function");
             std::string  name     = function.at("name");
             const auto   schema   = common_chat_tool_parameters(function);
@@ -89,11 +89,11 @@ common_chat_params common_chat_params_init_kimi_k2(const common_chat_template &
             auto tool_id = p.tool_id(p.literal("functions.") + p.tool_name(p.literal(name)) + p.literal(":") + p.chars("[0-9]", 1, -1));
             auto tool_parser = p.tool(
                 p.tool_open(tool_id + p.literal(ARGS_BEGIN)) +
-                p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)) +
+                p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", schema)) +
                 p.tool_close(p.optional((p.literal(CALL_END))))
             );
 
-            tool_choice |= p.rule("tool-" + name, tool_parser);
+            tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_parser);
         });
 
         // Tool calls section: <|tool_calls_section_begin|> tool_calls <|tool_calls_section_end|>
diff --git a/common/parsers/kimi-k3.cpp b/common/parsers/kimi-k3.cpp
index 989e39f95..cf153f6b5 100644
--- a/common/parsers/kimi-k3.cpp
+++ b/common/parsers/kimi-k3.cpp
@@ -95,7 +95,7 @@ common_chat_params common_chat_params_init_kimi_k3(const common_chat_template &
         }
 
         auto tool_choices = p.choice();
-        foreach_function(inputs.tools, [&](const json & tool) {
+        foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
             const auto & function = tool.at("function");
             std::string  name     = function.at("name");
             const json   schema   = common_chat_tool_parameters(function);
@@ -106,6 +106,7 @@ common_chat_params common_chat_params_init_kimi_k3(const common_chat_template &
             auto args = p.eps();
             if (schema.contains("properties") && !schema.at("properties").empty()) {
                 auto arg_choices = p.choice();
+                size_t param_index = 0;
                 for (const auto & prop : schema.at("properties").items()) {
                     const std::string & key = prop.key();
 
@@ -119,7 +120,7 @@ common_chat_params common_chat_params_init_kimi_k3(const common_chat_template &
                                                     p.tool_arg_value(p.until(ARG_END));
 
                     // skip the trailing type="..." attribute: anything up to <|sep|>
-                    arg_choices |= p.rule("kimi-k3-arg-" + name + "-" + key,
+                    arg_choices |= p.rule("kimi-k3-arg-" + std::to_string(tool_index) + "-" + std::to_string(param_index++),
                                           p.tool_arg(p.tool_arg_open(p.literal(ARG_START)) +
                                                      p.tool_arg_name(p.literal(key)) + p.literal("\"") +
                                                      p.until(SEP) + p.literal(SEP) + value +
@@ -133,7 +134,7 @@ common_chat_params common_chat_params_init_kimi_k3(const common_chat_template &
                                            p.until(SEP) + p.literal(SEP)) +
                                p.tool_args(args) + p.tool_close(p.literal(CALL_END)));
 
-            tool_choices |= p.rule("kimi-k3-tool-" + name, call);
+            tool_choices |= p.rule("kimi-k3-tool-" + std::to_string(tool_index), call);
         });
 
         // all calls go inside one tools section, then the message is closed. the
diff --git a/common/parsers/ling3.cpp b/common/parsers/ling3.cpp
index 8b60e40f4..67d44d47a 100644
--- a/common/parsers/ling3.cpp
+++ b/common/parsers/ling3.cpp
@@ -118,7 +118,7 @@ common_chat_params common_chat_params_init_ling3(const common_chat_template &
         auto arg_string   = p.rule("ling3-arg-string",
                                    p.tool_arg_string_value(p.until(ARG_VAL_END)) + arg_close);
 
-        foreach_function(inputs.tools, [&](const json & tool) {
+        foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
             const auto & function = tool.at("function");
             std::string  name     = function.at("name");
 
@@ -127,8 +127,8 @@ common_chat_params common_chat_params_init_ling3(const common_chat_template &
 
             // each argument may be preceded by whitespace: the model emits
             // newlines between arguments, the template history does not
-            foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
-                auto rule_name = "ling3-arg-" + name + "-" + param.name;
+            foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
+                auto rule_name = "ling3-arg-" + std::to_string(tool_index) + "-" + std::to_string(param_index);
 
                 auto types = param.schema->value_types();
 
@@ -159,7 +159,7 @@ common_chat_params common_chat_params_init_ling3(const common_chat_template &
 
             // required arguments in any order (as Qwen3-Coder does), then
             // optional ones in any order and number
-            auto args = p.permute("ling3-" + name + "-args", required_args);
+            auto args = p.permute("ling3-" + std::to_string(tool_index) + "-args", required_args);
             if (!optional_args.empty()) {
                 args = args + p.zero_or_more(p.choice(optional_args));
             }
@@ -169,7 +169,7 @@ common_chat_params common_chat_params_init_ling3(const common_chat_template &
                                p.tool_args(args) +
                                p.tool_close(p.optional(p.space()) + p.literal(CALL_END)));
 
-            tool_choices |= p.rule("ling3-tool-" + name, call);
+            tool_choices |= p.rule("ling3-tool-" + std::to_string(tool_index), call);
         });
 
         auto calls = inputs.parallel_tool_calls ?
diff --git a/common/parsers/llm-jp-harmony.cpp b/common/parsers/llm-jp-harmony.cpp
index b012f9a05..cafb82c12 100644
--- a/common/parsers/llm-jp-harmony.cpp
+++ b/common/parsers/llm-jp-harmony.cpp
@@ -109,13 +109,13 @@ common_chat_params common_chat_params_init_llm_jp_harmony(const common_chat_temp
         if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
             auto tool_choice = p.choice();
 
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
                 const auto   params   = common_chat_tool_parameters(function);
 
                 auto func_name = p.literal(" to=functions.") + p.tool_name(p.literal(name));
-                auto args      = p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", params));
+                auto args      = p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", params));
 
                 // recipient in role header
                 //   <|start|>assistant to=functions.NAME<|channel|>(commentary|analysis)[constraint]<|message|>ARGS
@@ -125,7 +125,7 @@ common_chat_params common_chat_params_init_llm_jp_harmony(const common_chat_temp
                 //   <|channel|>(commentary|analysis) to=functions.NAME[constraint]<|message|>ARGS
                 auto tool_in_channel = p.tool(p.tool_open(channel + func_name + constraint + message) + args);
 
-                tool_choice |= p.rule("tool-" + name, tool_in_role | tool_in_channel);
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_in_role | tool_in_channel);
             });
 
             // parallel calls are separated by <|end|>; inside the trigger rule so the lazy grammar covers all of them
diff --git a/common/parsers/minicpm5.cpp b/common/parsers/minicpm5.cpp
index 4d18d3d96..06f4cc8c8 100644
--- a/common/parsers/minicpm5.cpp
+++ b/common/parsers/minicpm5.cpp
@@ -68,18 +68,18 @@ common_chat_params common_chat_params_init_minicpm5(const common_chat_template &
             });
 
             auto tool_choice = p.choice();
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto &      function = tool.at("function");
                 const std::string name     = function.at("name");
 
                 std::vector arg_rules;
-                foreach_parameter(function, [&](const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) {
+                foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) {
                     auto value_parser = p.eps();
                     if (prop.schema->may_be_string()) {
                         value_parser = string_value;
                     } else {
                         value_parser = p.tool_arg_json_value(
-                                p.schema(p.json(), "tool-" + name + "-arg-" + prop.name + "-schema", doc, *prop.schema)
+                                p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index) + "-schema", doc, *prop.schema)
                             ) + p.tool_arg_close(p.literal(""));
                     }
 
@@ -99,7 +99,7 @@ common_chat_params common_chat_params_init_minicpm5(const common_chat_template &
                     << p.tool_args(args)
                     << p.tool_close(p.literal("")));
 
-                tool_choice |= p.rule("tool-" + name, tool_parser);
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_parser);
             });
 
             auto max_calls  = inputs.parallel_tool_calls ? -1 : 1;
diff --git a/common/parsers/minimax-m3.cpp b/common/parsers/minimax-m3.cpp
index 7ea9bfe5a..c38ab3e23 100644
--- a/common/parsers/minimax-m3.cpp
+++ b/common/parsers/minimax-m3.cpp
@@ -85,7 +85,7 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
         }
 
         auto tool_choice = p.choice();
-        foreach_function(inputs.tools, [&](const json & tool) {
+        foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
             const auto & function = tool.at("function");
             std::string  name     = function.at("name");
             auto         params   = common_chat_tool_parameters(function);
@@ -154,8 +154,9 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
             members_of = [&](const common_chat_schema_object & object, const std::string & rule_prefix) -> common_peg_parser {
                 std::vector required_elements;
                 std::vector optional_elements;
-                for (const auto & prop : object.properties) {
-                    auto element = element_of(prop.name, *prop.schema, rule_prefix + "-" + prop.name);
+                for (size_t i = 0; i < object.properties.size(); i++) {
+                    const auto & prop    = object.properties[i];
+                    auto         element = element_of(prop.name, *prop.schema, rule_prefix + "-" + std::to_string(i));
                     (prop.required ? required_elements : optional_elements).push_back(element);
                 }
 
@@ -180,7 +181,7 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
 
             common_peg_parser invoke_body = p.eps();
             if (doc->root->kind() == common_chat_schema::KIND_OBJECT) {
-                invoke_body = members_of(static_cast(*doc->root), "tool-" + name + "-arg");
+                invoke_body = members_of(static_cast(*doc->root), "tool-" + std::to_string(tool_index) + "-arg");
             }
 
             auto func_parser = p.tool(
@@ -189,7 +190,7 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
                 p.space() + invoke_body + p.space() +
                 p.tool_close(p.literal(INVOKE_END)));
 
-            tool_choice |= p.rule("tool-" + name, func_parser);
+            tool_choice |= p.rule("tool-" + std::to_string(tool_index), func_parser);
         });
 
         auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED;
diff --git a/common/parsers/ministral3.cpp b/common/parsers/ministral3.cpp
index 8bcd204c6..9cc5248b5 100644
--- a/common/parsers/ministral3.cpp
+++ b/common/parsers/ministral3.cpp
@@ -86,14 +86,14 @@ common_chat_params common_chat_params_init_ministral_3(const common_chat_templat
         // Tool call parser
         if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
             auto tool_choice = p.choice();
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
                 const auto   schema   = common_chat_tool_parameters(function);
 
                 tool_choice |=
-                    p.rule("tool-" + name, p.tool_open(p.tool_name(p.literal(name)) + "[ARGS]") +
-                                               p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)));
+                    p.rule("tool-" + std::to_string(tool_index), p.tool_open(p.tool_name(p.literal(name)) + "[ARGS]") +
+                                               p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", schema)));
             });
 
             auto min_calls  = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0;
diff --git a/common/parsers/muse-glimmer.cpp b/common/parsers/muse-glimmer.cpp
index 784928269..6c7303701 100644
--- a/common/parsers/muse-glimmer.cpp
+++ b/common/parsers/muse-glimmer.cpp
@@ -81,18 +81,18 @@ common_chat_params common_chat_params_init_muse_glimmer(const common_chat_templa
                 "");
 
             auto tool_choice = p.choice();
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto &      function = tool.at("function");
                 const std::string name     = function.at("name");
 
                 std::vector arg_rules;
-                foreach_parameter(function, [&](const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) {
+                foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) {
                     auto value_parser = p.eps();
                     if (prop.schema->may_be_string()) {
                         value_parser = string_value;
                     } else {
                         value_parser = p.tool_arg_json_value(
-                                p.schema(p.json(), "tool-" + name + "-arg-" + prop.name + "-schema", doc, *prop.schema))
+                                p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index) + "-schema", doc, *prop.schema))
                             + p.tool_arg_close(p.literal(""));
                     }
 
@@ -113,7 +113,7 @@ common_chat_params common_chat_params_init_muse_glimmer(const common_chat_templa
                     << p.tool_args(args)
                     << p.tool_close(p.literal("") + p.space() + p.literal("")));
 
-                tool_choice |= p.rule("tool-" + name, tool_parser);
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_parser);
             });
 
             auto tool_calls = inputs.parallel_tool_calls
diff --git a/common/parsers/parsers.cpp b/common/parsers/parsers.cpp
index 643186c34..2629fb3e9 100644
--- a/common/parsers/parsers.cpp
+++ b/common/parsers/parsers.cpp
@@ -2,24 +2,25 @@
 
 #include "log.h"
 
-void foreach_function(const json & tools, const std::function & fn) {
-    for (const auto & tool : tools) {
+void foreach_function(const json & tools, const std::function & fn) {
+    for (size_t i = 0; i < tools.size(); i++) {
+        const auto & tool = tools[i];
         if (!tool.contains("type") || tool.at("type") != "function" || !tool.contains("function")) {
             LOG_INF("Skipping tool without function: %s", tool.dump(2).c_str());
             continue;
         }
-        fn(tool);
+        fn(i, tool);
     }
 }
 
-void foreach_parameter(const json & function, const std::function & fn) {
+void foreach_parameter(const json & function, const std::function & fn) {
     auto         params = common_chat_tool_parameters(function);
     auto         doc    = std::make_shared(common_chat_schema_from_json(params));
     const auto * object = dynamic_cast(doc->root.get());
     if (!object) {
         return;
     }
-    for (const auto & prop : object->properties) {
-        fn(prop, doc);
+    for (size_t i = 0; i < object->properties.size(); i++) {
+        fn(i, object->properties[i], doc);
     }
 }
diff --git a/common/parsers/parsers.h b/common/parsers/parsers.h
index 7786b82f4..1b385264e 100644
--- a/common/parsers/parsers.h
+++ b/common/parsers/parsers.h
@@ -17,11 +17,11 @@
 
 using json = common_json;
 
-// iterate over the function tools of an OpenAI-style tools array
-void foreach_function(const json & tools, const std::function & fn);
+// iterate over the function tools of an OpenAI-style tools array, passing each tool with its index in the array
+void foreach_function(const json & tools, const std::function & fn);
 
-// iterate over the parameters of a function tool, with the document that owns them
-void foreach_parameter(const json & function, const std::function & fn);
+// iterate over the parameters of a function tool, passing each parameter with its index and the document that owns it
+void foreach_parameter(const json & function, const std::function & fn);
 
 // render a template; the override arguments let a parser feed in messages, tools or context it has rewritten
 std::string common_chat_template_direct_apply_impl(
diff --git a/common/parsers/qwen3-coder.cpp b/common/parsers/qwen3-coder.cpp
index 208f551f1..323b3da31 100644
--- a/common/parsers/qwen3-coder.cpp
+++ b/common/parsers/qwen3-coder.cpp
@@ -65,7 +65,7 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
         // Match complete  opener for Qwen3-Coder models that occasionally omit the
         // starting . The model may hallucinate a tool name, but it is preferable over
         // constraining on 
-        foreach_function(inputs.tools, [&](const json & tool) {
+        foreach_function(inputs.tools, [&](size_t, const json & tool) {
             const std::string name = tool.at("function").at("name");
             tool_call_starts.push_back("");
         });
@@ -93,15 +93,15 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
                 p.ac(p.tool_arg_string_value(p.until("\n\n")) + arg_close, "\n\n"));
 
             auto tool_choice = p.choice();
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
 
                 std::vector required_args;
                 std::vector optional_args;
 
-                foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
-                    auto rule_name = "tool-" + name + "-arg-" + param.name;
+                foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
+                    auto rule_name = "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index);
 
                     auto arg_open = p.tool_arg_open("\n");
 
@@ -141,7 +141,7 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
 
                 // Accept required arguments in any order, as Qwen does not always adhere to the
                 // order provided.
-                auto args = p.permute("tool-" + name + "-args", required_args);
+                auto args = p.permute("tool-" + std::to_string(tool_index) + "-args", required_args);
                 if (!optional_args.empty()) {
                     args = args + p.zero_or_more(p.choice(optional_args));
                 }
@@ -150,7 +150,7 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
                                    p.tool_args(args) +
                                    p.tool_close(p.literal("\n")));
 
-                tool_choice |= p.rule("tool-" + name, func);
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), func);
             });
 
             auto min_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0;
diff --git a/docs/development/parsing.md b/docs/development/parsing.md
index 0cb372eca..ab0ccad80 100644
--- a/docs/development/parsing.md
+++ b/docs/development/parsing.md
@@ -25,15 +25,16 @@ output from a model that emits arguments as JSON.
 auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
     // Build a choice of all available tools
     auto tool_choice = p.choice();
-    for (const auto & tool : tools) {
+    for (size_t i = 0; i < tools.size(); i++) {
+        const auto & tool = tools[i];
         const auto & function = tool.at("function");
         std::string name = function.at("name");
         const auto   schema = common_chat_tool_parameters(function);
 
         auto tool_name = p.json_member("name", "\"" + p.literal(name) + "\"");
-        auto tool_args = p.json_member("arguments", p.schema(p.json(), "tool-" + name + "-schema", schema));
+        auto tool_args = p.json_member("arguments", p.schema(p.json(), "tool-" + std::to_string(i) + "-schema", schema));
 
-        tool_choice |= p.rule("tool-" + name, "{" << tool_name << "," << tool_args << "}");
+        tool_choice |= p.rule("tool-" + std::to_string(i), "{" << tool_name << "," << tool_args << "}");
     }
 
     // Define the tool call structure: [{tool}]
diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp
index e31c093b0..f0aec5bbe 100644
--- a/tests/test-chat.cpp
+++ b/tests/test-chat.cpp
@@ -1087,6 +1087,7 @@ struct peg_test_case {
     common_chat_msg              expect;
     bool                         is_partial            = false;
     bool                         expect_reconstruction = false;
+    std::vector     expect_rules;
 };
 
 struct make_peg_parser {
@@ -1166,6 +1167,14 @@ static void test_peg_parser(common_chat_templates *                      tmpls,
         LOG_DBG("Generation prompt: '%s'\n", parser.params_.generation_prompt.c_str());
     }
 
+    for (const auto & rule : tc.expect_rules) {
+        if (!parser.arena_.has_rule(rule)) {
+            LOG_ERR("Missing rule: %s\n", rule.c_str());
+            common_log_flush(common_log_main());
+            throw std::runtime_error("Test failed");
+        }
+    }
+
     common_chat_msg msg_accum;
     common_chat_msg msg_prev;
     msg_accum.role = msg_prev.role = "assistant";
@@ -1573,6 +1582,11 @@ class peg_test_builder {
         return *this;
     }
 
+    peg_test_builder & expect_rules(std::vector rules) {
+        tc_.expect_rules = std::move(rules);
+        return *this;
+    }
+
     peg_test_builder & tool_choice(common_chat_tool_choice choice) {
         tc_.params.tool_choice = choice;
         return *this;
@@ -2174,6 +2188,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
             .tools({ special_function_tool })
             .expect(message_assist_call)
+            .expect_rules({ "tool-0", "tool-0-arg-0" })
             .run();
 
         tst.test(
@@ -2625,6 +2640,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .enable_thinking(true)
             .tools({ special_function_tool })
             .expect(message_assist_call)
+            .expect_rules({ "tool-0" })
             .run();
 
         tst.test(
@@ -3024,6 +3040,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
                 "<|tool_call>call:get_time{city:<|\"|>London<|\"|>}")
             .tools({ get_time_tool })
             .expect(message_with_tool_calls("get_time", R"({"city": "London"})"))
+            .expect_rules({ "tool-0" })
             .run();
 
         // Tool call with string argument containing special chars
@@ -3340,6 +3357,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
                "")
             .tools({ special_function_tool })
             .expect(message_assist_call)
+            .expect_rules({ "tool-0" })
             .run();
 
         tst.test(
@@ -3624,6 +3642,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
                 { "special_function", R"({"arg1": 1})", {} },
                 { "special_function_with_opt", R"({"arg1": 1, "arg2": 2})", {} },
             })
+            .expect_rules({ "tool-0", "tool-0-arg-0", "tool-1", "tool-1-arg-0", "tool-1-arg-1" })
             .run();
 
         tst.test(
@@ -4035,6 +4054,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
             .tools({ get_time_tool })
             .expect(message_with_tool_calls_and_reasoning("get_time", R"({"city": "Tokyo"})", "Let me check the time"))
+            .expect_rules({ "tool-0", "tool-0-arg-0" })
             .run();
 
         // Tool call without reasoning (non-thinking mode), integer param (string="false")
@@ -4673,6 +4693,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ get_time_tool })
             .expect_reasoning("I need to check the time first.\n")
             .expect_tool_calls({ { "get_time", R"({"city": "Paris"})", "" } })
+            .expect_rules({ "ling3-tool-0", "ling3-arg-0-0" })
             .run();
 
         // Closed think block, prose, then a tool call.
@@ -4911,6 +4932,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ special_function_tool })
             .expect(message_assist_call_thoughts)
             .expect_reconstruction()
+            .expect_rules({ "tool-0", "tool-0-arg-0" })
             .run();
 
         tst.test(
@@ -5123,6 +5145,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .expect_tool_calls({
                 { "special_function", R"({"arg1":1})", "" },
             })
+            .expect_rules({ "kimi-k3-tool-0", "kimi-k3-arg-0-0" })
             .run();
 
         // Tool call preceded by reasoning (no opening think tag) and content.
@@ -5411,6 +5434,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ special_function_tool })
             .expect(kimi_id_special_func_tool_call)
             .expect_reconstruction()
+            .expect_rules({ "tool-0" })
             .run();
 
         // Kimi-K2-Instruct
@@ -5465,6 +5489,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
             .tools({ special_function_tool })
             .expect(message_assist_call_thoughts)
+            .expect_rules({ "tool-0" })
             .run();
 
         // Tool call with reasoning and content
@@ -5742,6 +5767,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ get_time_tool })
             .expect(message_with_tool_calls_and_reasoning("get_time", R"({"city": "Tokyo"})", "Let me check the time"))
             .expect_reconstruction()
+            .expect_rules({ "tool-0", "tool-0-arg-0" })
             .run();
 
         // Tool call without reasoning, integer param
@@ -6140,6 +6166,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ special_function_tool })
             .expect(message_assist_call_id)
             .expect_reconstruction()
+            .expect_rules({ "tool-0" })
             .run();
 
         // Continuation tests
@@ -6175,6 +6202,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ special_function_tool })
             .expect(message_assist_call)
             .expect_reconstruction()
+            .expect_rules({ "tool-0" })
             .run();
 
         // Continuation tests
@@ -6547,7 +6575,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
     {
         // Llama 3.2
         auto tst = peg_tester("models/templates/meta-llama-Llama-3.2-3B-Instruct.jinja", detailed_debug);
-        tst.test("Hello, world!\nWhat's up?").tools({ special_function_tool }).expect(message_assist).expect_reconstruction().run();
+        tst.test("Hello, world!\nWhat's up?").tools({ special_function_tool }).expect(message_assist).expect_reconstruction().expect_rules({ "tool-0" }).run();
 
         // Continuation tests
         tst.test("world!\nWhat's up?")
@@ -6595,6 +6623,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ special_function_tool })
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
             .expect(message_assist_call)
+            .expect_rules({ "tool-0" })
             .run();
 
         // "Inform then act": the model answers the user and calls a tool in ONE generation,
@@ -6680,6 +6709,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
             .tools({ special_function_tool })
             .expect(message_assist_call)
+            .expect_rules({ "tool-0" })
             .run();
 
         // Tool call with recipient in channel header: "<|channel|>analysis to=functions.NAME<|message|>JSON"
@@ -6855,6 +6885,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
             .tools({ special_function_tool })
             .expect(message_assist_call_thoughts)
+            .expect_rules({ "tool-0" })
             .run();
 
         // Tool call, recipient in channel header
@@ -7111,6 +7142,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ special_function_tool })
             .expect(message_assist_call)
             .expect_reconstruction()
+            .expect_rules({ "tool-0" })
             .run();
 
         tst.test(
@@ -7139,6 +7171,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ special_function_tool })
             .expect(message_assist_call)
             .expect_reconstruction()
+            .expect_rules({ "tool-0" })
             .run();
 
         tst.test(
@@ -7174,6 +7207,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
             .tools({ python_tool })
             .expect_tool_calls({ { "python", R"#({"code": "print('Hello, World!')"})#", {} } })
+            .expect_rules({ "tool-0" })
             .run();
 
         tst.test(R"()")

From bd4eeaa047006cb1fe71999fbd11134b5836e167 Mon Sep 17 00:00:00 2001
From: Nicolas Budyn <16596345+nbud@users.noreply.github.com>
Date: Wed, 7 Oct 2026 22:48:06 +0200
Subject: [PATCH 19/40] chat: fix jinja parser for TranslateGemma (#30096)

* chat: fix jinja parser for translategemma

* log if missing soruce_lang_code or target_lang_code
---
 common/chat.cpp                               |   7 +
 common/parsers/parsers.h                      |   2 +
 common/parsers/sources.cmake                  |   1 +
 common/parsers/translate-gemma.cpp            |  63 ++
 .../google-translategemma-4b-it.jinja         | 641 ++++++++++++++++++
 tests/test-chat.cpp                           |  45 ++
 6 files changed, 759 insertions(+)
 create mode 100644 common/parsers/translate-gemma.cpp
 create mode 100644 models/templates/google-translategemma-4b-it.jinja

diff --git a/common/chat.cpp b/common/chat.cpp
index 2c795c281..c8f156a70 100644
--- a/common/chat.cpp
+++ b/common/chat.cpp
@@ -1223,6 +1223,13 @@ std::optional common_chat_try_specialized_template(
         return common_chat_params_init_minicpm5(tmpl, params);
     }
 
+    // TranslateGemma - user content must follow a custom schema with language codes
+    if (src.find("[source_lang_code]") != std::string::npos &&
+        src.find("[target_lang_code]") != std::string::npos) {
+        LOG_DBG("Using specialized template: TranslateGemma\n");
+        return common_chat_params_init_translate_gemma(tmpl, params);
+    }
+
     // Qwen3-Coder XML tool calls, also used by Nemotron Nano 3, Qwen3.5 and StepFun-3.5-Flash
     if (src.find("") != std::string::npos &&
         src.find("();
+        } else if (content.is_array()) {
+            for (const auto & part : content) {
+                if (!text.empty()) {
+                    text += "\n";
+                }
+                text += part.value("text", "");
+            }
+        }
+        message["content"] = json::array({
+            json{
+                {"type", "text"},
+                {"text", text},
+                {"source_lang_code", src_lang},
+                {"target_lang_code", tgt_lang},
+            }
+        });
+    }
+
+    data.prompt            = common_chat_template_direct_apply_impl(tmpl, inputs, messages);
+    data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, messages);
+    data.format            = COMMON_CHAT_FORMAT_PEG_NATIVE;
+    data.supports_thinking = false;
+
+    if (inputs.has_continuation()) {
+        data.generation_prompt = "model\n" + inputs.continue_msg.render_content();
+        data.prompt += data.generation_prompt;
+    }
+
+    auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
+        return p.literal(data.generation_prompt) << p.content(p.rest());
+    });
+    data.parser = parser.save();
+
+    return data;
+}
diff --git a/models/templates/google-translategemma-4b-it.jinja b/models/templates/google-translategemma-4b-it.jinja
new file mode 100644
index 000000000..d571968f9
--- /dev/null
+++ b/models/templates/google-translategemma-4b-it.jinja
@@ -0,0 +1,641 @@
+{%- set languages = {
+    "aa": "Afar",
+    "aa-DJ": "Afar",
+    "aa-ER": "Afar",
+    "ab": "Abkhazian",
+    "af": "Afrikaans",
+    "af-NA": "Afrikaans",
+    "ak": "Akan",
+    "am": "Amharic",
+    "an": "Aragonese",
+    "ar": "Arabic",
+    "ar-AE": "Arabic",
+    "ar-BH": "Arabic",
+    "ar-DJ": "Arabic",
+    "ar-DZ": "Arabic",
+    "ar-EG": "Arabic",
+    "ar-EH": "Arabic",
+    "ar-ER": "Arabic",
+    "ar-IL": "Arabic",
+    "ar-IQ": "Arabic",
+    "ar-JO": "Arabic",
+    "ar-KM": "Arabic",
+    "ar-KW": "Arabic",
+    "ar-LB": "Arabic",
+    "ar-LY": "Arabic",
+    "ar-MA": "Arabic",
+    "ar-MR": "Arabic",
+    "ar-OM": "Arabic",
+    "ar-PS": "Arabic",
+    "ar-QA": "Arabic",
+    "ar-SA": "Arabic",
+    "ar-SD": "Arabic",
+    "ar-SO": "Arabic",
+    "ar-SS": "Arabic",
+    "ar-SY": "Arabic",
+    "ar-TD": "Arabic",
+    "ar-TN": "Arabic",
+    "ar-YE": "Arabic",
+    "as": "Assamese",
+    "az": "Azerbaijani",
+    "az-Arab": "Azerbaijani",
+    "az-Arab-IQ": "Azerbaijani",
+    "az-Arab-TR": "Azerbaijani",
+    "az-Cyrl": "Azerbaijani",
+    "az-Latn": "Azerbaijani",
+    "ba": "Bashkir",
+    "be": "Belarusian",
+    "be-tarask": "Belarusian",
+    "bg": "Bulgarian",
+    "bg-BG": "Bulgarian",
+    "bm": "Bambara",
+    "bm-Nkoo": "Bambara",
+    "bn": "Bengali",
+    "bn-IN": "Bengali",
+    "bo": "Tibetan",
+    "bo-IN": "Tibetan",
+    "br": "Breton",
+    "bs": "Bosnian",
+    "bs-Cyrl": "Bosnian",
+    "bs-Latn": "Bosnian",
+    "ca": "Catalan",
+    "ca-AD": "Catalan",
+    "ca-ES": "Catalan",
+    "ca-FR": "Catalan",
+    "ca-IT": "Catalan",
+    "ce": "Chechen",
+    "co": "Corsican",
+    "cs": "Czech",
+    "cs-CZ": "Czech",
+    "cv": "Chuvash",
+    "cy": "Welsh",
+    "da": "Danish",
+    "da-DK": "Danish",
+    "da-GL": "Danish",
+    "de": "German",
+    "de-AT": "German",
+    "de-BE": "German",
+    "de-CH": "German",
+    "de-DE": "German",
+    "de-IT": "German",
+    "de-LI": "German",
+    "de-LU": "German",
+    "dv": "Divehi",
+    "dz": "Dzongkha",
+    "ee": "Ewe",
+    "ee-TG": "Ewe",
+    "el": "Greek",
+    "el-CY": "Greek",
+    "el-GR": "Greek",
+    "el-polyton": "Greek",
+    "en": "English",
+    "en-AE": "English",
+    "en-AG": "English",
+    "en-AI": "English",
+    "en-AS": "English",
+    "en-AT": "English",
+    "en-AU": "English",
+    "en-BB": "English",
+    "en-BE": "English",
+    "en-BI": "English",
+    "en-BM": "English",
+    "en-BS": "English",
+    "en-BW": "English",
+    "en-BZ": "English",
+    "en-CA": "English",
+    "en-CC": "English",
+    "en-CH": "English",
+    "en-CK": "English",
+    "en-CM": "English",
+    "en-CX": "English",
+    "en-CY": "English",
+    "en-CZ": "English",
+    "en-DE": "English",
+    "en-DG": "English",
+    "en-DK": "English",
+    "en-DM": "English",
+    "en-ER": "English",
+    "en-ES": "English",
+    "en-FI": "English",
+    "en-FJ": "English",
+    "en-FK": "English",
+    "en-FM": "English",
+    "en-FR": "English",
+    "en-GB": "English",
+    "en-GD": "English",
+    "en-GG": "English",
+    "en-GH": "English",
+    "en-GI": "English",
+    "en-GM": "English",
+    "en-GS": "English",
+    "en-GU": "English",
+    "en-GY": "English",
+    "en-HK": "English",
+    "en-HU": "English",
+    "en-ID": "English",
+    "en-IE": "English",
+    "en-IL": "English",
+    "en-IM": "English",
+    "en-IN": "English",
+    "en-IO": "English",
+    "en-IT": "English",
+    "en-JE": "English",
+    "en-JM": "English",
+    "en-KE": "English",
+    "en-KI": "English",
+    "en-KN": "English",
+    "en-KY": "English",
+    "en-LC": "English",
+    "en-LR": "English",
+    "en-LS": "English",
+    "en-MG": "English",
+    "en-MH": "English",
+    "en-MO": "English",
+    "en-MP": "English",
+    "en-MS": "English",
+    "en-MT": "English",
+    "en-MU": "English",
+    "en-MV": "English",
+    "en-MW": "English",
+    "en-MY": "English",
+    "en-NA": "English",
+    "en-NF": "English",
+    "en-NG": "English",
+    "en-NL": "English",
+    "en-NO": "English",
+    "en-NR": "English",
+    "en-NU": "English",
+    "en-NZ": "English",
+    "en-PG": "English",
+    "en-PH": "English",
+    "en-PK": "English",
+    "en-PL": "English",
+    "en-PN": "English",
+    "en-PR": "English",
+    "en-PT": "English",
+    "en-PW": "English",
+    "en-RO": "English",
+    "en-RW": "English",
+    "en-SB": "English",
+    "en-SC": "English",
+    "en-SD": "English",
+    "en-SE": "English",
+    "en-SG": "English",
+    "en-SH": "English",
+    "en-SI": "English",
+    "en-SK": "English",
+    "en-SL": "English",
+    "en-SS": "English",
+    "en-SX": "English",
+    "en-SZ": "English",
+    "en-TC": "English",
+    "en-TK": "English",
+    "en-TO": "English",
+    "en-TT": "English",
+    "en-TV": "English",
+    "en-TZ": "English",
+    "en-UG": "English",
+    "en-UM": "English",
+    "en-VC": "English",
+    "en-VG": "English",
+    "en-VI": "English",
+    "en-VU": "English",
+    "en-WS": "English",
+    "en-ZA": "English",
+    "en-ZM": "English",
+    "en-ZW": "English",
+    "eo": "Esperanto",
+    "es": "Spanish",
+    "es-AR": "Spanish",
+    "es-BO": "Spanish",
+    "es-BR": "Spanish",
+    "es-BZ": "Spanish",
+    "es-CL": "Spanish",
+    "es-CO": "Spanish",
+    "es-CR": "Spanish",
+    "es-CU": "Spanish",
+    "es-DO": "Spanish",
+    "es-EA": "Spanish",
+    "es-EC": "Spanish",
+    "es-ES": "Spanish",
+    "es-GQ": "Spanish",
+    "es-GT": "Spanish",
+    "es-HN": "Spanish",
+    "es-IC": "Spanish",
+    "es-MX": "Spanish",
+    "es-NI": "Spanish",
+    "es-PA": "Spanish",
+    "es-PE": "Spanish",
+    "es-PH": "Spanish",
+    "es-PR": "Spanish",
+    "es-PY": "Spanish",
+    "es-SV": "Spanish",
+    "es-US": "Spanish",
+    "es-UY": "Spanish",
+    "es-VE": "Spanish",
+    "et": "Estonian",
+    "et-EE": "Estonian",
+    "eu": "Basque",
+    "fa": "Persian",
+    "fa-AF": "Persian",
+    "fa-IR": "Persian",
+    "ff": "Fulah",
+    "ff-Adlm": "Fulah",
+    "ff-Adlm-BF": "Fulah",
+    "ff-Adlm-CM": "Fulah",
+    "ff-Adlm-GH": "Fulah",
+    "ff-Adlm-GM": "Fulah",
+    "ff-Adlm-GW": "Fulah",
+    "ff-Adlm-LR": "Fulah",
+    "ff-Adlm-MR": "Fulah",
+    "ff-Adlm-NE": "Fulah",
+    "ff-Adlm-NG": "Fulah",
+    "ff-Adlm-SL": "Fulah",
+    "ff-Adlm-SN": "Fulah",
+    "ff-Latn": "Fulah",
+    "ff-Latn-BF": "Fulah",
+    "ff-Latn-CM": "Fulah",
+    "ff-Latn-GH": "Fulah",
+    "ff-Latn-GM": "Fulah",
+    "ff-Latn-GN": "Fulah",
+    "ff-Latn-GW": "Fulah",
+    "ff-Latn-LR": "Fulah",
+    "ff-Latn-MR": "Fulah",
+    "ff-Latn-NE": "Fulah",
+    "ff-Latn-NG": "Fulah",
+    "ff-Latn-SL": "Fulah",
+    "fi": "Finnish",
+    "fi-FI": "Finnish",
+    "fil-PH": "Filipino",
+    "fo": "Faroese",
+    "fo-DK": "Faroese",
+    "fr": "French",
+    "fr-BE": "French",
+    "fr-BF": "French",
+    "fr-BI": "French",
+    "fr-BJ": "French",
+    "fr-BL": "French",
+    "fr-CA": "French",
+    "fr-CD": "French",
+    "fr-CF": "French",
+    "fr-CG": "French",
+    "fr-CH": "French",
+    "fr-CI": "French",
+    "fr-CM": "French",
+    "fr-DJ": "French",
+    "fr-DZ": "French",
+    "fr-FR": "French",
+    "fr-GA": "French",
+    "fr-GF": "French",
+    "fr-GN": "French",
+    "fr-GP": "French",
+    "fr-GQ": "French",
+    "fr-HT": "French",
+    "fr-KM": "French",
+    "fr-LU": "French",
+    "fr-MA": "French",
+    "fr-MC": "French",
+    "fr-MF": "French",
+    "fr-MG": "French",
+    "fr-ML": "French",
+    "fr-MQ": "French",
+    "fr-MR": "French",
+    "fr-MU": "French",
+    "fr-NC": "French",
+    "fr-NE": "French",
+    "fr-PF": "French",
+    "fr-PM": "French",
+    "fr-RE": "French",
+    "fr-RW": "French",
+    "fr-SC": "French",
+    "fr-SN": "French",
+    "fr-SY": "French",
+    "fr-TD": "French",
+    "fr-TG": "French",
+    "fr-TN": "French",
+    "fr-VU": "French",
+    "fr-WF": "French",
+    "fr-YT": "French",
+    "fy": "Western Frisian",
+    "ga": "Irish",
+    "ga-GB": "Irish",
+    "gd": "Scottish Gaelic",
+    "gl": "Galician",
+    "gn": "Guarani",
+    "gu": "Gujarati",
+    "gu-IN": "Gujarati",
+    "gv": "Manx",
+    "ha": "Hausa",
+    "ha-Arab": "Hausa",
+    "ha-Arab-SD": "Hausa",
+    "ha-GH": "Hausa",
+    "ha-NE": "Hausa",
+    "he": "Hebrew",
+    "he-IL": "Hebrew",
+    "hi": "Hindi",
+    "hi-IN": "Hindi",
+    "hi-Latn": "Hindi",
+    "hr": "Croatian",
+    "hr-BA": "Croatian",
+    "hr-HR": "Croatian",
+    "ht": "Haitian",
+    "hu": "Hungarian",
+    "hu-HU": "Hungarian",
+    "hy": "Armenian",
+    "ia": "Interlingua",
+    "id": "Indonesian",
+    "id-ID": "Indonesian",
+    "ie": "Interlingue",
+    "ig": "Igbo",
+    "ii": "Sichuan Yi",
+    "ik": "Inupiaq",
+    "io": "Ido",
+    "is": "Icelandic",
+    "it": "Italian",
+    "it-CH": "Italian",
+    "it-IT": "Italian",
+    "it-SM": "Italian",
+    "it-VA": "Italian",
+    "iu": "Inuktitut",
+    "iu-Latn": "Inuktitut",
+    "ja": "Japanese",
+    "ja-JP": "Japanese",
+    "jv": "Javanese",
+    "ka": "Georgian",
+    "ki": "Kikuyu",
+    "kk": "Kazakh",
+    "kk-Arab": "Kazakh",
+    "kk-Cyrl": "Kazakh",
+    "kk-KZ": "Kazakh",
+    "kl": "Kalaallisut",
+    "km": "Central Khmer",
+    "kn": "Kannada",
+    "kn-IN": "Kannada",
+    "ko": "Korean",
+    "ko-CN": "Korean",
+    "ko-KP": "Korean",
+    "ko-KR": "Korean",
+    "ks": "Kashmiri",
+    "ks-Arab": "Kashmiri",
+    "ks-Deva": "Kashmiri",
+    "ku": "Kurdish",
+    "kw": "Cornish",
+    "ky": "Kyrgyz",
+    "la": "Latin",
+    "lb": "Luxembourgish",
+    "lg": "Ganda",
+    "ln": "Lingala",
+    "ln-AO": "Lingala",
+    "ln-CF": "Lingala",
+    "ln-CG": "Lingala",
+    "lo": "Lao",
+    "lt": "Lithuanian",
+    "lt-LT": "Lithuanian",
+    "lu": "Luba-Katanga",
+    "lv": "Latvian",
+    "lv-LV": "Latvian",
+    "mg": "Malagasy",
+    "mi": "Maori",
+    "mk": "Macedonian",
+    "ml": "Malayalam",
+    "ml-IN": "Malayalam",
+    "mn": "Mongolian",
+    "mn-Mong": "Mongolian",
+    "mn-Mong-MN": "Mongolian",
+    "mr": "Marathi",
+    "mr-IN": "Marathi",
+    "ms": "Malay",
+    "ms-Arab": "Malay",
+    "ms-Arab-BN": "Malay",
+    "ms-BN": "Malay",
+    "ms-ID": "Malay",
+    "ms-SG": "Malay",
+    "mt": "Maltese",
+    "my": "Burmese",
+    "nb": "Norwegian Bokmål",
+    "nb-SJ": "Norwegian Bokmål",
+    "nd": "North Ndebele",
+    "ne": "Nepali",
+    "ne-IN": "Nepali",
+    "nl": "Dutch",
+    "nl-AW": "Dutch",
+    "nl-BE": "Dutch",
+    "nl-BQ": "Dutch",
+    "nl-CW": "Dutch",
+    "nl-NL": "Dutch",
+    "nl-SR": "Dutch",
+    "nl-SX": "Dutch",
+    "nn": "Norwegian Nynorsk",
+    "no": "Norwegian",
+    "no-NO": "Norwegian",
+    "nr": "South Ndebele",
+    "nv": "Navajo",
+    "ny": "Chichewa",
+    "oc": "Occitan",
+    "oc-ES": "Occitan",
+    "om": "Oromo",
+    "om-KE": "Oromo",
+    "or": "Oriya",
+    "os": "Ossetian",
+    "os-RU": "Ossetian",
+    "pa": "Punjabi",
+    "pa-IN": "Punjabi",
+    "pa-Arab": "Punjabi",
+    "pa-Guru": "Punjabi",
+    "pl": "Polish",
+    "pl-PL": "Polish",
+    "ps": "Pashto",
+    "ps-PK": "Pashto",
+    "pt": "Portuguese",
+    "pt-AO": "Portuguese",
+    "pt-BR": "Portuguese",
+    "pt-CH": "Portuguese",
+    "pt-CV": "Portuguese",
+    "pt-GQ": "Portuguese",
+    "pt-GW": "Portuguese",
+    "pt-LU": "Portuguese",
+    "pt-MO": "Portuguese",
+    "pt-MZ": "Portuguese",
+    "pt-PT": "Portuguese",
+    "pt-ST": "Portuguese",
+    "pt-TL": "Portuguese",
+    "qu": "Quechua",
+    "qu-BO": "Quechua",
+    "qu-EC": "Quechua",
+    "rm": "Romansh",
+    "rn": "Rundi",
+    "ro": "Romanian",
+    "ro-MD": "Romanian",
+    "ro-RO": "Romanian",
+    "ru": "Russian",
+    "ru-BY": "Russian",
+    "ru-KG": "Russian",
+    "ru-KZ": "Russian",
+    "ru-MD": "Russian",
+    "ru-RU": "Russian",
+    "ru-UA": "Russian",
+    "rw": "Kinyarwanda",
+    "sa": "Sanskrit",
+    "sc": "Sardinian",
+    "sd": "Sindhi",
+    "sd-Arab": "Sindhi",
+    "sd-Deva": "Sindhi",
+    "se": "Northern Sami",
+    "se-FI": "Northern Sami",
+    "se-SE": "Northern Sami",
+    "sg": "Sango",
+    "si": "Sinhala",
+    "sk": "Slovak",
+    "sk-SK": "Slovak",
+    "sl": "Slovenian",
+    "sl-SI": "Slovenian",
+    "sn": "Shona",
+    "so": "Somali",
+    "so-DJ": "Somali",
+    "so-ET": "Somali",
+    "so-KE": "Somali",
+    "sq": "Albanian",
+    "sq-MK": "Albanian",
+    "sq-XK": "Albanian",
+    "sr": "Serbian",
+    "sr-RS": "Serbian",
+    "sr-Cyrl": "Serbian",
+    "sr-Cyrl-BA": "Serbian",
+    "sr-Cyrl-ME": "Serbian",
+    "sr-Cyrl-XK": "Serbian",
+    "sr-Latn": "Serbian",
+    "sr-Latn-BA": "Serbian",
+    "sr-Latn-ME": "Serbian",
+    "sr-Latn-XK": "Serbian",
+    "ss": "Swati",
+    "ss-SZ": "Swati",
+    "st": "Southern Sotho",
+    "st-LS": "Southern Sotho",
+    "su": "Sundanese",
+    "su-Latn": "Sundanese",
+    "sv": "Swedish",
+    "sv-AX": "Swedish",
+    "sv-FI": "Swedish",
+    "sv-SE": "Swedish",
+    "sw": "Swahili",
+    "sw-CD": "Swahili",
+    "sw-KE": "Swahili",
+    "sw-TZ": "Swahili",
+    "sw-UG": "Swahili",
+    "ta": "Tamil",
+    "ta-IN": "Tamil",
+    "ta-LK": "Tamil",
+    "ta-MY": "Tamil",
+    "ta-SG": "Tamil",
+    "te": "Telugu",
+    "te-IN": "Telugu",
+    "tg": "Tajik",
+    "th": "Thai",
+    "th-TH": "Thai",
+    "ti": "Tigrinya",
+    "ti-ER": "Tigrinya",
+    "tk": "Turkmen",
+    "tl": "Tagalog",
+    "tn": "Tswana",
+    "tn-BW": "Tswana",
+    "to": "Tonga",
+    "tr": "Turkish",
+    "tr-CY": "Turkish",
+    "tr-TR": "Turkish",
+    "ts": "Tsonga",
+    "tt": "Tatar",
+    "ug": "Uyghur",
+    "uk": "Ukrainian",
+    "uk-UA": "Ukrainian",
+    "ur": "Urdu",
+    "ur-IN": "Urdu",
+    "ur-PK": "Urdu",
+    "uz": "Uzbek",
+    "uz-Arab": "Uzbek",
+    "uz-Cyrl": "Uzbek",
+    "uz-Latn": "Uzbek",
+    "ve": "Venda",
+    "vi": "Vietnamese",
+    "vi-VN": "Vietnamese",
+    "vo": "Volapük",
+    "wa": "Walloon",
+    "wo": "Wolof",
+    "xh": "Xhosa",
+    "yi": "Yiddish",
+    "yo": "Yoruba",
+    "yo-BJ": "Yoruba",
+    "za": "Zhuang",
+    "zh": "Chinese",
+    "zh-CH": "Chinese",
+    "zh-TW": "Chinese",
+    "zh-Hans": "Chinese",
+    "zh-Hans-HK": "Chinese",
+    "zh-Hans-MO": "Chinese",
+    "zh-Hans-MY": "Chinese",
+    "zh-Hans-SG": "Chinese",
+    "zh-Hant": "Chinese",
+    "zh-Hant-HK": "Chinese",
+    "zh-Hant-MO": "Chinese",
+    "zh-Hant-MY": "Chinese",
+    "zh-Latn": "Chinese",
+    "zu": "Zulu",
+    "zu-ZA": "Zulu",
+}
+-%}
+{{ bos_token }}
+{%- if (messages[0]['role'] != 'user') -%}
+    {{ raise_exception("Conversations must start with a user prompt.") }}
+{%- endif -%}
+{%- for message in messages -%}
+    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}
+        {{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }}
+    {%- endif -%}
+    {%- if (message['role'] == 'assistant') -%}
+        {%- if message['content'] is none or message['content'] is not string -%}
+            {{ raise_exception("Assistant role must provide content as a string") }}
+        {%- endif -%}
+        {{ 'model\n'}}
+        {{ message["content"] | trim }}
+    {%- elif (message['role'] == 'user') -%}
+        {%- if message['content'] is none or message['content'] is not iterable or message['content'] | length != 1 -%}
+            {{ raise_exception(
+                "User role must provide `content` as an iterable with exactly one item. That item must be a " +
+                "`mapping(type:'text' | 'image', source_lang_code:string, target_lang_code:string, " +
+                "text:string | none, image:string | none)`."
+            ) }}
+        {%- endif -%}
+        {%- set content = message["content"][0] -%}
+        {%- set source_lang_code = content["source_lang_code"] | replace("_", "-")  -%}
+        {%- set source_lang = languages[source_lang_code] -%}
+        {%- set target_lang_code = content["target_lang_code"] | replace("_", "-") -%}
+        {%- set target_lang = languages[target_lang_code] -%}
+        {{ 'user\nYou are a professional ' + source_lang + ' (' + source_lang_code + ') to ' +
+           target_lang + ' (' + target_lang_code + ') translator. Your goal is to accurately convey the meaning and '
+           'nuances of the original ' + source_lang + ' text while adhering to ' + target_lang + ' grammar, '
+           'vocabulary, and cultural sensitivities.\n'
+        }}
+        {%- if content["type"] == 'text' -%}
+            {{
+                'Produce only the ' + target_lang + ' translation, without any additional explanations or ' +
+                'commentary. Please translate the following ' + source_lang + ' text into ' + target_lang + ':\n\n\n' +
+                content["text"] | trim
+            }}
+        {%- elif content["type"] == 'image' -%}
+            {{
+                'Please translate the ' + source_lang + ' text in the provided image into ' + target_lang + '. ' +
+                'Produce only the ' + target_lang + ' translation, without any additional explanations, ' +
+                'alternatives or commentary. Focus only on the text, do not output where the text is located, ' +
+                'surrounding objects or any other explanation about the picture. Ignore symbols, pictogram, and ' +
+                'arrows!\n\n\n'
+            }}
+        {%- else -%}
+            {{ raise_exception("User role `content` can contain either `text` or `image` data.") }}
+        {%- endif -%}
+    {%- else -%}
+        {{ raise_exception("Conversations must only contain user or assistant roles.") }}
+    {%- endif -%}
+    {{ '\n' }}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{'model\n'}}
+{%- endif -%}
diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp
index f0aec5bbe..41e4948c4 100644
--- a/tests/test-chat.cpp
+++ b/tests/test-chat.cpp
@@ -7192,6 +7192,13 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .run();
     }
 
+    // TranslateGemma
+    {
+        // no reconstruction check, the template adds whitespace around assistant content
+        auto tst = peg_tester("models/templates/google-translategemma-4b-it.jinja", detailed_debug);
+        tst.test("Hello, world!\nWhat's up?").expect(message_assist).run();
+    }
+
     // MiniCPM5 - XML tool calls with ...
     {
         auto tst = peg_tester("models/templates/openbmb-MiniCPM5-1B.jinja", detailed_debug);
@@ -7609,6 +7616,43 @@ static void test_developer_role_to_system_workaround() {
     }
 }
 
+// TranslateGemma raises on plain string user content, the specialized handler must rewrite it
+static void test_translate_gemma() {
+    LOG_DBG("%s\n", __func__);
+
+    auto tmpls = read_templates("models/templates/google-translategemma-4b-it.jinja");
+
+    // server startup renders this example, it must not throw
+    auto example = common_chat_format_example(tmpls.get(), /* use_jinja= */ true, {});
+    assert_contains(example, "English (en-GB) to English (en-GB) translator");
+    assert_contains(example, "How are you?");
+
+    common_chat_templates_inputs inputs;
+    inputs.messages                                 = { message_user };
+    inputs.add_generation_prompt                    = true;
+    inputs.chat_template_kwargs["source_lang_code"] = R"("en")";
+    inputs.chat_template_kwargs["target_lang_code"] = R"("fr")";
+
+    auto params = common_chat_templates_apply(tmpls.get(), inputs);
+    assert_contains(params.prompt, "English (en) to French (fr) translator");
+    assert_contains(params.prompt, "into French:\n\n\nHey there!\n");
+    assert_equals(std::string("model\n"), params.generation_prompt);
+    assert_ends_with(params.prompt, params.generation_prompt);
+
+    // typed text parts are joined into one item
+    inputs.messages = { message_user_parts };
+    params = common_chat_templates_apply(tmpls.get(), inputs);
+    assert_contains(params.prompt, "into French:\n\n\nHey\nthere\n");
+
+    // assistant prefill is appended after the generation prompt
+    inputs.messages               = { message_user, message_assist_prefill_content };
+    inputs.add_generation_prompt  = false;
+    inputs.continue_final_message = COMMON_CHAT_CONTINUATION_CONTENT;
+    params = common_chat_templates_apply(tmpls.get(), inputs);
+    assert_equals(std::string("model\nHello, "), params.generation_prompt);
+    assert_ends_with(params.prompt, "Hey there!\nmodel\nHello, ");
+}
+
 // Verify reasoning-trace retention rules in the DeepSeek-V4 template:
 // all traces are retained unless drop_thinking is true AND the conversation
 // has no tool calls, in which case only the last (after-final-user) trace is
@@ -7986,6 +8030,7 @@ int main(int argc, char ** argv) {
         test_tools_oaicompat_json_conversion();
         test_convert_responses_to_chatcmpl();
         test_developer_role_to_system_workaround();
+        test_translate_gemma();
         test_deepseek_v4_thinking_retention();
         test_deepseek_v4_tool_result_ordering();
         test_template_generation_prompt();

From 70815103c8b251bdf5abba4eca12749eab1de686 Mon Sep 17 00:00:00 2001
From: kurquhar 
Date: Wed, 7 Oct 2026 16:57:09 -0700
Subject: [PATCH 20/40] hexagon: improved GELU accuracy (#30104)

Assisted-by: OpenCode
---
 ggml/src/ggml-hexagon/ggml-hexagon.cpp |  2 +-
 ggml/src/ggml-hexagon/htp/htp-ops.h    |  1 +
 ggml/src/ggml-hexagon/htp/hvx-erf.h    |  3 +-
 ggml/src/ggml-hexagon/htp/main.c       |  1 +
 ggml/src/ggml-hexagon/htp/unary-ops.c  | 97 ++++++++++++++++++++++++--
 ggml/src/ggml-hexagon/htp/unary-ops.h  |  1 +
 6 files changed, 96 insertions(+), 9 deletions(-)

diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
index 8c650b85c..f0d768388 100644
--- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp
+++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
@@ -7677,7 +7677,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) {
             switch (ggml_get_unary_op(t)) {
                 case GGML_UNARY_OP_SILU:       return HTP_OP_UNARY_SILU;
                 case GGML_UNARY_OP_GELU:       return HTP_OP_UNARY_GELU;
-                case GGML_UNARY_OP_GELU_QUICK: return HTP_OP_UNARY_GELU;
+                case GGML_UNARY_OP_GELU_QUICK: return HTP_OP_UNARY_GELU_QUICK;
                 case GGML_UNARY_OP_GELU_ERF:   return HTP_OP_UNARY_GELU_ERF;
                 case GGML_UNARY_OP_SIGMOID:    return HTP_OP_UNARY_SIGMOID;
                 case GGML_UNARY_OP_NEG:        return HTP_OP_UNARY_NEG;
diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops.h
index e47b42e64..5d2571fef 100644
--- a/ggml/src/ggml-hexagon/htp/htp-ops.h
+++ b/ggml/src/ggml-hexagon/htp/htp-ops.h
@@ -117,6 +117,7 @@ enum htp_op_code {
     HTP_OP_GLU_GEGLU_ERF,
     HTP_OP_POOL_2D,
     HTP_OP_POOL_1D,
+    HTP_OP_UNARY_GELU_QUICK,
 
     HTP_OP_INVALID
 };
diff --git a/ggml/src/ggml-hexagon/htp/hvx-erf.h b/ggml/src/ggml-hexagon/htp/hvx-erf.h
index 6eefc300f..b17c657a8 100644
--- a/ggml/src/ggml-hexagon/htp/hvx-erf.h
+++ b/ggml/src/ggml-hexagon/htp/hvx-erf.h
@@ -12,7 +12,7 @@ static __attribute__((noinline)) HVX_Vector hvx_vec_erf_f32(HVX_Vector x) {
     HVX_Vector t = hvx_vec_inverse_f32(hvx_vec_add_f32_f32(
         hvx_vec_splat_f32(1.0f), hvx_vec_mul_f32_f32(hvx_vec_splat_f32(0.3275911f), ax)));
 
-    HVX_Vector poly = hvx_vec_mul_f32_f32(hvx_vec_splat_f32(1.061405429f), t);
+    HVX_Vector poly = hvx_vec_splat_f32(1.061405429f);
     poly = hvx_vec_add_f32_f32(hvx_vec_splat_f32(-1.453152027f), hvx_vec_mul_f32_f32(poly, t));
     poly = hvx_vec_add_f32_f32(hvx_vec_splat_f32(1.421413741f), hvx_vec_mul_f32_f32(poly, t));
     poly = hvx_vec_add_f32_f32(hvx_vec_splat_f32(-0.284496736f), hvx_vec_mul_f32_f32(poly, t));
@@ -27,6 +27,7 @@ static __attribute__((noinline)) HVX_Vector hvx_vec_erf_f32(HVX_Vector x) {
     return result;
 }
 
+// GELU_ERF uses the erf definition.
 static inline HVX_Vector hvx_vec_gelu_erf_f32(HVX_Vector x) {
     const HVX_Vector scale = hvx_vec_splat_f32(0.7071067811865475f);
     const HVX_Vector half  = hvx_vec_splat_f32(0.5f);
diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c
index 41f78b387..2dd1b3fc1 100644
--- a/ggml/src/ggml-hexagon/htp/main.c
+++ b/ggml/src/ggml-hexagon/htp/main.c
@@ -851,6 +851,7 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_UNARY_SIGMOID:
         case HTP_OP_UNARY_SILU:
         case HTP_OP_UNARY_GELU:
+        case HTP_OP_UNARY_GELU_QUICK:
         case HTP_OP_UNARY_GELU_ERF:
         case HTP_OP_UNARY_NEG:
         case HTP_OP_UNARY_EXP:
diff --git a/ggml/src/ggml-hexagon/htp/unary-ops.c b/ggml/src/ggml-hexagon/htp/unary-ops.c
index b63fd4c29..c557b0907 100644
--- a/ggml/src/ggml-hexagon/htp/unary-ops.c
+++ b/ggml/src/ggml-hexagon/htp/unary-ops.c
@@ -497,7 +497,72 @@ static void silu_f32(const void * restrict src,
     }
 }
 
-// gelu(x) = x * sigmoid(1.702 * x)  (quick/sigmoid approximation, matches CPU GELU_QUICK reference)
+// GELU uses the tanh approximation.
+static __attribute__((noinline)) HVX_Vector hvx_vec_gelu_f32(HVX_Vector x) {
+    const HVX_Vector half = hvx_vec_splat_f32(0.5f);
+    const HVX_Vector one  = hvx_vec_splat_f32(1.0f);
+
+    HVX_Vector inner = hvx_vec_mul_f32_f32(x, x);
+    inner = hvx_vec_mul_f32_f32(inner, hvx_vec_splat_f32(0.044715f));
+    inner = hvx_vec_add_f32_f32(inner, one);
+    inner = hvx_vec_mul_f32_f32(inner, x);
+    inner = hvx_vec_mul_f32_f32(inner, hvx_vec_splat_f32(0.7978845608028654f));
+
+    return hvx_vec_mul_f32_f32(hvx_vec_mul_f32_f32(half, x),
+                               hvx_vec_add_f32_f32(one, hvx_vec_tanh_f32(inner)));
+}
+
+static inline void hvx_gelu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) dst % 128 == 0);
+    assert((unsigned long) src % 128 == 0);
+
+    HVX_Vector * restrict vdst = (HVX_Vector *) dst;
+    HVX_Vector * restrict vsrc = (HVX_Vector *) src;
+
+    const uint32_t nvec = n / VLEN_FP32;
+    const uint32_t nloe = n % VLEN_FP32;
+
+    uint32_t i = 0;
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        vdst[i] = hvx_vec_gelu_f32(vsrc[i]);
+    }
+    if (nloe) {
+        hvx_vec_store_a(&vdst[i], nloe * sizeof(float), hvx_vec_gelu_f32(vsrc[i]));
+    }
+}
+
+// GELU_QUICK uses x * sigmoid(1.702 * x).
+static __attribute__((noinline)) HVX_Vector hvx_vec_gelu_quick_f32(HVX_Vector x) {
+    const HVX_Vector one     = hvx_vec_splat_f32(1.0f);
+    const HVX_Vector max_exp = hvx_vec_splat_f32(87.0f);
+    const HVX_Vector min_exp = hvx_vec_splat_f32(-87.0f);
+    const HVX_Vector scaled  = hvx_vec_mul_f32_f32(x, hvx_vec_splat_f32(1.702f));
+    const HVX_Vector sigmoid = hvx_vec_fast_sigmoid_f32_guard(scaled, one, max_exp, min_exp);
+
+    return hvx_vec_mul_f32_f32(x, sigmoid);
+}
+
+static inline void hvx_gelu_quick_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) dst % 128 == 0);
+    assert((unsigned long) src % 128 == 0);
+
+    HVX_Vector * restrict vdst = (HVX_Vector *) dst;
+    HVX_Vector * restrict vsrc = (HVX_Vector *) src;
+
+    const uint32_t nvec = n / VLEN_FP32;
+    const uint32_t nloe = n % VLEN_FP32;
+
+    uint32_t i = 0;
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        vdst[i] = hvx_vec_gelu_quick_f32(vsrc[i]);
+    }
+    if (nloe) {
+        hvx_vec_store_a(&vdst[i], nloe * sizeof(float), hvx_vec_gelu_quick_f32(vsrc[i]));
+    }
+}
+
 static void gelu_f32(const void * restrict src,
                      void * restrict dst,
                      const uint32_t num_rows,
@@ -508,9 +573,21 @@ static void gelu_f32(const void * restrict src,
         const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
         uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
 
-        hvx_mul_scalar_f32(dst_local, src_local, 1.702f, ne0);
-        hvx_sigmoid_f32_aa(dst_local, dst_local, ne0);
-        hvx_mul_f32_aaa(dst_local, src_local, dst_local, ne0);
+        hvx_gelu_f32_aa(dst_local, src_local, ne0);
+    }
+}
+
+static void gelu_quick_f32(const void * restrict src,
+                           void * restrict dst,
+                           const uint32_t num_rows,
+                           const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_gelu_quick_f32_aa(dst_local, src_local, ne0);
     }
 }
 
@@ -774,9 +851,12 @@ static void tile_silu_f32(void * restrict dst, const void * restrict src, uint32
 
 static void tile_gelu_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) {
     (void) uctx;
-    hvx_mul_scalar_f32((uint8_t *) dst, (const uint8_t *) src, 1.702f, tw);
-    hvx_sigmoid_f32_aa((uint8_t *) dst, (uint8_t *) dst, tw);
-    hvx_mul_f32_aaa((uint8_t *) dst, (const uint8_t *) src, (uint8_t *) dst, tw);
+    hvx_gelu_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw);
+}
+
+static void tile_gelu_quick_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) {
+    (void) uctx;
+    hvx_gelu_quick_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw);
 }
 
 static void tile_gelu_erf_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) {
@@ -1533,6 +1613,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
         case HTP_OP_UNARY_SIGMOID:   op_type = "sigmoid-f32";                                break;
         case HTP_OP_UNARY_SILU:      op_type = "silu-f32";                                   break;
         case HTP_OP_UNARY_GELU:      op_type = "gelu-f32";                                   break;
+        case HTP_OP_UNARY_GELU_QUICK: op_type = "gelu-quick-f32";                             break;
         case HTP_OP_UNARY_GELU_ERF:  op_type = "gelu-erf-f32";                               break;
         case HTP_OP_UNARY_SOFTPLUS:  op_type = "softplus-f32";                               break;
         case HTP_OP_UNARY_TANH:      op_type = "tanh-f32";                                   break;
@@ -1684,6 +1765,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
             case HTP_OP_UNARY_SIGMOID:   compute_func = (void *) tile_sigmoid_f32;        break;
             case HTP_OP_UNARY_SILU:      compute_func = (void *) tile_silu_f32;           break;
             case HTP_OP_UNARY_GELU:      compute_func = (void *) tile_gelu_f32;           break;
+            case HTP_OP_UNARY_GELU_QUICK: compute_func = (void *) tile_gelu_quick_f32;     break;
             case HTP_OP_UNARY_GELU_ERF:  compute_func = (void *) tile_gelu_erf_f32;       break;
             case HTP_OP_UNARY_SOFTPLUS:  compute_func = (void *) tile_softplus_f32;       break;
             case HTP_OP_UNARY_TANH:      compute_func = (void *) tile_tanh_f32;           break;
@@ -1731,6 +1813,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
             case HTP_OP_UNARY_SIGMOID:   compute_func = (void *) sigmoid_f32;             break;
             case HTP_OP_UNARY_SILU:      compute_func = (void *) silu_f32;                break;
             case HTP_OP_UNARY_GELU:      compute_func = (void *) gelu_f32;                break;
+            case HTP_OP_UNARY_GELU_QUICK: compute_func = (void *) gelu_quick_f32;          break;
             case HTP_OP_UNARY_GELU_ERF:  compute_func = (void *) gelu_erf_f32;            break;
             case HTP_OP_UNARY_SOFTPLUS:  compute_func = (void *) softplus_f32;            break;
             case HTP_OP_UNARY_TANH:      compute_func = (void *) tanh_f32;                break;
diff --git a/ggml/src/ggml-hexagon/htp/unary-ops.h b/ggml/src/ggml-hexagon/htp/unary-ops.h
index 7b73cf10f..7e5779e5a 100644
--- a/ggml/src/ggml-hexagon/htp/unary-ops.h
+++ b/ggml/src/ggml-hexagon/htp/unary-ops.h
@@ -54,6 +54,7 @@ static inline bool htp_op_is_unary(uint32_t opcode) {
         case HTP_OP_UNARY_SIGMOID:
         case HTP_OP_UNARY_SILU:
         case HTP_OP_UNARY_GELU:
+        case HTP_OP_UNARY_GELU_QUICK:
         case HTP_OP_UNARY_GELU_ERF:
         case HTP_OP_UNARY_SOFTPLUS:
         case HTP_OP_UNARY_TANH:

From aa5e0092fdb8980b9e89498dfb73049b4fe204c2 Mon Sep 17 00:00:00 2001
From: kurquhar 
Date: Wed, 7 Oct 2026 17:04:45 -0700
Subject: [PATCH 21/40] hexagon: support tiled Q4_K and Q6_K GET_ROWS (#30115)

* hexagon: support tiled Q4_K GET_ROWS

Assisted-by: OpenCode

* properly reject Q4_K views

Assisted-by: OpenCode

* hexagon: support tiled Q6_K GET_ROWS

Assisted-by: OpenCode
---
 ggml/src/ggml-hexagon/ggml-hexagon.cpp   | 16 +++---
 ggml/src/ggml-hexagon/htp/get-rows-ops.c | 64 ++++++++++++++++++++----
 ggml/src/ggml-hexagon/htp/get-rows-ops.h |  2 +-
 3 files changed, 64 insertions(+), 18 deletions(-)

diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
index f0d768388..26b79582e 100644
--- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp
+++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
@@ -7232,7 +7232,7 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session *
     const struct ggml_tensor * src1 = op->src[1]; // indices
     const struct ggml_tensor * dst  = op;
 
-    if (src0->type == GGML_TYPE_Q4_0 && src0->view_src) {
+    if ((src0->type == GGML_TYPE_Q4_0 || src0->type == GGML_TYPE_Q4_K || src0->type == GGML_TYPE_Q6_K) && src0->view_src) {
         return false;
     }
 
@@ -7241,7 +7241,7 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session *
     if (src0_base->buffer && ggml_backend_buffer_is_hexagon(src0_base->buffer) && src0_base->extra) {
         const auto * extra = (const ggml_hexagon_tensor_extra *) src0_base->extra;
         is_repacked = (extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0;
-        if (is_repacked && src0->type != GGML_TYPE_Q4_0 && src0->type != GGML_TYPE_Q8_0) {
+        if (is_repacked && src0->type != GGML_TYPE_Q4_0 && src0->type != GGML_TYPE_Q4_K && src0->type != GGML_TYPE_Q6_K && src0->type != GGML_TYPE_Q8_0) {
             return false;
         }
     }
@@ -7252,7 +7252,7 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session *
         return false;
     }
 
-    if (src0->type == GGML_TYPE_Q4_0 && src0->buffer && !is_repacked) {
+    if ((src0->type == GGML_TYPE_Q4_0 || src0->type == GGML_TYPE_Q4_K || src0->type == GGML_TYPE_Q6_K) && src0->buffer && ggml_backend_buffer_get_size(src0->buffer) != 0 && !is_repacked) {
         return false;
     }
 
@@ -7261,7 +7261,11 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session *
     }
 
     if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16 &&
-        src0->type != GGML_TYPE_Q4_0 && src0->type != GGML_TYPE_Q8_0 && src0->type != GGML_TYPE_I32) {
+        src0->type != GGML_TYPE_Q4_0 && src0->type != GGML_TYPE_Q4_K && src0->type != GGML_TYPE_Q6_K && src0->type != GGML_TYPE_Q8_0 && src0->type != GGML_TYPE_I32) {
+        return false;
+    }
+
+    if ((src0->type == GGML_TYPE_Q4_K || src0->type == GGML_TYPE_Q6_K) && (!ggml_is_contiguous(src0) || ggml_is_permuted(src0) || src0->ne[0] % QK_K)) {
         return false;
     }
 
@@ -7290,8 +7294,8 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session *
         return false;
     }
 
-    // Q4_0 has no raw fallback. Mark only accepted tensors for repacking.
-    if (src0->type == GGML_TYPE_Q4_0 && !src0->buffer) {
+    // Tiled quantized weights have no raw fallback. Mark only accepted tensors for repacking.
+    if ((src0->type == GGML_TYPE_Q4_0 || src0->type == GGML_TYPE_Q4_K || src0->type == GGML_TYPE_Q6_K) && !src0->buffer) {
         sess->needs_repack.insert(src0);
     }
 
diff --git a/ggml/src/ggml-hexagon/htp/get-rows-ops.c b/ggml/src/ggml-hexagon/htp/get-rows-ops.c
index f354b813a..90f38fa56 100644
--- a/ggml/src/ggml-hexagon/htp/get-rows-ops.c
+++ b/ggml/src/ggml-hexagon/htp/get-rows-ops.c
@@ -217,7 +217,7 @@ GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int32_t, { compute_get_rows_q8_0((float
 GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int64_t, { compute_get_rows_q8_0((float *)dst_spad, src_spad, cur_elems); })
 
 
-static __attribute__((noinline)) void compute_get_rows_tiled(float * dst, const uint8_t * tile, uint32_t row, bool q4) {
+static __attribute__((noinline)) void compute_get_rows_tiled(float * dst, const uint8_t * tile, uint32_t row, bool q4, bool q4_k) {
     const HVX_VectorPred first2 = Q6_Q_vsetq_R(2);
     const HVX_VectorPred first4 = Q6_Q_vsetq_R(4);
     HVX_Vector vq = Q6_V_vzero();
@@ -235,7 +235,9 @@ static __attribute__((noinline)) void compute_get_rows_tiled(float * dst, const
         const HVX_Vector lo = Q6_V_vand_VV(vq, Q6_Vb_vsplat_R(0x0F));
         const HVX_Vector hi = Q6_Vub_vlsr_VubR(vq, 4);
         vq = Q6_V_lo_W(Q6_W_vshuff_VVR(hi, lo, -1));
-        vq = Q6_Vb_vsub_VbVb(vq, Q6_Vb_vsplat_R(8));
+        if (!q4_k) {
+            vq = Q6_Vb_vsub_VbVb(vq, Q6_Vb_vsplat_R(8));
+        }
     } else {
         for (int group = 7; group >= 0; --group) {
             const HVX_Vector v = Q6_V_vror_VR(hvx_vmem(tile + group * VLEN), 2 * row);
@@ -245,14 +247,46 @@ static __attribute__((noinline)) void compute_get_rows_tiled(float * dst, const
         }
     }
     const HVX_Vector scales = hvx_vmem(tile + (q4 ? 512 : 1024));
-    const HVX_Vector scale_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 2 * row));
+    const HVX_Vector scale_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, (q4_k ? 4 : 2) * row));
     const HVX_Vector scale = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hf));
     const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq);
     const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16));
-    const HVX_Vector values = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale);
+    HVX_Vector values = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale);
+    if (q4_k) {
+        const HVX_Vector offset_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 4 * row + 2));
+        const HVX_Vector offset = Q6_V_lo_W(hvx_vec_f16_to_f32(offset_hf));
+        values = hvx_vec_add_f32_f32(values, offset);
+    }
     *(HVX_Vector *) dst = values;
 }
 
+static __attribute__((noinline)) void compute_get_rows_q6_k(float * dst, const uint8_t * tile, uint32_t row) {
+    const HVX_VectorPred first4 = Q6_Q_vsetq_R(4);
+    const HVX_VectorPred first16 = Q6_Q_vsetq_R(16 * sizeof(float));
+    const HVX_Vector mask_0f = Q6_Vb_vsplat_R(0x0F);
+    const HVX_Vector mask_03 = Q6_Vb_vsplat_R(0x03);
+    HVX_Vector vq = Q6_V_vzero();
+
+    for (int group = 7; group >= 0; --group) {
+        const HVX_Vector lo_plane = Q6_V_vror_VR(hvx_vmem(tile + (group >> 1) * VLEN), 4 * row);
+        const HVX_Vector hi_plane = Q6_V_vror_VR(hvx_vmem(tile + 512 + (group >> 2) * VLEN), 4 * row);
+        const HVX_Vector lo = (group & 1) ? Q6_Vub_vlsr_VubR(lo_plane, 4) : Q6_V_vand_VV(lo_plane, mask_0f);
+        const HVX_Vector hi = Q6_Vub_vlsr_VubR(hi_plane, 2 * (group & 3));
+        const HVX_Vector packed = Q6_V_vor_VV(lo, Q6_Vw_vasl_VwR(Q6_V_vand_VV(hi, mask_03), 4));
+        vq = Q6_V_vmux_QVV(first4, packed, Q6_V_vror_VR(vq, VLEN - 4));
+    }
+
+    const HVX_Vector scales = hvx_vmem(tile + 768);
+    const HVX_Vector scale_lo_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 2 * row));
+    const HVX_Vector scale_hi_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 64 + 2 * row));
+    const HVX_Vector scale_lo = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_lo_hf));
+    const HVX_Vector scale_hi = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hi_hf));
+    const HVX_Vector scale = Q6_V_vmux_QVV(first16, scale_lo, scale_hi);
+    const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(Q6_Vb_vsub_VbVb(vq, Q6_Vb_vsplat_R(32)));
+    const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16));
+    *(HVX_Vector *) dst = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale);
+}
+
 struct get_rows_tiled_task {
     dma_addr_t tile_src_base;
     dma_addr_t dst_data;
@@ -315,7 +349,9 @@ static void get_rows_thread_tiled(unsigned int nth, unsigned int ith, void * dat
     const uint32_t tile_size   = grctx->tile_size;
     const uint32_t tile_stride = grctx->tile_stride;
     const uint32_t dst_bytes   = ne00 * sizeof(float);
-    const bool is_q4 = (octx->src[0]->type == HTP_TYPE_Q4_0);
+    const bool is_q4 = octx->src[0]->type == HTP_TYPE_Q4_0 || octx->src[0]->type == HTP_TYPE_Q4_K;
+    const bool is_q4_k = octx->src[0]->type == HTP_TYPE_Q4_K;
+    const bool is_q6_k = octx->src[0]->type == HTP_TYPE_Q6_K;
 
     for (uint32_t step = 0, spad_idx = 0; step < ir1 - ir0 && spad_idx < 2; ++step, ++spad_idx) {
         const uint32_t i = ir0 + step;
@@ -343,7 +379,11 @@ static void get_rows_thread_tiled(unsigned int nth, unsigned int ith, void * dat
         for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) {
             const uint8_t * tile = src_spad + k_tile * tile_stride;
             float * dst_block = dst_spad + k_tile * HTP_MM_HMX_TILE_N_COLS;
-            compute_get_rows_tiled(dst_block, tile, task.row, is_q4);
+            if (is_q6_k) {
+                compute_get_rows_q6_k(dst_block, tile, task.row);
+            } else {
+                compute_get_rows_tiled(dst_block, tile, task.row, is_q4, is_q4_k);
+            }
         }
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
 
@@ -369,10 +409,12 @@ int op_get_rows(struct htp_ops_context * octx) {
     const struct htp_get_rows_kernel_params * kparams = (const struct htp_get_rows_kernel_params *) octx->kernel_params;
 
     if (octx->src[0]->type != HTP_TYPE_F32 &&
-         octx->src[0]->type != HTP_TYPE_F16 &&
-         octx->src[0]->type != HTP_TYPE_Q4_0 &&
-         octx->src[0]->type != HTP_TYPE_Q8_0 &&
-         octx->src[0]->type != HTP_TYPE_I32) {
+        octx->src[0]->type != HTP_TYPE_F16 &&
+        octx->src[0]->type != HTP_TYPE_Q4_0 &&
+        octx->src[0]->type != HTP_TYPE_Q4_K &&
+        octx->src[0]->type != HTP_TYPE_Q6_K &&
+        octx->src[0]->type != HTP_TYPE_Q8_0 &&
+        octx->src[0]->type != HTP_TYPE_I32) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
@@ -426,7 +468,7 @@ int op_get_rows(struct htp_ops_context * octx) {
     grctx.task_start = task_start;
     grctx.tasks = tasks;
     grctx.tasks_per_thread = octx->ctx->mdev.count == 1 ? kparams->tasks_per_thread : fastdiv(tasks + n_threads - 1, &octx->n_threads_div);
-    grctx.tile_size = octx->src[0]->type == HTP_TYPE_Q4_0 ? HTP_MM_WEIGHT_TILE_SIZE_Q4_0 : HTP_MM_WEIGHT_TILE_SIZE_Q8_0;
+    grctx.tile_size = htp_mm_get_weight_tile_size(octx->src[0]->type);
     grctx.tile_stride = (grctx.tile_size + 127) & ~127;
     grctx.index_i32 = octx->src[1]->type == HTP_TYPE_I32;
 
diff --git a/ggml/src/ggml-hexagon/htp/get-rows-ops.h b/ggml/src/ggml-hexagon/htp/get-rows-ops.h
index 06ca1ea74..69fa54464 100644
--- a/ggml/src/ggml-hexagon/htp/get-rows-ops.h
+++ b/ggml/src/ggml-hexagon/htp/get-rows-ops.h
@@ -55,7 +55,7 @@ static inline void htp_get_rows_vtcm_layout_build(
     }
 
     if (kernel_type == HTP_GET_ROWS_KERNEL_TILED) {
-        const size_t tile_size   = type == HTP_TYPE_Q4_0 ? HTP_MM_WEIGHT_TILE_SIZE_Q4_0 : HTP_MM_WEIGHT_TILE_SIZE_Q8_0;
+        const size_t tile_size   = htp_mm_get_weight_tile_size(type);
         const size_t tile_stride = (tile_size + 127) & ~127;
         const uint32_t n_k_tiles = ne00 / HTP_MM_HMX_TILE_N_COLS;
         const size_t row_tiles_size = n_k_tiles > 0 ? (n_k_tiles * tile_stride) : tile_stride;

From a657f7e981ff8764d2ccce1e74ec3c7e9bf2cfa2 Mon Sep 17 00:00:00 2001
From: Tarek Dakhran 
Date: Thu, 8 Oct 2026 02:07:49 +0200
Subject: [PATCH 22/40] model : add LiquidAI/d1-omni-600M decision model
 (#30114)

* model : add LiquidAI/d1-omni-600M decision model

Assisted-by: Claude Opus 5.5

* mtmd : keep conformer GLU sigmoid on CUDA

Assisted-by: Claude Opus 5.5

* server : take d1omni audio through images and input_audio, scope memory-less lfm2 to non-causal

Assisted-by: Claude Opus 5.5

* common : rename decision type d1omni to lfm2-d1-omni, server : make images an alias of files

Assisted-by: Claude Opus 5.5
---
 common/common.cpp                |   4 +-
 common/common.h                  |   1 +
 conversion/__init__.py           |   2 +
 conversion/lfm2.py               | 167 +++++++++++++++++
 gguf-py/gguf/constants.py        |   7 +
 src/llama-model.cpp              |   5 +
 src/models/lfm2.cpp              | 295 ++++++++++++++++++++++++++++++-
 src/models/models.h              |  11 ++
 tools/mtmd/clip-impl.h           |   4 +
 tools/mtmd/clip-model.h          |   4 +
 tools/mtmd/clip.cpp              |  27 +++
 tools/mtmd/models/conformer.cpp  |  11 +-
 tools/mtmd/models/siglip.cpp     |   7 +-
 tools/mtmd/mtmd-audio.cpp        |  36 ++++
 tools/mtmd/mtmd-audio.h          |   6 +
 tools/mtmd/mtmd.cpp              |   9 +
 tools/server/README.md           |  18 +-
 tools/server/server-common.cpp   |   1 +
 tools/server/server-context.cpp  |   2 +-
 tools/server/server-decision.cpp | 233 ++++++++++++++++++++++--
 tools/server/server-decision.h   |  18 +-
 21 files changed, 833 insertions(+), 35 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index 0e891cea6..28ea8680e 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1171,6 +1171,7 @@ static const std::map COMMON_DECISION_TYPE_NA
     { COMMON_DECISION_TYPE_CLEF,           "clef"          },
     { COMMON_DECISION_TYPE_PPLX_DECIDER,   "pplx-decider"  },
     { COMMON_DECISION_TYPE_LFM2_D1,        "lfm2-d1"       },
+    { COMMON_DECISION_TYPE_LFM2_D1_OMNI,   "lfm2-d1-omni"  },
 };
 
 static common_decision_type common_decision_type_from_string(const std::string & str) {
@@ -1284,7 +1285,8 @@ common_init_result::common_init_result(common_params & params, bool model_only)
     // these decision models return a score for each token via the embeddings output
     // TODO: maybe improve this in the future
     const auto decision_type = common_get_decision_type(model);
-    if (decision_type == COMMON_DECISION_TYPE_LAYA || decision_type == COMMON_DECISION_TYPE_KEV || decision_type == COMMON_DECISION_TYPE_CLEF) {
+    if (decision_type == COMMON_DECISION_TYPE_LAYA || decision_type == COMMON_DECISION_TYPE_KEV || decision_type == COMMON_DECISION_TYPE_CLEF ||
+        decision_type == COMMON_DECISION_TYPE_LFM2_D1_OMNI) {
         params.embedding    = true;
         params.pooling_type = LLAMA_POOLING_TYPE_NONE;
 
diff --git a/common/common.h b/common/common.h
index 6f8acf31d..0a85f11f9 100644
--- a/common/common.h
+++ b/common/common.h
@@ -964,6 +964,7 @@ enum common_decision_type {
     COMMON_DECISION_TYPE_CLEF,    // all questions in one prompt, score of option i read from the embeddings output at row i
     COMMON_DECISION_TYPE_PPLX_DECIDER, // same as openjev, label codes of 1 or 2 letters
     COMMON_DECISION_TYPE_LFM2_D1, // same as openjev, the labels depend on the question type
+    COMMON_DECISION_TYPE_LFM2_D1_OMNI, // same as laya, other prompt layout
     COMMON_DECISION_TYPE_UNKNOWN, // a decision model of a type that is not supported
 };
 
diff --git a/conversion/__init__.py b/conversion/__init__.py
index 34c56bbee..72960be1a 100644
--- a/conversion/__init__.py
+++ b/conversion/__init__.py
@@ -162,6 +162,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "Lfm2BidirectionalModel": "lfm2",
     "Lfm2ForCausalLM": "lfm2",
     "D1Model": "lfm2",
+    "D1OmniModel": "lfm2",
     "Lfm2Model": "lfm2",
     "Lfm2MoeForCausalLM": "lfm2",
     "Llama4ForCausalLM": "llama",
@@ -335,6 +336,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
     "KimiK25ForConditionalGeneration": "kimivl",
     "KimiVLForConditionalGeneration": "kimivl",
     "Lfm2AudioForConditionalGeneration": "lfm2",
+    "D1OmniModel": "lfm2",
     "Lfm2VlForConditionalGeneration": "lfm2",
     "LightOnOCRForConditionalGeneration": "lighton_ocr",
     "Llama4ForConditionalGeneration": "llama4",
diff --git a/conversion/lfm2.py b/conversion/lfm2.py
index d4343675e..113c91116 100644
--- a/conversion/lfm2.py
+++ b/conversion/lfm2.py
@@ -161,6 +161,121 @@ class LFM2ColBertModel(LFM2Model):
         yield f"{self.dense_tensor_name}.weight", tensor.clone()
 
 
+def _is_d1_omni_checkpoint(dir_model: Path) -> bool:
+    if not (dir_model / "config.json").is_file():
+        return False
+    with open(dir_model / "config.json", encoding="utf-8") as f:
+        return json.load(f).get("model_type") == "d1_omni"
+
+
+@ModelBase.register_hparams_loader(_is_d1_omni_checkpoint)
+def _load_d1_omni_hparams(dir_model: Path) -> dict[str, Any]:
+    logger.info("gguf: detected d1-omni checkpoint")
+    hparams = ModelBase.load_hparams(dir_model, False, guess=False)
+    text = hparams["text_config"]
+    n_layer, n_layer_head = text["num_hidden_layers"], hparams["head_layers"]
+    # the trunk uses the LFM2 FFN sizing, the head blocks are appended with a plain 4x MLP
+    n_ff = int(text["block_ffn_dim_multiplier"] * int(2 * text["intermediate_size"] / 3))
+    n_ff = text["block_multiple_of"] * ((n_ff + text["block_multiple_of"] - 1) // text["block_multiple_of"])
+    text["num_hidden_layers"] = n_layer + n_layer_head
+    text["intermediate_size"] = [n_ff] * n_layer + [4 * text["hidden_size"]] * n_layer_head
+    text["block_auto_adjust_ff_dim"] = False
+    return hparams
+
+
+@ModelBase.register("D1OmniModel")
+@ModelBase.example("LiquidAI/d1-omni-600M")
+class D1OmniModel(LFM2Model):
+    model_arch = gguf.MODEL_ARCH.LFM2
+
+    # the server cuts the text to these lengths, see server-decision.cpp
+    _MAX_LENGTH = 16384
+    _IMAGE_TEXT_LENGTH = 896
+    _AUDIO_TEXT_LENGTH = 15360
+
+    def set_vocab(self):
+        super().set_vocab()
+        # the systemone template writes the BOS, after the media
+        self.gguf_writer.remove_key(gguf.Keys.Tokenizer.ADD_BOS)
+        self.gguf_writer.add_add_bos_token(False)
+        self.gguf_writer.add_token_type_count(3)  # choice, score, noul
+        self.gguf_writer.add_chat_template([{"name": "systemone", "template": self._systemone_template()}])
+
+    @staticmethod
+    def _systemone_template() -> str:
+        # follows prompt.py of the model repo, the server cuts each marked piece to its token budget
+        # the media (images, or an audio clip if audio is true) come first
+        description = jinja_str_or_json("o.description")
+        has_description = "o.description is not none and o.description != ''"
+        yes_no = "{{ 'yes' if o.key == 'true' else 'no' }}"
+        option_code = "{% if loop.index0 < 10 %}00{% elif loop.index0 < 100 %}0{% endif %}{{ loop.index0 }}"
+        option = (
+            "{% if type == 'choice' and audio %}option_" + option_code + ": "
+            "{% if " + has_description + " %}" + description + "{% else %}{{ o.key }}{% endif %}"
+            "{% elif type == 'choice' %}{{ o.key }}{% if " + has_description + " %}: " + description + "{% endif %}"
+            "{% elif type == 'score' %}level {{ o.key }}: " + description
+            + "{% elif audio %}{{ o.key }}: " + yes_no
+            + "{% else %}{{ o.key }}: {% if " + has_description + " %}" + description
+            + "{% elif images and not ns.criteria %}" + yes_no
+            + "{% elif o.key == 'true' %}yes, the statement holds"
+            "{% else %}no, the statement does not hold{% endif %}{% endif %}"
+        )
+        state = "{% if state is string %}{{ state }}{% elif state is not none %}{{ state | tojson }}{% elif audio %}{}{% endif %}"
+        return (
+            "{% set ns = namespace(criteria=false) %}"
+            "{% for o in options %}{% if o.description is not none %}{% set ns.criteria = true %}{% endif %}{% endfor %}"
+            "{% for image in images %}{{ image }}{% endfor %}{{ sep }}"
+            "<|startoftext|><|reserved_7|>{{ sep }}{{ mark_state }}" + state
+            + "{{ sep }}{{ mark_question }}<|reserved_8|>" + jinja_str_or_json("instructions")
+            + "{% for o in options %}{{ sep }}<|reserved_9|><|mask|>{{ sep }}{{ mark_option }} " + option
+            + "{{ sep }}<|reserved_10|>{% endfor %}{{ sep }}<|reserved_11|>"
+        )
+
+    def set_gguf_parameters(self):
+        lengths = (self.hparams["max_length"], self.hparams["image_text_length"], self.hparams["audio_text_length"])
+        if lengths != (self._MAX_LENGTH, self._IMAGE_TEXT_LENGTH, self._AUDIO_TEXT_LENGTH):
+            raise ValueError(f"unexpected text lengths: {lengths}")
+        n_head, n_layer_head = self.hparams["num_attention_heads"], self.hparams["head_layers"]
+        self.hparams["num_key_value_heads"] = [
+            self.hparams["num_key_value_heads"] if t != "conv" else 0 for t in self.hparams["layer_types"]
+        ] + [n_head] * n_layer_head
+
+        # the head needs per-layer sizes, LFM2Model writes a single feed forward length
+        TextModel.set_gguf_parameters(self)
+        self.gguf_writer.add_vocab_size(self.hparams["vocab_size"])
+        self.gguf_writer.add_shortconv_l_cache(self.hparams["conv_L_cache"])
+        self.gguf_writer.add_layer_norm_eps(1e-5)  # nn.LayerNorm of the head
+        self.gguf_writer.add_causal_attention(False)
+
+        self.gguf_writer.add_decision_type(gguf.DecisionType.LFM2_D1_OMNI)
+        self.gguf_writer.add_decision_block_count(n_layer_head)
+        # "choice:3-5" -> "choice.3_5", "choice:11+" -> "choice.11"
+        for name, value in self.hparams["temperatures"].items():
+            self.gguf_writer.add_decision_temperature(name.replace(":", ".").replace("-", "_").rstrip("+"), value)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+
+        if name.startswith(("vision.", "audio.")):
+            return None
+
+        name = name.replace("encoder.", "model.", 1) if name.startswith("encoder.") else name
+        name = name.replace("head.head.layers.", "head.layers.").replace("in_proj_", "in_proj.")
+        name = name.removeprefix("head.") if name.startswith(("head.type_emb", "head.scorer")) else name
+
+        return super().filter_tensors((name, gen))
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name.startswith("head.layers.") and bid is not None:
+            # the head blocks come after the trunk blocks
+            suffix = name.split(".", 3)[3]
+            bid += self.block_count - self.hparams["head_layers"]
+            name = f"head.layers.{bid}.{suffix}"
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+
 @ModelBase.register("Lfm2MoeForCausalLM")
 @ModelBase.example("LiquidAI/LFM2-8B-A1B")
 class LFM2MoeModel(TextModel):
@@ -276,6 +391,58 @@ class LFM2VLModel(MmprojModel):
         yield from super().modify_tensors(data_torch, name, bid)
 
 
+@ModelBase.register("D1OmniModel")
+@ModelBase.example("LiquidAI/d1-omni-600M")
+class D1OmniMmprojModel(ConformerAudioModel):
+    has_vision_encoder = True
+    has_audio_encoder  = True
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert self.hparams_vision is not None and self.hparams_audio is not None
+        # dynamic resolution, as LFM2VLModel
+        self.hparams_vision["image_size"] = 256
+        # the images are normalized to [-1, 1] (vision.py of the model repo)
+        self.preprocessor_config = {**self.preprocessor_config, "image_mean": [0.5] * 3, "image_std": [0.5] * 3}
+        self.hparams_audio["hidden_size"] = self.hparams_audio["d_model"]
+        self.hparams_audio["intermediate_size"] = self.hparams_audio["d_model"] * self.hparams_audio["ff_expansion_factor"]
+        self.hparams_audio["num_attention_heads"] = self.hparams_audio["n_heads"]
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.D1OMNI_V)
+        self.gguf_writer.add_vision_attention_layernorm_eps(self.find_vparam(["layer_norm_eps"]))
+        self.gguf_writer.add_vision_projector_scale_factor(self.global_config.get("downsample_factor", 2))
+        self.gguf_writer.add_vision_use_gelu(True)
+
+        assert self.hparams_audio is not None
+        self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.D1OMNI_A)
+        self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["feat_in"])
+        self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+
+        if name.startswith(("encoder.", "head.")):
+            return None
+
+        name = name.replace("vision.tower.", "vision_tower.").replace("vision.projector.", "multi_modal_projector.")
+        name = name.replace("audio.encoder.", "conformer.")
+        # the residual block continues the adapter: norm, linear, gelu, linear, then norm, down, up
+        for old, new in (("adapter.norm", 0), ("adapter.linear_1", 1), ("adapter.linear_2", 3),
+                         ("residual.ln", 4), ("residual.down", 5), ("residual.up", 6)):
+            name = name.replace(f"audio.{old}.", f"audio_adapter.model.{new}.")
+
+        return super().filter_tensors((name, gen))
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if "patch_embedding.weight" in name:
+            data_torch = data_torch.view(data_torch.shape[0], 16, 16, 3).permute(0, 3, 1, 2)
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+
 @ModelBase.register("Lfm2AudioForConditionalGeneration")
 @ModelBase.example("LiquidAI/LFM2.5-Audio-1.5B", "LiquidAI/LFM2-Audio-1.5B")
 class LFM2AudioModel(ConformerAudioModel):
diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index c181cb44a..fcf325fe9 100644
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -5172,6 +5172,10 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.ATTN_OUT,
         MODEL_TENSOR.OUTPUT,
         MODEL_TENSOR.DENSE_2_OUT, # LFM2-ColBert-350M
+        MODEL_TENSOR.TOKEN_TYPES, # decision head
+        MODEL_TENSOR.CLS,
+        MODEL_TENSOR.CLS_NORM,
+        MODEL_TENSOR.CLS_OUT,
     ],
     MODEL_ARCH.LFM2MOE: [
         MODEL_TENSOR.TOKEN_EMBD,
@@ -6072,6 +6076,7 @@ class DecisionType:
     CLEF    = "clef"     # joint head over all questions, one score per option
     PPLX_DECIDER = "pplx-decider"  # same as openjev, label codes of 1 or 2 letters
     LFM2_D1 = "lfm2-d1"  # same as openjev, the labels depend on the question type
+    LFM2_D1_OMNI = "lfm2-d1-omni"  # same head as laya on a bidirectional LFM2 trunk, other prompt layout
 
 
 class VisionProjectorType:
@@ -6133,6 +6138,8 @@ class VisionProjectorType:
     GRANITE4_VISION = "granite4_vision"
     MUSE_GLIMMER   = "muse-glimmer"
     COHERE2V       = "cohere2v"
+    D1OMNI_V       = "d1omni_v"  # lfm2 vision, without separator tokens
+    D1OMNI_A       = "d1omni_a"  # lfm2a audio, with a residual block after the projector
 
 
 # Items here are (block size, type size)
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
index 9f05a66e4..bf04945bf 100644
--- a/src/llama-model.cpp
+++ b/src/llama-model.cpp
@@ -2368,6 +2368,11 @@ ggml_tensor * llama_model::get_rope_factors(const llama_cparams & cparams, int i
 llama_memory_i * llama_model::create_memory(const llama_memory_params & params, const llama_cparams & cparams) const {
     llama_memory_i * res;
 
+    // the non-causal LFM2 decision graph reads the whole prompt in one batch, nothing is kept
+    if (arch == LLM_ARCH_LFM2 && !hparams.causal_attn && hparams.n_layer_decision > 0) {
+        return nullptr;
+    }
+
     switch (arch) {
         // Models that need specific instantiation should be handled in the
         // switch statement
diff --git a/src/models/lfm2.cpp b/src/models/lfm2.cpp
index 07b71ccd3..2488c0746 100644
--- a/src/models/lfm2.cpp
+++ b/src/models/lfm2.cpp
@@ -4,6 +4,9 @@
 
 #include 
 
+// question types of a decision model: choice, score, noul
+static const uint32_t N_DECISION_TYPES = 3;
+
 void llama_model_lfm2::load_arch_hparams(llama_model_loader & ml) {
     ml.get_key(LLM_KV_SHORTCONV_L_CACHE,           hparams.n_shortconv_l_cache);
     ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
@@ -23,6 +26,15 @@ void llama_model_lfm2::load_arch_hparams(llama_model_loader & ml) {
         default:    type = LLM_TYPE_UNKNOWN;
     }
 
+    ml.get_key(LLM_KV_DECISION_BLOCK_COUNT, hparams.n_layer_decision, false);
+    if (hparams.n_layer_decision > 0) {
+        if (hparams.n_layer_decision >= hparams.n_layer() || hparams.causal_attn) {
+            throw std::runtime_error("invalid decision head");
+        }
+        ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
+        hparams.n_embd_out_impl = N_DECISION_TYPES;
+    }
+
     if (const auto is_swa = ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false); is_swa && hparams.n_swa > 0) {
         hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
         for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
@@ -37,13 +49,49 @@ void llama_model_lfm2::load_arch_tensors(llama_model_loader &) {
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 
     output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM_LFM2, "weight"), {n_embd}, 0);
-    output      = create_tensor(tn(LLM_TENSOR_OUTPUT,           "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
 
-    if (output == NULL) {
-        output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
+    if (hparams.n_layer_decision > 0) {
+        // decision head: plain pre-norm blocks with biases
+        for (int i = n_layer - (int) hparams.n_layer_decision; i < n_layer; ++i) {
+            auto & layer = layers[i];
+            const int64_t n_ff_head = hparams.n_ff(i);
+
+            layer.attn_norm   = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+            layer.attn_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "bias",   i), {n_embd}, 0);
+
+            layer.wqkv   = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, 3 * n_embd}, 0);
+            layer.wqkv_b = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "bias",   i), {3 * n_embd}, 0);
+            layer.wo     = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}, 0);
+            layer.wo_b   = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias",   i), {n_embd}, 0);
+
+            layer.ffn_norm   = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+            layer.ffn_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias",   i), {n_embd}, 0);
+            layer.ffn_up     = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd, n_ff_head}, 0);
+            layer.ffn_up_b   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "bias",   i), {n_ff_head}, 0);
+            layer.ffn_down   = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff_head, n_embd}, 0);
+            layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias",   i), {n_embd}, 0);
+        }
+
+        if (n_token_types != N_DECISION_TYPES) {
+            throw std::runtime_error("decision model must have one token type per question type");
+        }
+        type_embd = create_tensor(tn(LLM_TENSOR_TOKEN_TYPES, "weight"), {n_embd, n_token_types}, 0);
+
+        cls_norm   = create_tensor(tn(LLM_TENSOR_CLS_NORM, "weight"), {n_embd},    0);
+        cls_norm_b = create_tensor(tn(LLM_TENSOR_CLS_NORM, "bias"),   {n_embd},    0);
+        cls        = create_tensor(tn(LLM_TENSOR_CLS,      "weight"), {n_embd, n_embd}, 0);
+        cls_b      = create_tensor(tn(LLM_TENSOR_CLS,      "bias"),   {n_embd},    0);
+        cls_out    = create_tensor(tn(LLM_TENSOR_CLS_OUT,  "weight"), {n_embd, 1}, 0);
+        cls_out_b  = create_tensor(tn(LLM_TENSOR_CLS_OUT,  "bias"),   {1},         0);
+    } else {
+        output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
+
+        if (output == NULL) {
+            output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
+        }
     }
 
-    for (int i = 0; i < n_layer; ++i) {
+    for (int i = 0; i < n_layer - (int) hparams.n_layer_decision; ++i) {
         auto & layer = layers[i];
 
         const bool is_moe_layer = i >= static_cast(hparams.n_layer_dense_lead);
@@ -87,6 +135,9 @@ void llama_model_lfm2::load_arch_tensors(llama_model_loader &) {
 }
 
 std::unique_ptr llama_model_lfm2::build_arch_graph(const llm_graph_params & params) const {
+    if (hparams.n_layer_decision > 0) {
+        return std::make_unique(*this, params);
+    }
     if (hparams.swa_type == LLAMA_SWA_TYPE_STANDARD) {
         return std::make_unique>(*this, params);
     } else {
@@ -294,6 +345,242 @@ llama_model_lfm2::graph::graph(const llama_model & model, const llm_graph_
     ggml_build_forward_expand(gf, cur);
 }
 
+// media entries (an image or audio prefix) are embeddings, text entries are tokens
+static bool lfm2_is_media(const llama_ubatch & ubatch, int64_t i) {
+    return ubatch.is_mixed() ? ubatch.type[i] != 0 : ubatch.token == nullptr;
+}
+
+// non-causal within a sequence, the media never reads the text, so it is a function of the media alone
+// in the head, the text and the media only read their own kind
+class llm_graph_input_attn_media : public llm_graph_input_attn_no_cache {
+public:
+    llm_graph_input_attn_media(const llama_hparams & hparams, const llama_cparams & cparams, bool is_head) :
+        llm_graph_input_attn_no_cache(hparams, cparams), is_head(is_head) {}
+
+    void set_input(const llama_ubatch * ubatch) override {
+        const int64_t n_tokens = ubatch->n_tokens;
+
+        std::vector is_media(n_tokens);
+        for (int64_t i = 0; i < n_tokens; ++i) {
+            is_media[i] = lfm2_is_media(*ubatch, i);
+        }
+
+        const auto fill_mask = [&](auto * data, auto zero, auto ninf) {
+            for (int64_t i1 = 0; i1 < n_tokens; ++i1) {
+                for (int64_t i0 = 0; i0 < n_tokens; ++i0) {
+                    bool visible = ubatch->seq_id[i0][0] == ubatch->seq_id[i1][0];
+                    if (is_head) {
+                        visible = visible && is_media[i0] == is_media[i1];
+                    } else {
+                        visible = visible && !(is_media[i1] && !is_media[i0]);
+                    }
+                    data[i1 * n_tokens + i0] = visible ? zero : ninf;
+                }
+            }
+        };
+
+        GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask->buffer));
+        if (self_kq_mask->type == GGML_TYPE_F16) {
+            fill_mask((ggml_fp16_t *) self_kq_mask->data, ggml_fp32_to_fp16(0.0f), ggml_fp32_to_fp16(-INFINITY));
+        } else {
+            fill_mask((float *) self_kq_mask->data, 0.0f, -INFINITY);
+        }
+    }
+
+    const bool is_head;
+};
+
+// 1 where the previous (next) token is the left (right) neighbor in the same sequence
+// the last media entry does not read the text on its right
+class llm_graph_input_conv_mask : public llm_graph_input_i {
+public:
+    void set_input(const llama_ubatch * ubatch) override {
+        const int64_t n_tokens = ubatch->n_tokens;
+
+        std::vector data_left(n_tokens, 0.0f);
+        std::vector data_right(n_tokens, 0.0f);
+        for (int64_t i = 0; i + 1 < n_tokens; ++i) {
+            const bool is_next = ubatch->seq_id[i][0] == ubatch->seq_id[i + 1][0] && ubatch->pos[i] + 1 == ubatch->pos[i + 1];
+            data_right[i]    = is_next && !(lfm2_is_media(*ubatch, i) && !lfm2_is_media(*ubatch, i + 1));
+            data_left[i + 1] = is_next;
+        }
+        ggml_backend_tensor_set(left,  data_left.data(),  0, ggml_nbytes(left));
+        ggml_backend_tensor_set(right, data_right.data(), 0, ggml_nbytes(right));
+    }
+
+    ggml_tensor * left  = nullptr; // F32 [1, n_tokens]
+    ggml_tensor * right = nullptr; // F32 [1, n_tokens]
+};
+
+llama_model_lfm2::graph_decision::graph_decision(const llama_model & model, const llm_graph_params & params) :
+    llm_graph_context(params) {
+    const int64_t n_embd_head = hparams.n_embd_head_v();
+    const int     n_layer_enc = n_layer - hparams.n_layer_decision;
+
+    ggml_tensor * cur = build_inp_embd(model.tok_embd);
+    cb(cur, "model.embed_tokens", -1);
+
+    ggml_tensor * inp_pos     = build_inp_pos();
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+    const auto type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
+
+    llm_graph_input_attn_no_cache * inp_attn[2];
+    for (bool is_head : {false, true}) {
+        auto inp = std::make_unique(hparams, cparams, is_head);
+        inp->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_tokens, n_tokens, 1, 1);
+        ggml_set_input(inp->self_kq_mask);
+        inp->self_kq_mask_cnv = inp->self_kq_mask;
+        inp_attn[is_head] = (llm_graph_input_attn_no_cache *) res->add_input(std::move(inp));
+    }
+
+    auto inp_conv = std::make_unique();
+    inp_conv->left  = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, n_tokens);
+    inp_conv->right = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, n_tokens);
+    ggml_set_input(inp_conv->left);
+    ggml_set_input(inp_conv->right);
+    ggml_tensor * conv_left  = inp_conv->left;
+    ggml_tensor * conv_right = inp_conv->right;
+    res->add_input(std::move(inp_conv));
+
+    for (int il = 0; il < n_layer_enc; ++il) {
+        const auto & layer = model.layers[il];
+
+        ggml_tensor * inpL = cur;
+        cur = build_norm(cur, layer.attn_norm, NULL, LLM_NORM_RMS, il);
+        cb(cur, "model.layers.{}.operator_norm", il);
+
+        if (hparams.is_recr(il)) {
+            ggml_tensor * bcx = build_lora_mm(layer.shortconv.in_proj, cur);
+            cb(bcx, "model.layers.{}.conv.in_proj", il);
+
+            ggml_tensor * b = ggml_view_2d(ctx0, bcx, n_embd, n_tokens, bcx->nb[1], 0 * n_embd * ggml_element_size(bcx));
+            ggml_tensor * c = ggml_view_2d(ctx0, bcx, n_embd, n_tokens, bcx->nb[1], 1 * n_embd * ggml_element_size(bcx));
+            ggml_tensor * x = ggml_view_2d(ctx0, bcx, n_embd, n_tokens, bcx->nb[1], 2 * n_embd * ggml_element_size(bcx));
+
+            // centred 3-tap conv, a tap outside the sequence reads 0
+            ggml_tensor * bx  = ggml_mul(ctx0, b, x);
+            ggml_tensor * bxp = ggml_pad_ext(ctx0, bx, 0, 0, 1, 1, 0, 0, 0, 0);
+            ggml_tensor * prv = ggml_view_2d(ctx0, bxp, n_embd, n_tokens, bxp->nb[1], 0);
+            ggml_tensor * nxt = ggml_view_2d(ctx0, bxp, n_embd, n_tokens, bxp->nb[1], 2 * bxp->nb[1]);
+
+            GGML_ASSERT(hparams.n_shortconv_l_cache == 3);
+            ggml_tensor * taps = ggml_cont(ctx0, ggml_transpose(ctx0, layer.shortconv.conv));
+            ggml_tensor * tap0 = ggml_view_1d(ctx0, taps, n_embd, 0 * taps->nb[1]);
+            ggml_tensor * tap1 = ggml_view_1d(ctx0, taps, n_embd, 1 * taps->nb[1]);
+            ggml_tensor * tap2 = ggml_view_1d(ctx0, taps, n_embd, 2 * taps->nb[1]);
+
+            ggml_tensor * y = ggml_mul(ctx0, bx, tap1);
+            y = ggml_add(ctx0, y, ggml_mul(ctx0, ggml_mul(ctx0, prv, tap0), conv_left));
+            y = ggml_add(ctx0, y, ggml_mul(ctx0, ggml_mul(ctx0, nxt, tap2), conv_right));
+            cb(y, "model.layers.{}.conv.conv", il);
+
+            cur = build_lora_mm(layer.shortconv.out_proj, ggml_mul(ctx0, c, y));
+            cb(cur, "model.layers.{}.conv.out_proj", il);
+        } else {
+            auto [q, k, v] = build_qkv(layer, cur, n_embd_head, n_head, hparams.n_head_kv(il), il);
+
+            q = build_norm(q, layer.attn_q_norm, NULL, LLM_NORM_RMS, il);
+            k = build_norm(k, layer.attn_k_norm, NULL, LLM_NORM_RMS, il);
+
+            q = ggml_rope_ext(ctx0, q, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor,
+                              attn_factor, beta_fast, beta_slow);
+            k = ggml_rope_ext(ctx0, k, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor,
+                              attn_factor, beta_fast, beta_slow);
+
+            cur = build_attn(inp_attn[0],
+                    layer.wo, NULL, layer.wo_s,
+                    q, k, v, nullptr, nullptr, nullptr, 1.0f / sqrtf(float(n_embd_head)), il);
+            cb(cur, "model.layers.{}.self_attn.out_proj", il);
+        }
+
+        cur = ggml_add(ctx0, cur, inpL);
+
+        ggml_tensor * ffn_out = build_norm(cur, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
+        ffn_out = build_ffn(ffn_out,
+                layer.ffn_up,   NULL, NULL,
+                layer.ffn_gate, NULL, NULL,
+                layer.ffn_down, NULL, NULL,
+                NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
+
+        cur = ggml_add(ctx0, cur, ffn_out);
+        cb(cur, "l_out", il);
+    }
+
+    cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
+    cb(cur, "result_norm", -1);
+
+    cur = build_decision_head(model, cur, inp_attn[1], inp_out_ids);
+
+    res->t_embd = cur;
+    ggml_build_forward_expand(gf, cur);
+}
+
+// same as llama_model_modern_bert::graph::build_decision_head(), with the head counts of the head layers
+ggml_tensor * llama_model_lfm2::graph_decision::build_decision_head(
+        const llama_model & model,
+        ggml_tensor * inp,
+        llm_graph_input_attn_no_cache * inp_attn,
+        ggml_tensor * inp_out_ids) {
+    const int64_t n_embd_head = hparams.n_embd_head_v();
+    const int     n_layer_enc = n_layer - hparams.n_layer_decision;
+
+    ggml_tensor * scores = nullptr;
+
+    // the question type is not a graph input, so the head is evaluated for each of them
+    for (uint32_t it = 0; it < N_DECISION_TYPES; ++it) {
+        ggml_tensor * type_row = ggml_view_1d(ctx0, model.type_embd, n_embd, it * model.type_embd->nb[1]);
+        ggml_tensor * inpL = ggml_add(ctx0, inp, type_row);
+
+        for (int il = n_layer_enc; il < n_layer; ++il) {
+            const auto & layer = model.layers[il];
+
+            ggml_tensor * cur = build_norm(inpL, layer.attn_norm, layer.attn_norm_b, LLM_NORM, il);
+            cb(cur, "attn_norm", il);
+
+            // no positional encoding in the head
+            auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il);
+
+            cur = build_attn(inp_attn,
+                        layer.wo, layer.wo_b, layer.wo_s,
+                        Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
+            cb(cur, "kqv_out", il);
+
+            if (il == n_layer - 1 && inp_out_ids) {
+                cur  = ggml_get_rows(ctx0,  cur, inp_out_ids);
+                inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
+            }
+
+            ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL);
+            cb(ffn_inp, "ffn_inp", il);
+
+            cur = build_norm(ffn_inp, layer.ffn_norm, layer.ffn_norm_b, LLM_NORM, il);
+            cb(cur, "ffn_norm", il);
+
+            cur = build_ffn(cur,
+                    layer.ffn_up,   layer.ffn_up_b,   NULL,
+                    NULL,           NULL,             NULL,
+                    layer.ffn_down, layer.ffn_down_b, NULL,
+                    NULL,
+                    LLM_FFN_RELU,
+                    LLM_FFN_SEQ, il);
+
+            inpL = ggml_add(ctx0, cur, ffn_inp);
+        }
+
+        // scorer
+        ggml_tensor * cur = build_norm(inpL, model.cls_norm, model.cls_norm_b, LLM_NORM, -1);
+        cur = ggml_add(ctx0, build_lora_mm(model.cls, cur), model.cls_b);
+        cur = ggml_gelu_erf(ctx0, cur);
+        cur = ggml_add(ctx0, build_lora_mm(model.cls_out, cur), model.cls_out_b);
+
+        scores = scores ? ggml_concat(ctx0, scores, cur, 0) : cur;
+    }
+    cb(scores, "decision_scores", -1);
+
+    return scores;
+}
+
 // Explicit template instantiations
 template struct llama_model_lfm2::graph;
 template struct llama_model_lfm2::graph;
diff --git a/src/models/models.h b/src/models/models.h
index 1ef0c5156..d6ddf9d16 100644
--- a/src/models/models.h
+++ b/src/models/models.h
@@ -2186,6 +2186,17 @@ struct llama_model_lfm2 : public llama_model_base {
         graph(const llama_model & model, const llm_graph_params & params);
     };
 
+    // non-causal trunk without memory, then the decision head
+    struct graph_decision : public llm_graph_context {
+        graph_decision(const llama_model & model, const llm_graph_params & params);
+
+        ggml_tensor * build_decision_head(
+                const llama_model & model,
+                ggml_tensor * inp,
+                llm_graph_input_attn_no_cache * inp_attn,
+                ggml_tensor * inp_out_ids);
+    };
+
     std::unique_ptr build_arch_graph(const llm_graph_params & params) const override;
 };
 
diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h
index abf75d998..de171b7a5 100644
--- a/tools/mtmd/clip-impl.h
+++ b/tools/mtmd/clip-impl.h
@@ -476,6 +476,7 @@ enum projector_type {
     PROJECTOR_TYPE_MERALION,
     PROJECTOR_TYPE_MUSIC_FLAMINGO,
     PROJECTOR_TYPE_LFM2,
+    PROJECTOR_TYPE_D1OMNI_V,
     PROJECTOR_TYPE_KIMIVL,
     PROJECTOR_TYPE_PADDLEOCR,
     PROJECTOR_TYPE_LIGHTONOCR,
@@ -488,6 +489,7 @@ enum projector_type {
     PROJECTOR_TYPE_DEEPSEEKOCR2,
     PROJECTOR_TYPE_DEEPSEEK4V,
     PROJECTOR_TYPE_LFM2A,
+    PROJECTOR_TYPE_D1OMNI_A,
     PROJECTOR_TYPE_GLM4V,
     PROJECTOR_TYPE_GLM5V,
     PROJECTOR_TYPE_YOUTUVL,
@@ -544,6 +546,7 @@ static std::map PROJECTOR_TYPE_NAMES = {
     { PROJECTOR_TYPE_MERALION,          "meralion"},
     { PROJECTOR_TYPE_MUSIC_FLAMINGO,    "musicflamingo"},
     { PROJECTOR_TYPE_LFM2,              "lfm2"},
+    { PROJECTOR_TYPE_D1OMNI_V,          "d1omni_v"},
     { PROJECTOR_TYPE_KIMIVL,            "kimivl"},
     { PROJECTOR_TYPE_PADDLEOCR,         "paddleocr"},
     { PROJECTOR_TYPE_LIGHTONOCR,        "lightonocr"},
@@ -556,6 +559,7 @@ static std::map PROJECTOR_TYPE_NAMES = {
     { PROJECTOR_TYPE_DEEPSEEKOCR2,      "deepseekocr2"},
     { PROJECTOR_TYPE_DEEPSEEK4V,        "deepseek4v"},
     { PROJECTOR_TYPE_LFM2A,             "lfm2a"},
+    { PROJECTOR_TYPE_D1OMNI_A,          "d1omni_a"},
     { PROJECTOR_TYPE_GLM4V,             "glm4v"},
     { PROJECTOR_TYPE_GLM5V,             "glm5v"},
     { PROJECTOR_TYPE_YOUTUVL,           "youtuvl"},
diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h
index 33f679fc0..a5ff38822 100644
--- a/tools/mtmd/clip-model.h
+++ b/tools/mtmd/clip-model.h
@@ -623,6 +623,10 @@ struct clip_model {
     ggml_tensor * mm_3_b = nullptr;
     ggml_tensor * mm_4_w = nullptr;
     ggml_tensor * mm_4_b = nullptr;
+    ggml_tensor * mm_5_w = nullptr;
+    ggml_tensor * mm_5_b = nullptr;
+    ggml_tensor * mm_6_w = nullptr;
+    ggml_tensor * mm_6_b = nullptr;
 
     // GLMV-Edge projection
     ggml_tensor * mm_model_adapter_conv_w = nullptr;
diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp
index b7302e4bf..ca38d76df 100644
--- a/tools/mtmd/clip.cpp
+++ b/tools/mtmd/clip.cpp
@@ -939,6 +939,7 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const
         case PROJECTOR_TYPE_IDEFICS3:
         case PROJECTOR_TYPE_COHERE2V:
         case PROJECTOR_TYPE_LFM2:
+        case PROJECTOR_TYPE_D1OMNI_V:
         case PROJECTOR_TYPE_JANUS_PRO:
         case PROJECTOR_TYPE_PHI4:
             {
@@ -1073,6 +1074,7 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const
                 builder = std::make_unique(ctx, img);
             } break;
         case PROJECTOR_TYPE_LFM2A:
+        case PROJECTOR_TYPE_D1OMNI_A:
             {
                 builder = std::make_unique(ctx, img);
             } break;
@@ -1525,6 +1527,7 @@ struct clip_model_loader {
                         }
                     } break;
                 case PROJECTOR_TYPE_LFM2:
+                case PROJECTOR_TYPE_D1OMNI_V:
                     {
                         // default for older GGUFs
                         std::string resize_algo = "bilinear";
@@ -1540,6 +1543,10 @@ struct clip_model_loader {
                         }
                         hparams.image_resize_algo_rf = hparams.image_resize_algo;
                         hparams.image_resize_algo_ov = hparams.image_resize_algo;
+                        // the tiles stretch the image to the grid (d1-omni vision.py)
+                        if (model.proj_type == PROJECTOR_TYPE_D1OMNI_V) {
+                            hparams.image_pad_rf = PAD_NONE;
+                        }
                         get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge, false);
                         // ref: https://huggingface.co/LiquidAI/LFM2.5-VL-1.6B/blob/main/processor_config.json
                         hparams.set_limit_image_tokens(64, 256);
@@ -1980,6 +1987,7 @@ struct clip_model_loader {
                         hparams.set_warmup_n_tokens(32*32);
                     } break;
                 case PROJECTOR_TYPE_LFM2A:
+                case PROJECTOR_TYPE_D1OMNI_A:
                     {
                         // audio preprocessing params
                         hparams.audio_chunk_len        = 1; // in seconds
@@ -2799,6 +2807,7 @@ struct clip_model_loader {
                     model.mm_2_b = get_tensor(string_format(TN_LLAVA_PROJ, 2, "bias"));
                 } break;
             case PROJECTOR_TYPE_LFM2:
+            case PROJECTOR_TYPE_D1OMNI_V:
                 {
                     model.mm_input_norm_w = get_tensor(TN_MM_INP_NORM, false);
                     model.mm_input_norm_b = get_tensor(TN_MM_INP_NORM_B, false);
@@ -3411,6 +3420,7 @@ struct clip_model_loader {
                     model.mm_input_proj_w = get_tensor(string_format(TN_A_MM_INP_PROJ, "weight"));
                 } break;
             case PROJECTOR_TYPE_LFM2A:
+            case PROJECTOR_TYPE_D1OMNI_A:
                 {
                     for (int i : {0, 2, 3, 5, 6}) {
                         model.pre_encode_conv_X_w[i] = get_tensor(string_format(TN_CONV1D, i, "weight"));
@@ -3426,6 +3436,16 @@ struct clip_model_loader {
                     model.mm_3_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 3, "weight"));
                     model.mm_3_b = get_tensor(string_format(TN_MM_AUDIO_MLP, 3, "bias"));
 
+                    // residual block after the projector: norm, down, up
+                    if (model.proj_type == PROJECTOR_TYPE_D1OMNI_A) {
+                        model.mm_4_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 4, "weight"));
+                        model.mm_4_b = get_tensor(string_format(TN_MM_AUDIO_MLP, 4, "bias"));
+                        model.mm_5_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 5, "weight"));
+                        model.mm_5_b = get_tensor(string_format(TN_MM_AUDIO_MLP, 5, "bias"));
+                        model.mm_6_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 6, "weight"));
+                        model.mm_6_b = get_tensor(string_format(TN_MM_AUDIO_MLP, 6, "bias"));
+                    }
+
                     for (int il = 0; il < hparams.n_layer; ++il) {
                         auto & layer = model.layers[il];
 
@@ -4277,6 +4297,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
                 n_patches = ctx->model.hparams.image_size / ctx->model.hparams.patch_size;
             } break;
         case PROJECTOR_TYPE_LFM2:
+        case PROJECTOR_TYPE_D1OMNI_V:
         case PROJECTOR_TYPE_KIMIVL:
         case PROJECTOR_TYPE_KIMIK25:
             {
@@ -4400,6 +4421,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
             }
         } break;
         case PROJECTOR_TYPE_LFM2A:
+        case PROJECTOR_TYPE_D1OMNI_A:
             {
                 n_patches = ((((img->nx() + 1) / 2) + 1) / 2 + 1) / 2;
             } break;
@@ -5358,6 +5380,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
         case PROJECTOR_TYPE_GLMA:
         case PROJECTOR_TYPE_ULTRAVOX:
         case PROJECTOR_TYPE_LFM2:
+        case PROJECTOR_TYPE_D1OMNI_V:
         case PROJECTOR_TYPE_VOXTRAL:
         case PROJECTOR_TYPE_MERALION:
         case PROJECTOR_TYPE_MUSIC_FLAMINGO:
@@ -5617,6 +5640,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
                 }
             } break;
         case PROJECTOR_TYPE_LFM2A:
+        case PROJECTOR_TYPE_D1OMNI_A:
             {
                 GGML_ASSERT(imgs.entries.size() == 1);
                 const auto n_frames = clip_n_output_tokens(ctx, &imgs.entries.front());
@@ -6080,6 +6104,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
             return ctx->model.mm_2_w->ne[1];
         case PROJECTOR_TYPE_GLMA:
         case PROJECTOR_TYPE_LFM2:
+        case PROJECTOR_TYPE_D1OMNI_V:
         case PROJECTOR_TYPE_KIMIVL:
         case PROJECTOR_TYPE_PADDLEOCR:
         case PROJECTOR_TYPE_KIMIK25:
@@ -6096,6 +6121,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
             return ctx->model.mm_fc_w->ne[1];
         case PROJECTOR_TYPE_LFM2A:
             return ctx->model.position_embeddings->ne[0];
+        case PROJECTOR_TYPE_D1OMNI_A:
+            return ctx->model.mm_3_w->ne[1];
         case PROJECTOR_TYPE_GRANITE_SPEECH:
             return ctx->model.qf_proj_blocks[0].qf_proj_linear_w->ne[1];
         case PROJECTOR_TYPE_GRANITE4_VISION:
diff --git a/tools/mtmd/models/conformer.cpp b/tools/mtmd/models/conformer.cpp
index 18c3d27bc..9463f9f1c 100644
--- a/tools/mtmd/models/conformer.cpp
+++ b/tools/mtmd/models/conformer.cpp
@@ -4,7 +4,7 @@ ggml_cgraph * clip_graph_conformer::build() {
     const int n_frames   = img.nx();
     const int n_pos      = n_frames / 2;
     const int n_pos_embd = (((((n_frames + 1) / 2) + 1) / 2 + 1) / 2) * 2 - 1;
-    GGML_ASSERT(model.position_embeddings->ne[1] >= n_pos);
+    GGML_ASSERT(!model.position_embeddings || model.position_embeddings->ne[1] >= n_pos);
 
     ggml_tensor * pos_emb = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 512, n_pos_embd);
     ggml_set_name(pos_emb, "pos_emb");
@@ -164,7 +164,7 @@ ggml_cgraph * clip_graph_conformer::build() {
             // TODO @ngxson : support this ops in ggml
             {
                 int64_t       d    = x->ne[0] / 2;
-                ggml_tensor * gate = ggml_sigmoid(ctx0, ggml_view_2d(ctx0, x, d, x->ne[1], x->nb[1], d * x->nb[0]));
+                ggml_tensor * gate = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, x, d, x->ne[1], x->nb[1], d * x->nb[0])));
                 x                  = ggml_mul(ctx0, ggml_view_2d(ctx0, x, d, x->ne[1], x->nb[1], 0), gate);
                 x                  = ggml_cont(ctx0, ggml_transpose(ctx0, x));
             }
@@ -207,6 +207,13 @@ ggml_cgraph * clip_graph_conformer::build() {
     cb(cur, "audio_adapter.model.{}", 0);
     cur = build_ffn(cur, model.mm_1_w, model.mm_1_b, nullptr, nullptr, model.mm_3_w, model.mm_3_b, FFN_GELU_ERF, -1);
 
+    // d1omni_a: residual block after the projector
+    if (model.mm_4_w) {
+        ggml_tensor * x = build_norm(cur, model.mm_4_w, model.mm_4_b, NORM_TYPE_NORMAL, 1e-5, -1);
+        x   = build_ffn(x, model.mm_5_w, model.mm_5_b, nullptr, nullptr, model.mm_6_w, model.mm_6_b, FFN_GELU_ERF, -1);
+        cur = ggml_add(ctx0, cur, x);
+    }
+
     cb(cur, "projected", -1);
 
     ggml_build_forward_expand(gf, cur);
diff --git a/tools/mtmd/models/siglip.cpp b/tools/mtmd/models/siglip.cpp
index 96d21ceee..8cc901a28 100644
--- a/tools/mtmd/models/siglip.cpp
+++ b/tools/mtmd/models/siglip.cpp
@@ -4,7 +4,7 @@ ggml_cgraph * clip_graph_siglip::build() {
     ggml_tensor * inp = build_inp();
 
     ggml_tensor * learned_pos_embd = model.position_embeddings;
-    if (proj_type == PROJECTOR_TYPE_LFM2 || proj_type == PROJECTOR_TYPE_PHI4) {
+    if (proj_type == PROJECTOR_TYPE_LFM2 || proj_type == PROJECTOR_TYPE_D1OMNI_V || proj_type == PROJECTOR_TYPE_PHI4) {
         learned_pos_embd = resize_position_embeddings();
     }
 
@@ -55,7 +55,7 @@ ggml_cgraph * clip_graph_siglip::build() {
         cur = build_mm(model.mm_2_w, cur);
         cur = ggml_add(ctx0, cur, model.mm_2_b);
 
-    } else if (proj_type == PROJECTOR_TYPE_LFM2) {
+    } else if (proj_type == PROJECTOR_TYPE_LFM2 || proj_type == PROJECTOR_TYPE_D1OMNI_V) {
         // pixel unshuffle block
         const int scale_factor = model.hparams.n_merge;
         cur = build_patch_merge_permute(cur, scale_factor);
@@ -70,11 +70,12 @@ ggml_cgraph * clip_graph_siglip::build() {
             cur = ggml_add(ctx0, cur, model.mm_input_norm_b);
         }
 
+        // d1-omni uses the exact gelu in the projector
         cur = build_ffn(cur,
             model.mm_1_w, model.mm_1_b,
             nullptr, nullptr,
             model.mm_2_w, model.mm_2_b,
-            FFN_GELU,
+            proj_type == PROJECTOR_TYPE_D1OMNI_V ? FFN_GELU_ERF : FFN_GELU,
             -1);
 
     } else if (proj_type == PROJECTOR_TYPE_JANUS_PRO) {
diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp
index 6c8b19441..1b4ea60bf 100644
--- a/tools/mtmd/mtmd-audio.cpp
+++ b/tools/mtmd/mtmd-audio.cpp
@@ -994,6 +994,42 @@ bool mtmd_audio_preprocessor_conformer::preprocess(const float *
     return true;
 }
 
+//
+// mtmd_audio_preprocessor_d1omni
+//
+
+bool mtmd_audio_preprocessor_d1omni::preprocess(const float *                 samples,
+                                                size_t                        n_samples,
+                                                std::vector & output) const {
+    if (n_samples == 0) {
+        return false;
+    }
+    const size_t n_max = 30 * hparams.audio_sample_rate;
+    const size_t n_min = hparams.audio_sample_rate / 2;
+
+    std::vector buf(samples, samples + std::min(n_samples, n_max));
+    buf.resize(std::max(buf.size(), n_min), 0.0f);
+
+    if (!mtmd_audio_preprocessor_conformer::preprocess(buf.data(), buf.size(), output)) {
+        return false;
+    }
+
+    // the encoder reads one frame per hop, not the extra frame of the centre padding (NeMo: seq_len)
+    const int64_t n_frames = buf.size() / hparams.audio_hop_len;
+    for (auto & mel : output) {
+        if (mel.n_len <= n_frames) {
+            continue;
+        }
+        std::vector data((size_t) mel.n_mel * n_frames);
+        for (int64_t j = 0; j < mel.n_mel; ++j) {
+            std::copy_n(mel.data.begin() + (size_t) j * mel.n_len, n_frames, data.begin() + (size_t) j * n_frames);
+        }
+        mel.n_len = n_frames;
+        mel.data  = std::move(data);
+    }
+    return true;
+}
+
 //
 // mtmd_audio_preprocessor_granite_speech
 //
diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h
index 36db50526..4ca6cbd55 100644
--- a/tools/mtmd/mtmd-audio.h
+++ b/tools/mtmd/mtmd-audio.h
@@ -80,6 +80,12 @@ struct mtmd_audio_preprocessor_conformer : mtmd_audio_preprocessor {
     mtmd_audio_cache cache;
 };
 
+// same as conformer, the audio is cut to 30 s and padded to 0.5 s (d1-omni audio.py)
+struct mtmd_audio_preprocessor_d1omni : mtmd_audio_preprocessor_conformer {
+    using mtmd_audio_preprocessor_conformer::mtmd_audio_preprocessor_conformer;
+    bool preprocess(const float * samples, size_t n_samples, std::vector & output) const override;
+};
+
 struct mtmd_audio_preprocessor_granite_speech : mtmd_audio_preprocessor {
     mtmd_audio_preprocessor_granite_speech(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {}
     void initialize() override;
diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp
index b2319b901..b6e331bc7 100644
--- a/tools/mtmd/mtmd.cpp
+++ b/tools/mtmd/mtmd.cpp
@@ -871,6 +871,11 @@ struct mtmd_context {
                     ov_img_first       = false;
                     image_preproc = std::make_unique(ctx_v);
                 } break;
+            case PROJECTOR_TYPE_D1OMNI_V:
+                {
+                    // same tiles and thumbnail as lfm2, without separator tokens
+                    image_preproc = std::make_unique(ctx_v);
+                } break;
             case PROJECTOR_TYPE_GLM4V:
                 {
                     // <|begin_of_image|> ... (image embeddings) ... <|end_of_image|>
@@ -982,6 +987,10 @@ struct mtmd_context {
                 {
                     audio_preproc = std::make_unique(ctx_a);
                 } break;
+            case PROJECTOR_TYPE_D1OMNI_A:
+                {
+                    audio_preproc = std::make_unique(ctx_a);
+                } break;
             case PROJECTOR_TYPE_GRANITE_SPEECH:
                 {
                     audio_preproc = std::make_unique(ctx_a);
diff --git a/tools/server/README.md b/tools/server/README.md
index 8b271395b..9231f06e8 100644
--- a/tools/server/README.md
+++ b/tools/server/README.md
@@ -1703,9 +1703,11 @@ Follows the [TypeSafe API](https://docs.typesafe.ai/api), streaming is not suppo
 
 *Options:*
 
-`state`: The content to evaluate. Can be a string, an object or an array. A value that is not a string is given to the model as JSON text. For lfm2-d1, it can be `null`, for example to ask about images only.
+`state`: The content to evaluate. Can be a string, an object or an array. A value that is not a string is given to the model as JSON text. For lfm2-d1 and lfm2-d1-omni, it can be `null`, for example to ask about images only.
 
-`images`: Optional. An array of images, the maximum number may be limited depending on the model. Each one is a data URL (`data:image/...;base64,...`). See the image input section below.
+`files`: Optional. An array of input files, the maximum number may be limited depending on the model. Each one is a data URL (`data:image/...;base64,...`). For audio-capable models, it can be audio clips (`data:audio/...;base64,...`). See the image input section below.
+
+`images`: Optional. An alias of `files`.
 
 `questions`: An object that maps a question id to a question. Each question has these fields:
 
@@ -1718,20 +1720,20 @@ Follows the [TypeSafe API](https://docs.typesafe.ai/api), streaming is not suppo
 
 The questions of a request are answered independently, an answer does not depend on the other questions. The exception is clef: it reads all the questions in one prompt and decides them jointly.
 
-The number of options of a `choice` question is limited by the model, for example: 52 for openjev, 255 for laya, clef, pplx-decider and lfm2-d1. For laya, long questions and options are truncated to the token budget the model was trained with.
+The number of options of a `choice` question is limited by the model, for example: 52 for openjev, 255 for laya, clef, pplx-decider, lfm2-d1 and lfm2-d1-omni. For laya, long questions and options are truncated to the token budget the model was trained with.
 
-For laya and clef, the whole prompt is evaluated in one batch: it must fit in `--ubatch-size`. A server that runs clef only serves this endpoint, text generation is not available.
+For laya, clef and lfm2-d1-omni, the whole prompt is evaluated in one batch: it must fit in `--ubatch-size`. An lfm2-d1-omni prompt is cut to 16384 tokens. A server that runs clef only serves this endpoint, text generation is not available.
 
 *Image input:*
 
-Image input needs a model that supports it (for example: openjev, clef, pplx-decider, lfm2-d1) and its multimodal projector, see `--mmproj`.
+Image input needs a model that supports it (for example: openjev, clef, pplx-decider, lfm2-d1, lfm2-d1-omni) and its multimodal projector, see `--mmproj`.
 
 Images can be given in two ways, and both can be used in the same request:
 
-- The `images` field.
-- A `state` made of chat messages, either an array of messages or an object with a `messages` array. An `image_url` part in the `content` of a message is taken as an image, in the same format as chat completions. Only data URLs are accepted.
+- The `files` field, or its alias `images`.
+- A `state` made of chat messages, either an array of messages or an object with a `messages` array. An `image_url` part in the `content` of a message is taken as an image, in the same format as chat completions. Only data URLs are accepted. For lfm2-d1-omni, an `input_audio` part is taken as an audio clip, as base64 data.
 
-All the images are placed before the state in the prompt, the ones from `images` first. The image parts are removed from the state.
+All the images are placed before the state in the prompt, the ones from `files` and `images` first. The image parts are removed from the state.
 
 *Response:*
 
diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp
index f56595231..d364ba898 100644
--- a/tools/server/server-common.cpp
+++ b/tools/server/server-common.cpp
@@ -157,6 +157,7 @@ std::vector server_model_output_modalities(common_decision_type dec
         case COMMON_DECISION_TYPE_CLEF:
         case COMMON_DECISION_TYPE_PPLX_DECIDER:
         case COMMON_DECISION_TYPE_LFM2_D1:
+        case COMMON_DECISION_TYPE_LFM2_D1_OMNI:
             return {"decisions"};
         default:
             // fallback when there is no decision type or the metadata is bad
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index 6f3519215..e3270747b 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -5660,7 +5660,7 @@ void server_routes::init_routes() {
                 scores.push_back(result->scores);
                 n_tokens += result->n_tokens;
             }
-            answers[question.id] = decision.format_answer(question, scores);
+            answers[question.id] = decision.format_answer(question, scores, !files.empty());
         }
 
         res->ok(json{
diff --git a/tools/server/server-decision.cpp b/tools/server/server-decision.cpp
index 71657a359..11ad66a19 100644
--- a/tools/server/server-decision.cpp
+++ b/tools/server/server-decision.cpp
@@ -126,6 +126,12 @@ void server_decision_context::init(const llama_model * model) {
     } else if (model_type == COMMON_DECISION_TYPE_LFM2_D1) {
         n_options_max   = 255;
         noul_true_first = true;
+    } else if (model_type == COMMON_DECISION_TYPE_LFM2_D1_OMNI) {
+        token_marker = llama_vocab_mask(vocab);
+        if (token_marker == LLAMA_TOKEN_NULL) {
+            throw std::runtime_error("decision model has no mask token");
+        }
+        n_options_max = 255;
     } else {
         throw std::runtime_error("unsupported decision model type: " + type_name);
     }
@@ -139,8 +145,8 @@ void server_decision_context::init(const llama_model * model) {
 //
 
 std::vector server_decision_context::parse_questions(const json & body) const {
-    // d1 accepts a null state (images only)
-    if (!body.contains("state") || (body.at("state").is_null() && type != COMMON_DECISION_TYPE_LFM2_D1)) {
+    // lfm2-d1 and lfm2-d1-omni accept a null state, for example to ask about images only
+    if (!body.contains("state") || (body.at("state").is_null() && type != COMMON_DECISION_TYPE_LFM2_D1 && type != COMMON_DECISION_TYPE_LFM2_D1_OMNI)) {
         throw std::invalid_argument("\"state\" must be provided");
     }
     if (!body.contains("questions") || !body.at("questions").is_object() || body.at("questions").empty()) {
@@ -193,7 +199,15 @@ std::vector server_decision_context::parse_questions(c
                 throw err("\"criteria\" must be an object");
             }
             for (const char * key : {"false", "true"}) {
-                question.options.push_back({key, criteria.is_object() && criteria.contains(key) ? criteria.at(key) : json()});
+                json description;
+                if (criteria.is_object() && criteria.contains(key)) {
+                    description = criteria.at(key);
+                } else if (criteria.is_object() && type == COMMON_DECISION_TYPE_LFM2_D1_OMNI) {
+                    // lfm2-d1-omni also reads the descriptions under "no" and "yes"
+                    const char * alias = std::string(key) == "true" ? "yes" : "no";
+                    description = criteria.contains(alias) ? criteria.at(alias) : json();
+                }
+                question.options.push_back({key, description});
             }
             if (noul_true_first) {
                 std::swap(question.options[0], question.options[1]);
@@ -217,9 +231,10 @@ std::vector server_decision_context::parse_questions(c
 
 static const size_t DECISION_MAX_IMAGES = 8;
 
+// any media, mtmd tells an audio clip from an image by its content
 static void decision_load_image(const json & url, std::vector & files) {
-    if (!url.is_string() || !string_starts_with(url.get(), "data:image/")) {
-        throw std::invalid_argument("images must be data URLs (data:image/...;base64,...)");
+    if (!url.is_string() || !string_starts_with(url.get(), "data:")) {
+        throw std::invalid_argument("images must be data URLs (data:image/...;base64,... or data:audio/...;base64,...)");
     }
     if (files.size() >= DECISION_MAX_IMAGES) {
         throw std::invalid_argument(string_format("too many images, the maximum is %zu", DECISION_MAX_IMAGES));
@@ -227,15 +242,30 @@ static void decision_load_image(const json & url, std::vector & file
     handle_media(files, url.get(), "");
 }
 
+// audio is base64 data, as a data URL or not (OpenAI input_audio)
+static void decision_load_audio(const json & data, std::vector & files) {
+    if (!data.is_string() || string_starts_with(data.get(), "http") || string_starts_with(data.get(), "file://")) {
+        throw std::invalid_argument("audio must be base64 data");
+    }
+    if (files.size() >= DECISION_MAX_IMAGES) {
+        throw std::invalid_argument(string_format("too many media files, the maximum is %zu", DECISION_MAX_IMAGES));
+    }
+    handle_media(files, data.get(), "");
+}
+
 json server_decision_context::parse_state(const json & body, std::vector & files) const {
     if (body.contains("videos") && !body.at("videos").is_null() && !body.at("videos").empty()) {
         throw std::invalid_argument("\"videos\" is not supported");
     }
-    if (body.contains("images") && !body.at("images").is_null()) {
-        if (!body.at("images").is_array()) {
-            throw std::invalid_argument("\"images\" must be an array");
+    // "images" is an alias of "files"
+    for (const char * key : {"files", "images"}) {
+        if (!body.contains(key) || body.at(key).is_null()) {
+            continue;
         }
-        for (const auto & url : body.at("images")) {
+        if (!body.at(key).is_array()) {
+            throw std::invalid_argument(string_format("\"%s\" must be an array", key));
+        }
+        for (const auto & url : body.at(key)) {
             decision_load_image(url, files);
         }
     }
@@ -247,7 +277,7 @@ json server_decision_context::parse_state(const json & body, std::vector");
 }
 
+// lfm2-d1-omni: special tokens written in the input must not be parsed as such, in keys too (d1-omni prompt.py: escape)
+static json decision_d1omni_escape(const json & val) {
+    static const std::regex re_special("<\\|([A-Za-z0-9_]+)\\|>");
+    if (val.is_string()) {
+        return std::regex_replace(val.get(), re_special, "<\xC2\xA6$1\xC2\xA6>");
+    }
+    if (val.is_array()) {
+        json out = json::array();
+        for (const auto & item : val) {
+            out.push_back(decision_d1omni_escape(item));
+        }
+        return out;
+    }
+    if (val.is_object()) {
+        json out = json::object();
+        for (const auto & [key, item] : val.items()) {
+            out[std::regex_replace(key, re_special, "<\xC2\xA6$1\xC2\xA6>")] = decision_d1omni_escape(item);
+        }
+        return out;
+    }
+    return val;
+}
+
+// given to the lfm2-d1-omni template: text between the pieces of the prompt, and at the start of the pieces that are cut to a token budget
+static const std::string D1OMNI_MARKER        = "< 1) {
@@ -515,7 +583,8 @@ std::string server_decision_context::render(
         const std::vector & questions,
         const server_decision_question & question,
         size_t variant,
-        size_t n_images) const {
+        size_t n_images,
+        bool is_audio) const {
     // the template is given raw JSON values, it serializes the ones that are not strings
     json inp = json{
         {"id",           question.id},
@@ -554,6 +623,15 @@ std::string server_decision_context::render(
         inp = decision_replace_text(inp, text_marker, " ");
     }
 
+    if (type == COMMON_DECISION_TYPE_LFM2_D1_OMNI) {
+        inp = decision_replace_text(decision_d1omni_escape(inp), D1OMNI_MARKER, "< 0) {
@@ -580,6 +658,11 @@ void server_decision_context::fill_task(
         mtmd_context * mctx,
         const mtmd_helper_init_opt & init_opt,
         server_task & task) const {
+    if (type == COMMON_DECISION_TYPE_LFM2_D1_OMNI) {
+        fill_task_d1omni(state, questions, question, variant, files, mctx, init_opt, task);
+        return;
+    }
+
     const std::string prompt = render(state, questions, question, variant, files.size());
 
     if (type == COMMON_DECISION_TYPE_OPENJEV || type == COMMON_DECISION_TYPE_LEV || type == COMMON_DECISION_TYPE_NIMBLE || type == COMMON_DECISION_TYPE_PPLX_DECIDER) {
@@ -678,6 +761,129 @@ void server_decision_context::fill_task_laya(llama_tokens & tokens, const server
     task.decision.column = question.type;
 }
 
+// each piece is cut to its budget as the model was trained (d1-omni prompt.py: encode), the state gets the room that is left
+// the media come first, the text after them has a budget of its own
+void server_decision_context::fill_task_d1omni(
+        const json & state,
+        const std::vector & questions,
+        const server_decision_question & question,
+        size_t variant,
+        const std::vector & files,
+        mtmd_context * mctx,
+        const mtmd_helper_init_opt & init_opt,
+        server_task & task) const {
+    const auto invalid = std::runtime_error("unexpected layout of the decision prompt");
+
+    // the prompt depends on the kind of media, mtmd tells an audio clip from an image by its content
+    std::string markers;
+    server_tokens media(llama_tokens(), false);
+    bool is_audio = false;
+    if (!files.empty()) {
+        for (size_t i = 0; i < files.size(); i++) {
+            markers += get_media_marker();
+        }
+        media = process_mtmd_prompt(mctx, markers, files, init_opt);
+        for (size_t i = 0; i < media.size(); i++) {
+            if (media[i] == LLAMA_TOKEN_NULL) {
+                const auto & chunk = media.find_chunk(i);
+                is_audio = is_audio || mtmd_input_chunk_get_type(chunk.get()) == MTMD_INPUT_CHUNK_TYPE_AUDIO;
+                i += mtmd_input_chunk_get_n_tokens(chunk.get()) - 1;
+            }
+        }
+    }
+    if (is_audio && files.size() > 1) {
+        throw std::invalid_argument("a request has images or one audio clip, not both");
+    }
+    const size_t n_media = media.size();
+
+    const std::string prompt = render(state, questions, question, variant, files.size(), is_audio);
+    std::vector pieces = string_split(prompt, D1OMNI_SEP);
+    if (pieces.empty() || pieces[0] != markers) {
+        throw invalid;
+    }
+
+    size_t n_max = D1OMNI_MAX_TOKENS;
+    if (!files.empty()) {
+        n_max = std::min(is_audio ? D1OMNI_MAX_TOKENS_AUDIO : D1OMNI_MAX_TOKENS_IMAGE, D1OMNI_MAX_TOKENS - std::min(D1OMNI_MAX_TOKENS, n_media));
+        if (n_max < 64) {
+            throw std::invalid_argument(string_format("the media take %zu of the %zu positions, send fewer images", n_media, D1OMNI_MAX_TOKENS));
+        }
+    }
+
+    // the options get max(96, min(24 n + 32, max / 2)) tokens, shared evenly
+    const int64_t n_options      = question.options.size();
+    const int64_t n_budget       = std::max(96, std::min(24 * n_options + 32, n_max / 2));
+    const int64_t n_option_max   = std::max(2, (n_budget - 3 * n_options) / n_options);
+    const int64_t n_question_max = std::max(16, n_budget);
+
+    llama_tokens head;  // before the state
+    llama_tokens body;  // the state
+    llama_tokens tail;  // after the state
+    bool has_state = false;
+    for (size_t i_piece = 1; i_piece < pieces.size(); i_piece++) {
+        std::string piece = pieces[i_piece];
+        int64_t n_piece_max = -1;
+        bool is_state = false;
+        if (string_starts_with(piece, D1OMNI_MARK_STATE)) {
+            piece    = piece.substr(D1OMNI_MARK_STATE.size());
+            is_state = true;
+        } else if (string_starts_with(piece, D1OMNI_MARK_QUESTION)) {
+            piece = piece.substr(D1OMNI_MARK_QUESTION.size());
+            n_piece_max = n_question_max;
+        } else if (string_starts_with(piece, D1OMNI_MARK_OPTION)) {
+            piece = piece.substr(D1OMNI_MARK_OPTION.size());
+            n_piece_max = n_option_max;
+        }
+
+        llama_tokens tokens = common_tokenize(vocab, piece, false, true);
+        if (n_piece_max >= 0 && (int64_t) tokens.size() > n_piece_max) {
+            tokens.resize(n_piece_max);
+        }
+
+        if (is_state) {
+            if (has_state) {
+                throw invalid;
+            }
+            body      = std::move(tokens);
+            has_state = true;
+        } else {
+            llama_tokens & dst = has_state ? tail : head;
+            dst.insert(dst.end(), tokens.begin(), tokens.end());
+        }
+    }
+    if (!has_state) {
+        throw invalid;
+    }
+
+    const size_t n_room = n_max - std::min(n_max, head.size() + tail.size());
+    body.resize(std::min(body.size(), n_room));
+
+    llama_tokens tokens = std::move(head);
+    tokens.insert(tokens.end(), body.begin(), body.end());
+    tokens.insert(tokens.end(), tail.begin(), tail.end());
+    tokens.resize(std::min(tokens.size(), n_max));
+
+    for (size_t i = 0; i < tokens.size(); i++) {
+        if (tokens[i] == token_marker) {
+            task.decision.markers.push_back(n_media + i);
+        }
+    }
+    if ((int64_t) task.decision.markers.size() != n_options) {
+        throw std::invalid_argument("the options do not fit in the context");
+    }
+
+    // the output has one score per question type
+    task.decision.column = question.type;
+    if (files.empty()) {
+        task.tokens = server_tokens(tokens, false);
+    } else {
+        task.tokens = std::move(media);
+        for (const llama_token token : tokens) {
+            task.tokens.push_back(token);
+        }
+    }
+}
+
 //
 // joint prompt (clef)
 //
@@ -850,14 +1056,15 @@ static double decision_confidence_score(const std::vector & probs) {
     return std::max(0.0, 1.0 - dist / dist_uniform);
 }
 
-json server_decision_context::format_answer(const server_decision_question & question, const std::vector> & scores) const {
+json server_decision_context::format_answer(const server_decision_question & question, const std::vector> & scores, bool has_media) const {
     const size_t n = n_outputs(question);
     if (scores.size() != n_variants(question)) {
         throw std::runtime_error("decision result does not match the number of variants");
     }
 
     // softmax over the outputs of each variant, then the average of the variants
-    const float temperature = get_temperature(question);
+    // lfm2-d1-omni: image and audio answers are not calibrated
+    const float temperature = has_media && type == COMMON_DECISION_TYPE_LFM2_D1_OMNI ? 1.0f : get_temperature(question);
     std::vector probs(n, 0.0);
     for (size_t v = 0; v < scores.size(); v++) {
         const auto & s = scores[v];
diff --git a/tools/server/server-decision.h b/tools/server/server-decision.h
index 8f357f374..5c7e44f37 100644
--- a/tools/server/server-decision.h
+++ b/tools/server/server-decision.h
@@ -62,6 +62,7 @@ struct server_decision_context {
             case COMMON_DECISION_TYPE_CLEF:
             case COMMON_DECISION_TYPE_PPLX_DECIDER:
             case COMMON_DECISION_TYPE_LFM2_D1:
+            case COMMON_DECISION_TYPE_LFM2_D1_OMNI:
                 return true;
             default:
                 return false;
@@ -72,7 +73,8 @@ struct server_decision_context {
     std::vector parse_questions(const json & body) const;
 
     // returns the state without its images, they are appended to files in order
-    // images come from "images" and from the image_url parts of a state made of chat messages
+    // images come from "files" (alias "images") and from the image_url and input_audio parts of a state made of chat messages
+    // an image can be an audio clip if the model supports it
     json parse_state(const json & body, std::vector & files) const;
 
     // number of prompts that are evaluated to answer this question, each one shows the options in a different order
@@ -101,7 +103,7 @@ struct server_decision_context {
             server_task & task) const;
 
     // scores: the raw model outputs of each variant
-    json format_answer(const server_decision_question & question, const std::vector> & scores) const;
+    json format_answer(const server_decision_question & question, const std::vector> & scores, bool has_media = false) const;
 
 private:
     const llama_vocab * vocab = nullptr;
@@ -128,12 +130,22 @@ private:
             const std::vector & questions,
             const server_decision_question & question,
             size_t variant,
-            size_t n_images) const;
+            size_t n_images,
+            bool is_audio = false) const;
     json render_options(const server_decision_question & question, size_t variant) const;
     size_t n_outputs(const server_decision_question & question) const;
     // LFM2_D1: label text and tokens of each option
     void d1_labels(const server_decision_question & question, std::vector & texts, std::vector & groups) const;
     void fill_task_laya(llama_tokens & tokens, const server_decision_question & question, server_task & task) const;
+    void fill_task_d1omni(
+            const json & state,
+            const std::vector & questions,
+            const server_decision_question & question,
+            size_t variant,
+            const std::vector & files,
+            mtmd_context * mctx,
+            const mtmd_helper_init_opt & init_opt,
+            server_task & task) const;
 
     float get_temperature(const server_decision_question & question) const;
 };

From 06cad0b9e77315bd2930bd4f70a6d7b37b2a01c1 Mon Sep 17 00:00:00 2001
From: kurquhar 
Date: Wed, 7 Oct 2026 17:13:21 -0700
Subject: [PATCH 23/40] hexagon: Q6_K weight dequant speedup (#30121)

* hexagon: Q6_K weight dequant speedup

Assisted-by: OpenCode

* unroll by another factor of 2

Assisted-by: OpenCode
---
 .../ggml-hexagon/htp/hmx-mm-kernels-tiled.h   | 39 +++++++++++++++----
 1 file changed, 31 insertions(+), 8 deletions(-)

diff --git a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h
index 698d6a33b..40404bf10 100644
--- a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h
+++ b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h
@@ -631,17 +631,40 @@ static void dequantize_tiled_weight_to_fp16_task_q6_k(
         HVX_Vector v_scale_k16 = Q6_V_lo_W(Q6_W_vshuff_VVR(v_sc_k16, v_sc_k16, -2));
 
         #pragma unroll
-        for (int g = 0; g < 8; g++) {
+        for (int g = 0; g < 8; g += 4) {
             const HVX_Vector v_scale = (g < 4) ? v_scale_k0 : v_scale_k16;
 
-            HVX_Vector     v_q   = unpack_q6_k_group(vptr, g, mask_0f, mask_03, i32);
-            HVX_VectorPair vp16  = Q6_Wh_vunpack_Vb(v_q);
-            HVX_VectorPair vp_k  = Q6_W_vdeal_VVR(Q6_V_hi_W(vp16), Q6_V_lo_W(vp16), -4);
+            HVX_Vector v_q0 = unpack_q6_k_group(vptr, g + 0, mask_0f, mask_03, i32);
+            HVX_Vector v_q1 = unpack_q6_k_group(vptr, g + 1, mask_0f, mask_03, i32);
+            HVX_Vector v_q2 = unpack_q6_k_group(vptr, g + 2, mask_0f, mask_03, i32);
+            HVX_Vector v_q3 = unpack_q6_k_group(vptr, g + 3, mask_0f, mask_03, i32);
 
-            hvx_vmem(dst_ptr + (2 * g + 0) * 64) =
-                Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_k)), v_scale));
-            hvx_vmem(dst_ptr + (2 * g + 1) * 64) =
-                Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_k)), v_scale));
+            HVX_VectorPair vp16_0 = Q6_Wh_vunpack_Vb(v_q0);
+            HVX_VectorPair vp16_1 = Q6_Wh_vunpack_Vb(v_q1);
+            HVX_VectorPair vp16_2 = Q6_Wh_vunpack_Vb(v_q2);
+            HVX_VectorPair vp16_3 = Q6_Wh_vunpack_Vb(v_q3);
+            HVX_VectorPair vp_k0  = Q6_W_vdeal_VVR(Q6_V_hi_W(vp16_0), Q6_V_lo_W(vp16_0), -4);
+            HVX_VectorPair vp_k1  = Q6_W_vdeal_VVR(Q6_V_hi_W(vp16_1), Q6_V_lo_W(vp16_1), -4);
+            HVX_VectorPair vp_k2  = Q6_W_vdeal_VVR(Q6_V_hi_W(vp16_2), Q6_V_lo_W(vp16_2), -4);
+            HVX_VectorPair vp_k3  = Q6_W_vdeal_VVR(Q6_V_hi_W(vp16_3), Q6_V_lo_W(vp16_3), -4);
+
+            HVX_Vector v_out00 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_k0)), v_scale));
+            HVX_Vector v_out01 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_k0)), v_scale));
+            HVX_Vector v_out10 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_k1)), v_scale));
+            HVX_Vector v_out11 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_k1)), v_scale));
+            HVX_Vector v_out20 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_k2)), v_scale));
+            HVX_Vector v_out21 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_k2)), v_scale));
+            HVX_Vector v_out30 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_k3)), v_scale));
+            HVX_Vector v_out31 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_k3)), v_scale));
+
+            hvx_vmem(dst_ptr + (2 * g + 0) * 64) = v_out00;
+            hvx_vmem(dst_ptr + (2 * g + 1) * 64) = v_out01;
+            hvx_vmem(dst_ptr + (2 * g + 2) * 64) = v_out10;
+            hvx_vmem(dst_ptr + (2 * g + 3) * 64) = v_out11;
+            hvx_vmem(dst_ptr + (2 * g + 4) * 64) = v_out20;
+            hvx_vmem(dst_ptr + (2 * g + 5) * 64) = v_out21;
+            hvx_vmem(dst_ptr + (2 * g + 6) * 64) = v_out30;
+            hvx_vmem(dst_ptr + (2 * g + 7) * 64) = v_out31;
         }
     }
 }

From 9c2e0e491a822adae1f0b1c831adb4160057d24f Mon Sep 17 00:00:00 2001
From: Max Krasnyansky 
Date: Wed, 7 Oct 2026 18:30:06 -0700
Subject: [PATCH 24/40] hexagon: enable alloc_buffer_n  (#30126)

* hex-bufs: add support for alloc_buffer_n

* hex-bufs: add support for splitting large tensors into separate buffers

* hex-bufs: update GGML_HEXAGON_MBUF to accept three values dyn,static,total

* hex-bufs: bump dyn. default to 512MB since 128MB causes perf regressions with big MOEs

* hex-run: add --no-embd-offload option to simplify command lines on devices that need it

* Update scripts/snapdragon/run.py

Co-authored-by: Jhen-Jie Hong 

---------

Co-authored-by: Jhen-Jie Hong 
---
 ggml/src/ggml-hexagon/ggml-hexagon.cpp | 251 +++++++++++++++++++++++--
 scripts/snapdragon/run.py              |   6 +-
 2 files changed, 239 insertions(+), 18 deletions(-)

diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
index 26b79582e..f3ca73124 100644
--- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp
+++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
@@ -53,6 +53,7 @@
 
 #define GGML_COMMON_IMPL_CPP
 #include "ggml-backend-impl.h"
+#include "ggml-alloc.h"
 #include "ggml-common.h"
 #include "ggml-hexagon.h"
 #include "ggml-impl.h"
@@ -101,13 +102,16 @@ static size_t opt_ndev    = 1;
 static size_t opt_nhvx    = 0; // use all
 static int    opt_nhmx    = 1; // when set, enable HMX; when 0, use HVX only
 static size_t opt_vmem    = HTP_OP_MAX_VMEM_DEFAULT;  // max available va space for buffer mappings
-static size_t opt_mbuf    = 1ul * 1024 * 1024 * 1024; // max buffer size
 static int    opt_etm     = 0;
 static int    opt_verbose = 0;
 static int    opt_profile = 0; // profiling mode (0-disabled, 1-basic, 2-pmu)
 static bool   opt_hostbuf = false;
 static bool   opt_dma64   = false;
 
+static size_t opt_mbuf_dyn    = 512ul * 1024 * 1024;      // max dynamic (compute) buffer size
+static size_t opt_mbuf_static = 1ul * 1024 * 1024 * 1024; // max static (weight/KV) buffer size
+static size_t opt_mbuf_total  = 0;                        // total buffer space limit (0 = unconstrained)
+
 static int    opt_mm_select  = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
 static int    opt_fa_select  = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
 static int    opt_fa_head_split = 1; // 1 = partition flash_attn by KV heads in multicore (default on), 0 = token-based (original)
@@ -121,7 +125,7 @@ static int    opt_ar_scatter = 1; // 1 = reduce-scatter the fused ALLREDUCE+ADD
 static u32vec opt_pmu_evt { 0x3, 0x111, 0x100, 0x105, 0x240, 0x256, 0x7D, 0x8C };
 
 static int opt_opbatch  = 1280; // max number of ops in a batch
-static int opt_opqueue  = 32;   // max number of pending batches
+static int opt_opqueue  = 8;   // max number of pending batches
 static int opt_optrace  = 0;    // trace buffer size per thread (0 means default)
 static int opt_oppoll   = 0;    // polling for batch completions
 static int opt_opfusion = 1;    // enable/disable op fusion
@@ -2820,8 +2824,30 @@ static size_t ggml_backend_hexagon_buffer_type_get_alloc_size(ggml_backend_buffe
     GGML_UNUSED(buft);
 }
 
+static size_t parse_size(const char * str, size_t default_unit = 1024 * 1024) {
+    if (!str || str[0] == '\0') {
+        return 0;
+    }
+    char * end = NULL;
+    double val = strtod(str, &end);
+    if (val < 0) {
+        return 0;
+    }
+    if (end && *end) {
+        while (*end == ' ') end++;
+        if (*end == 'k' || *end == 'K') {
+            return (size_t) (val * 1024);
+        } else if (*end == 'm' || *end == 'M') {
+            return (size_t) (val * 1024 * 1024);
+        } else if (*end == 'g' || *end == 'G') {
+            return (size_t) (val * 1024 * 1024 * 1024);
+        }
+    }
+    return (size_t) (val * default_unit);
+}
+
 static size_t ggml_backend_hexagon_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) {
-    return opt_mbuf;
+    return opt_mbuf_dyn;
     GGML_UNUSED(buft);
 }
 
@@ -2835,25 +2861,199 @@ static bool ggml_backend_hexagon_host_buffer_type_is_host(ggml_backend_buffer_ty
     GGML_UNUSED(buft);
 }
 
+struct ggml_backend_hexagon_alloc_buffer_n_plan_item {
+    size_t size;
+    int    first;
+    int    last;
+};
+
+using ggml_backend_hexagon_alloc_buffer_n_plan_t = std::vector;
+
+static const char * ggml_hexagon_kv_layer_suffix(const struct ggml_tensor * t) {
+    if (strncmp(t->name, "cache_", 6) != 0) {
+        return NULL;
+    }
+    const char * p = strstr(t->name, "_l");
+    if (!p || !isdigit((unsigned char)p[2])) {
+        return NULL;
+    }
+    return p;
+}
+
+struct ggml_backend_hexagon_alloc_unit {
+    size_t size;
+    int    first;
+    int    last;
+};
+
+static ggml_backend_hexagon_alloc_buffer_n_plan_t ggml_backend_hexagon_alloc_buffer_n_plan(
+        ggml_backend_buffer_type_t buft, struct ggml_tensor ** tensors, int n_tensors) {
+    ggml_backend_hexagon_alloc_buffer_n_plan_t plan;
+
+    const size_t alignment = ggml_backend_buft_get_alignment(buft);
+    const size_t max_size  = opt_mbuf_static > 0 ? opt_mbuf_static : SIZE_MAX;
+
+    std::vector units;
+
+    int i = 0;
+    while (i < n_tensors) {
+        struct ggml_tensor * t = tensors[i];
+        size_t unit_size = 0;
+        int unit_first = i;
+        int unit_last = i + 1;
+
+        if (t->data == NULL && t->view_src == NULL) {
+            unit_size += GGML_PAD(ggml_backend_buft_get_alloc_size(buft, t), alignment);
+        }
+
+        const char * layer_suffix = ggml_hexagon_kv_layer_suffix(t);
+
+        while (unit_last < n_tensors) {
+            struct ggml_tensor * next = tensors[unit_last];
+
+            if (next->view_src != NULL) {
+                unit_last++;
+                continue;
+            }
+
+            if (layer_suffix != NULL) {
+                const char * next_suffix = ggml_hexagon_kv_layer_suffix(next);
+                if (next_suffix != NULL && strcmp(layer_suffix, next_suffix) == 0) {
+                    if (next->data == NULL) {
+                        unit_size += GGML_PAD(ggml_backend_buft_get_alloc_size(buft, next), alignment);
+                    }
+                    unit_last++;
+                    continue;
+                }
+            }
+
+            break;
+        }
+
+        units.push_back({ unit_size, unit_first, unit_last });
+        i = unit_last;
+    }
+
+    size_t cur_buf_size  = 0;
+    int    cur_buf_first = 0;
+
+    for (const auto & unit : units) {
+        if (unit.size == 0) {
+            continue;
+        }
+
+        if (cur_buf_size > 0 && (cur_buf_size + unit.size) > max_size) {
+            plan.push_back({ cur_buf_size, cur_buf_first, unit.first });
+            cur_buf_size  = 0;
+            cur_buf_first = unit.first;
+        }
+
+        cur_buf_size += unit.size;
+    }
+
+    if (cur_buf_size > 0) {
+        plan.push_back({ cur_buf_size, cur_buf_first, n_tensors });
+    }
+
+    return plan;
+}
+
+static ggml_backend_buffer_t ggml_backend_hexagon_buffer_type_alloc_buffer_n(
+        ggml_backend_buffer_type_t buft, struct ggml_tensor ** tensors, int n_tensors) {
+    const ggml_backend_hexagon_alloc_buffer_n_plan_t plan = ggml_backend_hexagon_alloc_buffer_n_plan(buft, tensors, n_tensors);
+
+    std::vector buffers;
+    buffers.reserve(plan.size());
+
+    for (const auto & item : plan) {
+        ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, item.size);
+        if (buffer == NULL) {
+            GGML_LOG_ERROR("%s: failed to allocate %s buffer of size %zu\n", __func__, ggml_backend_buft_name(buft), item.size);
+            for (ggml_backend_buffer_t b : buffers) {
+                ggml_backend_buffer_free(b);
+            }
+            return NULL;
+        }
+
+        struct ggml_tallocr tallocr = ggml_tallocr_new(buffer);
+
+        struct ggml_tensor * t_failed = NULL;
+        for (int j = item.first; j < item.last; j++) {
+            struct ggml_tensor * t = tensors[j];
+            if (t->data == NULL) {
+                if (t->view_src == NULL) {
+                    if (ggml_tallocr_alloc(&tallocr, t) != GGML_STATUS_SUCCESS) {
+                        t_failed = t;
+                        break;
+                    }
+                } else if (t->buffer == NULL) {
+                    if (ggml_backend_view_init(t) != GGML_STATUS_SUCCESS) {
+                        t_failed = t;
+                        break;
+                    }
+                }
+            } else {
+                if (t->view_src != NULL && t->buffer == NULL) {
+                    if (ggml_backend_view_init(t) != GGML_STATUS_SUCCESS) {
+                        t_failed = t;
+                        break;
+                    }
+                }
+            }
+        }
+        if (t_failed != NULL) {
+            GGML_LOG_ERROR("%s: failed to initialize tensor %s\n", __func__, t_failed->name);
+            for (ggml_backend_buffer_t b : buffers) {
+                ggml_backend_buffer_free(b);
+            }
+            ggml_backend_buffer_free(buffer);
+            return NULL;
+        }
+
+        buffers.push_back(buffer);
+    }
+
+    if (buffers.empty()) {
+        return NULL;
+    }
+
+    if (buffers.size() == 1) {
+        return buffers[0];
+    }
+
+    return ggml_backend_multi_buffer_alloc_buffer(buffers.data(), buffers.size());
+}
+
+static size_t ggml_backend_hexagon_buffer_type_get_alloc_size_n(
+        ggml_backend_buffer_type_t buft, struct ggml_tensor ** tensors, int n_tensors) {
+    const ggml_backend_hexagon_alloc_buffer_n_plan_t plan = ggml_backend_hexagon_alloc_buffer_n_plan(buft, tensors, n_tensors);
+
+    size_t total = 0;
+    for (const auto & item : plan) {
+        total += item.size;
+    }
+    return total;
+}
+
 static ggml_backend_buffer_type_i ggml_backend_hexagon_buffer_type_interface = {
     /* .get_name            = */ ggml_backend_hexagon_buffer_type_name,
     /* .alloc_buffer        = */ ggml_backend_hexagon_buffer_type_alloc_buffer,
-    /* .alloc_buffer_n      = */ NULL,
+    /* .alloc_buffer_n      = */ ggml_backend_hexagon_buffer_type_alloc_buffer_n,
     /* .get_alignment       = */ ggml_backend_hexagon_buffer_type_get_alignment,
     /* .get_max_size        = */ ggml_backend_hexagon_buffer_type_get_max_size,
     /* .get_alloc_size      = */ ggml_backend_hexagon_buffer_type_get_alloc_size,
-    /* .get_alloc_size_n    = */ NULL,
+    /* .get_alloc_size_n    = */ ggml_backend_hexagon_buffer_type_get_alloc_size_n,
     /* .is_host             = */ ggml_backend_hexagon_buffer_type_is_host,
 };
 
 static ggml_backend_buffer_type_i ggml_backend_hexagon_host_buffer_type_interface = {
     /* .get_name            = */ ggml_backend_hexagon_buffer_type_name,
     /* .alloc_buffer        = */ ggml_backend_hexagon_host_buffer_type_alloc_buffer,
-    /* .alloc_buffer_n      = */ NULL,
+    /* .alloc_buffer_n      = */ ggml_backend_hexagon_buffer_type_alloc_buffer_n,
     /* .get_alignment       = */ ggml_backend_hexagon_buffer_type_get_alignment,
     /* .get_max_size        = */ ggml_backend_hexagon_buffer_type_get_max_size,
     /* .get_alloc_size      = */ ggml_backend_hexagon_buffer_type_get_alloc_size,
-    /* .get_alloc_size_n    = */ NULL,
+    /* .get_alloc_size_n    = */ ggml_backend_hexagon_buffer_type_get_alloc_size_n,
     /* .is_host             = */ ggml_backend_hexagon_host_buffer_type_is_host,
 };
 
@@ -8545,8 +8745,8 @@ static const char * ggml_backend_hexagon_device_get_description(ggml_backend_dev
 }
 
 static void ggml_backend_hexagon_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) {
-    *free  = 0;
-    *total = *free;
+    *free  = opt_mbuf_total;
+    *total = opt_mbuf_total;
 
     GGML_UNUSED(dev);
 }
@@ -9295,7 +9495,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
     size_t MiB = 1024 * 1024;
 
     // Update vmem default
-    opt_vmem = opt_arch >= 75 ? HTP_OP_MAX_VMEM_DEFAULT : 3000 * MiB;
+    opt_vmem  = opt_arch >= 75 ? HTP_OP_MAX_VMEM_DEFAULT : 3000 * MiB;
     opt_dma64 = opt_arch > 79 && (!str_dma64 || atoi(str_dma64) != 0);
 
     auto RE_ICASE = std::regex_constants::icase;
@@ -9313,13 +9513,30 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
     opt_nhmx      = str_nhmx     ? atoi(str_nhmx)                         : opt_nhmx;
     opt_mm_select = str_mm_select ? atoi(str_mm_select)                   : opt_mm_select;
     opt_fa_select = str_fa_select ? atoi(str_fa_select)                   : opt_fa_select;
-    opt_fa_head_split = str_fa_head_split ? atoi(str_fa_head_split)        : opt_fa_head_split;
-    opt_gdn_select = str_gdn_select ? atoi(str_gdn_select)                 : opt_gdn_select;
-    opt_ar_select = str_ar_select ? atoi(str_ar_select)                   : opt_ar_select;
-    opt_ar_scatter = str_ar_scatter ? atoi(str_ar_scatter)                : opt_ar_scatter;
-    opt_mbuf      = str_mbuf     ? strtoul(str_mbuf, NULL, 0) * MiB       : opt_mbuf;
-    opt_vmem      = str_vmem     ? strtoul(str_vmem, NULL, 0) * MiB       : opt_vmem;
-    opt_hostbuf   = str_hostbuf  ? atoi(str_hostbuf) != 0                 : opt_hostbuf;
+    opt_fa_head_split = str_fa_head_split ? atoi(str_fa_head_split)       : opt_fa_head_split;
+    opt_gdn_select    = str_gdn_select    ? atoi(str_gdn_select)          : opt_gdn_select;
+    opt_ar_select     = str_ar_select     ? atoi(str_ar_select)           : opt_ar_select;
+    opt_ar_scatter    = str_ar_scatter    ? atoi(str_ar_scatter)          : opt_ar_scatter;
+
+    if (str_mbuf) {
+        const char * p = str_mbuf;
+        for (int idx = 0; idx < 3 && p && *p; idx++) {
+            while (*p == ' ') p++;
+            const char * comma = strchr(p, ',');
+            size_t len = comma ? (size_t)(comma - p) : strlen(p);
+            while (len > 0 && p[len - 1] == ' ') len--;
+            if (len > 0) {
+                std::string token(p, len);
+                if (idx == 0) opt_mbuf_dyn    = parse_size(token.c_str());
+                if (idx == 1) opt_mbuf_static = parse_size(token.c_str());
+                if (idx == 2) opt_mbuf_total  = parse_size(token.c_str());
+            }
+            if (!comma) break;
+            p = comma + 1;
+        }
+    }
+    opt_vmem    = str_vmem    ? parse_size(str_vmem)   : opt_vmem;
+    opt_hostbuf = str_hostbuf ? atoi(str_hostbuf) != 0 : opt_hostbuf;
 
     // Parse device configuration
     const char * str_devices  = getenv("GGML_HEXAGON_DEVICES");
diff --git a/scripts/snapdragon/run.py b/scripts/snapdragon/run.py
index 14e53aee6..382879f7b 100755
--- a/scripts/snapdragon/run.py
+++ b/scripts/snapdragon/run.py
@@ -152,6 +152,7 @@ def main():
     parser.add_argument("--profile", help="Profiling flag (enables Hexagon profiling and OpenCL autotuning)")
     parser.add_argument("--sched-debug", action="store_true", help="Enable GGML/llama.cpp scheduler debug output (GGML_SCHED_DEBUG=2)")
     parser.add_argument("--mtmd-device", help="Specify the backend device ID for Multi-Threaded Multi-Device setup (MTMD_BACKEND_DEVICE)")
+    parser.add_argument("--no-embd-offload", action="store_true", help="Keep token embeddings and output projection on CPU (-ot token_embd.weight=CPU,output.weight=CPU)")
 
     # Hexagon specific parameters
     parser.add_argument("--hex-verbose", help="Enable verbose logging (GGML_HEXAGON_VERBOSE)")
@@ -166,7 +167,7 @@ def main():
     parser.add_argument("--hex-opfilter", help="Regex pattern to filter/select which operators are offloaded to NPU (GGML_HEXAGON_OPFILTER)")
     parser.add_argument("--hex-opfusion", help="NPU graph node fusion optimization level (0: disabled, 1: enabled) (GGML_HEXAGON_OPFUSION)")
     parser.add_argument("--hex-vmem", help="Maximum NPU VMEM size limit in MB to allocate (GGML_HEXAGON_VMEM)")
-    parser.add_argument("--hex-mbuf", help="Maximum host buffer size limit in MB to allocate (GGML_HEXAGON_MBUF)")
+    parser.add_argument("--hex-mbuf", help="Host buffer size limits in MB (supports K/M/G suffix): [,[,]] (default: 512,1024,0) (GGML_HEXAGON_MBUF)")
     parser.add_argument("--hex-mm-select", help="Select MUL_MAT and MUL_MAT_ID kernel (GGML_HEXAGON_MM_SELECT) 2:HMX,1:HVX,0:disable")
     parser.add_argument("--hex-fa-select", help="Select Flash Attention kernel (GGML_HEXAGON_FA_SELECT) 2:HMX,1:HVX,0:disable")
     parser.add_argument("--hex-fa-head-split", help="Enable (1) or disable (0) head-parallel flash_attn partitioning (GGML_HEXAGON_FA_HEAD_SPLIT)")
@@ -422,6 +423,9 @@ def main():
     if basename in ("llama-cli", "llama-completion", "llama-server", "llama-bench"):
         if "-t" not in cmd_args and "--threads" not in cmd_args:
             cmd_args += ["-t", "6"]
+        if getattr(args, "no_embd_offload", False):
+            if not any("token_embd" in arg or "output.weight" in arg for arg in cmd_args):
+                cmd_args += ["-ot", r"^(token_embd|output)\.weight$=CPU"]
 
     # Resolve target directory on device
     target_dir = args.target_dir

From 9b4ed0ca572be213f4440a490f341f810f7aa2cf Mon Sep 17 00:00:00 2001
From: Titaniumtown 
Date: Wed, 7 Oct 2026 22:39:25 -0700
Subject: [PATCH 25/40] sycl: remove duplicate block-size defines from op
 headers (#29507)

All were duplicates of the values in ggml/src/ggml-sycl/presets.hpp
---
 ggml/src/ggml-sycl/im2col.hpp  | 2 --
 ggml/src/ggml-sycl/pad.hpp     | 2 --
 ggml/src/ggml-sycl/rope.hpp    | 2 --
 ggml/src/ggml-sycl/upscale.hpp | 2 --
 4 files changed, 8 deletions(-)

diff --git a/ggml/src/ggml-sycl/im2col.hpp b/ggml/src/ggml-sycl/im2col.hpp
index 976d10946..c9eba49ef 100644
--- a/ggml/src/ggml-sycl/im2col.hpp
+++ b/ggml/src/ggml-sycl/im2col.hpp
@@ -15,8 +15,6 @@
 
 #include "common.hpp"
 
-#define SYCL_IM2COL_BLOCK_SIZE 256
-
 void ggml_sycl_op_im2col(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
 void ggml_sycl_op_im2col_3d(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
 
diff --git a/ggml/src/ggml-sycl/pad.hpp b/ggml/src/ggml-sycl/pad.hpp
index b099e9b73..4bc5a2cfd 100644
--- a/ggml/src/ggml-sycl/pad.hpp
+++ b/ggml/src/ggml-sycl/pad.hpp
@@ -15,8 +15,6 @@
 
 #include "common.hpp"
 
-#define SYCL_PAD_BLOCK_SIZE 256
-
 void ggml_sycl_pad(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
 
 void ggml_sycl_op_pad(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
diff --git a/ggml/src/ggml-sycl/rope.hpp b/ggml/src/ggml-sycl/rope.hpp
index b95a58580..94a769898 100644
--- a/ggml/src/ggml-sycl/rope.hpp
+++ b/ggml/src/ggml-sycl/rope.hpp
@@ -15,8 +15,6 @@
 
 #include "common.hpp"
 
-#define SYCL_ROPE_BLOCK_SIZE 256
-
 void ggml_sycl_rope(ggml_backend_sycl_context & ctx, ggml_tensor *dst);
 
 void ggml_sycl_rope_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
diff --git a/ggml/src/ggml-sycl/upscale.hpp b/ggml/src/ggml-sycl/upscale.hpp
index c36c1bdc9..2b72ea1b8 100644
--- a/ggml/src/ggml-sycl/upscale.hpp
+++ b/ggml/src/ggml-sycl/upscale.hpp
@@ -4,6 +4,4 @@
 #include "dpct/helper.hpp"
 #include "common.hpp"
 
-#define SYCL_UPSCALE_BLOCK_SIZE 256
-
 void ggml_sycl_upscale(ggml_backend_sycl_context & ctx, ggml_tensor * dst);

From 75118a3a5925d85e1850db67e87e447fa5890a47 Mon Sep 17 00:00:00 2001
From: Sam Malayek <12037535+SamMalayek@users.noreply.github.com>
Date: Thu, 8 Oct 2026 01:40:01 -0400
Subject: [PATCH 26/40] convert : support Qwen3.5 embedding models (#27920)

* Update convert_hf_to_gguf: support Qwen3.5 embedding models

* Behavior-preserving refactor for conventions.

* conversion: simplify pooling comment
---
 conversion/__init__.py         |  1 +
 conversion/base.py             | 10 ++++++++++
 conversion/qwen.py             | 15 ++++++++++++++-
 gguf-py/gguf/tensor_mapping.py |  9 +++++++++
 4 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/conversion/__init__.py b/conversion/__init__.py
index 72960be1a..5bf472889 100644
--- a/conversion/__init__.py
+++ b/conversion/__init__.py
@@ -255,6 +255,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "Qwen3_5ForConditionalGeneration": "qwen",
     "Qwen3_5MoeForCausalLM": "qwen",
     "Qwen3_5MoeForConditionalGeneration": "qwen",
+    "Qwen3_5TextModel": "qwen",
     "Qwen4ExpForCausalLM": "qwen4exp",
     "Qwen4ExpForConditionalGeneration": "qwen4exp",
     "RND1": "qwen",
diff --git a/conversion/base.py b/conversion/base.py
index 306a9f0f2..786e7faa1 100644
--- a/conversion/base.py
+++ b/conversion/base.py
@@ -2336,6 +2336,16 @@ class TextModel(ModelBase):
             else:
                 raise NotImplementedError("Only MEAN, CLS, and LAST pooling types supported")
             self.gguf_writer.add_pooling_type(pooling_type)
+        else:
+            embedding_config_path = self.dir_model / "embedding_config.json"
+            if embedding_config_path.is_file():
+                with open(embedding_config_path, encoding="utf-8") as f:
+                    embedding_config = json.load(f)
+                pooling = embedding_config.get("pooling")
+                if pooling == "last_token":
+                    self.gguf_writer.add_pooling_type(gguf.PoolingType.LAST)
+                elif pooling is not None:
+                    raise NotImplementedError(f"unsupported embedding_config.json pooling {pooling!r}")
 
         # pooling before a classification head (e.g. ModernBertForSequenceClassification)
         if (classifier_pooling := self.hparams.get("classifier_pooling")) is not None:
diff --git a/conversion/qwen.py b/conversion/qwen.py
index 0b3918733..3af7a035c 100644
--- a/conversion/qwen.py
+++ b/conversion/qwen.py
@@ -650,11 +650,24 @@ class _Qwen35MRopeMixin:
             self.gguf_writer.add_rope_dimension_sections(self._QWEN35_DEFAULT_MROPE_SECTION)
 
 
-@ModelBase.register("Qwen3_5ForConditionalGeneration", "Qwen3_5ForCausalLM")
+@ModelBase.register("Qwen3_5ForConditionalGeneration", "Qwen3_5ForCausalLM", "Qwen3_5TextModel")
 @ModelBase.example("Qwen/Qwen3.5-9B")
 class Qwen3_5TextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
     model_arch = gguf.MODEL_ARCH.QWEN35
 
+    def __init__(self, dir_model, *args, **kwargs):
+        # Inner TextModel does not own mtp.*. Set no_mtp before mixin bumps block_count.
+        hparams = kwargs.pop("hparams", None)
+        if hparams is None:
+            hparams = ModelBase.load_hparams(dir_model, self.is_mistral_format)
+        if get_model_architecture(hparams, ModelType.TEXT) == "Qwen3_5TextModel":
+            self.no_mtp = True
+        super().__init__(dir_model, *args, hparams=hparams, **kwargs)
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self._try_set_pooling_type()
+
 
 def _is_openjev_checkpoint(dir_model: Path) -> bool:
     return (dir_model / "helper" / "shim.py").is_file() and (dir_model / "config.json").is_file()
diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py
index a45e8e7d3..566b7348a 100644
--- a/gguf-py/gguf/tensor_mapping.py
+++ b/gguf-py/gguf/tensor_mapping.py
@@ -255,6 +255,7 @@ class TensorNameMap:
             "model.layers.{bid}.self_attn.language_expert_query_key_value",        # cogvlm
             "model.layers.{bid}.linear_attn.in_proj_qkv",                          # qwen3.5
             "head.layers.{bid}.self_attn.in_proj",  # laya
+            "layers.{bid}.linear_attn.in_proj_qkv",                                # qwen3.5 text
         ),
 
         # Attention query
@@ -397,6 +398,7 @@ class TensorNameMap:
         MODEL_TENSOR.ATTN_GATE: (
             "model.layers.{bid}.self_attn.gate_proj", # afmoe muse-glimmer
             "model.layers.{bid}.linear_attn.in_proj_z",  # qwen3.5
+            "layers.{bid}.linear_attn.in_proj_z",        # qwen3.5 text
             "model.layers.{bid}.self_attn.g_proj",    # step3.5 head-wise attention gate
             "model.layers.{bid}.self_attn.output_gate",  # minimax-01
             "model.layers.{bid}.self_attn.linear_gate",  # hy-v4
@@ -842,6 +844,7 @@ class TensorNameMap:
             "model.layers.{bid}.mamba.conv1d",         # jamba falcon-h1 granite-hybrid
             "model.layers.layers.{bid}.mixer.conv1d",  # plamo2
             "model.layers.{bid}.linear_attn.conv1d",   # qwen3next
+            "layers.{bid}.linear_attn.conv1d",         # qwen3.5 text
         ),
 
         MODEL_TENSOR.SSM_X: (
@@ -857,6 +860,7 @@ class TensorNameMap:
             "model.layers.{bid}.mamba.dt_proj",         # jamba falcon-h1 granite-hybrid
             "model.layers.layers.{bid}.mixer.dt_proj",  # plamo2
             "model.layers.{bid}.linear_attn.dt_proj",   # qwen3next
+            "layers.{bid}.linear_attn.dt_proj",         # qwen3.5 text
             "backbone.layers.{bid}.mixer.dt",           # nemotron-h-moe
             "model.layers.{bid}.self_attn.dt_proj",     # kimi
             "model.layers.{bid}.attention.dt_proj",     # bailingmoe3
@@ -873,6 +877,7 @@ class TensorNameMap:
             "model.layers.{bid}.mamba.A_log",         # jamba falcon-h1 granite-hybrid
             "model.layers.layers.{bid}.mixer.A_log",  # plamo2
             "model.layers.{bid}.linear_attn.A_log",   # qwen3next
+            "layers.{bid}.linear_attn.A_log",         # qwen3.5 text
             "model.layers.{bid}.self_attn.A_log",     # kimi
             "model.layers.{bid}.attention.A_log",     # bailingmoe3
         ),
@@ -899,6 +904,7 @@ class TensorNameMap:
         MODEL_TENSOR.SSM_NORM: (
             "model.layers.{bid}.mamba.norm",        # falcon-h1 granite-hybrid
             "model.layers.{bid}.linear_attn.norm",  # qwen3next
+            "layers.{bid}.linear_attn.norm",        # qwen3.5 text
             "backbone.layers.{bid}.mixer.norm",     # mamba2
             "model.layers.{bid}.self_attn.o_norm",  # kimi
             "model.layers.{bid}.attention.o_norm",  # bailingmoe3
@@ -909,11 +915,13 @@ class TensorNameMap:
             "backbone.layers.{bid}.mixer.out_proj",      # mamba
             "model.layers.{bid}.mamba.out_proj",         # jamba falcon-h1 granite-hybrid
             "model.layers.{bid}.linear_attn.out_proj",   # qwen3next
+            "layers.{bid}.linear_attn.out_proj",         # qwen3.5 text
             "model.layers.layers.{bid}.mixer.out_proj",  # plamo2
         ),
 
         MODEL_TENSOR.SSM_ALPHA: (
             "model.layers.{bid}.linear_attn.in_proj_a",  # qwen3.5
+            "layers.{bid}.linear_attn.in_proj_a",        # qwen3.5 text
         ),
 
         MODEL_TENSOR.SSM_BETA_ALPHA: (
@@ -941,6 +949,7 @@ class TensorNameMap:
         ),
         MODEL_TENSOR.SSM_BETA: (
             "model.layers.{bid}.linear_attn.in_proj_b",  # qwen3.5
+            "layers.{bid}.linear_attn.in_proj_b",        # qwen3.5 text
             "model.layers.{bid}.self_attn.b_proj",       # Kimi Linear
             "model.layers.{bid}.attention.b_proj",       # bailingmoe3
         ),

From 847f447c310a3784d623a953e6f0bb07f04c0f27 Mon Sep 17 00:00:00 2001
From: cwriter 
Date: Thu, 8 Oct 2026 07:41:13 +0200
Subject: [PATCH 27/40] sycl: add grouped MoE XMX GEMM (#29245)

Co-authored-by: cwriter 
---
 docs/backend/SYCL.md              |    4 +
 ggml/src/ggml-sycl/CMakeLists.txt |   23 +
 ggml/src/ggml-sycl/common.hpp     |   62 ++
 ggml/src/ggml-sycl/fused-gemm.cpp | 1093 +++++++++++++++++++++++++++++
 ggml/src/ggml-sycl/fused-gemm.hpp |   88 +++
 ggml/src/ggml-sycl/ggml-sycl.cpp  |  104 ++-
 6 files changed, 1358 insertions(+), 16 deletions(-)
 create mode 100644 ggml/src/ggml-sycl/fused-gemm.cpp
 create mode 100644 ggml/src/ggml-sycl/fused-gemm.hpp

diff --git a/docs/backend/SYCL.md b/docs/backend/SYCL.md
index 70f5f2eff..c2ae0c1ba 100644
--- a/docs/backend/SYCL.md
+++ b/docs/backend/SYCL.md
@@ -816,6 +816,10 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_MKL_FA_DIAG | 0 (default) or 1 | Enable output fingerprinting for MKL flash attention. Dumps the first 64 float output values for the first 6 FA calls with n_kv ≥ 1024, labeled with kernel type (MKL/TILE/VEC) for cross-kernel comparison. |
 | GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute. Unsupported types and layouts fall back to the standalone op kernels. See `ggml_sycl_can_fuse()`. |
 | GGML_SYCL_ENABLE_ESIMD | 0 or 1 (default)| Enable ESIMD kernels when available. |
+| GGML_SYCL_XMX_GATHER_TYPES | decimal bitmask, all bits set (default) | Weight formats that may use the XMX dequant-GEMM paths, which dequantize weights straight into the XMX tiles. This speeds up prompt processing of MoE models on GPUs with XMX units (Arc A- and B-series, Arc Pro, Data Center GPU Max), for example pp512 of Qwen3-30B-A3B UD-IQ3_XXS by about 50% on an Arc Pro B60. Bits:
* 1: IQ4_NL, 2: IQ3_S, 4: IQ4_XS, 8: IQ3_XXS, 16: IQ2_XXS, 32: IQ2_XS, 64: IQ2_S, 128: IQ1_S, 256: IQ1_M
* 512: Q8_0, 1024: Q4_K, 2048: Q5_K, 4096: Q6_K (MoE `MUL_MAT_ID` only)
Add values to combine them, for example `3` for IQ4_NL and IQ3_S; `0` disables the paths. A set bit does not force the path: batches of more than 64 tokens per expert or row lengths that are not a multiple of 256 (32 for IQ4_NL and Q8_0) use the library GEMM. |
+| GGML_SYCL_XMX_GATHER_SHAPES | decimal bitmask, 255 (default) | XMX `joint_matrix` combinations the paths of `GGML_SYCL_XMX_GATHER_TYPES` may use; the operand type comes from `GGML_SYCL_DYNAMIC_PRECISION` and the best supported combination is picked automatically (logged as `fg_pick_combo`). Bits:
* Xe2, Xe3, Xe-HPC: 1: f16 8x16x16, 2: f16 16x16x16, 4: f16 32x64x16, 8: f16 32x64x32, 32: tf32 8x16x8, 64: bf16 8x16x16
* Xe-HPG (Arc A770, ARL-H): 16: f16 8x8x16, 128: bf16 8x8x16
Clear a bit to exclude a combination, or set a single bit to force one for testing. |
+| GGML_SYCL_DYNAMIC_PRECISION | `F16` (default with `GGML_SYCL_F16=ON`), `BF16`, `TF32` or `F32` (default otherwise) | Operand type of the XMX dequant-GEMM paths (`GGML_SYCL_XMX_GATHER_TYPES`); accumulation is always f32. `F16` is the fastest, but activations above 65504 overflow. `BF16` keeps the f32 range at a 7-bit mantissa, `TF32` keeps the range and the f16 mantissa but is about 30% slower and needs Xe2, Xe3 or Xe-HPC, and `F32` turns the XMX paths off. Ops that request a higher src1 precision ([TAG_GGML_PREC]) get it regardless of this setting. |
+| GGML_SYCL_DYNAMIC_REQUIRED_PRECISION | `F32` (default), `TF32`, `BF16` or `F16` | Lowest type the XMX paths may use for an op that requests an F32 src1, such as Mistral 4 `ffn_down_exps`. The default runs such ops on the library f32 GEMM; `TF32` or `BF16` trade mantissa for speed while keeping the f32 range. `F16` ignores the request and can overflow; it is meant for testing only. |
 | GGML_SYCL_MMVQ_WIDE | 0 or 1 (default) | Use the wide-load variant of the reordered Q8_0 mat-vec kernel, which reads four contiguous dwords per operand instead of one value at a time. Set to 0 to fall back to the per-value loads. Only affects Q8_0 weights in the reordered layout. |
 | GGML_SYCL_SPARSE_FA | 0 (default) or 1 | Enable Sparse Flash-attention.|
 | GGML_SYCL_SPARSE_FA_DEBUG | 0 (default) or 1 | Enable to debug for Sparse Flash-attention.|
diff --git a/ggml/src/ggml-sycl/CMakeLists.txt b/ggml/src/ggml-sycl/CMakeLists.txt
index d2196f74d..1697ecd95 100644
--- a/ggml/src/ggml-sycl/CMakeLists.txt
+++ b/ggml/src/ggml-sycl/CMakeLists.txt
@@ -221,4 +221,27 @@ if (GGML_SYCL_DEVICE_ARCH)
         "SHELL:-Xsycl-target-backend=spir64_gen \"-device ${GGML_SYCL_DEVICE_ARCH}\""
         -fsycl-max-parallel-link-jobs=${GGML_SYCL_MAX_PARALLEL_LINK_JOBS}
     )
+
+    # The XMX dequant-GEMM tiles need the sub-group size of the target: 8 on Xe-HPG (DG2, ARL-H),
+    # 16 on Xe-HPC and Xe2 or newer. ocloc fails on the other size, so build only the one that fits.
+    # 0 (unknown name, mixed list, or no XMX) builds no XMX tile and the path stays off.
+    set(_ggml_sycl_xmx_sg "")
+    string(TOLOWER "${GGML_SYCL_DEVICE_ARCH}" _ggml_sycl_archs)
+    string(REPLACE "," ";" _ggml_sycl_archs "${_ggml_sycl_archs}")
+    foreach(_arch IN LISTS _ggml_sycl_archs)
+        if (_arch MATCHES "^(dg2|acm|ats-m|arl-h|xe-hpg|12\\.5[567]\\.|12\\.74\\.)")
+            set(_sg 8)
+        elseif (_arch MATCHES "^(pvc|bmg|lnl|ptl|wcl|nvl|cri|xe2|xe3|xe-hpc|12\\.60\\.|20\\.|30\\.)")
+            set(_sg 16)
+        else()
+            set(_sg 0)
+        endif()
+        if (_ggml_sycl_xmx_sg STREQUAL "" OR _ggml_sycl_xmx_sg EQUAL _sg)
+            set(_ggml_sycl_xmx_sg ${_sg})
+        else()
+            set(_ggml_sycl_xmx_sg 0)
+        endif()
+    endforeach()
+    message(STATUS "GGML_SYCL_DEVICE_ARCH: XMX dequant-GEMM sub-group size ${_ggml_sycl_xmx_sg} (0 = off)")
+    target_compile_definitions(ggml-sycl PRIVATE GGML_SYCL_XMX_AOT_SG=${_ggml_sycl_xmx_sg})
 endif()
diff --git a/ggml/src/ggml-sycl/common.hpp b/ggml/src/ggml-sycl/common.hpp
index 5e904bd2a..475509372 100644
--- a/ggml/src/ggml-sycl/common.hpp
+++ b/ggml/src/ggml-sycl/common.hpp
@@ -65,6 +65,61 @@ extern int g_ggml_sycl_enable_fusion;
 extern int g_ggml_sycl_enable_esimd;
 extern int g_ggml_sycl_mmvq_wide;
 extern int g_ggml_sycl_prioritize_dmmv;
+
+// Which quantized weight formats may take the XMX dequant-GEMM paths. A bitmask rather than one
+// flag per path, so a format can be enabled or measured on its own and adding a format is one bit.
+enum ggml_sycl_xmx_gather_type {
+    GGML_SYCL_XMX_GATHER_IQ4_NL   = 1 << 0,
+    GGML_SYCL_XMX_GATHER_IQ3_S    = 1 << 1,
+    GGML_SYCL_XMX_GATHER_IQ4_XS   = 1 << 2,
+    GGML_SYCL_XMX_GATHER_IQ3_XXS  = 1 << 3,
+    GGML_SYCL_XMX_GATHER_IQ2_XXS  = 1 << 4,
+    GGML_SYCL_XMX_GATHER_IQ2_XS   = 1 << 5,
+    GGML_SYCL_XMX_GATHER_IQ2_S    = 1 << 6,
+    GGML_SYCL_XMX_GATHER_IQ1_S    = 1 << 7,
+    GGML_SYCL_XMX_GATHER_IQ1_M    = 1 << 8,
+    GGML_SYCL_XMX_GATHER_Q8_0     = 1 << 9,
+    GGML_SYCL_XMX_GATHER_Q4_K     = 1 << 10,
+    GGML_SYCL_XMX_GATHER_Q5_K     = 1 << 11,
+    GGML_SYCL_XMX_GATHER_Q6_K     = 1 << 12,
+};
+static constexpr int GGML_SYCL_XMX_GATHER_TYPES_DEFAULT = ~0;
+extern int g_ggml_sycl_xmx_gather_types;
+// Which joint_matrix combinations the XMX dequant-GEMM paths may use, one bit each (see fused-gemm.cpp).
+// GGML_SYCL_DYNAMIC_PRECISION picks the operand type, this mask the combinations of that type.
+static constexpr int GGML_SYCL_XMX_GATHER_SHAPES_DEFAULT = 0xff;
+extern int g_ggml_sycl_xmx_gather_shapes;
+
+// GGML_SYCL_DYNAMIC_PRECISION: operand type of the XMX dequant-GEMM paths. F32 turns them off and
+// keeps the library GEMM in f32. A src1 precision request of an op [TAG_GGML_PREC] is always met.
+enum ggml_sycl_dynamic_precision {
+    GGML_SYCL_DYNAMIC_PRECISION_F16,
+    GGML_SYCL_DYNAMIC_PRECISION_BF16,
+    GGML_SYCL_DYNAMIC_PRECISION_TF32,
+    GGML_SYCL_DYNAMIC_PRECISION_F32,
+};
+#ifdef GGML_SYCL_F16
+static constexpr int GGML_SYCL_DYNAMIC_PRECISION_DEFAULT = GGML_SYCL_DYNAMIC_PRECISION_F16;
+#else
+static constexpr int GGML_SYCL_DYNAMIC_PRECISION_DEFAULT = GGML_SYCL_DYNAMIC_PRECISION_F32;
+#endif
+extern int g_ggml_sycl_dynamic_precision;
+// GGML_SYCL_DYNAMIC_REQUIRED_PRECISION: the XMX type an F32 src1 request may run on instead of f32
+// (TF32, or BF16 which also allows tf32). F32 (default): none. F16: src1 requests are ignored.
+extern int g_ggml_sycl_dynamic_required_precision;
+
+// [TAG_GGML_PREC] src1 precision request of the MUL_MAT/MUL_MAT_ID op dst
+static inline int32_t ggml_sycl_src1_prec(const ggml_tensor * dst) {
+    return g_ggml_sycl_dynamic_required_precision == GGML_SYCL_DYNAMIC_PRECISION_F16 ? GGML_PREC_UNDEFINED :
+                                                                                       dst->op_params[3];
+}
+
+// [TAG_GGML_PREC] the library GEMM and dmmv may convert src1 of the MUL_MAT/MUL_MAT_ID op dst to f16
+static inline bool ggml_sycl_src1_f16_ok(const ggml_tensor * dst) {
+    const int32_t src1_prec = ggml_sycl_src1_prec(dst);
+    return g_ggml_sycl_dynamic_precision != GGML_SYCL_DYNAMIC_PRECISION_F32 &&
+           (src1_prec == GGML_PREC_UNDEFINED || src1_prec >= GGML_PREC_F16);
+}
 extern int g_ggml_sycl_enable_flash_attention;
 extern int g_ggml_sycl_dev2dev_memcpy;
 extern int g_ggml_sycl_fa_onednn;
@@ -333,6 +388,12 @@ struct mmid_row_mapping {
     int32_t i2;
 };
 
+struct ggml_sycl_gg_tile {
+    int32_t expert;
+    int32_t n0;
+    int32_t n1;
+};
+
 namespace sycl_ex = sycl::ext::oneapi::experimental;
 struct ggml_backend_sycl_context {
     int device;
@@ -410,6 +471,7 @@ struct ggml_backend_sycl_context {
     std::unique_ptr host_pools[GGML_SYCL_MAX_DEVICES];
 
     std::vector mmid_row_mapping_host;
+    std::vector mmid_tile_schedule_host;
 
     static std::unique_ptr new_pool_for_device(queue_ptr qptr, int device);
 
diff --git a/ggml/src/ggml-sycl/fused-gemm.cpp b/ggml/src/ggml-sycl/fused-gemm.cpp
new file mode 100644
index 000000000..38c69bc6b
--- /dev/null
+++ b/ggml/src/ggml-sycl/fused-gemm.cpp
@@ -0,0 +1,1093 @@
+#include "fused-gemm.hpp"
+
+#include 
+
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+
+namespace mx = sycl::ext::oneapi::experimental::matrix;
+
+// FG_ / fg_ is short for fused GEMM: the weights are dequantized inside the GEMM, into the XMX tiles.
+
+// A k step is one 32-value weight sub-block; iq3_s and the other superblock formats split their
+// superblock into steps of this width. The sub-groups of a work-group each walk their own K range
+// and are summed at the end.
+static constexpr int FG_BK     = QK4_NL;
+static constexpr int FG_KSPLIT = 4;
+
+// Element traits of one joint_matrix operand type. The A stage and the B pack compute in f32 and
+// convert once, in registers, when they write the element, so any type costs the same one pass.
+//   store: storage in SLM (A) and in the packed B buffer
+//   mtype: matrix_type in matrix_combinations
+//   mode:  GGML_SYCL_DYNAMIC_PRECISION value that selects this type
+//   src:   ggml type that needs no conversion into this type (GGML_TYPE_COUNT: none)
+//   slow:  XMX throughput class, 0 is fastest. f16 and bf16 share the DPAS rate; tf32 does half the
+//          K per instruction. B60, Qwen3-30B-A3B pp512: f16 1108, bf16 1000, tf32 751 t/s
+template  struct fg_elem;
+
+template <> struct fg_elem {
+    using store = sycl::half;
+    using pair  = sycl::half2;
+    static constexpr mx::matrix_type mtype = mx::matrix_type::fp16;
+    static constexpr int             mode  = GGML_SYCL_DYNAMIC_PRECISION_F16;
+    static constexpr ggml_type       src   = GGML_TYPE_F16;
+    static constexpr int             mant  = 10;
+    static constexpr int             slow  = 0;
+    static store cvt(float x) { return (store) x; }
+    static pair make(float x, float y) { return pair((store) x, (store) y); }
+};
+
+// tf32 rounds to nearest even with plain bit ops: round_to_tf32 needs a SPIR-V extension that the
+// DG2 AOT target rejects
+static inline uint32_t fg_round_bits(float x, int drop) {
+    const uint32_t u = sycl::bit_cast(x);
+    if ((u & 0x7f800000u) == 0x7f800000u) {
+        return (u & 0x7fffffu) ? u | (1u << drop) : u; // nan stays nan
+    }
+    return u + ((1u << (drop - 1)) - 1) + ((u >> drop) & 1);
+}
+
+struct alignas(4) fg_bf16x2 {
+    sycl::ext::oneapi::bfloat16 x, y;
+};
+
+template <> struct fg_elem {
+    using store = sycl::ext::oneapi::bfloat16;
+    using pair  = fg_bf16x2;
+    static constexpr mx::matrix_type mtype = mx::matrix_type::bf16;
+    static constexpr int             mode  = GGML_SYCL_DYNAMIC_PRECISION_BF16;
+    static constexpr ggml_type       src   = GGML_TYPE_BF16;
+    static constexpr int             mant  = 7;
+    static constexpr int             slow  = 0;
+    static store cvt(float x) { return store(x); }
+    static pair make(float x, float y) { return { cvt(x), cvt(y) }; }
+};
+
+// tf32 keeps f32 range and f16 mantissa, in f32 storage
+template <> struct fg_elem {
+    using store = float;
+    using pair  = sycl::float2;
+    static constexpr mx::matrix_type mtype = mx::matrix_type::tf32;
+    static constexpr int             mode  = GGML_SYCL_DYNAMIC_PRECISION_TF32;
+    static constexpr ggml_type       src   = GGML_TYPE_COUNT;
+    static constexpr int             mant  = 10;
+    static constexpr int             slow  = 1;
+    static store cvt(float x) { return sycl::bit_cast(fg_round_bits(x, 13) & ~0x1fffu); }
+    static pair make(float x, float y) { return pair(cvt(x), cvt(y)); }
+};
+
+// One joint_matrix combination (A type, B type, TM x TN x TK, sub-group size; C and D are f32) and
+// the tiling built on it. A sub-group owns SG_ROWS rows of A (at least 16) and BN (at least 32)
+// columns of B. A and B may differ: the device lists the pairs it supports.
+template  struct fg_combo {
+    using ta  = TA;
+    using tb  = TB;
+    using EA  = fg_elem;
+    using EB  = fg_elem;
+    using tsa = typename EA::store;
+    using tsb = typename EB::store;
+    static constexpr int TM = TM_;
+    static constexpr int TN = TN_;
+    static constexpr int TK = TK_;
+    static constexpr int SG = SG_;
+    static constexpr int VNNI    = 4 / sizeof(tsb);  // K rows of B packed in one 32-bit word
+    static constexpr int SG_ROWS = TM > 16 ? TM : 16;
+    static constexpr int RPL     = SG_ROWS / SG;     // A rows one lane decodes per k step
+    static constexpr int MT      = SG_ROWS / TM;
+    static constexpr int BN      = TN > 32 ? TN : 32;
+    static constexpr int NT      = BN / TN;
+    static constexpr int WG_SIZE = FG_KSPLIT * SG;
+    static constexpr mx::layout b_layout = VNNI == 1 ? mx::layout::row_major : mx::layout::ext_intel_packed;
+    // a 64-wide N is mostly padding here and a 32x64 f32 accumulator needs 128 registers per lane,
+    // so it spills: 13x slower on B60
+    static constexpr bool efficient = TN <= 32;
+    static_assert(SG_ROWS % SG == 0 && SG_ROWS % TM == 0 && BN % TN == 0 && FG_BK % TK == 0, "bad tile");
+    static_assert(BN <= GGML_SYCL_FG_MAX_N, "header gate must cover the tile width");
+};
+
+using fg_half = sycl::half;
+using fg_bf16 = sycl::ext::oneapi::bfloat16;
+using fg_tf32 = mx::precision::tf32;
+
+// One bit of GGML_SYCL_XMX_GATHER_SHAPES per combination. Only combinations some device lists in
+// matrix_combinations are built (appendix of sycl_ext_oneapi_matrix and the runtime's own list).
+template  static void fg_visit_combo(int idx, F && f);
+static constexpr int FG_N_COMBOS = 8;
+
+// A spir64_gen AOT build (GGML_SYCL_XMX_AOT_SG) drops the combinations of the other sub-group size
+// entirely: ocloc rejects even an empty kernel that asks for a sub-group size it lacks.
+template  static constexpr bool fg_listed() {
+#if defined(GGML_SYCL_XMX_AOT_SG)
+    return SG == GGML_SYCL_XMX_AOT_SG;
+#else
+    return true;
+#endif
+}
+
+template  static void fg_call_combo(F && f) {
+    if constexpr (fg_listed()) {
+        f(S{});
+    }
+}
+
+template  static void fg_visit_combo(int idx, F && f) {
+    switch (idx) {
+        case 0: fg_call_combo>(f);  break; // Xe2, Xe3, Xe-HPC
+        case 1: fg_call_combo>(f); break; // Xe2, Xe3, Xe-HPC
+        case 2: fg_call_combo>(f); break; // Xe2, Xe3, Xe-HPC
+        case 3: fg_call_combo>(f); break; // Xe2, Xe3, Xe-HPC
+        case 4: fg_call_combo>(f);    break; // Xe-HPG (Arc A), ARL-H
+        case 5: fg_call_combo>(f);   break; // Xe2, Xe3, Xe-HPC
+        case 6: fg_call_combo>(f);  break; // Xe2, Xe3, Xe-HPC
+        case 7: fg_call_combo>(f);    break; // Xe-HPG (Arc A), ARL-H
+        default: GGML_ABORT("bad XMX combination %d", idx);
+    }
+}
+
+// AOT with -fsycl-targets=intel_gpu_*: compile each tile body only for targets with its sub-group
+// size, since IGC fails on the other ones. A JIT build keeps them all, but each combination lands in
+// its own device image (joint_matrix is an optional kernel feature) and only a combination the
+// device reports is launched, so the runtime never asks IGC for the others.
+#if defined(__SYCL_DEVICE_ONLY__)
+#    if __SYCL_TARGET_INTEL_GPU_ACM_G10__ || __SYCL_TARGET_INTEL_GPU_ACM_G11__ || __SYCL_TARGET_INTEL_GPU_ACM_G12__ || \
+        __SYCL_TARGET_INTEL_GPU_ARL_H__
+#        define FG_AOT_SG 8
+#    elif __SYCL_TARGET_INTEL_GPU_PVC__ || __SYCL_TARGET_INTEL_GPU_PVC_VG__ || __SYCL_TARGET_INTEL_GPU_BMG_G21__ || \
+        __SYCL_TARGET_INTEL_GPU_BMG_G31__ || __SYCL_TARGET_INTEL_GPU_LNL_M__ || __SYCL_TARGET_INTEL_GPU_PTL_H__ ||   \
+        __SYCL_TARGET_INTEL_GPU_PTL_U__ || __SYCL_TARGET_INTEL_GPU_WCL__ || __SYCL_TARGET_INTEL_GPU_NVL_S__ ||       \
+        __SYCL_TARGET_INTEL_GPU_NVL_U__ || __SYCL_TARGET_INTEL_GPU_NVL_P__
+#        define FG_AOT_SG 16
+#    elif __SYCL_TARGET_INTEL_GPU_TGLLP__ || __SYCL_TARGET_INTEL_GPU_RKL__ || __SYCL_TARGET_INTEL_GPU_ADL_S__ || \
+        __SYCL_TARGET_INTEL_GPU_ADL_P__ || __SYCL_TARGET_INTEL_GPU_ADL_N__ || __SYCL_TARGET_INTEL_GPU_DG1__ ||   \
+        __SYCL_TARGET_INTEL_GPU_MTL_U__ || __SYCL_TARGET_INTEL_GPU_MTL_H__
+#        define FG_AOT_SG 0 // no XMX
+#    endif
+#endif
+
+template  static constexpr bool fg_built() {
+#if defined(FG_AOT_SG)
+    return SG == FG_AOT_SG;
+#else
+    return true;
+#endif
+}
+
+// Upper bound on the tile count when total_rows rows are routed to n_as experts: the worst case gives
+// each expert one row and fills whole tiles with the rest. The bound depends only on the shape, not
+// on the routing, so the pool reuses one buffer every ubatch instead of keeping one per size seen.
+static constexpr int64_t grouped_gemm_max_tiles(int64_t total_rows, int64_t n_as, int64_t BN) {
+    return total_rows <= n_as ? total_rows : n_as + (total_rows - n_as) / BN;
+}
+// Tiles do not cross experts, so the bound is not ceil(total_rows / BN): 34 rows over 2 experts with
+// BN = 16 split 17 + 17 need 2 + 2 tiles, where the ceil gives 3.
+static_assert(grouped_gemm_max_tiles(34, 2, 16) == 4);
+
+// the device lists S with an f32 accumulator and output
+template  static bool fg_device_has_combo(const std::vector & combinations) {
+    for (const auto & c : combinations) {
+        if (c.atype == S::EA::mtype && c.btype == S::EB::mtype && c.ctype == mx::matrix_type::fp32 &&
+            c.dtype == mx::matrix_type::fp32 &&
+            (c.max_msize >= (size_t) S::TM || c.msize == (size_t) S::TM) &&
+            (c.max_nsize >= (size_t) S::TN || c.nsize == (size_t) S::TN) &&
+            (c.max_ksize >= (size_t) S::TK || c.ksize == (size_t) S::TK)) {
+            return true;
+        }
+    }
+    return false;
+}
+
+template  static const char * fg_type_name() {
+    return std::is_same_v ? "f16" : std::is_same_v ? "bf16" : "tf32";
+}
+
+static std::string fg_combo_name(int idx) {
+    std::string name;
+    fg_visit_combo(idx, [&](auto s) {
+        using S = decltype(s);
+        name = std::string(fg_type_name()) + "x" + fg_type_name() + " " +
+               std::to_string(S::TM) + "x" + std::to_string(S::TN) + "x" + std::to_string(S::TK) + " sg" +
+               std::to_string(S::SG);
+    });
+    return name;
+}
+
+// Combinations this build has kernels for and the device lists, one bit each. Cached per device:
+// on a mixed box the first caller's verdict is not the others'.
+static int fg_device_combos(const sycl::device & dev) {
+    static std::mutex                            mtx;
+    static std::unordered_map known;
+    std::lock_guard                  lock(mtx);
+    const auto                                   it = known.find(dev);
+    if (it != known.end()) {
+        return it->second;
+    }
+    int available = 0;
+    try {
+        const auto combinations = dev.get_info();
+        const auto sg_sizes     = dev.get_info();
+        for (int idx = 0; idx < FG_N_COMBOS; ++idx) {
+            fg_visit_combo(idx, [&](auto s) {
+                using S = decltype(s);
+                const bool sg_ok = std::find(sg_sizes.begin(), sg_sizes.end(), (size_t) S::SG) != sg_sizes.end();
+                if (sg_ok && fg_device_has_combo(combinations)) {
+                    available |= 1 << idx;
+                }
+            });
+        }
+    } catch (const sycl::exception &) {
+        available = 0;
+    }
+    GGML_LOG_INFO("%s: %s: XMX dequant-GEMM combinations available 0x%x, allowed 0x%x\n", __func__,
+                  dev.get_info().c_str(), available, g_ggml_sycl_xmx_gather_shapes);
+    known.emplace(dev, available);
+    return available;
+}
+
+// Rank of combination S for a src1 of type src1_type, lower is better. Order:
+//  1. throughput: a tile that does not spill, then the fastest type class of A and B
+//  2. B type equal to the src1 type, so the pack is a plain copy
+//  3. B at least as precise as f16
+//  4. the device's native DPAS tile (8 x SG x 32 bytes of K), then the largest M x K
+// A costs nothing to convert: the A stage emits any type at the same cost.
+template  static int64_t fg_rank(ggml_type src1_type) {
+    const int64_t spills  = !S::efficient;
+    const int64_t slow    = std::max(S::EA::slow, S::EB::slow);
+    const int64_t convert = S::EB::src != src1_type;
+    const int64_t lossy   = S::EB::mant < 10;
+    const int64_t foreign = !(S::TM == 8 && S::TN == S::SG);
+    const int64_t mk      = 1024 - S::TM * S::TK;
+    return ((((spills * 2 + slow) * 2 + convert) * 2 + lossy) * 2 + foreign) * 2048 + mk;
+}
+
+// whether XMX operands of type mode (a GGML_SYCL_DYNAMIC_PRECISION value) meet the src1 request
+// [TAG_GGML_PREC]. f16 lacks the f32 range that BF16 and F32 ask for; an F32 request goes only as far
+// down as GGML_SYCL_DYNAMIC_REQUIRED_PRECISION allows.
+static bool fg_mode_meets(int mode, int32_t src1_prec) {
+    if (src1_prec == GGML_PREC_UNDEFINED || src1_prec >= GGML_PREC_F16) {
+        return true;
+    }
+    if (src1_prec >= GGML_PREC_BF16) {
+        return mode != GGML_SYCL_DYNAMIC_PRECISION_F16;
+    }
+    switch (g_ggml_sycl_dynamic_required_precision) {
+        case GGML_SYCL_DYNAMIC_PRECISION_TF32: return mode == GGML_SYCL_DYNAMIC_PRECISION_TF32;
+        case GGML_SYCL_DYNAMIC_PRECISION_BF16: return mode != GGML_SYCL_DYNAMIC_PRECISION_F16;
+        default:                               return false;
+    }
+}
+
+// Best allowed combination for this call, or -1 if none. The type is GGML_SYCL_DYNAMIC_PRECISION if it
+// meets the src1 request; if not, bf16 then tf32 for a BF16 request (fastest first) and tf32 then bf16
+// for an F32 request (most mantissa first).
+static int fg_pick_combo(dpct::queue_ptr stream, ggml_type src1_type, int32_t src1_prec) {
+    const sycl::device dev     = stream->get_device();
+    const int          allowed = fg_device_combos(dev) & g_ggml_sycl_xmx_gather_shapes;
+    const bool         f32_req = src1_prec != GGML_PREC_UNDEFINED && src1_prec < GGML_PREC_BF16;
+    const int          modes[] = {
+        g_ggml_sycl_dynamic_precision,
+        f32_req ? GGML_SYCL_DYNAMIC_PRECISION_TF32 : GGML_SYCL_DYNAMIC_PRECISION_BF16,
+        f32_req ? GGML_SYCL_DYNAMIC_PRECISION_BF16 : GGML_SYCL_DYNAMIC_PRECISION_TF32,
+    };
+    int     best      = -1;
+    int64_t best_rank = 0;
+    for (int i = 0; i < 3 && best < 0; ++i) {
+        const int mode = modes[i];
+        if (!fg_mode_meets(mode, src1_prec)) {
+            continue;
+        }
+        for (int idx = 0; idx < FG_N_COMBOS; ++idx) {
+            if (!(allowed & (1 << idx))) {
+                continue;
+            }
+            fg_visit_combo(idx, [&](auto s) {
+                using S = decltype(s);
+                if (S::EA::mode != mode || S::EB::mode != mode) {
+                    return;
+                }
+                const int64_t rank = fg_rank(src1_type);
+                if (best < 0 || rank < best_rank) {
+                    best      = idx;
+                    best_rank = rank;
+                }
+            });
+        }
+    }
+    // log each distinct decision once
+    static std::mutex                                      mtx;
+    static std::set> seen;
+    std::lock_guard                            lock(mtx);
+    if (seen.emplace(std::hash{}(dev), (int) src1_type, src1_prec, best).second) {
+        GGML_LOG_INFO("%s: src1 %s, src1 prec %d -> %s\n", __func__, ggml_type_name(src1_type), src1_prec,
+                      best >= 0 ? fg_combo_name(best).c_str() : "none (library GEMM)");
+    }
+    return best;
+}
+
+// src1 [N][K] -> packed [K/V][Npad][V] so B tiles load straight from global memory
+template 
+static void fused_gemm_pack_b(const T_src * y, typename E::store * packed, int N, int Npad, int K,
+                              dpct::queue_ptr stream) {
+    constexpr int V   = 4 / sizeof(typename E::store);
+    const int     kqs = K / V;
+    stream->parallel_for(sycl::range<1>((size_t) Npad * kqs), [=](sycl::id<1> id) {
+        const int idx = id[0];
+        const int n   = idx / kqs;
+        const int kq  = idx - n * kqs;
+        typename E::store vals[V] = {};
+        if (n < N) {
+            const T_src * src = y + (size_t) n * K + V * kq;
+#pragma unroll
+            for (int v = 0; v < V; ++v) {
+                vals[v] = E::cvt((float) src[v]);
+            }
+        }
+        typename E::store * out = packed + ((size_t) kq * Npad + n) * V;
+#pragma unroll
+        for (int v = 0; v < V; ++v) {
+            out[v] = vals[v];
+        }
+    });
+}
+
+// A stage: one lane owns one row and decodes FG_BK values of it per k step, with every scale
+// folded into the value so the mad below sees plain A elements. One overload per weight format.
+template 
+static __dpct_inline__ void fg_stage_a(const block_iq4_nl * __restrict__ xrow, const int kb, typename E::pair * a) {
+    const block_iq4_nl blk = xrow[kb];
+    const float        d   = (float) blk.d;
+#pragma unroll
+    for (int j = 0; j < QK4_NL / 2; j += 2) {
+        const uint8_t q0 = blk.qs[j];
+        const uint8_t q1 = blk.qs[j + 1];
+        a[j / 2]     = E::make(d * kvalues_iq4nl[q0 & 0xf], d * kvalues_iq4nl[q1 & 0xf]);
+        a[j / 2 + 8] = E::make(d * kvalues_iq4nl[q0 >> 4], d * kvalues_iq4nl[q1 >> 4]);
+    }
+}
+
+// iq3_s: k step kb is sub-block kb % 8 of superblock kb / 8. The superblock is 110 bytes, so read
+// only the fields of that sub-block instead of copying the block. Same decode as
+// dequantize_block_iq3_s: grid entries are taken as dwords and the sign bit is a plain shift.
+template 
+static __dpct_inline__ void fg_stage_a(const block_iq3_s * __restrict__ xrow, const int kb, typename E::pair * a) {
+    static_assert(QK_K == 256, "the iq3_s A stage assumes 8 sub-blocks per superblock");
+    const block_iq3_s * blk = xrow + kb / (QK_K / 32);
+    const int           ib8 = kb % (QK_K / 32);
+    const uint8_t *     qs  = blk->qs + 8 * ib8;
+    const int           qh  = blk->qh[ib8];
+    const float         d   = (float) blk->d * (1 + 2 * ((blk->scales[ib8 / 2] >> (4 * (ib8 % 2))) & 0xf));
+#pragma unroll
+    for (int il = 0; il < 4; ++il) {
+        const uint32_t grid1 = iq3s_grid[qs[2 * il + 0] | ((qh << (8 - 2 * il)) & 256)];
+        const uint32_t grid2 = iq3s_grid[qs[2 * il + 1] | ((qh << (7 - 2 * il)) & 256)];
+        const int      signs = blk->signs[4 * ib8 + il];
+#pragma unroll
+        for (int j = 0; j < 2; ++j) {
+            const float g1a = (float) ((grid1 >> (16 * j + 0)) & 0xff);
+            const float g1b = (float) ((grid1 >> (16 * j + 8)) & 0xff);
+            const float g2a = (float) ((grid2 >> (16 * j + 0)) & 0xff);
+            const float g2b = (float) ((grid2 >> (16 * j + 8)) & 0xff);
+            const int   s   = 2 * j;
+            a[4 * il + j]     = E::make(d * ((signs & (1 << (s + 0))) ? -g1a : g1a),
+                                        d * ((signs & (1 << (s + 1))) ? -g1b : g1b));
+            a[4 * il + j + 2] = E::make(d * ((signs & (1 << (s + 4))) ? -g2a : g2a),
+                                        d * ((signs & (1 << (s + 5))) ? -g2b : g2b));
+        }
+    }
+}
+
+// values per stored block, so a row of K values is K/qk blocks
+template  struct fg_block_traits;
+template <> struct fg_block_traits { static constexpr int qk = QK4_NL; };
+template <> struct fg_block_traits  { static constexpr int qk = QK_K; };
+template <> struct fg_block_traits { static constexpr int qk = QK_K; };
+template <> struct fg_block_traits  { static constexpr int qk = QK_K; };
+template <> struct fg_block_traits { static constexpr int qk = QK_K; };
+template <> struct fg_block_traits  { static constexpr int qk = QK_K; };
+template <> struct fg_block_traits   { static constexpr int qk = QK_K; };
+template <> struct fg_block_traits   { static constexpr int qk = QK_K; };
+template <> struct fg_block_traits   { static constexpr int qk = QK_K; };
+template <> struct fg_block_traits    { static constexpr int qk = QK8_0; };
+template <> struct fg_block_traits    { static constexpr int qk = QK_K; };
+template <> struct fg_block_traits    { static constexpr int qk = QK_K; };
+template <> struct fg_block_traits    { static constexpr int qk = QK_K; };
+
+// The A stages below are the dequantize_block_iq* kernels rewritten for one k step. There a
+// work-item handled one quarter (il) of one 32-wide sub-block (ib); here one lane produces the
+// whole step, so il becomes a loop and ib is kb inside the superblock. Each quarter yields 8
+// consecutive values, i.e. 4 pairs at a[4*il], so nothing larger than 8 floats is ever live.
+#define FG_SUPERBLOCK(T)                                                         \
+    static_assert(QK_K == 256, "the " #T " A stage assumes 8 sub-blocks per superblock"); \
+    const T * blk = xrow + kb / (QK_K / 32);                                     \
+    const int ib  = kb % (QK_K / 32)
+
+template 
+static __dpct_inline__ void fg_pack_quarter(const float * __restrict__ t, typename E::pair * a, int il) {
+#pragma unroll
+    for (int j = 0; j < 4; ++j) {
+        a[4 * il + j] = E::make(t[2 * j], t[2 * j + 1]);
+    }
+}
+
+template 
+static __dpct_inline__ void fg_stage_a(const block_iq4_xs * __restrict__ xrow, const int kb, typename E::pair * a) {
+    FG_SUPERBLOCK(block_iq4_xs);
+    // low nibbles fill the first half of the step, high nibbles the second, so the two halves
+    // land at a[0..7] and a[8..15] and no quarter loop is needed
+    const float d = (float) blk->d *
+        ((((blk->scales_l[ib / 2] >> (4 * (ib % 2))) & 0xf) | (((blk->scales_h >> (2 * ib)) & 3) << 4)) - 32);
+    const uint8_t * q4 = blk->qs + 16 * ib;
+#pragma unroll
+    for (int j = 0; j < 8; ++j) {
+        a[j]     = E::make(d * kvalues_iq4nl[q4[2 * j] & 0xf], d * kvalues_iq4nl[q4[2 * j + 1] & 0xf]);
+        a[8 + j] = E::make(d * kvalues_iq4nl[q4[2 * j] >> 4],  d * kvalues_iq4nl[q4[2 * j + 1] >> 4]);
+    }
+}
+
+template 
+static __dpct_inline__ void fg_stage_a(const block_iq3_xxs * __restrict__ xrow, const int kb, typename E::pair * a) {
+    FG_SUPERBLOCK(block_iq3_xxs);
+    const uint8_t *  q3    = blk->qs + 8 * ib;
+    const uint16_t * gas   = (const uint16_t *) (blk->qs + QK_K / 4) + 2 * ib;
+    const uint32_t   aux32 = gas[0] | (gas[1] << 16);
+    const float      d     = (float) blk->d * (0.5f + (aux32 >> 28)) * 0.5f;
+#pragma unroll
+    for (int il = 0; il < 4; ++il) {
+        const uint8_t * grid1 = (const uint8_t *) (iq3xxs_grid + q3[2 * il + 0]);
+        const uint8_t * grid2 = (const uint8_t *) (iq3xxs_grid + q3[2 * il + 1]);
+        const uint8_t   signs = ksigns_iq2xs[(aux32 >> (7 * il)) & 127];
+        float t[8];
+#pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            t[j + 0] = d * grid1[j] * (signs & kmask_iq2xs[j + 0] ? -1.f : 1.f);
+            t[j + 4] = d * grid2[j] * (signs & kmask_iq2xs[j + 4] ? -1.f : 1.f);
+        }
+        fg_pack_quarter(t, a, il);
+    }
+}
+
+template 
+static __dpct_inline__ void fg_stage_a(const block_iq2_xxs * __restrict__ xrow, const int kb, typename E::pair * a) {
+    FG_SUPERBLOCK(block_iq2_xxs);
+    const uint16_t * q2    = blk->qs + 4 * ib;
+    const uint8_t *  aux8  = (const uint8_t *) q2;
+    const uint32_t   aux32 = q2[2] | (q2[3] << 16);
+    const float      d     = (float) blk->d * (0.5f + (aux32 >> 28)) * 0.25f;
+#pragma unroll
+    for (int il = 0; il < 4; ++il) {
+        const uint8_t * grid  = (const uint8_t *) (iq2xxs_grid + aux8[il]);
+        const uint8_t   signs = ksigns_iq2xs[(aux32 >> (7 * il)) & 127];
+        float t[8];
+#pragma unroll
+        for (int j = 0; j < 8; ++j) {
+            t[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
+        }
+        fg_pack_quarter(t, a, il);
+    }
+}
+
+template 
+static __dpct_inline__ void fg_stage_a(const block_iq2_xs * __restrict__ xrow, const int kb, typename E::pair * a) {
+    FG_SUPERBLOCK(block_iq2_xs);
+    const uint16_t * q2 = blk->qs + 4 * ib;
+#pragma unroll
+    for (int il = 0; il < 4; ++il) {
+        const uint8_t * grid  = (const uint8_t *) (iq2xs_grid + (q2[il] & 511));
+        const float     d     = (float) blk->d * (0.5f + ((blk->scales[ib] >> (4 * (il / 2))) & 0xf)) * 0.25f;
+        const uint8_t   signs = ksigns_iq2xs[q2[il] >> 9];
+        float t[8];
+#pragma unroll
+        for (int j = 0; j < 8; ++j) {
+            t[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
+        }
+        fg_pack_quarter(t, a, il);
+    }
+}
+
+template 
+static __dpct_inline__ void fg_stage_a(const block_iq2_s * __restrict__ xrow, const int kb, typename E::pair * a) {
+    FG_SUPERBLOCK(block_iq2_s);
+#pragma unroll
+    for (int il = 0; il < 4; ++il) {
+        const uint8_t * grid =
+            (const uint8_t *) (iq2s_grid + (blk->qs[4 * ib + il] | ((blk->qh[ib] << (8 - 2 * il)) & 0x300)));
+        const float   d     = (float) blk->d * (0.5f + ((blk->scales[ib] >> (4 * (il / 2))) & 0xf)) * 0.25f;
+        const uint8_t signs = blk->qs[QK_K / 8 + 4 * ib + il];
+        float t[8];
+#pragma unroll
+        for (int j = 0; j < 8; ++j) {
+            t[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
+        }
+        fg_pack_quarter(t, a, il);
+    }
+}
+
+template 
+static __dpct_inline__ void fg_stage_a(const block_iq1_s * __restrict__ xrow, const int kb, typename E::pair * a) {
+    FG_SUPERBLOCK(block_iq1_s);
+    const float delta = blk->qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA;
+    const float d     = (float) blk->d * (2 * ((blk->qh[ib] >> 12) & 7) + 1);
+#pragma unroll
+    for (int il = 0; il < 4; ++il) {
+        uint32_t       grid32[2];
+        const int8_t * q = (const int8_t *) grid32;
+        grid32[0] = iq1s_grid_gpu[blk->qs[4 * ib + il] | (((blk->qh[ib] >> (3 * il)) & 7) << 8)];
+        grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
+        grid32[0] &= 0x0f0f0f0f;
+        float t[8];
+#pragma unroll
+        for (int j = 0; j < 8; ++j) {
+            t[j] = d * (q[j] + delta);
+        }
+        fg_pack_quarter(t, a, il);
+    }
+}
+
+template 
+static __dpct_inline__ void fg_stage_a(const block_iq1_m * __restrict__ xrow, const int kb, typename E::pair * a) {
+    FG_SUPERBLOCK(block_iq1_m);
+    const uint16_t * sc = (const uint16_t *) blk->scales;
+    iq1m_scale_t     scale;
+    scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
+#pragma unroll
+    for (int il = 0; il < 4; ++il) {
+        const int   ib16  = 2 * ib + il / 2;
+        const float d     = (float) scale.f16 * (2 * ((sc[ib16 / 4] >> (3 * (ib16 % 4))) & 0x7) + 1);
+        const float delta = blk->qh[2 * ib + il / 2] & (0x08 << (4 * (il % 2))) ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA;
+        uint32_t       grid32[2];
+        const int8_t * q = (const int8_t *) grid32;
+        grid32[0] = iq1s_grid_gpu[blk->qs[4 * ib + il] |
+                                  (((blk->qh[2 * ib + il / 2] >> (4 * (il % 2))) & 7) << 8)];
+        grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
+        grid32[0] &= 0x0f0f0f0f;
+        float t[8];
+#pragma unroll
+        for (int j = 0; j < 8; ++j) {
+            t[j] = d * (q[j] + delta);
+        }
+        fg_pack_quarter(t, a, il);
+    }
+}
+
+
+// q8_0 and the k-quants. Each decode takes the fields of one 32-value sub-block, so the canonical
+// layout and the reorder (SoA) layout of reorder_qw() share it and differ only in where the fields
+// live, as in dequantize.hpp.
+template 
+static __dpct_inline__ void fg_decode_q8_0(const int8_t * __restrict__ qs, const float d, typename E::pair * a) {
+#pragma unroll
+    for (int j = 0; j < QK8_0 / 2; ++j) {
+        a[j] = E::make(d * qs[2 * j], d * qs[2 * j + 1]);
+    }
+}
+
+// same unpack as get_scale_min_k4() in dequantize.hpp
+static __dpct_inline__ void fg_scale_min_k4(const int j, const uint8_t * __restrict__ q, uint8_t & d, uint8_t & m) {
+    if (j < 4) {
+        d = q[j] & 63;
+        m = q[j + 4] & 63;
+    } else {
+        d = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4);
+        m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4);
+    }
+}
+
+// q4_K sub-block ib (0..7) uses scale/min pair ib and the low (even ib) or high (odd ib) nibbles of
+// qs[32 * (ib / 2) ...], as in dequantize_row_q4_K
+template 
+static __dpct_inline__ void fg_decode_q4_K(const uint8_t * __restrict__ qs, const uint8_t * __restrict__ scales,
+                                           const sycl::half2 dm, const int ib, typename E::pair * a) {
+    uint8_t sc, mb;
+    fg_scale_min_k4(ib, scales, sc, mb);
+    const float     d     = (float) dm[0] * sc;
+    const float     m     = (float) dm[1] * mb;
+    const uint8_t * q     = qs + 32 * (ib / 2);
+    const int       shift = 4 * (ib % 2);
+#pragma unroll
+    for (int j = 0; j < 16; ++j) {
+        a[j] = E::make(d * ((q[2 * j] >> shift) & 0xF) - m, d * ((q[2 * j + 1] >> shift) & 0xF) - m);
+    }
+}
+
+// q5_K: q4_K plus one high bit per value, bit ib of qh[l]
+template 
+static __dpct_inline__ void fg_decode_q5_K(const uint8_t * __restrict__ qs, const uint8_t * __restrict__ qh,
+                                           const uint8_t * __restrict__ scales, const sycl::half2 dm, const int ib,
+                                           typename E::pair * a) {
+    uint8_t sc, mb;
+    fg_scale_min_k4(ib, scales, sc, mb);
+    const float     d     = (float) dm[0] * sc;
+    const float     m     = (float) dm[1] * mb;
+    const uint8_t * q     = qs + 32 * (ib / 2);
+    const int       shift = 4 * (ib % 2);
+#pragma unroll
+    for (int j = 0; j < 16; ++j) {
+        const int l = 2 * j;
+        a[j] = E::make(d * (((q[l] >> shift) & 0xF) | (((qh[l] >> ib) & 1) << 4)) - m,
+                       d * (((q[l + 1] >> shift) & 0xF) | (((qh[l + 1] >> ib) & 1) << 4)) - m);
+    }
+}
+
+// q6_K: sub-block ib is quarter r = ib % 4 of half h = ib / 4. The half selects ql + 64h, qh + 32h
+// and scales + 8h; the quarter selects ql + 32(r & 1), the ql nibble r / 2, the qh bit pair r and
+// scales + 2r, as in dequantize_row_q6_K. The scale changes at value 16 of the sub-block.
+template 
+static __dpct_inline__ void fg_decode_q6_K(const uint8_t * __restrict__ ql, const uint8_t * __restrict__ qh,
+                                           const int8_t * __restrict__ scales, const float d, const int ib,
+                                           typename E::pair * a) {
+    const int       h  = ib / 4;
+    const int       r  = ib % 4;
+    const uint8_t * q  = ql + 64 * h + 32 * (r & 1);
+    const uint8_t * hb = qh + 32 * h;
+    const int8_t *  sc = scales + 8 * h + 2 * r;
+#pragma unroll
+    for (int j = 0; j < 16; ++j) {
+        const int   l  = 2 * j;
+        const float dl = d * sc[j / 8];
+        const int   q0 = (((q[l] >> (4 * (r / 2))) & 0xF) | (((hb[l] >> (2 * r)) & 3) << 4)) - 32;
+        const int   q1 = (((q[l + 1] >> (4 * (r / 2))) & 0xF) | (((hb[l + 1] >> (2 * r)) & 3) << 4)) - 32;
+        a[j] = E::make(dl * q0, dl * q1);
+    }
+}
+
+template 
+static __dpct_inline__ void fg_stage_a(const block_q8_0 * __restrict__ xrow, const int kb, typename E::pair * a) {
+    fg_decode_q8_0(xrow[kb].qs, (float) xrow[kb].d, a);
+}
+
+template 
+static __dpct_inline__ void fg_stage_a(const block_q4_K * __restrict__ xrow, const int kb, typename E::pair * a) {
+    FG_SUPERBLOCK(block_q4_K);
+    fg_decode_q4_K(blk->qs, blk->scales, blk->dm, ib, a);
+}
+
+template 
+static __dpct_inline__ void fg_stage_a(const block_q5_K * __restrict__ xrow, const int kb, typename E::pair * a) {
+    FG_SUPERBLOCK(block_q5_K);
+    fg_decode_q5_K(blk->qs, blk->qh, blk->scales, blk->dm, ib, a);
+}
+
+template 
+static __dpct_inline__ void fg_stage_a(const block_q6_K * __restrict__ xrow, const int kb, typename E::pair * a) {
+    FG_SUPERBLOCK(block_q6_K);
+    fg_decode_q6_K(blk->ql, blk->qh, blk->scales, (float) blk->d, ib, a);
+}
+
+// Reorder (SoA) layout: each block field is one stream over the nblocks of the matrix (of the expert
+// slice for MUL_MAT_ID), in the order reorder_qw() writes them. Block ib holds k step kb.
+template  struct fg_soa {
+    static constexpr bool supported = false;
+};
+
+template <> struct fg_soa {
+    static constexpr bool supported = true;
+    // [qs][d]
+    template 
+    static __dpct_inline__ void stage(const uint8_t * x, const size_t nblocks, const size_t ib, const int,
+                                      typename E::pair * a) {
+        const float d = (float) ((const sycl::half *) (x + nblocks * QK8_0))[ib];
+        fg_decode_q8_0((const int8_t *) x + ib * QK8_0, d, a);
+    }
+};
+
+template <> struct fg_soa {
+    static constexpr bool supported = true;
+    // [qs][scales][dm]
+    template 
+    static __dpct_inline__ void stage(const uint8_t * x, const size_t nblocks, const size_t ib, const int kb,
+                                      typename E::pair * a) {
+        const uint8_t *   scales = x + nblocks * (QK_K / 2);
+        const sycl::half2 dm     = ((const sycl::half2 *) (scales + nblocks * K_SCALE_SIZE))[ib];
+        fg_decode_q4_K(x + ib * (QK_K / 2), scales + ib * K_SCALE_SIZE, dm, kb % (QK_K / 32), a);
+    }
+};
+
+template <> struct fg_soa {
+    static constexpr bool supported = true;
+    // [qs][qh][scales][dm]
+    template 
+    static __dpct_inline__ void stage(const uint8_t * x, const size_t nblocks, const size_t ib, const int kb,
+                                      typename E::pair * a) {
+        const uint8_t *   qh     = x + nblocks * (QK_K / 2);
+        const uint8_t *   scales = qh + nblocks * (QK_K / 8);
+        const sycl::half2 dm     = ((const sycl::half2 *) (scales + nblocks * K_SCALE_SIZE))[ib];
+        fg_decode_q5_K(x + ib * (QK_K / 2), qh + ib * (QK_K / 8), scales + ib * K_SCALE_SIZE, dm,
+                          kb % (QK_K / 32), a);
+    }
+};
+
+template <> struct fg_soa {
+    static constexpr bool supported = true;
+    // [ql][qh][scales][d]
+    template 
+    static __dpct_inline__ void stage(const uint8_t * x, const size_t nblocks, const size_t ib, const int kb,
+                                      typename E::pair * a) {
+        const uint8_t * qh     = x + nblocks * (QK_K / 2);
+        const uint8_t * scales = qh + nblocks * (QK_K / 4);
+        const float     d      = (float) ((const sycl::half *) (scales + nblocks * (QK_K / 16)))[ib];
+        fg_decode_q6_K(x + ib * (QK_K / 2), qh + ib * (QK_K / 4), (const int8_t *) scales + ib * (QK_K / 16), d,
+                          kb % (QK_K / 32), a);
+    }
+};
+
+
+// one SG_ROWS x BN output tile: B columns [b0, b0 + BN) of packed_b go to dst columns [n0, n1),
+// n1 - n0 <= BN
+template 
+static void fused_dequant_gemm_tile(
+    const block_q_t * __restrict__ x,
+    const typename S::tsb * __restrict__ packed_b,
+    float * __restrict__ dst,
+    const int M, const int Npad, const int K, const int ldd,
+    const int b0, const int n0, const int n1,
+    sycl::local_accessor tile_a,
+    sycl::local_accessor tile_c,
+    const sycl::nd_item<2> & item) {
+    if constexpr (fg_built()) {
+        using EA = typename S::EA;
+        using TA = typename S::ta;
+        using TB = typename S::tb;
+        const auto sg     = item.get_sub_group();
+        const int  sg_id  = sg.get_group_id()[0];
+        const int  lane   = sg.get_local_id()[0];
+        const int  m0     = item.get_group(1) * S::SG_ROWS;
+        const int  nstep  = K / FG_BK;
+        const int  a_base = sg_id * S::SG_ROWS * FG_BK;
+        const int  c_base = sg_id * S::SG_ROWS * S::BN;
+
+        mx::joint_matrix acc[S::MT][S::NT];
+#pragma unroll
+        for (int mt = 0; mt < S::MT; ++mt) {
+#pragma unroll
+            for (int nt = 0; nt < S::NT; ++nt) {
+                mx::joint_matrix_fill(sg, acc[mt][nt], 0.0f);
+            }
+        }
+
+        // lane decodes rows lane, lane + SG, ... of the sub-group's SG_ROWS
+        constexpr int         KPB     = fg_block_traits::qk / FG_BK; // k steps per block
+        const size_t          bpr     = K / fg_block_traits::qk;
+        const size_t          nblocks = (size_t) M * bpr;
+        size_t                row_blk[S::RPL];
+        bool                  row_ok[S::RPL];
+        typename EA::pair *   a[S::RPL];
+#pragma unroll
+        for (int r = 0; r < S::RPL; ++r) {
+            const int row = m0 + r * S::SG + lane;
+            row_ok[r]  = row < M;
+            row_blk[r] = (size_t) (row_ok[r] ? row : 0) * bpr;
+            a[r]       = (typename EA::pair *) &tile_a[a_base + (r * S::SG + lane) * FG_BK];
+        }
+
+        const auto b_ptr = sycl::address_space_cast(packed_b);
+        const int b_stride = Npad * S::VNNI;
+
+        const int kb_begin = (sg_id * nstep) / FG_KSPLIT;
+        const int kb_end   = ((sg_id + 1) * nstep) / FG_KSPLIT;
+        for (int kb = kb_begin; kb < kb_end; ++kb) {
+#pragma unroll
+            for (int r = 0; r < S::RPL; ++r) {
+                if (row_ok[r]) {
+                    if constexpr (reordered) {
+                        fg_soa::template stage((const uint8_t *) x, nblocks, row_blk[r] + kb / KPB, kb,
+                                                              a[r]);
+                    } else {
+                        fg_stage_a(x + row_blk[r], kb, a[r]);
+                    }
+                } else {
+#pragma unroll
+                    for (int j = 0; j < FG_BK / 2; ++j) {
+                        a[r][j] = EA::make(0.0f, 0.0f);
+                    }
+                }
+            }
+            sycl::group_barrier(sg);
+
+#pragma unroll
+            for (int kt = 0; kt < FG_BK / S::TK; ++kt) {
+                const int kq0 = (kb * FG_BK + kt * S::TK) / S::VNNI;
+                mx::joint_matrix sub_b[S::NT];
+#pragma unroll
+                for (int nt = 0; nt < S::NT; ++nt) {
+                    mx::joint_matrix_load(sg, sub_b[nt], b_ptr + (size_t) kq0 * b_stride + (b0 + nt * S::TN) * S::VNNI, b_stride);
+                }
+#pragma unroll
+                for (int mt = 0; mt < S::MT; ++mt) {
+                    mx::joint_matrix sub_a;
+                    mx::joint_matrix_load(sg, sub_a,
+                        tile_a.template get_multi_ptr() + a_base + (mt * S::TM) * FG_BK + kt * S::TK,
+                        FG_BK);
+#pragma unroll
+                    for (int nt = 0; nt < S::NT; ++nt) {
+                        mx::joint_matrix_mad(sg, acc[mt][nt], sub_a, sub_b[nt], acc[mt][nt]);
+                    }
+                }
+            }
+            // the next step overwrites tile_a
+            sycl::group_barrier(sg);
+        }
+
+#pragma unroll
+        for (int mt = 0; mt < S::MT; ++mt) {
+#pragma unroll
+            for (int nt = 0; nt < S::NT; ++nt) {
+                mx::joint_matrix_store(sg, acc[mt][nt],
+                    tile_c.template get_multi_ptr() + c_base + (mt * S::TM) * S::BN + nt * S::TN,
+                    S::BN, mx::layout::row_major);
+            }
+        }
+        sycl::group_barrier(item.get_group());
+
+        // sum the K splits; consecutive lanes write consecutive rows of one dst column
+        for (int idx = item.get_local_linear_id(); idx < S::SG_ROWS * S::BN; idx += S::WG_SIZE) {
+            const int r = idx % S::SG_ROWS;
+            const int c = idx / S::SG_ROWS;
+            const int m = m0 + r;
+            const int n = n0 + c;
+            if (m < M && n < n1) {
+                float sum = 0.0f;
+#pragma unroll
+                for (int s = 0; s < FG_KSPLIT; ++s) {
+                    sum += tile_c[s * S::SG_ROWS * S::BN + r * S::BN + c];
+                }
+                dst[(size_t) n * ldd + m] = sum;
+            }
+        }
+    }
+}
+
+template 
+static void fused_dequant_gemm_launch(const void * src0, const typename S::tsb * packed, float * dst, const int M,
+                                      const int N, const int Npad, const int K, const int ldd,
+                                      const int64_t groups_n, const int64_t groups_m, dpct::queue_ptr stream) {
+    stream->submit([&](sycl::handler & cgh) {
+        sycl::local_accessor tile_a(FG_KSPLIT * S::SG_ROWS * FG_BK, cgh);
+        sycl::local_accessor          tile_c(FG_KSPLIT * S::SG_ROWS * S::BN, cgh);
+        cgh.parallel_for(
+            sycl::nd_range<2>(sycl::range<2>(groups_n, groups_m * S::WG_SIZE), sycl::range<2>(1, S::WG_SIZE)),
+            [=](sycl::nd_item<2> item) [[sycl::reqd_sub_group_size(S::SG)]] {
+                const int n0 = item.get_group(0) * S::BN;
+                fused_dequant_gemm_tile((const block_q_t *) src0, packed, dst, M, Npad, K, ldd,
+                                                      n0, n0, N, tile_a, tile_c, item);
+            });
+    });
+}
+
+// grouped: work-group (t, mt) is tile t of the schedule; its B columns sit at t * BN
+template 
+static void grouped_dequant_gemm_launch(const char * src0_dd, const size_t expert_stride,
+                                        const ggml_sycl_gg_tile * tiles_ptr, const typename S::tsb * packed, float * dst,
+                                        const int M, const int Npad, const int K, const int64_t n_tiles,
+                                        const int64_t groups_m, dpct::queue_ptr stream) {
+    stream->submit([&](sycl::handler & cgh) {
+        sycl::local_accessor tile_a(FG_KSPLIT * S::SG_ROWS * FG_BK, cgh);
+        sycl::local_accessor          tile_c(FG_KSPLIT * S::SG_ROWS * S::BN, cgh);
+        cgh.parallel_for(
+            sycl::nd_range<2>(sycl::range<2>(n_tiles, groups_m * S::WG_SIZE), sycl::range<2>(1, S::WG_SIZE)),
+            [=](sycl::nd_item<2> item) [[sycl::reqd_sub_group_size(S::SG)]] {
+                const int               t    = item.get_group(0);
+                const ggml_sycl_gg_tile tile = tiles_ptr[t];
+                const block_q_t *       x    = (const block_q_t *) (src0_dd + (size_t) tile.expert * expert_stride);
+                fused_dequant_gemm_tile(x, packed, dst, M, Npad, K, M, t * S::BN, tile.n0,
+                                                                 tile.n1, tile_a, tile_c, item);
+            });
+    });
+}
+
+// src1 f32 rows -> packed [K/V][n_tiles*BN][V], tile t holds its rows [n0, n1) at columns t*BN..,
+// zero past n1. The column runs fastest so a sub-group writes one contiguous run.
+template 
+static void grouped_gemm_pack_b(const float * y, typename S::tsb * packed, const ggml_sycl_gg_tile * tiles, int Npad,
+                                int K, dpct::queue_ptr stream) {
+    using E        = typename S::EB;
+    constexpr int V = S::VNNI;
+    const int kqs   = K / V;
+    stream->parallel_for(sycl::range<1>((size_t) Npad * kqs), [=](sycl::id<1> id) {
+        const size_t idx = id[0];
+        const int    kq  = idx / Npad;
+        const int    n   = idx - (size_t) kq * Npad;
+        const ggml_sycl_gg_tile tile = tiles[n / S::BN];
+        const int    row = tile.n0 + n % S::BN;
+        // one guarded load run per work-item, as a per-element select costs ~1.5% prefill
+        typename S::tsb vals[V] = {};
+        if (row < tile.n1) {
+            const float * src = y + (size_t) row * K + V * kq;
+#pragma unroll
+            for (int v = 0; v < V; ++v) {
+                vals[v] = E::cvt(src[v]);
+            }
+        }
+        typename S::tsb * out = packed + ((size_t) kq * Npad + n) * V;
+#pragma unroll
+        for (int v = 0; v < V; ++v) {
+            out[v] = vals[v];
+        }
+    });
+}
+
+// q8_0 and the k-quants take only the grouped path. The plain kernel decodes A again for every BN columns
+// of a dense batch, and for these formats that costs more than the one dequantization of the library GEMM.
+template  static constexpr bool fg_plain_ok() {
+    return !std::is_same_v && !std::is_same_v && !std::is_same_v &&
+           !std::is_same_v;
+}
+
+template  struct fg_tag {
+    using type                     = T;
+    static constexpr bool reordered = R;
+};
+
+// calls f(fg_tag{}) for the weight format and layout; false if it has no A stage
+template  static bool fg_visit_layout(bool reordered, F && f) {
+    if (!reordered) {
+        f(fg_tag{});
+        return true;
+    }
+    if constexpr (fg_soa::supported) {
+        f(fg_tag{});
+        return true;
+    }
+    return false;
+}
+
+template  static bool fg_visit_type(ggml_type type, bool reordered, F && f) {
+    switch (type) {
+        case GGML_TYPE_IQ4_NL:  return fg_visit_layout(reordered, f);
+        case GGML_TYPE_IQ3_S:   return fg_visit_layout(reordered, f);
+        case GGML_TYPE_IQ4_XS:  return fg_visit_layout(reordered, f);
+        case GGML_TYPE_IQ3_XXS: return fg_visit_layout(reordered, f);
+        case GGML_TYPE_IQ2_XXS: return fg_visit_layout(reordered, f);
+        case GGML_TYPE_IQ2_XS:  return fg_visit_layout(reordered, f);
+        case GGML_TYPE_IQ2_S:   return fg_visit_layout(reordered, f);
+        case GGML_TYPE_IQ1_S:   return fg_visit_layout(reordered, f);
+        case GGML_TYPE_IQ1_M:   return fg_visit_layout(reordered, f);
+        case GGML_TYPE_Q8_0:    return fg_visit_layout(reordered, f);
+        case GGML_TYPE_Q4_K:    return fg_visit_layout(reordered, f);
+        case GGML_TYPE_Q5_K:    return fg_visit_layout(reordered, f);
+        case GGML_TYPE_Q6_K:    return fg_visit_layout(reordered, f);
+        default:                return false;
+    }
+}
+
+template 
+static bool fg_fused_run(ggml_type src0_type, const void * src0, const void * src1, ggml_type src1_type,
+                         float * dst, int64_t M, int64_t N, int64_t K, int64_t ldd, ggml_sycl_pool & pool,
+                         dpct::queue_ptr stream) {
+    const int64_t groups_n = (N + S::BN - 1) / S::BN;
+    const int64_t groups_m = (M + S::SG_ROWS - 1) / S::SG_ROWS;
+    const int     Npad     = (int) (groups_n * S::BN);
+
+    // src1 is read in its own type: one pass, converted in registers only if B differs
+    ggml_sycl_pool_alloc packed_b(pool, (size_t) K * Npad);
+    if (src1_type == GGML_TYPE_F16) {
+        fused_gemm_pack_b((const sycl::half *) src1, packed_b.get(), (int) N, Npad, (int) K, stream);
+    } else if (src1_type == GGML_TYPE_BF16) {
+        fused_gemm_pack_b((const fg_bf16 *) src1, packed_b.get(), (int) N, Npad, (int) K, stream);
+    } else {
+        fused_gemm_pack_b((const float *) src1, packed_b.get(), (int) N, Npad, (int) K, stream);
+    }
+
+    const typename S::tsb * packed = packed_b.get();
+    return fg_visit_type(src0_type, false, [&](auto tag) {
+        using block_q_t = typename decltype(tag)::type;
+        if constexpr (fg_plain_ok()) {
+            fused_dequant_gemm_launch(src0, packed, dst, (int) M, (int) N, Npad, (int) K, (int) ldd,
+                                                    groups_n, groups_m, stream);
+        }
+    });
+}
+
+bool ggml_sycl_fused_dequant_gemm(ggml_type src0_type, const void * src0, const void * src1, ggml_type src1_type,
+                                  int32_t src1_prec, float * dst, int64_t M, int64_t N, int64_t K, int64_t ldd,
+                                  ggml_sycl_pool & pool, dpct::queue_ptr stream) {
+    // every BN columns dequantize A again, so wide N is left to the library GEMM
+    bool plain_ok = false;
+    fg_visit_type(src0_type, false, [&](auto tag) { plain_ok = fg_plain_ok(); });
+    if (g_ggml_sycl_dynamic_precision == GGML_SYCL_DYNAMIC_PRECISION_F32 ||
+        !ggml_sycl_xmx_gather_type_enabled(src0_type) || !plain_ok) {
+        return false;
+    }
+    if (src1_type != GGML_TYPE_F32 && src1_type != GGML_TYPE_F16 && src1_type != GGML_TYPE_BF16) {
+        return false;
+    }
+    if (!ggml_sycl_fused_dequant_gemm_shape_ok(src0_type, M, N, K, ldd)) {
+        return false;
+    }
+    const int combo = fg_pick_combo(stream, src1_type, src1_prec);
+    if (combo < 0) {
+        return false;
+    }
+    bool launched = false;
+    fg_visit_combo(combo, [&](auto s) {
+        launched = fg_fused_run(src0_type, src0, src1, src1_type, dst, M, N, K, ldd, pool, stream);
+    });
+    return launched;
+}
+
+template 
+static bool fg_grouped_run(ggml_type src0_type, bool reordered, const void * src0_base, size_t expert_stride,
+                           const float * src1, float * dst, const int64_t * expert_row_offsets, int64_t n_as, int64_t M,
+                           int64_t K, std::vector & tiles, ggml_sycl_pool & pool,
+                           dpct::queue_ptr stream) {
+    // the host knows every slice, so it lays out the work-groups: no search on the device
+    tiles.clear();
+    for (int64_t e = 0; e < n_as; ++e) {
+        const int64_t end = expert_row_offsets[e + 1];
+        for (int64_t n0 = expert_row_offsets[e]; n0 < end; n0 += S::BN) {
+            tiles.push_back({ (int32_t) e, (int32_t) n0, (int32_t) std::min(n0 + S::BN, end) });
+        }
+    }
+    const int64_t n_tiles  = tiles.size();
+    const int64_t groups_m = (M + S::SG_ROWS - 1) / S::SG_ROWS;
+    const int     Npad     = (int) (n_tiles * S::BN);
+
+    const int64_t max_tiles = grouped_gemm_max_tiles(expert_row_offsets[n_as], n_as, S::BN);
+    GGML_ASSERT(n_tiles <= max_tiles);
+    ggml_sycl_pool_alloc tiles_dev(pool, max_tiles);
+    SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy(tiles_dev.get(), tiles.data(), n_tiles * sizeof(ggml_sycl_gg_tile))));
+
+    ggml_sycl_pool_alloc packed_b(pool, (size_t) K * max_tiles * S::BN);
+    grouped_gemm_pack_b(src1, packed_b.get(), tiles_dev.get(), Npad, (int) K, stream);
+
+    const typename S::tsb *   packed    = packed_b.get();
+    const ggml_sycl_gg_tile * tiles_ptr = tiles_dev.get();
+    const char *              src0_dd   = (const char *) src0_base;
+    return fg_visit_type(src0_type, reordered, [&](auto tag) {
+        using T = decltype(tag);
+        grouped_dequant_gemm_launch(src0_dd, expert_stride, tiles_ptr, packed, dst,
+                                                                       (int) M, Npad, (int) K, n_tiles, groups_m,
+                                                                       stream);
+    });
+}
+
+bool ggml_sycl_grouped_dequant_gemm(ggml_type src0_type, bool reordered, const void * src0_base, size_t expert_stride,
+                                    const float * src1, int32_t src1_prec, float * dst,
+                                    const int64_t * expert_row_offsets, int64_t n_as, int64_t M, int64_t K,
+                                    int64_t total_rows, std::vector & tiles,
+                                    ggml_sycl_pool & pool, dpct::queue_ptr stream) {
+    int64_t n_active = 0;
+    for (int64_t e = 0; e < n_as; ++e) {
+        n_active += expert_row_offsets[e + 1] > expert_row_offsets[e];
+    }
+    if (g_ggml_sycl_dynamic_precision == GGML_SYCL_DYNAMIC_PRECISION_F32 ||
+        !ggml_sycl_xmx_gather_type_enabled(src0_type) || !fg_visit_type(src0_type, reordered, [](auto) {})) {
+        return false;
+    }
+    if (!ggml_sycl_grouped_dequant_gemm_shape_ok(src0_type, M, K, total_rows, n_active)) {
+        return false;
+    }
+    const int combo = fg_pick_combo(stream, GGML_TYPE_F32, src1_prec);
+    if (combo < 0) {
+        return false;
+    }
+    bool launched = false;
+    fg_visit_combo(combo, [&](auto s) {
+        launched = fg_grouped_run(src0_type, reordered, src0_base, expert_stride, src1, dst,
+                                               expert_row_offsets, n_as, M, K, tiles, pool, stream);
+    });
+    return launched;
+}
diff --git a/ggml/src/ggml-sycl/fused-gemm.hpp b/ggml/src/ggml-sycl/fused-gemm.hpp
new file mode 100644
index 000000000..298989204
--- /dev/null
+++ b/ggml/src/ggml-sycl/fused-gemm.hpp
@@ -0,0 +1,88 @@
+#ifndef GGML_SYCL_FUSED_GEMM_HPP
+#define GGML_SYCL_FUSED_GEMM_HPP
+
+#include "common.hpp"
+
+
+// Shape and type gates for the kernels below. Device capability is separate: it needs a queue to ask.
+static constexpr int GGML_SYCL_FG_MAX_N = 64; // widest N taken; each shape covers it in BN-wide tiles
+
+// weight formats the fused A stage decodes; K must cover whole stored blocks
+constexpr bool ggml_sycl_fused_dequant_gemm_type_ok(ggml_type src0_type, int64_t K) {
+    // iq4_nl and q8_0 store 32 values per block; every other format here is a 256-value superblock
+    // that the A stage walks in steps of 32, so K must cover whole superblocks.
+    if (src0_type == GGML_TYPE_IQ4_NL || src0_type == GGML_TYPE_Q8_0) {
+        return K % 32 == 0;
+    }
+    const bool superblock =
+           src0_type == GGML_TYPE_Q4_K ||
+           src0_type == GGML_TYPE_Q5_K ||
+           src0_type == GGML_TYPE_Q6_K ||
+           src0_type == GGML_TYPE_IQ3_S ||
+           src0_type == GGML_TYPE_IQ4_XS ||
+           src0_type == GGML_TYPE_IQ3_XXS ||
+           src0_type == GGML_TYPE_IQ2_XXS ||
+           src0_type == GGML_TYPE_IQ2_XS ||
+           src0_type == GGML_TYPE_IQ2_S ||
+           src0_type == GGML_TYPE_IQ1_S ||
+           src0_type == GGML_TYPE_IQ1_M;
+    return superblock && QK_K == 256 && K % QK_K == 0;
+}
+
+constexpr bool ggml_sycl_fused_dequant_gemm_shape_ok(ggml_type src0_type, int64_t M, int64_t N, int64_t K,
+                                                     int64_t ldd) {
+    return ggml_sycl_fused_dequant_gemm_type_ok(src0_type, K) && M > 0 && N > 0 && K > 0 &&
+           N <= GGML_SYCL_FG_MAX_N &&
+           M <= INT32_MAX && N <= INT32_MAX && K <= INT32_MAX && ldd <= INT32_MAX;
+}
+
+// grouped variant: the per-expert fused kernel is only worth it while each expert is narrow,
+// so wider average slices are left to the per-expert library GEMM loop
+constexpr bool ggml_sycl_grouped_dequant_gemm_shape_ok(ggml_type src0_type, int64_t M, int64_t K,
+                                                       int64_t total_rows, int64_t n_active) {
+    return ggml_sycl_fused_dequant_gemm_shape_ok(src0_type, M, 1, K, M) && total_rows > 0 &&
+           total_rows <= INT32_MAX && total_rows <= n_active * GGML_SYCL_FG_MAX_N;
+}
+
+// Runtime type gate, kept out of the constexpr predicates above so those stay pure.
+inline bool ggml_sycl_xmx_gather_type_enabled(ggml_type src0_type) {
+    switch (src0_type) {
+        case GGML_TYPE_IQ4_NL:  return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_IQ4_NL  ) != 0;
+        case GGML_TYPE_IQ3_S:   return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_IQ3_S   ) != 0;
+        case GGML_TYPE_IQ4_XS:  return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_IQ4_XS  ) != 0;
+        case GGML_TYPE_IQ3_XXS: return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_IQ3_XXS ) != 0;
+        case GGML_TYPE_IQ2_XXS: return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_IQ2_XXS ) != 0;
+        case GGML_TYPE_IQ2_XS:  return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_IQ2_XS  ) != 0;
+        case GGML_TYPE_IQ2_S:   return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_IQ2_S   ) != 0;
+        case GGML_TYPE_IQ1_S:   return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_IQ1_S   ) != 0;
+        case GGML_TYPE_IQ1_M:   return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_IQ1_M   ) != 0;
+        case GGML_TYPE_Q8_0:    return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_Q8_0    ) != 0;
+        case GGML_TYPE_Q4_K:    return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_Q4_K    ) != 0;
+        case GGML_TYPE_Q5_K:    return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_Q5_K    ) != 0;
+        case GGML_TYPE_Q6_K:    return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_Q6_K    ) != 0;
+        default:          return false;
+    }
+}
+
+// dst[n*ldd + m] = sum_k dequant(src0)[m*K + k] * src1[n*K + k], src1 is F32, F16 or BF16.
+// The XMX combination is picked per call from the src1 type and its precision request src1_prec
+// (op_params[3], [TAG_GGML_PREC]); the accumulator is f32, which meets any request.
+// q8_0 and the k-quants are not handled here, only in the grouped path below.
+// Returns false when the case is not handled (type, device, precision, or shape).
+bool ggml_sycl_fused_dequant_gemm(ggml_type src0_type, const void * src0, const void * src1, ggml_type src1_type,
+                                  int32_t src1_prec, float * dst, int64_t M, int64_t N, int64_t K, int64_t ldd,
+                                  ggml_sycl_pool & pool, dpct::queue_ptr stream);
+
+// One launch for every expert of a MUL_MAT_ID: rows of src1/dst are grouped by expert, expert e
+// owns rows [expert_row_offsets[e], expert_row_offsets[e+1]) and reads its weights at
+// src0_base + e*expert_stride. tiles is host scratch that must stay alive until the queue drains.
+// reordered: each expert slice is in the reorder (SoA) layout of reorder_qw().
+// dst[n*M + m] = sum_k dequant(src0_e)[m*K + k] * src1[n*K + k]
+// Returns false when the case is not handled (type, layout, device, precision, or shape).
+bool ggml_sycl_grouped_dequant_gemm(ggml_type src0_type, bool reordered, const void * src0_base, size_t expert_stride,
+                                    const float * src1, int32_t src1_prec, float * dst,
+                                    const int64_t * expert_row_offsets, int64_t n_as, int64_t M, int64_t K,
+                                    int64_t total_rows, std::vector & tiles,
+                                    ggml_sycl_pool & pool, dpct::queue_ptr stream);
+
+#endif // GGML_SYCL_FUSED_GEMM_HPP
diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp
index b4d200526..ebca92886 100644
--- a/ggml/src/ggml-sycl/ggml-sycl.cpp
+++ b/ggml/src/ggml-sycl/ggml-sycl.cpp
@@ -14,6 +14,7 @@
 #include 
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -60,6 +61,7 @@
 #include "ggml-sycl/common.hpp"
 #include "ggml-sycl/element_wise.hpp"
 #include "ggml-sycl/fwht.hpp"
+#include "ggml-sycl/fused-gemm.hpp"
 #include "ggml-sycl/gemm.hpp"
 #include "ggml-sycl/getrows.hpp"
 #include "ggml-sycl/mem.hpp"
@@ -105,6 +107,30 @@ int g_ggml_sycl_enable_fusion = 1;
 int g_ggml_sycl_enable_esimd = 1;
 int g_ggml_sycl_mmvq_wide = 1;
 int g_ggml_sycl_prioritize_dmmv = 0;
+int g_ggml_sycl_xmx_gather_types = GGML_SYCL_XMX_GATHER_TYPES_DEFAULT;
+int g_ggml_sycl_xmx_gather_shapes = GGML_SYCL_XMX_GATHER_SHAPES_DEFAULT;
+int g_ggml_sycl_dynamic_precision = GGML_SYCL_DYNAMIC_PRECISION_DEFAULT;
+int g_ggml_sycl_dynamic_required_precision = GGML_SYCL_DYNAMIC_PRECISION_F32;
+static const char * ggml_sycl_dynamic_precision_names[] = { "F16", "BF16", "TF32", "F32" };
+
+// value of a GGML_SYCL_DYNAMIC_PRECISION-style variable; def if unset or invalid
+static int ggml_sycl_get_env_precision(const char * name, int def) {
+    const char * env = getenv(name);
+    if (!env) {
+        return def;
+    }
+    std::string mode(env);
+    for (char & c : mode) {
+        c = (char) std::toupper((unsigned char) c);
+    }
+    for (int i = GGML_SYCL_DYNAMIC_PRECISION_F16; i <= GGML_SYCL_DYNAMIC_PRECISION_F32; i++) {
+        if (mode == ggml_sycl_dynamic_precision_names[i]) {
+            return i;
+        }
+    }
+    GGML_LOG_WARN("%s: unknown %s=%s, using %s\n", __func__, name, env, ggml_sycl_dynamic_precision_names[def]);
+    return def;
+}
 int g_ggml_sycl_use_async_mem_op = 0;
 int g_ggml_sycl_use_async_mem_op_requested = 1;
 int g_ggml_sycl_use_level_zero_api = 0;
@@ -401,6 +427,12 @@ static void ggml_check_sycl() try {
         g_ggml_sycl_enable_esimd = ggml_sycl_get_env("GGML_SYCL_ENABLE_ESIMD", 1);
         g_ggml_sycl_mmvq_wide = ggml_sycl_get_env("GGML_SYCL_MMVQ_WIDE", 1);
         g_ggml_sycl_prioritize_dmmv = ggml_sycl_get_env("GGML_SYCL_PRIORITIZE_DMMV", 0);
+        g_ggml_sycl_xmx_gather_types = ggml_sycl_get_env("GGML_SYCL_XMX_GATHER_TYPES", GGML_SYCL_XMX_GATHER_TYPES_DEFAULT);
+        g_ggml_sycl_xmx_gather_shapes = ggml_sycl_get_env("GGML_SYCL_XMX_GATHER_SHAPES", GGML_SYCL_XMX_GATHER_SHAPES_DEFAULT);
+        g_ggml_sycl_dynamic_precision =
+            ggml_sycl_get_env_precision("GGML_SYCL_DYNAMIC_PRECISION", GGML_SYCL_DYNAMIC_PRECISION_DEFAULT);
+        g_ggml_sycl_dynamic_required_precision =
+            ggml_sycl_get_env_precision("GGML_SYCL_DYNAMIC_REQUIRED_PRECISION", GGML_SYCL_DYNAMIC_PRECISION_F32);
 
 #ifdef GGML_SYCL_SUPPORT_LEVEL_ZERO_API
         g_ggml_sycl_use_level_zero_api = ggml_sycl_get_env("GGML_SYCL_USE_LEVEL_ZERO_API", 1);
@@ -509,6 +541,12 @@ static void ggml_check_sycl() try {
 #endif
 
         GGML_LOG_INFO("  GGML_SYCL_ENABLE_OPT: %d\n", g_ggml_sycl_enable_optimize);
+        GGML_LOG_INFO("  GGML_SYCL_XMX_GATHER_TYPES: %d\n", g_ggml_sycl_xmx_gather_types);
+        GGML_LOG_INFO("  GGML_SYCL_XMX_GATHER_SHAPES: %d\n", g_ggml_sycl_xmx_gather_shapes);
+        GGML_LOG_INFO("  GGML_SYCL_DYNAMIC_PRECISION: %s\n",
+                      ggml_sycl_dynamic_precision_names[g_ggml_sycl_dynamic_precision]);
+        GGML_LOG_INFO("  GGML_SYCL_DYNAMIC_REQUIRED_PRECISION: %s\n",
+                      ggml_sycl_dynamic_precision_names[g_ggml_sycl_dynamic_required_precision]);
 
 #if defined(GGML_SYCL_SUPPORT_VMM)
         GGML_LOG_INFO("  GGML_SYCL_ENABLE_VMM: %d\n", g_ggml_sycl_enable_vmm);
@@ -3027,22 +3065,18 @@ inline void ggml_sycl_op_mul_mat_sycl(
     }
 #endif
 
+    // dequantize inside the GEMM instead of writing the f16 weights out and reading them back; src1
+    // goes in its own type, so there is no separate conversion pass
+    if (ggml_is_quantized(src0->type) && ggml_is_contiguous(src0) && row_diff == src0->ne[1] &&
+        ggml_sycl_fused_dequant_gemm(src0->type, src0_dd_i, src1_ddf_i, src1->type, ggml_sycl_src1_prec(dst), dst_dd_i,
+                                     row_diff, src1_ncols, ne10, ldc, ctx.pool(), stream)) {
+        return;
+    }
+
+    // the f16 route converts src1 to f16 [TAG_GGML_PREC]
+    use_fp16 = use_fp16 && ggml_sycl_src1_f16_ok(dst);
     if ((src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && use_fp16 && ggml_is_contiguous(src0) &&
         row_diff == src0->ne[1] && dst->op_params[0] == GGML_PREC_DEFAULT) {
-        ggml_sycl_pool_alloc src0_as_f16(ctx.pool());
-        if (src0->type != GGML_TYPE_F16) {
-            scope_op_debug_print scope_dbg_print(__func__, "/to_fp16_sycl", dst, /*num_src=*/2,
-                                                 " : converting src0 to fp16");
-            const to_fp16_sycl_t to_fp16_sycl = ggml_get_to_fp16_sycl(src0->type, dst);
-            GGML_ASSERT(to_fp16_sycl != nullptr);
-            size_t ne = row_diff*ne00;
-            src0_as_f16.alloc(ne);
-            to_fp16_sycl(src0_dd_i, src0_as_f16.get(), ne, stream);
-        }
-        const sycl::half *src0_ptr = src0->type == GGML_TYPE_F16
-                                         ? (const sycl::half *)src0_dd_i
-                                         : src0_as_f16.get();
-
         ggml_sycl_pool_alloc src1_as_f16(ctx.pool());
         if (src1->type != GGML_TYPE_F16) {
             scope_op_debug_print scope_dbg_print(__func__, "/to_fp16_sycl", dst, /*num_src=*/2,
@@ -3057,6 +3091,20 @@ inline void ggml_sycl_op_mul_mat_sycl(
                 ? (const sycl::half *)src1->data + src1_padded_row_size
                                          : src1_as_f16.get();
 
+        ggml_sycl_pool_alloc src0_as_f16(ctx.pool());
+        if (src0->type != GGML_TYPE_F16) {
+            scope_op_debug_print scope_dbg_print(__func__, "/to_fp16_sycl", dst, /*num_src=*/2,
+                                                 " : converting src0 to fp16");
+            const to_fp16_sycl_t to_fp16_sycl = ggml_get_to_fp16_sycl(src0->type, dst);
+            GGML_ASSERT(to_fp16_sycl != nullptr);
+            size_t ne = row_diff*ne00;
+            src0_as_f16.alloc(ne);
+            to_fp16_sycl(src0_dd_i, src0_as_f16.get(), ne, stream);
+        }
+        const sycl::half *src0_ptr = src0->type == GGML_TYPE_F16
+                                         ? (const sycl::half *)src0_dd_i
+                                         : src0_as_f16.get();
+
 #if GGML_SYCL_DNNL
         if (g_ggml_sycl_enable_dnn && ggml_sycl_dnnl_has_optimized_gemm(ggml_sycl_get_device())) {
                 DnnlGemmWrapper::row_gemm(ctx,row_diff, src1_ncols , ne10, src0_ptr,
@@ -4859,6 +4907,10 @@ static void ggml_sycl_mul_mat(ggml_backend_sycl_context & ctx, const ggml_tensor
 
     // check data types and tensor shapes for custom matrix multiplication kernels:
     bool use_dequantize_mul_mat_vec = can_use_dequantize_mul_mat_vec(src0, src1, dst);
+#ifdef GGML_SYCL_F16
+    // dmmv may convert src1 to f16 in this build [TAG_GGML_PREC]
+    use_dequantize_mul_mat_vec = use_dequantize_mul_mat_vec && ggml_sycl_src1_f16_ok(dst);
+#endif
 
     bool use_mul_mat_vec_q = can_use_mul_mat_vec_q(src0, src1, dst);
 
@@ -5272,7 +5324,9 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
     SYCL_CHECK(CHECK_TRY_ERROR(
         stream->memcpy(ids_host.data(), ids_dev, ggml_nbytes(ids))));
 
-    // also ensures ctx.mmid_row_mapping_host is drained before we use it again
+    // also ensures ctx.mmid_row_mapping_host and ctx.mmid_tile_schedule_host are drained before we
+    // refill them: the grouped GEMM enqueues an async copy out of the tile schedule, so removing
+    // this wait would let the next node overwrite a buffer the device is still reading
     SYCL_CHECK(CHECK_TRY_ERROR(stream->wait()));
 
     ggml_tensor src0_row = *src0;
@@ -5363,7 +5417,25 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
             });
         }
 
-        for (int64_t i02 = 0; i02 < n_as; i02++) {
+        bool grouped = false;
+        if (ggml_is_contiguous(src0) && src1->type == GGML_TYPE_F32 &&
+            dst->type == GGML_TYPE_F32 && nb11 == sizeof(float)*ne10 && nb1 == sizeof(float)*ne0) {
+            // the grouped GEMM reads the reorder (SoA) layout faster, and the first decode step installs it
+            // anyway: install it here already, so prefill does not depend on whether a decode ran before
+            if (g_ggml_sycl_dynamic_precision != GGML_SYCL_DYNAMIC_PRECISION_F32 &&
+                ggml_sycl_xmx_gather_type_enabled(src0->type)) {
+                opt_for_reorder_id(&ctx, src0);
+            }
+            const bool src0_reordered =
+                src0->extra && ((const ggml_tensor_extra_gpu *) src0->extra)->optimized_feature.reorder;
+            grouped = ggml_sycl_grouped_dequant_gemm(src0->type, src0_reordered, src0_original, nb02,
+                                                     (const float *) src1_contiguous.get(), ggml_sycl_src1_prec(dst),
+                                                     (float *) dst_contiguous.get(),
+                                                     expert_row_offsets.data(), n_as, ne01, ne10, n_routed_rows,
+                                                     ctx.mmid_tile_schedule_host, ctx.pool(), stream);
+        }
+
+        for (int64_t i02 = 0; i02 < n_as && !grouped; i02++) {
             const int64_t num_src1_rows = expert_row_counts[i02];
 
             if (num_src1_rows == 0) {

From 24e41838e0c834f67db214e6d280af51f266056e Mon Sep 17 00:00:00 2001
From: SXX 
Date: Thu, 8 Oct 2026 14:07:34 +0800
Subject: [PATCH 28/40] ggml-cuda: assign four GDN state columns per warp
 (#30087)

* ggml-cuda: assign two GDN state columns per warp

* ggml-cuda: use 4 GDN state columns per warp at S_v=128

* ggml-cuda: default cols_per_warp=4

* ggml-cuda: address GDN review nits
---
 ggml/src/ggml-cuda/gated_delta_net.cu | 59 +++++++++++++++++----------
 1 file changed, 38 insertions(+), 21 deletions(-)

diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu
index 1b431a724..b58257067 100644
--- a/ggml/src/ggml-cuda/gated_delta_net.cu
+++ b/ggml/src/ggml-cuda/gated_delta_net.cu
@@ -1,7 +1,9 @@
 #include "gated_delta_net.cuh"
 #include "ggml-cuda/common.cuh"
 
-template 
+constexpr int gdn_cols_per_warp = 4;
+
+template 
 __global__ void __launch_bounds__((ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v) * 4, 2)
 gated_delta_net_cuda(const float * q,
                                      const float * k,
@@ -30,9 +32,19 @@ gated_delta_net_cuda(const float * q,
                                      int           K) {
     const uint32_t h_idx    = blockIdx.x;
     const uint32_t sequence = blockIdx.y;
-    // each warp owns one column, using warp-level primitives to reduce across rows
-    const int      lane     = threadIdx.x;
-    const int      col      = blockIdx.z * blockDim.y + threadIdx.y;
+
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v;
+    static_assert(S_v % warp_size == 0, "S_v must be a multiple of warp_size");
+    // the warp is split into cols_per_warp segments of lanes_per_col lanes; each segment owns
+    // one state column and reduces within itself
+    constexpr int lanes_per_col = warp_size / cols_per_warp;
+    constexpr int rows_per_lane = S_v / lanes_per_col;
+    static_assert(S_v % lanes_per_col == 0, "S_v must be a multiple of lanes_per_col");
+
+    const int lane        = threadIdx.x;
+    const int col_in_warp = lane / lanes_per_col;              // column slot within the warp
+    const int lane_in_col = lane - col_in_warp * lanes_per_col;  // lane within the column's reduction segment
+    const int col = (blockIdx.z * blockDim.y + threadIdx.y) * cols_per_warp + col_in_warp;
 
     const uint32_t iq1 = fastmodulo(h_idx, neqk1_magic);
     const uint32_t iq3 = fastdiv(sequence, rq3_magic);
@@ -47,16 +59,13 @@ gated_delta_net_cuda(const float * q,
     curr_state += state_in_offset + col * S_v;
     attn_data += (sequence * n_tokens * H + h_idx) * S_v;
 
-    constexpr int warp_size = ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v;
-    static_assert(S_v % warp_size == 0, "S_v must be a multiple of warp_size");
-    constexpr int rows_per_lane = (S_v + warp_size - 1) / warp_size;
     float         s_shard[rows_per_lane];
     // state is stored transposed: M[col][i] = S[i][col], row col is contiguous
 
     ggml_cuda_pdl_sync();
 #pragma unroll
     for (int r = 0; r < rows_per_lane; r++) {
-        const int i = r * warp_size + lane;
+        const int i = r * lanes_per_col + lane_in_col;
         s_shard[r]  = curr_state[i];
     }
 
@@ -76,7 +85,7 @@ gated_delta_net_cuda(const float * q,
         float q_reg[rows_per_lane];
 #pragma unroll
         for (int r = 0; r < rows_per_lane; r++) {
-            const int i = r * warp_size + lane;
+            const int i = r * lanes_per_col + lane_in_col;
             k_reg[r] = k_t[i];
             q_reg[r] = q_t[i];
         }
@@ -90,7 +99,7 @@ gated_delta_net_cuda(const float * q,
             for (int r = 0; r < rows_per_lane; r++) {
                 kv_shard += s_shard[r] * k_reg[r];
             }
-            float kv_col = warp_reduce_sum(kv_shard);
+            float kv_col = warp_reduce_sum(kv_shard);
 
             // delta[col] = (v[col] - g * kv[col]) * beta
             float delta_col = (v_t[col] - g_val * kv_col) * beta_val;
@@ -104,9 +113,9 @@ gated_delta_net_cuda(const float * q,
                 attn_partial += s_shard[r] * q_reg[r];
             }
 
-            float attn_col = warp_reduce_sum(attn_partial);
+            float attn_col = warp_reduce_sum(attn_partial);
 
-            if (lane == 0) {
+            if (lane_in_col == 0) {
                 attn_data[col] = attn_col * scale;
             }
         } else {
@@ -114,11 +123,11 @@ gated_delta_net_cuda(const float * q,
             float kv_shard = 0.0f;
 #pragma unroll
             for (int r = 0; r < rows_per_lane; r++) {
-                const int i = r * warp_size + lane;
+                const int i = r * lanes_per_col + lane_in_col;
                 kv_shard += expf(g_t[i]) * s_shard[r] * k_reg[r];
             }
 
-            float kv_col = warp_reduce_sum(kv_shard);
+            float kv_col = warp_reduce_sum(kv_shard);
 
             // delta[col] = (v[col] - kv[col]) * beta
             float delta_col = (v_t[col] - kv_col) * beta_val;
@@ -128,14 +137,14 @@ gated_delta_net_cuda(const float * q,
             float attn_partial = 0.0f;
 #pragma unroll
             for (int r = 0; r < rows_per_lane; r++) {
-                const int i = r * warp_size + lane;
+                const int i = r * lanes_per_col + lane_in_col;
                 s_shard[r]  = expf(g_t[i]) * s_shard[r] + k_reg[r] * delta_col;
                 attn_partial += s_shard[r] * q_reg[r];
             }
 
-            float attn_col = warp_reduce_sum(attn_partial);
+            float attn_col = warp_reduce_sum(attn_partial);
 
-            if (lane == 0) {
+            if (lane_in_col == 0) {
                 attn_data[col] = attn_col * scale;
             }
         }
@@ -150,7 +159,7 @@ gated_delta_net_cuda(const float * q,
                 float * curr_state = state + target_slot * state_slot_stride;
 #pragma unroll
                 for (int r = 0; r < rows_per_lane; r++) {
-                    const int i = r * warp_size + lane;
+                    const int i = r * lanes_per_col + lane_in_col;
                     curr_state[col * S_v + i] = s_shard[r];
                 }
             }
@@ -160,7 +169,7 @@ gated_delta_net_cuda(const float * q,
     if constexpr (!keep_rs_t) {
 #pragma unroll
         for (int r = 0; r < rows_per_lane; r++) {
-            const int i          = r * warp_size + lane;
+            const int i          = r * lanes_per_col + lane_in_col;
             state[col * S_v + i] = s_shard[r];
         }
     }
@@ -179,8 +188,16 @@ static void launch_gated_delta_net(
         float scale, int64_t state_slot_stride, int K, cudaStream_t stream) {
     //TODO: Add chunked kernel for even faster pre-fill
     const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
-    const int num_warps = 4;
-    dim3      grid_dims(H, n_seqs, (S_v + num_warps - 1) / num_warps);
+    // four columns per warp (see the kernel); shrink the CTA when the wider CTA would leave
+    // SMs without a CTA, so small head counts keep the device filled
+    const int nsm = ggml_cuda_info().devices[ggml_cuda_get_device()].nsm;
+    const int cols_per_warp = gdn_cols_per_warp;
+    int num_warps = 4;
+    while (num_warps > 1 && H*n_seqs*(S_v / (cols_per_warp * num_warps)) < nsm) {
+        num_warps /= 2;
+    }
+    // one CTA covers cols_per_warp*num_warps columns (see the kernel)
+    dim3      grid_dims(H, n_seqs, (S_v + cols_per_warp * num_warps - 1) / (cols_per_warp * num_warps));
     dim3      block_dims(warp_size <= S_v ? warp_size : S_v, num_warps, 1);
 
     const uint3 neqk1_magic = init_fastdiv_values(neqk1);

From 37ac634566439e680959b90e1c960714ca6b8b92 Mon Sep 17 00:00:00 2001
From: bosh <98094229+boshjerns@users.noreply.github.com>
Date: Thu, 8 Oct 2026 13:29:20 +0700
Subject: [PATCH 29/40] model : support classifier_activation for rerankers
 (#29692)

* model : support classifier_activation for rerankers

Assisted-by: Claude Opus 5.5

* model : map classifier gelu to gelu_erf and accept tanh

Assisted-by: Claude Opus 5.5

* model : default act_cls to tanh, ModernBERT falls back to gelu_erf

Assisted-by: Claude Opus 5.5
---
 conversion/base.py          |  4 ++++
 gguf-py/gguf/constants.py   |  1 +
 gguf-py/gguf/gguf_writer.py |  3 +++
 src/llama-arch.cpp          |  1 +
 src/llama-arch.h            |  1 +
 src/llama-graph.cpp         |  6 +-----
 src/llama-hparams.h         |  1 +
 src/llama-model.cpp         | 13 +++++++++++++
 src/models/modern-bert.cpp  |  6 ++++++
 9 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/conversion/base.py b/conversion/base.py
index 786e7faa1..7a93855c8 100644
--- a/conversion/base.py
+++ b/conversion/base.py
@@ -2352,6 +2352,10 @@ class TextModel(ModelBase):
             if classifier_pooling not in ("cls", "mean"):
                 raise NotImplementedError(f"Unsupported classifier_pooling: {classifier_pooling}")
             self.gguf_writer.add_classifier_pooling_type(mode_mapping[classifier_pooling])
+        if (classifier_activation := self.hparams.get("classifier_activation")) is not None:
+            if classifier_activation not in ("gelu", "silu", "tanh"):
+                raise NotImplementedError(f"Unsupported classifier_activation: {classifier_activation}")
+            self.gguf_writer.add_classifier_activation(classifier_activation)
 
     def _set_vocab_glmedge(self):
         from transformers import AutoTokenizer
diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index fcf325fe9..c003ed2b6 100644
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -318,6 +318,7 @@ class Keys:
     class Classifier:
         OUTPUT_LABELS = "{arch}.classifier.output_labels"
         POOLING_TYPE  = "{arch}.classifier.pooling_type"
+        ACTIVATION    = "{arch}.classifier.activation"
 
     class ShortConv:
         L_CACHE = "{arch}.shortconv.l_cache"
diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py
index f33a8a551..3d063bb3a 100644
--- a/gguf-py/gguf/gguf_writer.py
+++ b/gguf-py/gguf/gguf_writer.py
@@ -1340,6 +1340,9 @@ class GGUFWriter:
     def add_classifier_pooling_type(self, value: PoolingType) -> None:
         self.add_uint32(Keys.Classifier.POOLING_TYPE.format(arch=self.arch), value.value)
 
+    def add_classifier_activation(self, value: str) -> None:
+        self.add_string(Keys.Classifier.ACTIVATION.format(arch=self.arch), value)
+
     def add_decision_type(self, value: str) -> None:
         self.add_string(Keys.Decision.TYPE.format(arch=self.arch), value)
 
diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
index eea2ef589..821479064 100644
--- a/src/llama-arch.cpp
+++ b/src/llama-arch.cpp
@@ -370,6 +370,7 @@ static const std::map LLM_KV_NAMES = {
 
     { LLM_KV_CLASSIFIER_OUTPUT_LABELS, "%s.classifier.output_labels" },
     { LLM_KV_CLASSIFIER_POOLING_TYPE,  "%s.classifier.pooling_type"  },
+    { LLM_KV_CLASSIFIER_ACTIVATION,    "%s.classifier.activation"    },
 
     { LLM_KV_DECISION_BLOCK_COUNT,         "%s.decision.block_count"         },
     { LLM_KV_DECISION_ROUTING_BLOCK_COUNT, "%s.decision.routing_block_count" },
diff --git a/src/llama-arch.h b/src/llama-arch.h
index c8abee234..1bf4744ff 100644
--- a/src/llama-arch.h
+++ b/src/llama-arch.h
@@ -416,6 +416,7 @@ enum llm_kv {
 
     LLM_KV_CLASSIFIER_OUTPUT_LABELS,
     LLM_KV_CLASSIFIER_POOLING_TYPE,
+    LLM_KV_CLASSIFIER_ACTIVATION,
 
     LLM_KV_DECISION_BLOCK_COUNT,
     LLM_KV_DECISION_ROUTING_BLOCK_COUNT,
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index e0b7a47a9..f030eb9e6 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -3903,11 +3903,7 @@ void llm_graph_context::build_pooling(
                     if (cls_b) {
                         cur = ggml_add(ctx0, cur, cls_b);
                     }
-                    if (arch == LLM_ARCH_MODERN_BERT) {
-                        cur = ggml_gelu(ctx0, cur);
-                    } else {
-                        cur = ggml_tanh(ctx0, cur);
-                    }
+                    cur = ggml_unary(ctx0, cur, hparams.act_cls);
                     if (cls_norm) {
                         // head norm
                         cur = build_norm(cur, cls_norm, NULL, LLM_NORM, -1);
diff --git a/src/llama-hparams.h b/src/llama-hparams.h
index 848428706..98afe8a62 100644
--- a/src/llama-hparams.h
+++ b/src/llama-hparams.h
@@ -371,6 +371,7 @@ struct llama_hparams {
     // llm_ffn_op_type_from_string() in llama-model.cpp, mirroring how
     // rope_scaling_type_train is handled.
     enum llm_ffn_op_type llm_ffn_op;
+    enum ggml_unary_op   act_cls = GGML_UNARY_OP_TANH; // activation of the classifier head (RANK)
 
     // Step35: optional per-layer clamps for (Swi)GLU
     std::array swiglu_clamp_exp; // clamping for expert FFN
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
index bf04945bf..2e86914bc 100644
--- a/src/llama-model.cpp
+++ b/src/llama-model.cpp
@@ -1078,6 +1078,13 @@ static const std::map LLM_FFN_OP_TYPES_FROM_STRING
     { "reglu",  LLM_FFN_REGLU  },
 };
 
+// transformers names, "gelu" is the exact (erf) variant
+static const std::map LLM_CLS_ACT_TYPES_FROM_STRING = {
+    { "gelu", GGML_UNARY_OP_GELU_ERF },
+    { "silu", GGML_UNARY_OP_SILU     },
+    { "tanh", GGML_UNARY_OP_TANH     },
+};
+
 llm_ffn_op_type llm_ffn_op_type_from_string(const std::string & name, llm_ffn_op_type fallback) {
     const auto it = LLM_FFN_OP_TYPES_FROM_STRING.find(name);
     if (it != LLM_FFN_OP_TYPES_FROM_STRING.end()) {
@@ -1336,6 +1343,12 @@ void llama_model_base::load_hparams(llama_model_loader & ml) {
     ml.get_key(LLM_KV_ATTENTION_CAUSAL,        hparams.causal_attn,     false);
     ml.get_key(LLM_KV_POOLING_TYPE,            hparams.pooling_type,    false);
     ml.get_key(LLM_KV_CLASSIFIER_POOLING_TYPE, hparams.pooling_type_cls, false);
+    std::string act_cls;
+    if (ml.get_key(LLM_KV_CLASSIFIER_ACTIVATION, act_cls, false)) {
+        const auto it = LLM_CLS_ACT_TYPES_FROM_STRING.find(act_cls);
+        GGML_ASSERT(it != LLM_CLS_ACT_TYPES_FROM_STRING.end() && "unsupported classifier activation");
+        hparams.act_cls = it->second;
+    }
     ml.get_key(LLM_KV_BLOCK_COUNT,             hparams.n_layer_all);
     GGML_ASSERT(hparams.n_layer_all > 0 && hparams.n_layer_all <= LLAMA_MAX_LAYERS);
     ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS,    hparams.n_layer_nextn,   false);
diff --git a/src/models/modern-bert.cpp b/src/models/modern-bert.cpp
index 45455acbb..d265b03bf 100644
--- a/src/models/modern-bert.cpp
+++ b/src/models/modern-bert.cpp
@@ -28,6 +28,12 @@ void llama_model_modern_bert::load_arch_hparams(llama_model_loader & ml) {
         hparams.pooling_type_cls = LLAMA_POOLING_TYPE_MEAN;
     }
 
+    // GGUFs without a classifier activation use gelu, the transformers default
+    std::string act_cls;
+    if (!ml.get_key(LLM_KV_CLASSIFIER_ACTIVATION, act_cls, false)) {
+        hparams.act_cls = GGML_UNARY_OP_GELU_ERF;
+    }
+
     ml.get_key(LLM_KV_DECISION_BLOCK_COUNT, hparams.n_layer_decision, false);
     if (hparams.n_layer_decision > 0) {
         if (hparams.n_layer_decision >= hparams.n_layer()) {

From 000bee54a544070356b44d68af299b973d391d26 Mon Sep 17 00:00:00 2001
From: cwriter 
Date: Thu, 8 Oct 2026 08:30:00 +0200
Subject: [PATCH 30/40] sycl: stage bulk uploads (model loading) through a
 pinned ring buffer (#29608)

Co-authored-by: cwriter 
---
 docs/backend/SYCL.md             |  1 +
 ggml/src/ggml-sycl/ggml-sycl.cpp | 63 ++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/docs/backend/SYCL.md b/docs/backend/SYCL.md
index c2ae0c1ba..19dc69851 100644
--- a/docs/backend/SYCL.md
+++ b/docs/backend/SYCL.md
@@ -803,6 +803,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still on development, no better performance. |
 | GGML_SYCL_ENABLE_HOST_PINNED_MEM | 0 or 1 (default) | Enable host pinned memory to speed up copy data from host to device. When disable it, host memory will common malloc() on CPU. Disable it when use `--load-model mlock`.|
 | GGML_SYCL_HOST_PINNED_MEM_2G | 0 (default) or 1 | Limit the max memory allocation to be no more than 2GB when enable host pinned memory. USM allocations above 2 GiB take the relaxed/large-allocation path, which serializes H2D copies with compute and prevents copy/compute overlap. It will impact the startup time. Need more test. Depend on `GGML_SYCL_ENABLE_HOST_PINNED_MEM=1`.|
+| GGML_SYCL_UPLOAD_STAGING_SLOTS | 4 (default) or non-negative integer | Number of 8 MiB pinned host slots used to stage tensor uploads (model loading), so the host copy of one slot overlaps the transfer of the previous one. Set to 0 to use the old path: a malloc'd bounce buffer and a blocking copy per tensor. |
 | GGML_SYCL_GET_MEM_API | 0 (default) or 1  | Set to get memory info (free, total) by Level Zero or SYCL API:
0 - Level Zero API: support more GPUs, only run on Level Zero running time. When there is an error, fallback to call SYCL API. Depend on GGML_SYCL_SUPPORT_LEVEL_ZERO_API.
1 - SYCL API: legacy, support more running time, it can't get the free size of some GPUs (like Arc770). In such case, return the free size as value of total size.|
 | GGML_SYCL_USE_LEVEL_ZERO_API | 1 (default) or 0 | Use Level Zero API for device memory allocation instead of SYCL. Reduces system RAM usage on Intel dGPUs by avoiding DMA-buf/TTM host memory staging. Requires GGML_SYCL_SUPPORT_LEVEL_ZERO_API=ON at build time. SYCL backend always runs on Level Zero running time even if it's set as OFF (The SYCL api will be usage for memory allocation).|
 | GGML_SYCL_ENABLE_DNN | 0 or 1 (default)| Enable running computations through oneDNN and always use oneMKL. |
diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp
index ebca92886..13fc02325 100644
--- a/ggml/src/ggml-sycl/ggml-sycl.cpp
+++ b/ggml/src/ggml-sycl/ggml-sycl.cpp
@@ -139,6 +139,7 @@ int g_ggml_sycl_dev2dev_memcpy = DEV2DEV_MEMCPY_SYCL;
 int g_ggml_sycl_usm_system = 0;
 int g_ggml_sycl_enable_host_pinned_mem = 1;
 int g_ggml_sycl_host_pinned_mem_2g = 0;
+int g_ggml_sycl_upload_staging_slots = 4;
 int g_ggml_sycl_get_mem_api = MEMORY_API_TYPE_LEVEL_ZERO;
 int g_ggml_sycl_enable_sparse_fa = 0;
 int g_ggml_sycl_debug_sparse_fa = 0;
@@ -458,6 +459,7 @@ static void ggml_check_sycl() try {
 
         g_ggml_sycl_host_pinned_mem_2g =
             ggml_sycl_get_env("GGML_SYCL_HOST_PINNED_MEM_2G", 0) & g_ggml_sycl_enable_host_pinned_mem;
+        g_ggml_sycl_upload_staging_slots = std::max(0, ggml_sycl_get_env("GGML_SYCL_UPLOAD_STAGING_SLOTS", 4));
 
         g_ggml_sycl_enable_sparse_fa = ggml_sycl_get_env("GGML_SYCL_SPARSE_FA", 0);
         g_ggml_sycl_debug_sparse_fa = ggml_sycl_get_env("GGML_SYCL_SPARSE_FA_DEBUG", 0);
@@ -555,6 +557,7 @@ static void ggml_check_sycl() try {
 #endif
 
         GGML_LOG_INFO("  GGML_SYCL_ENABLE_FUSION: %d\n", g_ggml_sycl_enable_fusion);
+        GGML_LOG_INFO("  GGML_SYCL_UPLOAD_STAGING_SLOTS: %d\n", g_ggml_sycl_upload_staging_slots);
 
 #if defined(__INTEL_LLVM_COMPILER)
         GGML_LOG_INFO("  GGML_SYCL_ENABLE_ESIMD: %d\n", g_ggml_sycl_enable_esimd);
@@ -667,12 +670,23 @@ inline void free_aligned_mem_host(void * memblock) {
 // sycl buffer
 
 struct ggml_backend_sycl_buffer_context {
+    // pinned staging for uploads; the host fills one slot while the previous one transfers
+    static constexpr size_t staging_slot_size = 8*1024*1024;
+
+    struct host_staging {
+        void * data = nullptr;
+        std::vector events;
+        std::vector submitted;
+        int next = 0;
+    };
+
     int device;
     void * dev_ptr = nullptr;
     queue_ptr stream;
     std::string name;
     optimize_feature opt_feature;
     std::vector tensor_extras;
+    host_staging staging;
     bool is_usm_system;
 
     ggml_backend_sycl_buffer_context(int device, void * dev_ptr, queue_ptr stream, bool is_usm_system) :
@@ -682,7 +696,22 @@ struct ggml_backend_sycl_buffer_context {
             opt_feature = ggml_sycl_info().devices[device].opt_feature;
         }
 
+    // waits for every queued upload, then releases the pinned block
+    void drop_host_staging() {
+        for (size_t i = 0; i < staging.submitted.size(); ++i) {
+            if (staging.submitted[i]) {
+                staging.events[i].wait_and_throw();
+                staging.submitted[i] = false;
+            }
+        }
+        if (staging.data != nullptr) {
+            sycl::free(staging.data, *stream);
+            staging.data = nullptr;
+        }
+    }
+
     ~ggml_backend_sycl_buffer_context() {
+        drop_host_staging();
         if (dev_ptr != nullptr) {
             ggml_sycl_set_device(device);
             if (is_usm_system)
@@ -783,6 +812,40 @@ static void ggml_backend_sycl_buffer_set_tensor(ggml_backend_buffer_t buffer,
     GGML_SYCL_DEBUG(" size=%zu offset=%zu\n", size, offset);
     ggml_backend_sycl_buffer_context * ctx = ( ggml_backend_sycl_buffer_context *)buffer->context;
     ggml_sycl_set_device(ctx->device);
+
+    // 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) {
+        queue_ptr    stream    = ctx->stream;
+        char *       dst       = (char *) tensor->data + offset;
+        const char * src       = (const char *) data;
+        size_t       remaining = size;
+        while (remaining > 0) {
+            const size_t chunk = std::min(remaining, ctx->staging_slot_size);
+            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();
+            }
+            void * stage = (char *) ctx->staging.data + slot * ctx->staging_slot_size;
+            memcpy(stage, src, chunk);
+            ctx->staging.events[slot] = stream->memcpy(dst, stage, chunk);
+            ctx->staging.submitted[slot] = true;
+            src       += chunk;
+            dst       += chunk;
+            remaining -= chunk;
+        }
+        return;
+    }
+
     auto stream = &(dpct::dev_mgr::instance().get_device(ctx->device).default_queue());
     SYCL_CHECK(CHECK_TRY_ERROR(dpct::dev_mgr::instance().get_device(ctx->device).queues_wait_and_throw()));
 #ifndef _WIN32

From ff30363a0e3e2630828309ef5a37da5d1d40396c Mon Sep 17 00:00:00 2001
From: Titaniumtown 
Date: Wed, 7 Oct 2026 23:45:21 -0700
Subject: [PATCH 31/40] sycl: fuse the delta-net alpha gate (add + unary + mul)
 (#29687)

* sycl: fuse the delta-net alpha gate (add + unary + mul)

* tests: cover the fused add + unary + mul chain

* sycl: give the fused alpha gate a flat path and pin the node skip
---
 ggml/src/ggml-sycl/element_wise.cpp |  95 +++++++++++++++++++++++--
 ggml/src/ggml-sycl/element_wise.hpp |   5 ++
 ggml/src/ggml-sycl/fusion.cpp       |  59 +++++++++++++++-
 ggml/src/ggml-sycl/ggml-sycl.cpp    |  10 +++
 tests/test-backend-ops.cpp          | 105 ++++++++++++++++++++++++++++
 5 files changed, 264 insertions(+), 10 deletions(-)

diff --git a/ggml/src/ggml-sycl/element_wise.cpp b/ggml/src/ggml-sycl/element_wise.cpp
index 2e926abea..0d263aa76 100644
--- a/ggml/src/ggml-sycl/element_wise.cpp
+++ b/ggml/src/ggml-sycl/element_wise.cpp
@@ -452,6 +452,46 @@ static void unary_mul_sycl(const T * x, const T * g, T * dst, const int64_t k, c
     });
 }
 
+// ADD(bias) + UNARY + MUL(scale) with both broadcast over dim 0, the delta-net alpha gate:
+// dst[i] = op(a[i] + bias[i % ne0]) * scale[i % ne0]. k == ne0 makes that the flat index.
+template
+static void add_unary_mul_flat_kernel(const float * a, const float * bias, const float * scale, float * dst,
+                                      const int64_t k, const sycl::nd_item<1> &item_ct1, F op) {
+    SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
+        dst[i] = op(a[i] + bias[i]) * scale[i];
+    }
+}
+
+template
+static void add_unary_mul_bcast_kernel(const float * a, const float * bias, const float * scale, float * dst,
+                                       const int64_t k, const sycl::uint3 ne0_fd, const sycl::nd_item<1> &item_ct1, F op) {
+    SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
+        const uint32_t h = fastmodulo((uint32_t) i, ne0_fd);
+        dst[i] = op(a[i] + bias[h]) * scale[h];
+    }
+}
+
+template
+static void add_unary_mul_sycl(const float * a, const float * bias, const float * scale, float * dst,
+                               const int64_t k, const int64_t ne0, queue_ptr main_stream, F op) {
+    const size_t            num_blocks = ceil_div((size_t) k, (size_t) SYCL_GLU_BLOCK_SIZE);
+    const sycl::nd_range<1> range(num_blocks * sycl::range<1>(SYCL_GLU_BLOCK_SIZE), sycl::range<1>(SYCL_GLU_BLOCK_SIZE));
+
+    if (k == ne0) {
+        main_stream->parallel_for(range, [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+            add_unary_mul_flat_kernel(a, bias, scale, dst, k, item_ct1, op);
+        });
+        return;
+    }
+
+    // 32-bit fastdiv, exact only below 2^31; ggml_sycl_can_fuse() already declined past that
+    GGML_ASSERT(k < ((int64_t) 1 << 31));
+    const sycl::uint3 ne0_fd = init_fastdiv_values((uint32_t) ne0);
+    main_stream->parallel_for(range, [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+        add_unary_mul_bcast_kernel(a, bias, scale, dst, k, ne0_fd, item_ct1, op);
+    });
+}
+
 namespace ggml_sycl_detail {
 static void acc_f32_sycl(const char *x, const char *y, float *dst,
                          const int64_t n_elements,
@@ -995,6 +1035,19 @@ static inline void ggml_sycl_op_swiglu(ggml_backend_sycl_context & ctx, ggml_ten
     });
 }
 
+// Hands `launch` the functor for the unary op of a fused unary chain. Anything else
+// ggml_sycl_can_fuse() has already declined, so the default is a dispatcher bug.
+template
+static void dispatch_fused_unary_op(ggml_unary_op uop, F && launch) {
+    switch (uop) {
+        case GGML_UNARY_OP_SILU:     launch([](float v) { return op_silu(v); });     break;
+        case GGML_UNARY_OP_SIGMOID:  launch([](float v) { return op_sigmoid(v); });  break;
+        case GGML_UNARY_OP_SOFTPLUS: launch([](float v) { return op_softplus(v); }); break;
+        default:
+            GGML_ABORT("fused unary chain: unsupported unary op %s", ggml_unary_op_name(uop));
+    }
+}
+
 // dst = op(unary_node->src[0]) * other, written straight to the MUL output, saving the
 // standalone unary launch. Preconditions come from ggml_sycl_can_fuse(); re-asserted here.
 void ggml_sycl_op_unary_mul_fused(ggml_backend_sycl_context & ctx, ggml_tensor * unary_node, ggml_tensor * mul_node) {
@@ -1032,13 +1085,41 @@ void ggml_sycl_op_unary_mul_fused(ggml_backend_sycl_context & ctx, ggml_tensor *
         }
     };
 
-    switch (ggml_get_unary_op(unary_node)) {
-        case GGML_UNARY_OP_SILU:     dispatch_type([](float v) { return op_silu(v); });     break;
-        case GGML_UNARY_OP_SIGMOID:  dispatch_type([](float v) { return op_sigmoid(v); });  break;
-        case GGML_UNARY_OP_SOFTPLUS: dispatch_type([](float v) { return op_softplus(v); }); break;
-        default:
-            GGML_ABORT("fused unary+mul: unsupported unary op %s", ggml_unary_op_name(ggml_get_unary_op(unary_node)));
-    }
+    dispatch_fused_unary_op(ggml_get_unary_op(unary_node), dispatch_type);
+}
+
+// dst = op(a + bias) * scale for an ADD + UNARY + MUL chain whose bias and scale broadcast
+// over dim 0. Preconditions come from ggml_sycl_can_fuse(); re-asserted here.
+void ggml_sycl_op_add_unary_mul_fused(ggml_backend_sycl_context & ctx, ggml_tensor * add_node,
+                                      ggml_tensor * unary_node, ggml_tensor * mul_node) {
+    // the dst-arity convention the other fusions follow; a and bias live on add_node
+    scope_op_debug_print scope_dbg_print(__func__, mul_node, /*num_src=*/2);
+
+    const ggml_tensor * a     = add_node->src[0];
+    const ggml_tensor * bias  = add_node->src[1];
+    const ggml_tensor * scale = (mul_node->src[0] == unary_node) ? mul_node->src[1] : mul_node->src[0];
+
+    // scale is picked by elimination; ggml_can_fuse()'s single-use rule rules out MUL(unary, unary)
+    GGML_ASSERT(scale != unary_node);
+    GGML_ASSERT(a->type == GGML_TYPE_F32 && bias->type == GGML_TYPE_F32);
+    GGML_ASSERT(scale->type == GGML_TYPE_F32 && mul_node->type == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_are_same_shape(a, mul_node));
+    // a and dst are indexed flat
+    GGML_ASSERT(ggml_is_contiguous(a) && ggml_is_contiguous(mul_node));
+    // bias and scale are one contiguous ne0-length row each, broadcast over the outer dims
+    GGML_ASSERT(bias->ne[0] == a->ne[0] && scale->ne[0] == a->ne[0]);
+    GGML_ASSERT(ggml_nrows(bias) == 1 && ggml_nrows(scale) == 1);
+    GGML_ASSERT(ggml_is_contiguous(bias) && ggml_is_contiguous(scale));
+
+    queue_ptr main_stream = ctx.stream();
+    SYCL_CHECK(ggml_sycl_set_device(ctx.device));
+
+    const auto dispatch_op = [&](auto op) {
+        add_unary_mul_sycl((const float *) a->data, (const float *) bias->data, (const float *) scale->data,
+                           (float *) mul_node->data, ggml_nelements(mul_node), mul_node->ne[0], main_stream, op);
+    };
+
+    dispatch_fused_unary_op(ggml_get_unary_op(unary_node), dispatch_op);
 }
 
 __dpct_inline__ float ggml_sycl_op_swiglu_oai_single(float x, float g, float alpha = 1.702f, float limit = 7.0f) {
diff --git a/ggml/src/ggml-sycl/element_wise.hpp b/ggml/src/ggml-sycl/element_wise.hpp
index d280066ef..0a5bfa956 100644
--- a/ggml/src/ggml-sycl/element_wise.hpp
+++ b/ggml/src/ggml-sycl/element_wise.hpp
@@ -132,4 +132,9 @@ void ggml_sycl_arange(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
 // fused UNARY(silu|sigmoid|softplus) + MUL; see ggml_sycl_can_fuse() for the accepted shapes
 void ggml_sycl_op_unary_mul_fused(ggml_backend_sycl_context & ctx, ggml_tensor * unary_node, ggml_tensor * mul_node);
 
+// fused f32 ADD + UNARY(silu|sigmoid|softplus) + MUL with the bias and the scale broadcast
+// over dim 0; see ggml_sycl_can_fuse() for the accepted shapes
+void ggml_sycl_op_add_unary_mul_fused(ggml_backend_sycl_context & ctx, ggml_tensor * add_node,
+                                      ggml_tensor * unary_node, ggml_tensor * mul_node);
+
 #endif // GGML_SYCL_ELEMENTWISE_HPP
diff --git a/ggml/src/ggml-sycl/fusion.cpp b/ggml/src/ggml-sycl/fusion.cpp
index d3e995233..79fe13a1d 100644
--- a/ggml/src/ggml-sycl/fusion.cpp
+++ b/ggml/src/ggml-sycl/fusion.cpp
@@ -64,6 +64,12 @@ static bool ggml_sycl_should_fuse_mul_mat_glu(const ggml_tensor * gate, const gg
     return true;
 }
 
+// the unary ops the fused unary chains in element_wise.cpp have a functor for
+static bool ggml_sycl_fused_unary_has_kernel(ggml_unary_op unary_op) {
+    return unary_op == GGML_UNARY_OP_SILU || unary_op == GGML_UNARY_OP_SIGMOID ||
+           unary_op == GGML_UNARY_OP_SOFTPLUS;
+}
+
 bool ggml_sycl_can_fuse(const ggml_cgraph * cgraph, int node_idx, std::initializer_list ops,
                         std::initializer_list unary_ops) {
 #ifndef NDEBUG
@@ -184,9 +190,7 @@ bool ggml_sycl_can_fuse(const ggml_cgraph * cgraph, int node_idx, std::initializ
             return false;
         }
 
-        // the ops ggml_sycl_op_unary_mul_fused() has a kernel for
-        if (unary_op != GGML_UNARY_OP_SILU && unary_op != GGML_UNARY_OP_SIGMOID &&
-            unary_op != GGML_UNARY_OP_SOFTPLUS) {
+        if (!ggml_sycl_fused_unary_has_kernel(unary_op)) {
             return false;
         }
 
@@ -233,6 +237,55 @@ bool ggml_sycl_can_fuse(const ggml_cgraph * cgraph, int node_idx, std::initializ
         return true;
     }
 
+    // ADD(bias) + UNARY + MUL(scale): the delta-net alpha gate, softplus(alpha + dt) * a.
+    // The broadcast is what stops the same-shape UNARY + MUL branch above firing past one token.
+    if (ops.size() == 3 && ops.begin()[0] == GGML_OP_ADD && ops.begin()[1] == GGML_OP_UNARY &&
+        ops.begin()[2] == GGML_OP_MUL && unary_ops.size() == 1) {
+        const ggml_tensor * add   = cgraph->nodes[node_idx];
+        const ggml_tensor * unary = cgraph->nodes[node_idx + 1];
+        const ggml_tensor * mul   = cgraph->nodes[node_idx + 2];
+
+        const ggml_unary_op unary_op = ggml_get_unary_op(unary);
+        if (unary_op != unary_ops.begin()[0]) {
+            return false;
+        }
+
+        if (!ggml_sycl_fused_unary_has_kernel(unary_op)) {
+            return false;
+        }
+
+        // ggml_can_fuse() has already pinned the chain: unary consumes add, mul consumes
+        // unary, add and unary have one use each, and all three have the same shape
+        const ggml_tensor * a     = add->src[0];
+        const ggml_tensor * bias  = add->src[1];
+        const ggml_tensor * scale = (mul->src[0] == unary) ? mul->src[1] : mul->src[0];
+
+        if (a->type != GGML_TYPE_F32 || bias->type != GGML_TYPE_F32 ||
+            scale->type != GGML_TYPE_F32 || mul->type != GGML_TYPE_F32) {
+            return false;
+        }
+
+        // the activation and the destination are indexed flat
+        if (!ggml_is_contiguous(a) || !ggml_is_contiguous(mul) || !ggml_are_same_shape(a, mul)) {
+            return false;
+        }
+
+        // the kernel reads the bias and the scale as v[col], so each must be a single
+        // contiguous row spanning ne0
+        if (bias->ne[0] != a->ne[0] || scale->ne[0] != a->ne[0] ||
+            ggml_nrows(bias) != 1 || ggml_nrows(scale) != 1 ||
+            !ggml_is_contiguous(bias) || !ggml_is_contiguous(scale)) {
+            return false;
+        }
+
+        // the 32-bit fastdiv is inexact past 2^31; decline, the unfused path handles it
+        if (ggml_nelements(mul) >= ((int64_t) 1 << 31)) {
+            return false;
+        }
+
+        return true;
+    }
+
     if (ops.size() == 3 && ops.begin()[0] == GGML_OP_SSM_CONV && ops.begin()[1] == GGML_OP_ADD &&
         ops.begin()[2] == GGML_OP_UNARY && unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_SILU) {
         const ggml_tensor * ssm_conv = cgraph->nodes[node_idx];
diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp
index 13fc02325..2a5184e01 100644
--- a/ggml/src/ggml-sycl/ggml-sycl.cpp
+++ b/ggml/src/ggml-sycl/ggml-sycl.cpp
@@ -6312,6 +6312,16 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc
             i++;
             continue;
         }
+        // ADD(bias) + UNARY + MUL(scale) with both broadcast over dim 0, the form the branch
+        // above cannot take; ggml_get_unary_op() asserts, so check the op first.
+        if (node->op == GGML_OP_ADD && i + 2 < cgraph->n_nodes &&
+            cgraph->nodes[i + 1]->op == GGML_OP_UNARY &&
+            ggml_sycl_can_fuse(cgraph, i, { GGML_OP_ADD, GGML_OP_UNARY, GGML_OP_MUL },
+                               { ggml_get_unary_op(cgraph->nodes[i + 1]) })) {
+            ggml_sycl_op_add_unary_mul_fused(*sycl_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]);
+            i += 2;
+            continue;
+        }
 
         // Batch consecutive independent same-shape F32 L2_NORM siblings (the GDN q/k
         // norms) into one launch; sources are strided views of the fused qkv buffer, so
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index cfe5b73e7..f31016bbb 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -4159,6 +4159,93 @@ struct test_unary_mul : public test_case {
     }
 };
 
+// GGML_OP_ADD + GGML_OP_UNARY(SILU|SIGMOID|SOFTPLUS) + GGML_OP_MUL with the ADD's bias and
+// the MUL's scale broadcast over dim 0: the delta-net alpha gate, softplus(alpha + dt) * a_coeff.
+struct test_add_unary_mul : public test_case {
+    const ggml_unary_op op;
+    const ggml_type type;
+    const std::array ne;
+    const bool swap;          // unary result is the second MUL operand
+    const std::string layout; // bias/scale layout, see build_graph()
+    const std::string tail;   // extra consumer past the MUL, see build_graph()
+
+    std::string op_desc(ggml_tensor * t) override {
+        GGML_UNUSED(t);
+        return "ADD_" + std::string(ggml_unary_op_name(op)) + "_MUL";
+    }
+    bool run_whole_graph() override { return true; }
+
+    double max_nmse_err() override {
+        switch (type) {
+            // f16 never fuses (the kernel is f32-only), so this bound is the unfused
+            // chain's own f16 rounding drift, as in test_unary_mul
+            case GGML_TYPE_F16: return 5e-5;
+            // gelu never fuses either, and the backends' exp form drifts from the CPU's tanhf
+            default:            return op == GGML_UNARY_OP_GELU ? 5e-7 : 1e-7;
+        }
+    }
+
+    std::string vars() override {
+        return VARS_TO_STR5(type, ne, swap, layout, tail);
+    }
+
+    test_add_unary_mul(ggml_unary_op op,
+            ggml_type type = GGML_TYPE_F32,
+            std::array ne = {32, 7, 1, 1},
+            bool swap = false,
+            std::string layout = "bcast",
+            std::string tail = "")
+        : op(op), type(type), ne(ne), swap(swap), layout(std::move(layout)), tail(std::move(tail)) {}
+
+    ggml_tensor * build_graph(ggml_context * ctx) override {
+        ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data());
+        ggml_set_name(a, "a");
+
+        std::array ne_v = { ne[0], 1, 1, 1 };
+        if (layout == "bcast") {
+            // one ne0 row each, broadcast over the outer dims, which is the alpha-gate form
+        } else if (layout == "same_shape") {
+            // no broadcast at all; fuses only while the activation is a single row
+            ne_v = ne;
+        } else if (layout == "rep_ne0") {
+            // repeat on dim 0, which bias[col] cannot address, so this must not fuse
+            ne_v[0] = ne[0] / 4;
+        } else {
+            GGML_ABORT("unknown layout %s", layout.c_str());
+        }
+
+        ggml_tensor * bias = ggml_new_tensor(ctx, type, 4, ne_v.data());
+        ggml_set_name(bias, "bias");
+
+        ggml_tensor * scale = ggml_new_tensor(ctx, type, 4, ne_v.data());
+        ggml_set_name(scale, "scale");
+
+        ggml_tensor * s = ggml_add(ctx, a, bias);
+        ggml_set_name(s, "add");
+
+        ggml_tensor * u = ggml_unary(ctx, s, op);
+        ggml_set_name(u, "unary");
+
+        // a broadcasting operand can only be the second one, so swap needs same-shape operands
+        ggml_tensor * out = swap ? ggml_mul(ctx, scale, u) : ggml_mul(ctx, u, scale);
+
+        if (tail == "reuse") {
+            // a second read of the add result must block the fusion
+            ggml_set_name(out, "mul");
+            out = ggml_add(ctx, out, s);
+        } else if (tail == "consumer") {
+            // fusion still applies; catches a dispatcher that skips one node too many
+            ggml_set_name(out, "mul");
+            out = ggml_add(ctx, out, scale);
+        } else if (!tail.empty()) {
+            GGML_ABORT("unknown tail %s", tail.c_str());
+        }
+        ggml_set_name(out, "out");
+
+        return out;
+    }
+};
+
 // SNAKE activation fusion: y = x + sin(a*x)^2 * inv_b
 // CUDA backend matches the naive 5-op chain (mul, sin, sqr, mul, add)
 // and dispatches a single fused kernel.
@@ -9393,6 +9480,24 @@ static std::vector> make_test_cases_eval() {
         }
     }
 
+    // fused add + unary + mul: the delta-net alpha gate, bias and scale broadcast over dim 0
+    for (ggml_unary_op op : { GGML_UNARY_OP_SILU, GGML_UNARY_OP_SIGMOID, GGML_UNARY_OP_SOFTPLUS }) {
+        test_cases.emplace_back(new test_add_unary_mul(op, GGML_TYPE_F32, { 32, 7, 1, 1 }));
+        test_cases.emplace_back(new test_add_unary_mul(op, GGML_TYPE_F32, { 32, 512, 1, 1 }));
+        test_cases.emplace_back(new test_add_unary_mul(op, GGML_TYPE_F32, { 5, 7, 11, 13 }));
+        // one token: no broadcast left, and the unary result may be either MUL operand
+        test_cases.emplace_back(new test_add_unary_mul(op, GGML_TYPE_F32, { 32, 1, 1, 1 }, false, "same_shape"));
+        test_cases.emplace_back(new test_add_unary_mul(op, GGML_TYPE_F32, { 32, 1, 1, 1 }, true, "same_shape"));
+        test_cases.emplace_back(new test_add_unary_mul(op, GGML_TYPE_F32, { 32, 7, 1, 1 }, false, "bcast", "consumer"));
+        // must not fuse
+        test_cases.emplace_back(new test_add_unary_mul(op, GGML_TYPE_F32, { 32, 7, 1, 1 }, false, "same_shape"));
+        test_cases.emplace_back(new test_add_unary_mul(op, GGML_TYPE_F32, { 32, 7, 1, 1 }, false, "rep_ne0"));
+        test_cases.emplace_back(new test_add_unary_mul(op, GGML_TYPE_F32, { 32, 7, 1, 1 }, false, "bcast", "reuse"));
+        test_cases.emplace_back(new test_add_unary_mul(op, GGML_TYPE_F16, { 32, 7, 1, 1 }));
+    }
+    // a unary op with no fused kernel must fall back to the three-op chain
+    test_cases.emplace_back(new test_add_unary_mul(GGML_UNARY_OP_GELU, GGML_TYPE_F32, { 32, 7, 1, 1 }));
+
     // SNAKE activation fusion: x + sin(a*x)^2 * inv_b
     for (ggml_type type : { GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16 }) {
         test_cases.emplace_back(new test_snake_fuse(type, {   5,   7, 1, 1}));   // primes sub-block

From fda186661304735bb9027abd472c1403a9902081 Mon Sep 17 00:00:00 2001
From: Eden 
Date: Thu, 8 Oct 2026 14:46:21 +0800
Subject: [PATCH 32/40] CUDA: fix norm family kernels when ne[2]/ne[3] exceed
 grid dim limits (#28175)

---
 ggml/src/ggml-cuda/norm.cu | 249 ++++++++++++++++++++-----------------
 ggml/src/ggml-cuda/rope.cu | 130 ++++++++++---------
 tests/test-backend-ops.cpp |  18 ++-
 3 files changed, 225 insertions(+), 172 deletions(-)

diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu
index 5543307be..d804b933b 100644
--- a/ggml/src/ggml-cuda/norm.cu
+++ b/ggml/src/ggml-cuda/norm.cu
@@ -3,38 +3,46 @@
 
 template 
 static __global__ void norm_f32(
-        const float * x, float * dst, const int ncols, const int64_t stride_row, const int64_t stride_channel,
-        const int64_t stride_sample, const float eps) {
-    const int nrows     = gridDim.x;
-    const int nchannels = gridDim.y;
+        const float * x, float * dst, const int ncols, const int nchannels, const int nsamples,
+        const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps) {
+    const int nrows = gridDim.x;
+    const int row   = blockIdx.x;
+    const int tid   = threadIdx.x;
 
-    const int row       = blockIdx.x;
-    const int channel   = blockIdx.y;
-    const int sample    = blockIdx.z;
-    const int tid       = threadIdx.x;
-
-    x   += sample*stride_sample + channel*stride_channel + row*stride_row;
-    dst += ((sample*nchannels + channel)*nrows + row)*ncols;
-
-    float2 mean_var = make_float2(0.0f, 0.0f);
+    extern __shared__ float2 s_sum2[];
 
     ggml_cuda_pdl_sync();
-    for (int col = tid; col < ncols; col += block_size) {
-        const float xi = x[col];
-        mean_var.x += xi;
-        mean_var.y += xi * xi;
-    }
 
-    // sum up partial sums
-    extern __shared__ float2 s_sum2[];
-    mean_var = block_reduce(mean_var, s_sum2);
+    // grid.y and grid.z are clamped to the CUDA limit, iterate over the excess channels/samples
+    for (int sample = blockIdx.z; sample < nsamples; sample += gridDim.z) {
+        for (int channel = blockIdx.y; channel < nchannels; channel += gridDim.y) {
+            const float * xc   = x   + sample*stride_sample + channel*stride_channel + row*stride_row;
+            float       * dstc = dst + ((sample*nchannels + channel)*nrows + row)*ncols;
 
-    const float mean = mean_var.x / ncols;
-    const float var = mean_var.y / ncols - mean * mean;
-    const float inv_std = rsqrtf(var + eps);
+            float2 mean_var = make_float2(0.0f, 0.0f);
 
-    for (int col = tid; col < ncols; col += block_size) {
-        dst[col] = (x[col] - mean) * inv_std;
+            for (int col = tid; col < ncols; col += block_size) {
+                const float xi = xc[col];
+                mean_var.x += xi;
+                mean_var.y += xi * xi;
+            }
+
+            // sum up partial sums
+            mean_var = block_reduce(mean_var, s_sum2);
+
+            const float mean = mean_var.x / ncols;
+            const float var = mean_var.y / ncols - mean * mean;
+            const float inv_std = rsqrtf(var + eps);
+
+            for (int col = tid; col < ncols; col += block_size) {
+                dstc[col] = (xc[col] - mean) * inv_std;
+            }
+
+            if constexpr (block_size > WARP_SIZE) {
+                // sync is needed as we reuse s_sum2 across block_reduce invocations, see #26385
+                __syncthreads();
+            }
+        }
     }
 }
 
@@ -77,6 +85,8 @@ template (tmp, s_sum);
+    // grid.y and grid.z are clamped to the CUDA limit, iterate over the excess channels/samples
+    for (int sample = blockIdx.z; sample < nsamples; sample += gridDim.z) {
+        for (int channel = blockIdx.y; channel < nchannels; channel += gridDim.y) {
+            const float * xc   = x   + sample*stride_sample + channel*stride_channel + row*stride_row;
+            float       * dstc = dst + ((sample*nchannels + channel)*nrows + row)*ncols;
 
-    const float mean = tmp / ncols;
-    const float scale = rsqrtf(mean + eps);
+            [[maybe_unused]] const float * mulc = nullptr;
+            if constexpr (do_multiply) {
+                const uint32_t mul_row     = fastmodulo(row, mul_nrows_packed);
+                const uint32_t mul_channel = fastmodulo(channel, mul_nchannels_packed);
+                const uint32_t mul_sample  = fastmodulo(sample, mul_nsamples_packed);
+                mulc = mul + mul_sample * mul_stride_sample + mul_channel * mul_stride_channel + mul_row * mul_stride_row;
+            }
 
-    for (int col = tid; col < ncols; col += block_size) {
-        if constexpr (do_multiply && do_add) {
-            const int mul_col = fastmodulo(col, mul_ncols_packed);
-            const int add_col = fastmodulo(col, add_ncols_packed);
-            dst[col]          = scale * x[col] * mul[mul_col] + add[add_col];
-        } else if constexpr (do_multiply) {
-            const int mul_col = fastmodulo(col, mul_ncols_packed);
-            dst[col]          = scale * x[col] * mul[mul_col];
-        } else if constexpr (do_scale) {
-            dst[col] = scale_out * (scale * x[col]);
-        } else {
-            dst[col] = scale * x[col];
+            [[maybe_unused]] const float * addc = nullptr;
+            if constexpr (do_add) {
+                const int add_row     = fastmodulo(row, add_nrows_packed);
+                const int add_channel = fastmodulo(channel, add_nchannels_packed);
+                const int add_sample  = fastmodulo(sample, add_nsamples_packed);
+                addc = add + add_sample * add_stride_sample + add_channel * add_stride_channel + add_row * add_stride_row;
+            }
+
+            float tmp = 0.0f; // partial sum for thread in warp
+
+            for (int col = tid; col < ncols; col += block_size) {
+                const float xi = xc[col];
+                tmp += xi * xi;
+            }
+
+            // sum up partial sums
+            tmp = block_reduce(tmp, s_sum);
+
+            const float mean = tmp / ncols;
+            const float scale = rsqrtf(mean + eps);
+
+            for (int col = tid; col < ncols; col += block_size) {
+                if constexpr (do_multiply && do_add) {
+                    const int mul_col = fastmodulo(col, mul_ncols_packed);
+                    const int add_col = fastmodulo(col, add_ncols_packed);
+                    dstc[col]         = scale * xc[col] * mulc[mul_col] + addc[add_col];
+                } else if constexpr (do_multiply) {
+                    const int mul_col = fastmodulo(col, mul_ncols_packed);
+                    dstc[col]         = scale * xc[col] * mulc[mul_col];
+                } else if constexpr (do_scale) {
+                    dstc[col] = scale_out * (scale * xc[col]);
+                } else {
+                    dstc[col] = scale * xc[col];
+                }
+            }
+
+            if constexpr (block_size > WARP_SIZE) {
+                // sync is needed as we reuse s_sum across block_reduce invocations, see #26385
+                __syncthreads();
+            }
         }
     }
 }
@@ -247,50 +267,57 @@ static __global__ void rms_norm_back_f32(
 
 template 
 static __global__ void l2_norm_f32(
-        const float * x, float * dst, const int ncols, const int64_t stride_row, const int64_t stride_channel,
-        const int64_t stride_sample, const float eps) {
-    const int nrows     = gridDim.x;
-    const int nchannels = gridDim.y;
+        const float * x, float * dst, const int ncols, const int nchannels, const int nsamples,
+        const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps) {
+    const int nrows = gridDim.x;
+    const int row   = blockIdx.x;
+    const int tid   = threadIdx.x;
 
-    const int row       = blockIdx.x;
-    const int channel   = blockIdx.y;
-    const int sample    = blockIdx.z;
-    const int tid       = threadIdx.x;
-
-    x   += sample*stride_sample + channel*stride_channel + row*stride_row;
-    dst += ((sample*nchannels + channel)*nrows + row)*ncols;
-
-    float tmp = 0.0f; // partial sum for thread in warp
+    extern __shared__ float s_sum[];
 
     ggml_cuda_pdl_sync();
-    for (int col = tid; col < ncols; col += block_size) {
-        const float xi = x[col];
-        tmp += xi * xi;
-    }
 
-    // sum up partial sums
-    extern __shared__ float s_sum[];
-    tmp = block_reduce(tmp, s_sum);
-    ggml_cuda_pdl_lc();
+    // grid.y and grid.z are clamped to the CUDA limit, iterate over the excess channels/samples
+    for (int sample = blockIdx.z; sample < nsamples; sample += gridDim.z) {
+        for (int channel = blockIdx.y; channel < nchannels; channel += gridDim.y) {
+            const float * xc   = x   + sample*stride_sample + channel*stride_channel + row*stride_row;
+            float       * dstc = dst + ((sample*nchannels + channel)*nrows + row)*ncols;
 
-    // from https://pytorch.org/docs/stable/generated/torch.nn.functional.normalize.html
-    const float scale = rsqrtf(fmaxf(tmp, eps * eps));
+            float tmp = 0.0f; // partial sum for thread in warp
 
-    for (int col = tid; col < ncols; col += block_size) {
-        dst[col] = scale * x[col];
+            for (int col = tid; col < ncols; col += block_size) {
+                const float xi = xc[col];
+                tmp += xi * xi;
+            }
+
+            // sum up partial sums
+            tmp = block_reduce(tmp, s_sum);
+
+            // from https://pytorch.org/docs/stable/generated/torch.nn.functional.normalize.html
+            const float scale = rsqrtf(fmaxf(tmp, eps * eps));
+
+            for (int col = tid; col < ncols; col += block_size) {
+                dstc[col] = scale * xc[col];
+            }
+
+            if constexpr (block_size > WARP_SIZE) {
+                // sync is needed as we reuse s_sum across block_reduce invocations, see #26385
+                __syncthreads();
+            }
+        }
     }
 }
 
 static void norm_f32_cuda(
         const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples,
         const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream) {
-    const dim3 blocks_num(nrows, nchannels, nsamples);
+    const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX));
     if (ncols < 1024) {
         const dim3 block_dims(WARP_SIZE, 1, 1);
-        norm_f32<<>>(x, dst, ncols, stride_row, stride_channel, stride_sample, eps);
+        norm_f32<<>>(x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps);
     } else {
         const dim3 block_dims(1024, 1, 1);
-        norm_f32<1024><< WARP_SIZE ? 32 * sizeof(float2): 0, stream>>>(x, dst, ncols, stride_row, stride_channel, stride_sample, eps);
+        norm_f32<1024><< WARP_SIZE ? 32 * sizeof(float2): 0, stream>>>(x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps);
     }
 }
 
@@ -310,19 +337,19 @@ static void rms_norm_f32_cuda(
         const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples,
         const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream,
         const float scale_out = 1.0f) {
-    const dim3 blocks_num(nrows, nchannels, nsamples);
+    const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX));
     if (ncols < 1024) {
         const dim3 block_dims(256, 1, 1);
         const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
         ggml_cuda_kernel_launch(rms_norm_f32<256, false, false, do_scale>, launch_params,
-            x, dst, ncols, stride_row, stride_channel, stride_sample, eps,
+            x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
         // underlying cudaLaunchKernelEx does not support default params
         nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0),
         nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), scale_out);
     } else {
         const dim3 block_dims(1024, 1, 1);
         const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
-        ggml_cuda_kernel_launch(rms_norm_f32<1024, false, false, do_scale>, launch_params, x, dst, ncols, stride_row, stride_channel, stride_sample, eps,
+        ggml_cuda_kernel_launch(rms_norm_f32<1024, false, false, do_scale>, launch_params, x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
         // underlying cudaLaunchKernelEx does not support default params
         nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0),
         nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), scale_out);
@@ -356,7 +383,7 @@ static void rms_norm_mul_f32_cuda(const float *  x,
                                   const uint32_t add_nsamples,
                                   const float    eps,
                                   cudaStream_t   stream) {
-    const dim3 blocks_num(nrows, nchannels, nsamples);
+    const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX));
     if (mul == nullptr) {
         rms_norm_f32_cuda(x, dst, ncols, nrows, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, stream);
         return;
@@ -370,7 +397,7 @@ static void rms_norm_mul_f32_cuda(const float *  x,
             const dim3 block_dims(256, 1, 1);
             const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
             ggml_cuda_kernel_launch(rms_norm_f32<256, true>, launch_params,
-                x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
+                x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
                 mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
                 // underlying cudaLaunchKernelEx does not support default params
             nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), 1.0f);
@@ -378,7 +405,7 @@ static void rms_norm_mul_f32_cuda(const float *  x,
             const dim3 block_dims(1024, 1, 1);
             const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
             ggml_cuda_kernel_launch(rms_norm_f32<1024, true>, launch_params,
-                x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
+                x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
                 mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
                 // underlying cudaLaunchKernelEx does not support default params
             nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), 1.0f);
@@ -397,7 +424,7 @@ static void rms_norm_mul_f32_cuda(const float *  x,
             const dim3 block_dims(256, 1, 1);
             const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims,block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
             ggml_cuda_kernel_launch(rms_norm_f32<256, true, true>, launch_params,
-                x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
+                x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
                 mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add,
                 add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed,
                 add_nchannels_packed, add_nsamples_packed, 1.0f);
@@ -405,7 +432,7 @@ static void rms_norm_mul_f32_cuda(const float *  x,
             const dim3 block_dims(1024, 1, 1);
             const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
             ggml_cuda_kernel_launch(rms_norm_f32<1024, true, true>, launch_params,
-                x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
+                x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
                 mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add,
                 add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed,
                 add_nchannels_packed, add_nsamples_packed, 1.0f);
@@ -426,15 +453,15 @@ static void rms_norm_back_f32_cuda(const float * grad, const float * xf, float *
 static void l2_norm_f32_cuda(
         const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples,
         const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream) {
-    const dim3 blocks_num(nrows, nchannels, nsamples);
+    const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX));
     if (ncols < 1024) {
         const dim3 block_dims(WARP_SIZE, 1, 1);
         const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, 0, stream};
-        ggml_cuda_kernel_launch(l2_norm_f32, launch_params, x, dst, ncols, stride_row, stride_channel, stride_sample, eps);
+        ggml_cuda_kernel_launch(l2_norm_f32, launch_params, x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps);
     } else {
         const dim3 block_dims(1024, 1, 1);
         const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
-        ggml_cuda_kernel_launch(l2_norm_f32<1024>, launch_params, x, dst, ncols, stride_row, stride_channel, stride_sample, eps);
+        ggml_cuda_kernel_launch(l2_norm_f32<1024>, launch_params, x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps);
     }
 }
 
diff --git a/ggml/src/ggml-cuda/rope.cu b/ggml/src/ggml-cuda/rope.cu
index e546fb655..b4fca778f 100644
--- a/ggml/src/ggml-cuda/rope.cu
+++ b/ggml/src/ggml-cuda/rope.cu
@@ -709,7 +709,7 @@ void ggml_cuda_op_rope_fused(ggml_backend_cuda_context & ctx, ggml_tensor * rope
 // one block per row: block_reduce gives the norm scale, then each thread applies mul and rope to the elements it owns
 template 
 static __global__ void rms_norm_mul_rope_f32(
-        const float * x, D * dst, const int ncols,
+        const float * x, D * dst, const int ncols, const int nchannels, const int nsamples,
         const int64_t s01, const int64_t s02, const int64_t s03,
         const int64_t s1, const int64_t s2, const int64_t s3,
         const float eps,
@@ -724,66 +724,76 @@ static __global__ void rms_norm_mul_rope_f32(
         const int64_t * row_indices, const int set_rows_stride,
         const bool is_neox) {
     ggml_cuda_pdl_lc();
-    const int row     = blockIdx.x;
-    const int channel = blockIdx.y;
-    const int sample  = blockIdx.z;
-    const int tid     = threadIdx.x;
-
-    x += sample*s03 + channel*s02 + row*s01;
-
-    const uint32_t mul_row     = fastmodulo(row,     mul_nrows_packed);
-    const uint32_t mul_channel = fastmodulo(channel, mul_nchannels_packed);
-    const uint32_t mul_sample  = fastmodulo(sample,  mul_nsamples_packed);
-    mul += mul_sample*mul_s03 + mul_channel*mul_s02 + mul_row*mul_s01;
-
-    float tmp = 0.0f;
-
-    ggml_cuda_pdl_sync();
-    for (int col = tid; col < ncols; col += block_size) {
-        const float xi = x[col];
-        tmp += xi * xi;
-    }
+    const int row = blockIdx.x;
+    const int tid = threadIdx.x;
 
     extern __shared__ float s_sum[];
-    tmp = block_reduce(tmp, s_sum);
 
-    const float scale = rsqrtf(tmp/ncols + eps);
+    ggml_cuda_pdl_sync();
 
-    int64_t idst = sample*s3 + channel*s2 + row*s1;
-    if (set_rows_stride != 0) {
-        idst = row*s1 + row_indices[channel]*set_rows_stride;
-    }
-    dst += idst;
+    // grid.y and grid.z are clamped to the CUDA limit, iterate over the excess channels/samples
+    for (int sample = blockIdx.z; sample < nsamples; sample += gridDim.z) {
+        for (int channel = blockIdx.y; channel < nchannels; channel += gridDim.y) {
+            const float * xc = x + sample*s03 + channel*s02 + row*s01;
 
-    for (int i0 = 2*tid; i0 < ncols; i0 += 2*block_size) {
-        int ix0;
-        int ix1;
-        if (is_neox && i0 < n_dims) {
-            ix0 = i0/2;
-            ix1 = i0/2 + n_dims/2;
-        } else {
-            ix0 = i0 + 0;
-            ix1 = i0 + 1;
+            const uint32_t mul_row     = fastmodulo(row,     mul_nrows_packed);
+            const uint32_t mul_channel = fastmodulo(channel, mul_nchannels_packed);
+            const uint32_t mul_sample  = fastmodulo(sample,  mul_nsamples_packed);
+            const float * mulc = mul + mul_sample*mul_s03 + mul_channel*mul_s02 + mul_row*mul_s01;
+
+            float tmp = 0.0f;
+
+            for (int col = tid; col < ncols; col += block_size) {
+                const float xi = xc[col];
+                tmp += xi * xi;
+            }
+
+            tmp = block_reduce(tmp, s_sum);
+
+            const float scale = rsqrtf(tmp/ncols + eps);
+
+            int64_t idst = sample*s3 + channel*s2 + row*s1;
+            if (set_rows_stride != 0) {
+                idst = row*s1 + row_indices[channel]*set_rows_stride;
+            }
+            D * dstc = dst + idst;
+
+            for (int i0 = 2*tid; i0 < ncols; i0 += 2*block_size) {
+                int ix0;
+                int ix1;
+                if (is_neox && i0 < n_dims) {
+                    ix0 = i0/2;
+                    ix1 = i0/2 + n_dims/2;
+                } else {
+                    ix0 = i0 + 0;
+                    ix1 = i0 + 1;
+                }
+
+                const float x0 = scale * xc[ix0] * mulc[fastmodulo(ix0, mul_ncols_packed)];
+                const float x1 = scale * xc[ix1] * mulc[fastmodulo(ix1, mul_ncols_packed)];
+
+                if (i0 >= n_dims) {
+                    dstc[ix0] = ggml_cuda_cast(x0);
+                    dstc[ix1] = ggml_cuda_cast(x1);
+                    continue;
+                }
+
+                const float theta_base  = pos[channel]*powf(theta_scale, i0/2.0f);
+                const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f;
+
+                float cos_theta;
+                float sin_theta;
+                rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta);
+
+                dstc[ix0] = ggml_cuda_cast(x0*cos_theta - x1*sin_theta);
+                dstc[ix1] = ggml_cuda_cast(x0*sin_theta + x1*cos_theta);
+            }
+
+            if constexpr (block_size > WARP_SIZE) {
+                // sync is needed as we reuse s_sum across block_reduce invocations, see #26385
+                __syncthreads();
+            }
         }
-
-        const float x0 = scale * x[ix0] * mul[fastmodulo(ix0, mul_ncols_packed)];
-        const float x1 = scale * x[ix1] * mul[fastmodulo(ix1, mul_ncols_packed)];
-
-        if (i0 >= n_dims) {
-            dst[ix0] = ggml_cuda_cast(x0);
-            dst[ix1] = ggml_cuda_cast(x1);
-            continue;
-        }
-
-        const float theta_base  = pos[channel]*powf(theta_scale, i0/2.0f);
-        const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f;
-
-        float cos_theta;
-        float sin_theta;
-        rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta);
-
-        dst[ix0] = ggml_cuda_cast(x0*cos_theta - x1*sin_theta);
-        dst[ix1] = ggml_cuda_cast(x0*sin_theta + x1*cos_theta);
     }
 }
 
@@ -806,7 +816,7 @@ static void rms_norm_mul_rope_cuda(
         const bool is_neox, cudaStream_t stream) {
     GGML_ASSERT(ncols % 2 == 0);
 
-    const dim3 blocks_num(nrows, nchannels, nsamples);
+    const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX));
 
     const float theta_scale = powf(freq_base, -2.0f/n_dims);
 
@@ -820,13 +830,13 @@ static void rms_norm_mul_rope_cuda(
         const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, 32*sizeof(float), stream};
         if (freq_factors == nullptr) {
             ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<256, false, D>, launch_params,
-                x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
+                x, dst, ncols, nchannels, nsamples, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
                 mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
                 n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
                 freq_factors, row_indices, set_rows_stride, is_neox);
         } else {
             ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<256, true, D>, launch_params,
-                x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
+                x, dst, ncols, nchannels, nsamples, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
                 mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
                 n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
                 freq_factors, row_indices, set_rows_stride, is_neox);
@@ -836,13 +846,13 @@ static void rms_norm_mul_rope_cuda(
         const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, 32*sizeof(float), stream};
         if (freq_factors == nullptr) {
             ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<1024, false, D>, launch_params,
-                x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
+                x, dst, ncols, nchannels, nsamples, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
                 mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
                 n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
                 freq_factors, row_indices, set_rows_stride, is_neox);
         } else {
             ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<1024, true, D>, launch_params,
-                x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
+                x, dst, ncols, nchannels, nsamples, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
                 mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
                 n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
                 freq_factors, row_indices, set_rows_stride, is_neox);
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index f31016bbb..a38766cba 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -2826,7 +2826,14 @@ struct test_rms_norm_mul_rope : public test_case {
     }
 
     double max_nmse_err() override {
-        return ne[0] == 8192 ? 5e-6 : test_case::max_nmse_err();
+        if (ne[0] == 8192) {
+            return 5e-6;
+        }
+        // large positions amplify the difference between the GPU and CPU trig functions
+        if (ne[2] > 8192) {
+            return 1e-5;
+        }
+        return test_case::max_nmse_err();
     }
 };
 
@@ -10273,6 +10280,12 @@ static std::vector> make_test_cases_eval() {
 
     // in-place tests
     test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, {64, 5, 4, 3}, false, 1e-6f, true));
+    // shapes below exceed the CUDA gridDim.y/gridDim.z limit of 65535 (#27901)
+    test_cases.emplace_back(new test_norm        (GGML_TYPE_F32, {4, 1, 65536, 1}, false, 1e-6f));
+    test_cases.emplace_back(new test_rms_norm    (GGML_TYPE_F32, {4, 1, 65536, 1}, false, 1e-6f, false));
+    test_cases.emplace_back(new test_rms_norm    (GGML_TYPE_F32, {4, 1, 1, 65536}, false, 1e-6f, false));
+    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 (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));
@@ -10331,6 +10344,9 @@ static std::vector> make_test_cases_eval() {
             }
         }
     }
+    // ne[2] > 65535 exceeds the CUDA gridDim.y limit (#27901)
+    test_cases.emplace_back(new test_rms_norm_mul_rope({4, 1, 65536, 1}, 1e-6f, false, false, false, GGML_ROPE_TYPE_NORMAL));
+    test_cases.emplace_back(new test_rms_norm_mul_rope({4, 1, 65536, 1}, 1e-6f, false, true,  false, GGML_ROPE_TYPE_NEOX));
     for (int64_t d_conv : {3, 4, 9}) {
         for (int64_t d_inner: {1024, 1536, 2048}) {
             test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {d_conv, d_inner, 1, 1}, {d_conv, d_inner, 1, 1}));

From d8880160413fc77f63d6a73a5b1ac21ac65e36e1 Mon Sep 17 00:00:00 2001
From: Bertay Eren <39909689+bertaye@users.noreply.github.com>
Date: Thu, 8 Oct 2026 09:46:38 +0300
Subject: [PATCH 33/40] cuda : Use byte strides for roll to allow
 non-contiguous ROLL operations (#29547)

---
 ggml/src/ggml-cuda/ggml-cuda.cu | 2 +-
 ggml/src/ggml-cuda/roll.cu      | 8 ++++++--
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index 33ebfdfa0..20775a921 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -5658,7 +5658,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
             return max_bias == 0.0f;
         }
         case GGML_OP_ROLL:
-            if(op->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[0])) {
+            if(op->src[0]->type == GGML_TYPE_F32) {
                 return true;
             }
             return false;
diff --git a/ggml/src/ggml-cuda/roll.cu b/ggml/src/ggml-cuda/roll.cu
index a339dfc1a..c9bbe76f0 100644
--- a/ggml/src/ggml-cuda/roll.cu
+++ b/ggml/src/ggml-cuda/roll.cu
@@ -17,6 +17,10 @@ static __global__ void roll_f32_cuda(const float * __restrict__ src,
                                      const int64_t ne01,
                                      const int64_t ne02,
                                      const int64_t ne03,
+                                     const int64_t nb00,
+                                     const int64_t nb01,
+                                     const int64_t nb02,
+                                     const int64_t nb03,
                                      const int     s0,
                                      const int     s1,
                                      const int     s2,
@@ -39,7 +43,7 @@ static __global__ void roll_f32_cuda(const float * __restrict__ src,
     const int64_t d3 = wrap_index(i3 - s3, ne03);
 
     dst[i3 * (ne00 * ne01 * ne02) + i2 * (ne01 * ne00) + i1 * ne00 + i0] =
-        src[d3 * (ne00 * ne01 * ne02) + d2 * (ne01 * ne00) + d1 * ne00 + d0];
+        src[(d3 * nb03 + d2 * nb02 + d1 * nb01 + d0 * nb00) / sizeof(float)];
 }
 
 void ggml_cuda_op_roll(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
@@ -63,5 +67,5 @@ void ggml_cuda_op_roll(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     int64_t num_blocks = (sz + CUDA_ROLL_BLOCK_SIZE - 1) / CUDA_ROLL_BLOCK_SIZE;
 
     roll_f32_cuda<<>>(
-        src0_d, dst_d, ne00, ne01, ne02, ne03, s0, s1, s2, s3);
+        src0_d, dst_d, ne00, ne01, ne02, ne03, nb00, nb01, nb02, nb03, s0, s1, s2, s3);
 }

From 097f5b5332f46559f1ee11d9d64e48a5065b2989 Mon Sep 17 00:00:00 2001
From: Max Krasnyansky 
Date: Wed, 7 Oct 2026 23:48:06 -0700
Subject: [PATCH 34/40] hex-mmadd: do not assume aligned read/write when
 bias-add is fused (#30133)

---
 ggml/src/ggml-hexagon/htp/matmul-ops.c | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c
index a589e0a20..1a513d160 100644
--- a/ggml/src/ggml-hexagon/htp/matmul-ops.c
+++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c
@@ -425,7 +425,7 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void
     if (copy_cnt > 0) {                                                                                                  \
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct_end);                                                       \
         if (src2) {                                                                                                      \
-            hvx_add_f32_uaa((uint8_t *) &dst_col[src0_start_row],                                                        \
+            hvx_add_f32_uuu((uint8_t *) &dst_col[src0_start_row],                                                        \
                             (const uint8_t *) tmp,                                                                       \
                             (const uint8_t *) ((const float *) mmctx->vtcm_bias + src0_start_row),                       \
                             copy_cnt);                                                                                   \
@@ -1108,7 +1108,7 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) {
     if (copy_cnt > 0) {
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, src0_end_row);
         if (src2) {
-            hvx_add_f32_uaa((uint8_t *) &dst_col[src0_start_row],
+            hvx_add_f32_uuu((uint8_t *) &dst_col[src0_start_row],
                             (const uint8_t *) tmp,
                             (const uint8_t *) ((const float *) mmctx->vtcm_bias + src0_start_row),
                             copy_cnt);

From 46baf1f1fec5a06d1e52122a9207ca978b720f95 Mon Sep 17 00:00:00 2001
From: bri-prism <288398250+bri-prism@users.noreply.github.com>
Date: Wed, 7 Oct 2026 23:50:22 -0700
Subject: [PATCH 35/40] sycl: FWHT optimizations (#29605)

---
 ggml/src/ggml-sycl/fwht.cpp | 72 +++++++++++++++++++++----------------
 tests/test-backend-ops.cpp  | 11 ++++++
 2 files changed, 53 insertions(+), 30 deletions(-)

diff --git a/ggml/src/ggml-sycl/fwht.cpp b/ggml/src/ggml-sycl/fwht.cpp
index fb48d7fec..5af549c23 100644
--- a/ggml/src/ggml-sycl/fwht.cpp
+++ b/ggml/src/ggml-sycl/fwht.cpp
@@ -46,8 +46,8 @@ static constexpr float H20[20][20] = {
 #undef P
 #undef N
 
-template 
-static void fwht_kernel(const float * __restrict__ src, float * __restrict__ dst, const int64_t n_rows,
+template 
+static void fwht_kernel(const T * __restrict__ src, float * __restrict__ dst, const int64_t n_rows,
                         const float scale, const sycl::nd_item<2> & item) {
     const sycl::sub_group sg = item.get_sub_group();
 
@@ -67,7 +67,7 @@ static void fwht_kernel(const float * __restrict__ src, float * __restrict__ dst
 
 #pragma unroll
     for (int i = 0; i < el_w; ++i) {
-        reg[i] = src[i * WARP_SIZE + lane] * scale;
+        reg[i] = static_cast(src[i * WARP_SIZE + lane]) * scale;
     }
 
     // Butterflies inside the sub-group. The partner of a lane with bit h clear is the
@@ -107,8 +107,8 @@ static void fwht_kernel(const float * __restrict__ src, float * __restrict__ dst
     }
 }
 
-template 
-static void launch_fwht(const float * src, float * dst, const int64_t n_rows, const float scale,
+template 
+static void launch_fwht(const T * src, float * dst, const int64_t n_rows, const float scale,
                         dpct::queue_ptr stream) {
     constexpr int rows_per_block = 4;
 
@@ -120,7 +120,7 @@ static void launch_fwht(const float * src, float * dst, const int64_t n_rows, co
 
     stream->parallel_for(sycl::nd_range<2>(global, local),
                          [=](sycl::nd_item<2> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                             fwht_kernel(src, dst, n_rows, scale, item);
+                             fwht_kernel(src, dst, n_rows, scale, item);
                          });
 }
 
@@ -128,8 +128,8 @@ static void launch_fwht(const float * src, float * dst, const int64_t n_rows, co
 // keeps N/NT values rather than N/WARP_SIZE. Butterflies below the sub-group width
 // still shuffle; those up to NT go through work-group local memory; the rest stay
 // in registers.
-template 
-static void fwht_kernel_wide(const float * __restrict__ src,
+template 
+static void fwht_kernel_wide(const T * __restrict__ src,
                              float * __restrict__ dst,
                              const int64_t            n_rows,
                              const float              scale,
@@ -151,7 +151,7 @@ static void fwht_kernel_wide(const float * __restrict__ src,
     float reg[el_w];
 #pragma unroll
     for (int i = 0; i < el_w; ++i) {
-        reg[i] = src[i * NT + tid] * scale;
+        reg[i] = static_cast(src[i * NT + tid]) * scale;
     }
 
     const sycl::sub_group sg   = item.get_sub_group();
@@ -207,8 +207,8 @@ static void fwht_kernel_wide(const float * __restrict__ src,
     }
 }
 
-template 
-static void launch_fwht_wide(const float *   src,
+template 
+static void launch_fwht_wide(const T *       src,
                              float *         dst,
                              const int64_t   n_rows,
                              const float     scale,
@@ -220,13 +220,13 @@ static void launch_fwht_wide(const float *   src,
         sycl::local_accessor smem(sycl::range<1>(N), cgh);
         cgh.parallel_for(sycl::nd_range<2>(global, local),
                          [=](sycl::nd_item<2> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                             fwht_kernel_wide(src, dst, n_rows, scale, item, get_pointer(smem));
+                             fwht_kernel_wide(src, dst, n_rows, scale, item, get_pointer(smem));
                          });
     });
 }
 
-template 
-static void kronecker_kernel(const float * __restrict__ src,
+template 
+static void kronecker_kernel(const T * __restrict__ src,
                              float * __restrict__ dst,
                              const int64_t            n_rows,
                              const float              scale,
@@ -255,7 +255,7 @@ static void kronecker_kernel(const float * __restrict__ src,
 
 #pragma unroll
         for (int j = 0; j < m; ++j) {
-            reg[i * m + j] = src[b_idx * m + j] * scale;
+            reg[i * m + j] = static_cast(src[b_idx * m + j]) * scale;
         }
     }
 
@@ -321,8 +321,8 @@ static void kronecker_kernel(const float * __restrict__ src,
     }
 }
 
-template 
-static void launch_kronecker(const float *   src,
+template 
+static void launch_kronecker(const T *       src,
                              float *         dst,
                              const int64_t   n_rows,
                              const float     scale,
@@ -337,25 +337,16 @@ static void launch_kronecker(const float *   src,
 
     stream->parallel_for(sycl::nd_range<2>(global, local),
                          [=](sycl::nd_item<2> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                             kronecker_kernel(src, dst, n_rows, scale, item);
+                             kronecker_kernel(src, dst, n_rows, scale, item);
                          });
 }
 
-bool ggml_sycl_op_fwht(ggml_backend_sycl_context & ctx, const ggml_tensor * src, ggml_tensor * dst) {
-    if (src->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
-        return false;
-    }
-    if (!ggml_are_same_shape(src, dst)) {
-        return false;
-    }
-    if (!ggml_is_contiguous(src) || !ggml_is_contiguous(dst)) {
-        return false;
-    }
-
+template 
+static bool ggml_sycl_op_fwht_impl(ggml_backend_sycl_context & ctx, const ggml_tensor * src, ggml_tensor * dst) {
     const int     n    = (int) src->ne[0];
     const int64_t rows = ggml_nrows(src);
 
-    const float *   src_d  = (const float *) src->data;
+    const T *       src_d  = (const T *) src->data;
     float *         dst_d  = (float *) dst->data;
     dpct::queue_ptr stream = ctx.stream();
 
@@ -402,3 +393,24 @@ bool ggml_sycl_op_fwht(ggml_backend_sycl_context & ctx, const ggml_tensor * src,
             return false;
     }
 }
+
+bool ggml_sycl_op_fwht(ggml_backend_sycl_context & ctx, const ggml_tensor * src, ggml_tensor * dst) {
+    if (dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+    if (!ggml_are_same_shape(src, dst)) {
+        return false;
+    }
+    if (!ggml_is_contiguous(src) || !ggml_is_contiguous(dst)) {
+        return false;
+    }
+
+    switch (src->type) {
+        case GGML_TYPE_F32:
+            return ggml_sycl_op_fwht_impl(ctx, src, dst);
+        case GGML_TYPE_F16:
+            return ggml_sycl_op_fwht_impl(ctx, src, dst);
+        default:
+            return false;
+    }
+}
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index a38766cba..453fb78d7 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -11923,6 +11923,17 @@ static std::vector> make_test_cases_perf() {
     test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 128, 2048, 128));
     test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 256, 2048, 256));
     test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 512, 2048, 512));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 128, 1, 128));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 64, 1, 64));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 256, 1, 256));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 128, 32, 128));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 64, 2048, 64));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 128, 2048, 128));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 256, 2048, 256));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 512, 2048, 512));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 1024, 2048, 1024));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 4096, 2048, 4096));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 8192, 2048, 8192));
 
     test_cases.emplace_back(new test_solve_tri(GGML_TYPE_F32, { 64, 64, 4, 4 }, { 32, 64, 4, 4 }));
     test_cases.emplace_back(new test_solve_tri(GGML_TYPE_F32, { 128, 128, 4, 2 }, { 32, 128, 4, 2 }));

From 08246a28f6000100433d297c4e037c02e9d2d464 Mon Sep 17 00:00:00 2001
From: Muhammad Saad 
Date: Thu, 8 Oct 2026 03:21:42 -0400
Subject: [PATCH 36/40] cuda : support arbitrary striding for unary ops on f16,
 f32, and bf16 (#29781)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* cuda : support arbitrary striding for unary ops on f16, f32, and bf16

* Remove added newline

---------

Co-authored-by: Johannes Gäßler 
---
 ggml/src/ggml-cuda/ggml-cuda.cu |  7 ++--
 ggml/src/ggml-cuda/unary.cu     | 57 +++++++++++++++++++++++++++++----
 2 files changed, 54 insertions(+), 10 deletions(-)

diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index 20775a921..995037f1b 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -5314,9 +5314,10 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                 case GGML_UNARY_OP_CEIL:
                 case GGML_UNARY_OP_ROUND:
                 case GGML_UNARY_OP_TRUNC:
-                    // TODO: should become:
-                    //return ggml_is_contiguous_rows(op->src[0]);
-                    return ggml_is_contiguous(op->src[0]);
+                    if (op->src[0]->type == GGML_TYPE_BF16 && ggml_get_unary_op(op) == GGML_UNARY_OP_XIELU) {
+                        return false;
+                    }
+                    return op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_BF16;
                 default:
                     return false;
             }
diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu
index c9a8cf6b0..436fdfffb 100644
--- a/ggml/src/ggml-cuda/unary.cu
+++ b/ggml/src/ggml-cuda/unary.cu
@@ -134,24 +134,67 @@ static void unary_cuda(const T * x, T * dst, const int k, cudaStream_t stream) {
     ggml_cuda_kernel_launch(unary_op_kernel, launch_params, x, dst, k);
 }
 
+template 
+static __global__ void unary_op_kernel_strided(const T * x, T * dst, const int k,const int64_t ne00,const int64_t ne01,const int64_t ne02,const size_t nb00,const size_t nb01,const size_t nb02,const size_t nb03) {
+    ggml_cuda_pdl_lc();
+    const int i = blockDim.x*blockIdx.x + threadIdx.x;
+
+    if (i >= k) {
+        return;
+    }
+
+    int64_t rem = i;
+    const int64_t i0 = rem % ne00; rem /= ne00;
+    const int64_t i1 = rem % ne01; rem /= ne01;
+    const int64_t i2 = rem % ne02;
+    const int64_t i3 = rem / ne02;
+    const size_t src_byte_offset = i0 * nb00 + i1 * nb01 + i2 * nb02 + i3 * nb03;
+    const T * src_ptr = (const T *)((const char *)x + src_byte_offset);
+
+    ggml_cuda_pdl_sync();
+    dst[i] = ggml_cuda_cast(op(ggml_cuda_cast(*src_ptr)));
+}
+
+template 
+static void unary_cuda_strided(const T * x, T * dst, const int k,const int64_t ne00,const int64_t ne01,const int64_t ne02,const size_t nb00,const size_t nb01,const size_t nb02,const size_t nb03, cudaStream_t stream) {
+    const int num_blocks = (k + CUDA_NEG_BLOCK_SIZE - 1) / CUDA_NEG_BLOCK_SIZE;
+    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params((dim3)num_blocks, CUDA_NEG_BLOCK_SIZE, 0, stream);
+    ggml_cuda_kernel_launch(unary_op_kernel_strided, launch_params, x, dst, k, ne00,ne01,ne02,nb00,nb01,nb02,nb03);
+}
+
 template 
 void ggml_cuda_op_unary(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const ggml_tensor * src0 = dst->src[0];
     const void * src0_d = src0->data;
     void * dst_d = dst->data;
-    cudaStream_t stream = ctx.stream();
 
-    GGML_ASSERT(ggml_is_contiguous(src0));
+    cudaStream_t stream = ctx.stream();
 
     GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16);
     GGML_ASSERT(src0->type == dst->type);
 
-    if (src0->type == GGML_TYPE_F16) {
-        unary_cuda((const half *)src0_d, (half *)dst_d, ggml_nelements(src0), stream);
-    } else if (src0->type == GGML_TYPE_BF16) {
-        unary_cuda((const nv_bfloat16 *)src0_d, (nv_bfloat16 *)dst_d, ggml_nelements(src0), stream);
+    if (ggml_is_contiguous(src0)) {
+        if (src0->type == GGML_TYPE_F16) {
+            unary_cuda((const half *)src0_d, (half *)dst_d, ggml_nelements(src0), stream);
+        } else if (src0->type == GGML_TYPE_BF16) {
+            unary_cuda((const nv_bfloat16 *)src0_d, (nv_bfloat16 *)dst_d, ggml_nelements(src0), stream);
+        } else {
+            unary_cuda((const float *)src0_d, (float *)dst_d, ggml_nelements(src0), stream);
+        }
     } else {
-        unary_cuda((const float *)src0_d, (float *)dst_d, ggml_nelements(src0), stream);
+        if (src0->type == GGML_TYPE_F16) {
+            unary_cuda_strided((const half *)src0_d, (half *)dst_d, ggml_nelements(src0),
+                src0->ne[0], src0->ne[1], src0->ne[2],
+                src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3], stream);
+        } else if (src0->type == GGML_TYPE_BF16) {
+            unary_cuda_strided((const nv_bfloat16 *)src0_d, (nv_bfloat16 *)dst_d, ggml_nelements(src0),
+                src0->ne[0], src0->ne[1], src0->ne[2],
+                src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3], stream);
+        } else {
+            unary_cuda_strided((const float *)src0_d, (float *)dst_d, ggml_nelements(src0),
+                src0->ne[0], src0->ne[1], src0->ne[2],
+                src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3], stream);
+        }
     }
 }
 

From 03aa006acb547549162a00dddeab3d3b38ebf7b2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrien=20Gallou=C3=ABt?= 
Date: Thu, 8 Oct 2026 10:57:47 +0200
Subject: [PATCH 37/40] vendor : update cpp-httplib to 0.60.1 (#30134)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Adrien Gallouët 
---
 scripts/sync_vendor.py         |  2 +-
 vendor/cpp-httplib/httplib.cpp | 10 +++-------
 vendor/cpp-httplib/httplib.h   |  8 +++++---
 3 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/scripts/sync_vendor.py b/scripts/sync_vendor.py
index c667fcb0a..df1ea1eaf 100755
--- a/scripts/sync_vendor.py
+++ b/scripts/sync_vendor.py
@@ -5,7 +5,7 @@ import os
 import sys
 import subprocess
 
-HTTPLIB_VERSION = "refs/tags/v0.60.0"
+HTTPLIB_VERSION = "refs/tags/v0.60.1"
 
 # used by examples/gguf-hash, these repos have no release tag, so we pin a commit
 XXHASH_COMMIT      = "9f465f1ea932d6ad9a26cd77496311ffa544cd68"
diff --git a/vendor/cpp-httplib/httplib.cpp b/vendor/cpp-httplib/httplib.cpp
index ac6f6cda7..2898454d9 100644
--- a/vendor/cpp-httplib/httplib.cpp
+++ b/vendor/cpp-httplib/httplib.cpp
@@ -5192,8 +5192,10 @@ bool range_error(Request &req, Response &res) {
         last_pos = content_len;
       }
 
+      // RFC 9110 14.1.2: a suffix-length longer than the representation
+      // selects the entire representation.
       if (first_pos == -1) {
-        first_pos = content_len - last_pos;
+        first_pos = (std::max)(static_cast(0), content_len - last_pos);
         last_pos = content_len - 1;
       }
 
@@ -6136,12 +6138,6 @@ bool verify_cert_with_windows_schannel(
   auto chain_guard =
       scope_exit([&] { CertFreeCertificateChain(chain_context); });
 
-  // Check if chain has errors
-  if (chain_context->TrustStatus.dwErrorStatus != CERT_TRUST_NO_ERROR) {
-    out_error = chain_context->TrustStatus.dwErrorStatus;
-    return false;
-  }
-
   // Verify SSL policy
   SSL_EXTRA_CERT_CHAIN_POLICY_PARA extra_policy_para = {};
   extra_policy_para.cbSize = sizeof(extra_policy_para);
diff --git a/vendor/cpp-httplib/httplib.h b/vendor/cpp-httplib/httplib.h
index 553924015..95432d01e 100644
--- a/vendor/cpp-httplib/httplib.h
+++ b/vendor/cpp-httplib/httplib.h
@@ -8,8 +8,8 @@
 #ifndef CPPHTTPLIB_HTTPLIB_H
 #define CPPHTTPLIB_HTTPLIB_H
 
-#define CPPHTTPLIB_VERSION "0.60.0"
-#define CPPHTTPLIB_VERSION_NUM "0x003c00"
+#define CPPHTTPLIB_VERSION "0.60.1"
+#define CPPHTTPLIB_VERSION_NUM "0x003c01"
 
 #ifdef _WIN32
 #if defined(_WIN32_WINNT) && _WIN32_WINNT < 0x0A00
@@ -4412,6 +4412,8 @@ private:
 
 namespace ws {
 
+class WebSocketClient;
+
 enum class Opcode : uint8_t {
   Continuation = 0x0,
   Text = 0x1,
@@ -4513,7 +4515,7 @@ public:
 
 private:
   friend class httplib::Server;
-  friend class WebSocketClient;
+  friend class httplib::ws::WebSocketClient;
 
   WebSocket(
       Stream &strm, const Request &req, bool is_server,

From dac308739429d60ccf99c46ecc1fbb1dff20ca77 Mon Sep 17 00:00:00 2001
From: Jeff Bolz 
Date: Thu, 8 Oct 2026 04:36:20 -0500
Subject: [PATCH 38/40] vulkan: extend sparse FA support to coopmat2 (#30003)

---
 ggml/src/ggml-vulkan/ggml-vulkan.cpp          |  8 +-
 .../vulkan-shaders/flash_attn_cm2.comp        | 99 +++++++++++--------
 tests/test-backend-ops.cpp                    | 11 +++
 3 files changed, 77 insertions(+), 41 deletions(-)

diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index 740de5b5a..998cc693c 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -8154,10 +8154,14 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
     // cm2 dense is fast, so it needs a larger reduction to win.
     // With quantized K/V, sparse only breaks even around 16x (measured on RDNA3/RDNA4).
     const int64_t min_ratio = tuning_params.path == FA_COOPMAT2 ? 4 : (kv_f16 ? 2 : 16);
+    // coopmat2 vector decode requires 8B strides.
+    auto sparse_gather_aligned = [](const ggml_tensor * t) {
+        return (t->type != GGML_TYPE_F16 && t->type != GGML_TYPE_BF16) ||
+               (t->nb[1] | t->nb[2] | t->nb[3]) % (4 * sizeof(ggml_fp16_t)) == 0;
+    };
     const bool use_sparse = !disable_sparse && n_kv_max > 0 && mask &&
                             max_bias == 0.0f && logit_softcap == 0.0f &&
-                            // the cm2 sparse gather only reads f16
-                            (kv_f16 || tuning_params.path != FA_COOPMAT2) &&
+                            (tuning_params.path != FA_COOPMAT2 || (sparse_gather_aligned(k) && sparse_gather_aligned(v))) &&
                             nem0 == KV &&
                             (int64_t)KV >= std::max(4096, min_ratio * (int64_t)n_kv_max) &&
                             (gqa_ratio > 1 || (tuning_params.path == FA_SCALAR && N == 1));
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp
index c6ed63dd4..c7253784c 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp
@@ -18,7 +18,8 @@
 #ifdef GL_NV_cooperative_matrix_decode_vector
 #extension GL_NV_cooperative_matrix_decode_vector : enable
 #endif
-#extension GL_EXT_buffer_reference : enable
+#extension GL_EXT_buffer_reference2 : enable
+#extension GL_EXT_shader_explicit_arithmetic_types_int64 : enable
 #extension GL_KHR_shader_subgroup_ballot : enable
 #extension GL_KHR_shader_subgroup_vote : enable
 #extension GL_EXT_null_initializer : enable
@@ -35,6 +36,10 @@
 #define FA_GATHER_BS 1u
 #endif
 
+layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufFA_Byte {
+    uint8_t raw;
+};
+
 // buffer_reference stride = sizeof(struct) = FaBlockBytesK/V.
 layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufFA_K {
     uint8_t raw[FaBlockBytesK];
@@ -113,48 +118,71 @@ layout (binding = 1) readonly buffer K {uint8_t data_k[];};
 layout (binding = 2) readonly buffer V {uint8_t data_v[];};
 layout (binding = 3) readonly buffer M {uint8_t data_m[];};
 
-// f16 aliases for the sparse gather callbacks.
-layout (binding = 1) readonly buffer KF16 {float16_t data_kf16[];};
-layout (binding = 2) readonly buffer VF16 {float16_t data_vf16[];};
+// Native 16-bit aliases for the sparse gather callbacks.
+layout (binding = 1) readonly buffer K16 {FLOAT_TYPE data_k16[];};
+layout (binding = 2) readonly buffer V16 {FLOAT_TYPE data_v16[];};
 layout (binding = 3) readonly buffer MF16 {float16_t data_mf16[];};
 #ifdef GL_NV_cooperative_matrix_decode_vector
-layout (binding = 1) readonly buffer KF16V4 {f16vec4 data_kf16v4[];};
-layout (binding = 2) readonly buffer VF16V4 {f16vec4 data_vf16v4[];};
+layout (binding = 1) readonly buffer K16V4 {FLOAT_TYPEV4 data_k16v4[];};
+layout (binding = 2) readonly buffer V16V4 {FLOAT_TYPEV4 data_v16v4[];};
 #endif
 
-// K/V/mask f16-element offsets for the current head/batch, set in main().
+// K/V/mask offsets in 16-bit elements for the current head/batch, set in main().
 uint32_t g_k_off_elem, g_v_off_elem, g_m_off_elem;
 
-#if !defined(BFLOAT16)
-// blockCoords are in block units: KV slot = blockCoords[0],
-// head dim = blockCoords[1]*FA_GATHER_BS + coordInBlock[1].
-float16_t faGatherK(const decodeBufFA_K unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
-    if (blockCoords[0] >= p.split_kv) { return float16_t(0); }
+FLOAT_TYPE faGatherK(const decodeBufFA_K bl_in, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
+    if (blockCoords[0] >= p.split_kv) { return FLOAT_TYPE(0.0); }
     const int r = data_sparse[sparse_base + blockCoords[0]];
-    return r < 0 ? float16_t(0) : data_kf16[g_k_off_elem + uint(r) * k_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]];
+    if (r < 0) { return FLOAT_TYPE(0.0); }
+#if !defined(BFLOAT16)
+    if (USE_DECODE_K) {
+        decodeBufFA_K block = decodeBufFA_K(decodeBufFA_Byte(bl_in) + uint64_t(uint(r) - blockCoords[0]) * k_stride * FaBlockBytesK);
+        return faDecodeK(block, blockCoords, coordInBlock);
+    }
+#endif
+    return data_k16[g_k_off_elem + uint(r) * k_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]];
 }
 
-float16_t faGatherV(const decodeBufFA_V unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
-    if (blockCoords[0] >= p.split_kv) { return float16_t(0); }
+FLOAT_TYPE faGatherV(const decodeBufFA_V bl_in, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
+    if (blockCoords[0] >= p.split_kv) { return FLOAT_TYPE(0.0); }
     const int r = data_sparse[sparse_base + blockCoords[0]];
-    return r < 0 ? float16_t(0) : data_vf16[g_v_off_elem + uint(r) * v_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]];
+    if (r < 0) { return FLOAT_TYPE(0.0); }
+#if !defined(BFLOAT16)
+    if (USE_DECODE_V) {
+        decodeBufFA_V block = decodeBufFA_V(decodeBufFA_Byte(bl_in) + uint64_t(uint(r) - blockCoords[0]) * v_stride * FaBlockBytesV);
+        return faDecodeV(block, blockCoords, coordInBlock);
+    }
+#endif
+    return data_v16[g_v_off_elem + uint(r) * v_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]];
 }
 
 #ifdef GL_NV_cooperative_matrix_decode_vector
-f16vec4 faGatherKVector(const decodeBufFA_K unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
-    if (blockCoords[0] >= p.split_kv) { return f16vec4(0); }
+FLOAT_TYPEV4 faGatherKVector(const decodeBufFA_K bl_in, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
+    if (blockCoords[0] >= p.split_kv) { return FLOAT_TYPEV4(0.0); }
     const int r = data_sparse[sparse_base + blockCoords[0]];
-    if (r < 0) { return f16vec4(0); }
+    if (r < 0) { return FLOAT_TYPEV4(0.0); }
+#if !defined(BFLOAT16)
+    if (USE_DECODE_K) {
+        decodeBufFA_K block = decodeBufFA_K(decodeBufFA_Byte(bl_in) + uint64_t(uint(r) - blockCoords[0]) * k_stride * FaBlockBytesK);
+        return faDecodeKVector(block, blockCoords, coordInBlock);
+    }
+#endif
     const uint32_t o = g_k_off_elem + uint(r) * k_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1];
-    return data_kf16v4[o / 4];
+    return data_k16v4[o / 4];
 }
 
-f16vec4 faGatherVVector(const decodeBufFA_V unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
-    if (blockCoords[0] >= p.split_kv) { return f16vec4(0); }
+FLOAT_TYPEV4 faGatherVVector(const decodeBufFA_V bl_in, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
+    if (blockCoords[0] >= p.split_kv) { return FLOAT_TYPEV4(0.0); }
     const int r = data_sparse[sparse_base + blockCoords[0]];
-    if (r < 0) { return f16vec4(0); }
+    if (r < 0) { return FLOAT_TYPEV4(0.0); }
+#if !defined(BFLOAT16)
+    if (USE_DECODE_V) {
+        decodeBufFA_V block = decodeBufFA_V(decodeBufFA_Byte(bl_in) + uint64_t(uint(r) - blockCoords[0]) * v_stride * FaBlockBytesV);
+        return faDecodeVVector(block, blockCoords, coordInBlock);
+    }
+#endif
     const uint32_t o = g_v_off_elem + uint(r) * v_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1];
-    return data_vf16v4[o / 4];
+    return data_v16v4[o / 4];
 }
 
 #define FAGATHERK , faGatherK, faGatherKVector
@@ -163,7 +191,6 @@ f16vec4 faGatherVVector(const decodeBufFA_V unused, const uint32_t blockCoords[2
 #define FAGATHERK , faGatherK
 #define FAGATHERV , faGatherV
 #endif
-#endif
 
 // Add gathered mask to S (slope==1 since sparse requires max_bias==0). col = slot in block jblk.
 ACC_TYPE faAddSparseMask(const uint32_t row, const uint32_t col, const ACC_TYPE elem, const uint32_t jblk) {
@@ -252,8 +279,8 @@ void main() {
 
     tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0);
 
-    const uint bs_k = USE_SPARSE ? FA_GATHER_BS : fa_block_elems(FaTypeK);
-    const uint bs_v = USE_SPARSE ? FA_GATHER_BS : fa_block_elems(FaTypeV);
+    const uint bs_k = USE_SPARSE ? max(FA_GATHER_BS, BLOCK_SIZE_K) : BLOCK_SIZE_K;
+    const uint bs_v = USE_SPARSE ? max(FA_GATHER_BS, BLOCK_SIZE_V) : BLOCK_SIZE_V;
     tensorLayoutK = setTensorLayoutBlockSizeNV(tensorLayoutK, 1, bs_k);
     tensorLayoutV = setTensorLayoutBlockSizeNV(tensorLayoutV, 1, bs_v);
 
@@ -384,18 +411,15 @@ void main() {
 
         uint32_t k_offset = ik2*p.nb12 + ik3*p.nb13;
         // F16: bs_k==1 (direct load). F32: bs_k==4 (vec4 / dequantFuncF32). Quantized types: bs_k==32.
-#if defined(BFLOAT16)
-        coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose);
-#else
-        const bool k_use_decode = (bs_k > 1u);
         if (USE_SPARSE) {
             coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose FAGATHERK);
-        } else if (k_use_decode) {
+#if !defined(BFLOAT16)
+        } else if (USE_DECODE_K) {
             coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose FADECODEK);
+#endif
         } else {
             coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose);
         }
-#endif
         S = coopMatMulAdd(Qf16, K_T, S);
 
         if (LOGIT_SOFTCAP) {
@@ -458,18 +482,15 @@ void main() {
 
         coopmat V;
         uint32_t v_offset = iv2*p.nb22 + iv3*p.nb23;
-#if defined(BFLOAT16)
-        coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad));
-#else
-        const bool v_use_decode = (bs_v > 1u);
         if (USE_SPARSE) {
             coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad) FAGATHERV);
-        } else if (v_use_decode) {
+#if !defined(BFLOAT16)
+        } else if (USE_DECODE_V) {
             coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad) FADECODEV);
+#endif
         } else {
             coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad));
         }
-#endif
 
         L = eM*L + rowsum;
 
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 453fb78d7..709997680 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -11526,6 +11526,15 @@ static std::vector> make_test_cases_eval() {
     // KV not a multiple of the compaction workgroup size.
     test_cases.emplace_back(new test_flash_attn_ext(128, 128, 1, { 8, 1}, 5003, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false,  512));
 
+    // Sparse gather: native block sizes, padded slots, and head/batch strides.
+    for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0, GGML_TYPE_IQ4_NL}) {
+        test_cases.emplace_back(new test_flash_attn_ext(128, 96, 2, {8, 2}, 5003, 3, true, false, 0, 0, GGML_PREC_F32, type, type, {0, 1, 2, 3}, true, false, 257));
+    }
+    test_cases.emplace_back(new test_flash_attn_ext(128, 96, 2, {8, 2}, 5003, 3, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_F16, {0, 2, 1, 3}, true, false, 257));
+    test_cases.emplace_back(new test_flash_attn_ext(128, 96, 2, {8, 2}, 5003, 3, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_Q4_0, {0, 2, 1, 3}, true, false, 257));
+    test_cases.emplace_back(new test_flash_attn_ext(128, 96, 2, {8, 2}, 5003, 3, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_F32, {0, 1, 2, 3}, true, false, 257));
+    test_cases.emplace_back(new test_flash_attn_ext(128, 96, 2, {8, 2}, 5003, 3, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F32, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, false, 257));
+
     // more V-is-sub-view-of-K cases: other head shapes, and full views with equal head sizes
     test_cases.emplace_back(new test_flash_attn_ext(320, 256, 1, {32, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true));
     test_cases.emplace_back(new test_flash_attn_ext(192, 128, 4, {8, 1},  512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true));
@@ -12017,6 +12026,8 @@ static std::vector> make_test_cases_perf() {
         test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true,   512));
         test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {12, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2048));
         test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {12, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, false, 2048));
+        test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {12, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0, {0, 1, 2, 3}, true, false, 2048));
+        test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {12, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16, {0, 1, 2, 3}, true, false, 2048));
         test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {12, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, false, 0));
     }
 

From ff5888f999f24e2f9eaf6537b4eb6b85fd559188 Mon Sep 17 00:00:00 2001
From: gianni-cor 
Date: Thu, 8 Oct 2026 12:14:35 +0200
Subject: [PATCH 39/40] vulkan : fix TOP_K for +inf/NaN inputs and k = 1 on
 negative values (#30107)

The bucket search in topk_nary_search.comp started from the range
[0, 0xFF800000), which ends just below the ordered-uint mapping of +inf,
so +inf and NaN were never counted. A workgroup block with fewer than k
countable values left the ballot empty and the shader read uninitialized
shared state (hang/device lost on NVIDIA, wrong indices on AMD), and a few
+inf in a block were selected without being counted, dropping real top
values.

Map NaN to -inf on input, start from [0, 0xFFFFFFFF) so every value is
counted, and clamp the top bucket's end (2^32) instead of wrapping to 0.

The k = 1 path compared float bits as signed integers, which orders
negative values backwards; compare floats instead.

Add test_top_k_inf to test-backend-ops: negative values, fewer than k
+inf and many -inf, for k = 1, 10, 40.

Assisted-by: Claude Opus 5.5
---
 .../vulkan-shaders/topk_nary_search.comp      | 26 ++++++++---
 tests/test-backend-ops.cpp                    | 45 +++++++++++++++++++
 2 files changed, 65 insertions(+), 6 deletions(-)

diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/topk_nary_search.comp b/ggml/src/ggml-vulkan/vulkan-shaders/topk_nary_search.comp
index 0b757f38e..3f6296bda 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/topk_nary_search.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/topk_nary_search.comp
@@ -60,7 +60,13 @@ void topk(const uint row) {
     if (gl_GlobalInvocationID.x < p.ncols_input) {
         if (p.first_pass != 0) {
             const uint row_offset = row * p.ncols_input;
-            dst_row[tid] = ivec2(gl_GlobalInvocationID.x, floatBitsToInt(data_a[row_offset + gl_GlobalInvocationID.x]));
+            // NaN ranks lowest, like -inf, so that every value has a place in
+            // the ordering the search below counts
+            float a = float(data_a[row_offset + gl_GlobalInvocationID.x]);
+            if (isnan(a)) {
+                a = uintBitsToFloat(0xFF800000);
+            }
+            dst_row[tid] = ivec2(gl_GlobalInvocationID.x, floatBitsToInt(a));
         } else {
             const uint row_offset = row * p.ncols_input;
             dst_row[tid] = data_s[row_offset + gl_GlobalInvocationID.x];
@@ -76,8 +82,10 @@ void topk(const uint row) {
             if (tid < s) {
                 ivec2 a = dst_row[tid];
                 ivec2 b = dst_row[tid + s];
+                // compare as floats: the bit patterns of negative values
+                // order the other way as integers
                 if (a.x >= p.orig_ncols ||
-                    b.x < p.orig_ncols && b.y > a.y) {
+                    b.x < p.orig_ncols && intBitsToFloat(b.y) > intBitsToFloat(a.y)) {
                     dst_row[tid] = b;
                 }
             }
@@ -95,9 +103,11 @@ void topk(const uint row) {
         int shift = 32 - SUBGROUP_SIZE_LOG2;
         uint mask = ((1 << SUBGROUP_SIZE_LOG2) - 1) << shift;
 
-        // The current range.
+        // The current range, [range_min, range_max). It starts as every value
+        // (+inf maps to 0xFF800000 and NaN was replaced by -inf), so the
+        // buckets always hold at least limit values.
         uint range_min = 0;
-        uint range_max = 0xFF800000;
+        uint range_max = 0xFFFFFFFF;
         // How many are above the current range, and how many we need to find.
         uint total = 0;
         uint limit = min(p.k, p.ncols_input - gl_WorkGroupID.x * BLOCK_SIZE);
@@ -138,8 +148,12 @@ void topk(const uint row) {
             total = sh_total;
 
             // Update the range, and break if we've found the K-th largest.
-            range_max = range_min + ((min_idx + 1) << shift);
-            range_min = range_min + (min_idx << shift);
+            // The end of the top bucket wraps past 2^32, clamp it instead.
+            range_min = range_min + (uint(min_idx) << shift);
+            range_max = range_min + (1u << shift);
+            if (range_max < range_min) {
+                range_max = 0xFFFFFFFF;
+            }
 
             if (total == p.k) {
                 break;
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 709997680..c5b1cd935 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -7042,6 +7042,46 @@ struct test_top_k : public test_case {
     }
 };
 
+// top_k over rows like log-probabilities: distinct negative values, fewer
+// than k +inf (none for k = 1, so the expected indices are unique) and many
+// -inf (masked tokens)
+struct test_top_k_inf : public test_top_k {
+    test_top_k_inf(std::array ne, int k)
+        : test_top_k(GGML_TYPE_F32, ne, k, false) {}
+
+    std::string vars() override {
+        return test_top_k::vars() + ",inf=1";
+    }
+
+    // compare only the output: the input holds infinities, which err() would
+    // read as indices
+    bool run_whole_graph() override { return true; }
+
+    void initialize_tensors(ggml_context * ctx) override {
+        std::random_device rd;
+        std::default_random_engine rng(rd());
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+            for (int64_t r = 0; r < ggml_nrows(t); r++) {
+                std::vector data(t->ne[0]);
+                for (int i = 0; i < t->ne[0]; i++) {
+                    data[i] = -1.0f - i;
+                }
+                std::shuffle(data.begin(), data.end(), rng);
+                const int n_pinf = k / 2;
+                for (int i = 0; i < t->ne[0]; i++) {
+                    if (i < n_pinf) {
+                        data[i] = INFINITY;
+                    } else if (i % 3 == 0) {
+                        data[i] = -INFINITY;
+                    }
+                }
+                std::shuffle(data.begin(), data.end(), rng);
+                ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(float));
+            }
+        }
+    }
+};
+
 // qwen4exp QSA indexer top-k fusion: expand per-block scores to cells, add the f16 mask, top-k.
 struct test_topk_qsa : public test_case {
     const int64_t n_blocks;
@@ -11179,6 +11219,11 @@ static std::vector> make_test_cases_eval() {
             test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {n, 2, 1, 3}, k, true));
         }
     }
+    for (int k : {1, 10, 40}) {
+        test_cases.emplace_back(new test_top_k_inf({4096, 2, 1, 1}, k));
+        test_cases.emplace_back(new test_top_k_inf({248320, 1, 1, 1}, k));
+    }
+
     for (int i = 0; i < 20; ++i) {
         for (int k : {1, 2, 3, 7, 15, 100, 500, 1023, 9999}) {
             if (k <= 1<
Date: Thu, 8 Oct 2026 12:18:20 +0200
Subject: [PATCH 40/40] server : preserve context checkpoints across slot
 save/restore (#26004)

* server : preserve context checkpoints across slot save/restore

Append the checkpoints after the packed server_tokens payload added in #26640
and count them in n_written / n_read, so a restored slot can still roll back to
a checkpoint instead of re-processing the whole prompt.

* server : drop draft checkpoint data that does not match the draft context

Restoring a slot saved with a different draft KV cache type aborted in
load_dft(). Test-load one draft checkpoint on restore and drop the draft
data if it does not fit, instead of crashing. Adds a regression test.

Co-authored-by: Igor Okulist 

* server : harden the checkpoint appendix of slot save files

Bound each blob size by the bytes left in the file before allocating, open the
file with UTF-8 paths on Windows like the llama state payload, fall back to full
prompt re-processing when a checkpoint restored from a slot file fails to load,
and replace the 1024 count cap by keeping the last n_ctx_checkpoints while reading.

* server : report an incomplete checkpoint appendix as a failed slot save

Return an error to the client when the appendix cannot be written, like a
failed payload write, and make the oversized-blob test declare a size that
cannot be allocated, so an unbounded allocation fails the test.

* server : reject an empty target state in the checkpoint appendix

A saved checkpoint always holds a target state, an empty blob would roll back
without restoring anything. Also log with the slot id, and load the draft test
model from the HF cache instead of a second download.

* common : return bool from checkpoint load_tgt / load_dft

A checkpoint restored from a slot file falls back to full prompt re-processing
when it fails to load, a checkpoint created in memory still aborts.

---------

Co-authored-by: Igor Okulist 
---
 common/common.cpp                             |  20 +-
 common/common.h                               |   5 +-
 .../speculative-simple/speculative-simple.cpp |   6 +-
 tools/server/server-context.cpp               | 163 +++++++++++-
 tools/server/tests/unit/test_slot_save.py     | 243 ++++++++++++++++++
 tools/server/tests/utils.py                   |   6 +
 6 files changed, 419 insertions(+), 24 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index 28ea8680e..48caa80c1 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -2388,40 +2388,36 @@ void common_prompt_checkpoint::update_dft(
     }
 }
 
-void common_prompt_checkpoint::load_tgt(
+bool common_prompt_checkpoint::load_tgt(
         llama_context * ctx,
         llama_seq_id seq_id,
         llama_state_seq_flags flags) const {
     if (ctx == nullptr) {
-        return;
+        return true;
     }
 
     if (data_tgt.empty()) {
-        return;
+        return true;
     }
 
     const size_t n = llama_state_seq_set_data_ext(ctx, data_tgt.data(), data_tgt.size(), seq_id, flags);
-    if (n != data_tgt.size()) {
-        GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", data_tgt.size(), n);
-    }
+    return n == data_tgt.size();
 }
 
-void common_prompt_checkpoint::load_dft(
+bool common_prompt_checkpoint::load_dft(
         llama_context * ctx,
         llama_seq_id seq_id,
         llama_state_seq_flags flags) const {
     if (ctx == nullptr) {
-        return;
+        return true;
     }
 
     if (data_dft.empty()) {
-        return;
+        return true;
     }
 
     const size_t n = llama_state_seq_set_data_ext(ctx, data_dft.data(), data_dft.size(), seq_id, flags);
-    if (n != data_dft.size()) {
-        GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", data_dft.size(), n);
-    }
+    return n == data_dft.size();
 }
 
 void common_prompt_checkpoint::clear_tgt() {
diff --git a/common/common.h b/common/common.h
index 0a85f11f9..de88dfb9b 100644
--- a/common/common.h
+++ b/common/common.h
@@ -1296,12 +1296,13 @@ struct common_prompt_checkpoint {
             llama_seq_id seq_id,
             llama_state_seq_flags flags);
 
-    void load_tgt(
+    // return false if the state could not be restored
+    bool load_tgt(
             llama_context * ctx,
             llama_seq_id seq_id,
             llama_state_seq_flags flags) const;
 
-    void load_dft(
+    bool load_dft(
             llama_context * ctx,
             llama_seq_id seq_id,
             llama_state_seq_flags flags) const;
diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp
index 08a1f2a88..e4a878b82 100644
--- a/examples/speculative-simple/speculative-simple.cpp
+++ b/examples/speculative-simple/speculative-simple.cpp
@@ -206,7 +206,7 @@ int main(int argc, char ** argv) {
             // reset the draft context to the checkpoint before verification
             if (ctx_dft) {
                 if (use_ckpt_dft) {
-                    ckpt.load_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+                    GGML_ASSERT(ckpt.load_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY));
                 }
 
                 llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, ckpt.pos_max + 1, -1);
@@ -269,13 +269,13 @@ int main(int argc, char ** argv) {
             draft = std::move(ids);
 
             {
-                ckpt.load_tgt(ctx_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+                GGML_ASSERT(ckpt.load_tgt(ctx_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY));
 
                 llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, ckpt.pos_max + 1, -1);
             }
 
             if (ctx_dft) {
-                ckpt.load_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+                GGML_ASSERT(ckpt.load_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY));
 
                 llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, ckpt.pos_max + 1, -1);
             }
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index e3270747b..1e3021d4a 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -2583,6 +2583,136 @@ private:
                 cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024);
     }
 
+    // checkpoints are appended to the slot save file, after the llama state payload
+    // they cannot be recreated from the final state alone (a recurrent state cannot be rewound)
+    static constexpr uint32_t SLOT_CKPT_MAGIC   = 0x504b4353; // "SCKP"
+    static constexpr uint32_t SLOT_CKPT_VERSION = 1;
+
+    static bool ckpt_read(std::ifstream & ifs, void * dst, size_t size, size_t & n_read) {
+        if (!ifs.read((char *) dst, size)) {
+            return false;
+        }
+        n_read += size;
+        return true;
+    }
+
+    static bool ckpt_read_buf(std::ifstream & ifs, std::vector & buf, size_t n_avail, size_t & n_read) {
+        uint64_t n = 0;
+        // check the size against the bytes left in the file before allocating, the size field may be corrupted
+        if (!ckpt_read(ifs, &n, sizeof(n), n_read) || n > n_avail - n_read) {
+            return false;
+        }
+        buf.resize(n);
+        return n == 0 || ckpt_read(ifs, buf.data(), n, n_read);
+    }
+
+    static void ckpt_write(std::ofstream & ofs, const void * src, size_t size, size_t & n_written) {
+        ofs.write((const char *) src, size);
+        n_written += size;
+    }
+
+    static void ckpt_write_buf(std::ofstream & ofs, const std::vector & buf, size_t & n_written) {
+        const uint64_t n = buf.size();
+        ckpt_write(ofs, &n, sizeof(n), n_written);
+        if (n > 0) {
+            ckpt_write(ofs, buf.data(), n, n_written);
+        }
+    }
+
+    // returns false if the appendix could not be written completely
+    bool save_slot_checkpoints(const std::string & filepath, const server_slot & slot, size_t & n_written) const {
+        n_written = 0;
+        if (slot.prompt.checkpoints.empty()) {
+            return true;
+        }
+        std::ofstream ofs(std::filesystem::u8path(filepath), std::ios::binary | std::ios::app);
+        if (!ofs) {
+            SLT_WRN(slot, "failed to append context checkpoints to '%s'\n", filepath.c_str());
+            return false;
+        }
+        const uint32_t magic   = SLOT_CKPT_MAGIC;
+        const uint32_t version = SLOT_CKPT_VERSION;
+        const uint32_t count   = (uint32_t) slot.prompt.checkpoints.size();
+        ckpt_write(ofs, &magic,   sizeof(magic),   n_written);
+        ckpt_write(ofs, &version, sizeof(version), n_written);
+        ckpt_write(ofs, &count,   sizeof(count),   n_written);
+        for (const auto & cur : slot.prompt.checkpoints) {
+            ckpt_write(ofs, &cur.n_tokens, sizeof(cur.n_tokens), n_written);
+            ckpt_write(ofs, &cur.pos_min,  sizeof(cur.pos_min),  n_written);
+            ckpt_write(ofs, &cur.pos_max,  sizeof(cur.pos_max),  n_written);
+            ckpt_write_buf(ofs, cur.data_tgt,  n_written);
+            ckpt_write_buf(ofs, cur.data_dft,  n_written);
+            ckpt_write_buf(ofs, cur.data_spec, n_written);
+        }
+        ofs.flush();
+        if (!ofs) {
+            SLT_WRN(slot, "failed to append context checkpoints to '%s' - the appendix is incomplete\n", filepath.c_str());
+            return false;
+        }
+        SLT_INF(slot, "appended %u context checkpoint(s) (%.3f MiB) to '%s'\n",
+                count, (float) n_written / 1024 / 1024, filepath.c_str());
+        return true;
+    }
+
+    // returns the number of bytes consumed, 0 if there is no usable appendix
+    size_t load_slot_checkpoints(const std::string & filepath, size_t offset, server_slot & slot) const {
+        std::ifstream ifs(std::filesystem::u8path(filepath), std::ios::binary | std::ios::ate);
+        const size_t file_size = ifs ? (size_t) ifs.tellg() : 0;
+        if (!ifs || file_size < offset || !ifs.seekg(offset)) {
+            return 0;
+        }
+        const size_t n_avail = file_size - offset; // bytes after the llama state payload
+        size_t n_read = 0;
+        uint32_t magic   = 0;
+        uint32_t version = 0;
+        uint32_t count   = 0;
+        if (!ckpt_read(ifs, &magic, sizeof(magic), n_read) || magic != SLOT_CKPT_MAGIC) {
+            return 0;
+        }
+        if (!ckpt_read(ifs, &version, sizeof(version), n_read) || version != SLOT_CKPT_VERSION ||
+            !ckpt_read(ifs, &count,   sizeof(count),   n_read)) {
+            SLT_WRN(slot, "invalid context checkpoint appendix in '%s' - ignored\n", filepath.c_str());
+            return 0;
+        }
+        std::list checkpoints;
+        for (uint32_t i = 0; i < count; ++i) {
+            common_prompt_checkpoint cur;
+            cur.id_task = -1; // not created by a task - marks a checkpoint restored from a slot file
+            if (!ckpt_read(ifs, &cur.n_tokens, sizeof(cur.n_tokens), n_read) ||
+                !ckpt_read(ifs, &cur.pos_min,  sizeof(cur.pos_min),  n_read) ||
+                !ckpt_read(ifs, &cur.pos_max,  sizeof(cur.pos_max),  n_read) ||
+                !ckpt_read_buf(ifs, cur.data_tgt,  n_avail, n_read) ||
+                !ckpt_read_buf(ifs, cur.data_dft,  n_avail, n_read) ||
+                !ckpt_read_buf(ifs, cur.data_spec, n_avail, n_read)) {
+                SLT_WRN(slot, "truncated context checkpoint appendix in '%s' - ignored\n", filepath.c_str());
+                return 0;
+            }
+            // a saved checkpoint always holds a target state - an empty blob would roll back without restoring anything
+            if (cur.data_tgt.empty()) {
+                SLT_WRN(slot, "invalid context checkpoint appendix in '%s' - ignored\n", filepath.c_str());
+                return 0;
+            }
+            checkpoints.push_back(std::move(cur));
+            if (checkpoints.size() > (size_t) params_base.n_ctx_checkpoints) {
+                checkpoints.pop_front();
+            }
+        }
+        // the slot file does not check the draft context - test-load one draft checkpoint, drop the draft data if it does not fit
+        if (ctx_dft != nullptr && !checkpoints.empty() && !checkpoints.back().data_dft.empty()) {
+            const bool ok = checkpoints.back().load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+            llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, -1, -1);
+            if (!ok) {
+                SLT_WRN(slot, "draft context checkpoint data in '%s' does not match the draft context - dropped\n", filepath.c_str());
+                for (auto & cur : checkpoints) {
+                    cur.clear_dft();
+                }
+            }
+        }
+        slot.prompt.checkpoints = std::move(checkpoints);
+        SLT_INF(slot, "restored %zu context checkpoint(s) from '%s'\n", slot.prompt.checkpoints.size(), filepath.c_str());
+        return n_read;
+    }
+
     // returns false to decline the task, it is offered again after the decode is done
     bool process_single_task(server_task && task, bool is_yielding) {
         // while yielding, an encode / decode is running and only reading the server state is safe
@@ -2792,6 +2922,12 @@ private:
                         break;
                     }
 
+                    size_t nwrite_ckpt = 0;
+                    if (!save_slot_checkpoints(filepath, *slot, nwrite_ckpt)) {
+                        send_error(task, "Unable to save slot: incomplete context checkpoints", ERROR_TYPE_SERVER);
+                        break;
+                    }
+
                     const int64_t t_end = ggml_time_us();
                     const double t_save_ms = (t_end - t_start) / 1000.0;
 
@@ -2801,7 +2937,7 @@ private:
                     res->filename = filename;
                     res->is_save  = true;
                     res->n_tokens = slot->prompt.tokens.size();
-                    res->n_bytes  = nwrite;
+                    res->n_bytes  = nwrite + nwrite_ckpt;
                     res->t_ms     = t_save_ms;
                     queue_results.send(std::move(res));
                 } break;
@@ -2857,6 +2993,9 @@ private:
                         break;
                     }
 
+                    // nread is the end offset of the llama state payload within the file
+                    const size_t nread_ckpt = load_slot_checkpoints(filepath, nread, *slot);
+
                     const int64_t t_end = ggml_time_us();
                     const double t_restore_ms = (t_end - t_start) / 1000.0;
 
@@ -2866,7 +3005,7 @@ private:
                     res->filename = filename;
                     res->is_save  = false;
                     res->n_tokens = slot->prompt.tokens.size();
-                    res->n_bytes  = nread;
+                    res->n_bytes  = nread + nread_ckpt;
                     res->t_ms     = t_restore_ms;
                     queue_results.send(std::move(res));
                 } break;
@@ -3278,7 +3417,7 @@ private:
 
             if (ctx_dft) {
                 if (use_ckpt_dft) {
-                    ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+                    GGML_ASSERT(ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY));
                 }
 
                 if (!llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, ckpt.pos_max + 1, -1)) {
@@ -3604,8 +3743,18 @@ private:
 
                                     if (!do_reset) {
                                         // restore the context checkpoint
-                                        it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
-                                        it->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+                                        if (!it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) ||
+                                            !it->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY)) {
+                                            if (it->id_task != -1) {
+                                                GGML_ABORT("failed to restore context checkpoint\n");
+                                            }
+                                            // restored from a slot file, not guaranteed to load - fall back to full prompt re-processing
+                                            SLT_WRN(slot, "%s", "failed to load context checkpoint restored from a slot file\n");
+                                            do_reset = true;
+                                        }
+                                    }
+
+                                    if (!do_reset) {
                                         // restore the draft's speculative state
                                         common_speculative_set_state(spec.get(), slot.id, it->data_spec);
 
@@ -4300,10 +4449,10 @@ private:
 
                         SLT_DBG(slot, "restoring speculative checkpoint (pos_min = %d, pos_max = %d, size = %zu)\n", ckpt.pos_min, ckpt.pos_max, ckpt.size());
 
-                        ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+                        GGML_ASSERT(ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY));
 
                         if (slot.ctx_dft) {
-                            ckpt.load_dft(slot.ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+                            GGML_ASSERT(ckpt.load_dft(slot.ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY));
                         }
 
                         slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1);
diff --git a/tools/server/tests/unit/test_slot_save.py b/tools/server/tests/unit/test_slot_save.py
index 5eca46cb2..33bcede93 100644
--- a/tools/server/tests/unit/test_slot_save.py
+++ b/tools/server/tests/unit/test_slot_save.py
@@ -37,6 +37,8 @@ def test_slot_save_restore():
     })
     assert res.status_code == 200
     assert res.body["n_saved"] == 84
+    slot_file = os.path.join(server.slot_save_path, "slot1.bin")
+    assert res.body["n_written"] == os.path.getsize(slot_file)
 
     # Since we have cache, this should only process the last tokens
     res = server.make_request("POST", "/completion", data={
@@ -54,6 +56,7 @@ def test_slot_save_restore():
     })
     assert res.status_code == 200
     assert res.body["n_restored"] == 84
+    assert res.body["n_read"] == os.path.getsize(slot_file)
 
     # Since we have cache, slot 0 should only process the last tokens
     res = server.make_request("POST", "/completion", data={
@@ -546,3 +549,243 @@ def test_slot_restore_media_file_without_mmproj(mmproj_server):
     assert res.status_code == 200
     assert res.body["timings"]["cache_n"] == 0
     assert res.body["content"] == content
+
+
+@pytest.fixture
+def swa_server():
+    swa = ServerPreset.tinygemma3()
+    swa.slot_save_path = "./tmp"
+    swa.temperature = 0.0
+    swa.cache_ram = 0
+    # Keep the first prompt checkpoint before the divergence point.
+    swa.n_ubatch = 32
+    return swa
+
+
+# the non-ASCII name checks that the appendix lands in the same file as the llama state on Windows
+@pytest.mark.parametrize("filename", ["ckpt_slot1.bin", "ckpt_slot1_é.bin"])
+def test_slot_restore_preserves_context_checkpoints(swa_server, filename):
+    server = swa_server
+    server.start()
+
+    base = "The quick brown fox jumps over the lazy dog. " * 20
+
+    res = server.make_request("POST", "/completion", data={
+        "prompt": base + "The first ending of this story is a happy one.",
+        "id_slot": 1,
+        "cache_prompt": True,
+    })
+    assert res.status_code == 200
+    n_full = res.body["timings"]["prompt_n"]
+
+    res = server.make_request("POST", "/completion", data={
+        "prompt": base + "But the second ending was different and sad.",
+        "id_slot": 1,
+        "cache_prompt": True,
+    })
+    assert res.status_code == 200
+    n_live = res.body["timings"]["prompt_n"]
+    assert n_live < n_full
+
+    res = server.make_request("POST", "/slots/1?action=erase")
+    assert res.status_code == 200
+
+    res = server.make_request("POST", "/completion", data={
+        "prompt": base + "The first ending of this story is a happy one.",
+        "id_slot": 1,
+        "cache_prompt": True,
+    })
+    assert res.status_code == 200
+
+    res = server.make_request("POST", "/slots/1?action=save", data={
+        "filename": filename,
+    })
+    assert res.status_code == 200
+    assert res.body["n_saved"] > 0
+    ckpt_file = os.path.join(server.slot_save_path, filename)
+    assert res.body["n_written"] == os.path.getsize(ckpt_file)
+
+    res = server.make_request("POST", "/completion", data={
+        "prompt": "Unrelated text with no common prefix occupies the slot now.",
+        "id_slot": 1,
+        "cache_prompt": True,
+    })
+    assert res.status_code == 200
+
+    res = server.make_request("POST", "/slots/1?action=restore", data={
+        "filename": filename,
+    })
+    assert res.status_code == 200
+    assert res.body["n_read"] == os.path.getsize(ckpt_file)
+
+    res = server.make_request("POST", "/completion", data={
+        "prompt": base + "But the second ending was different and sad.",
+        "id_slot": 1,
+        "cache_prompt": True,
+    })
+    assert res.status_code == 200
+    assert res.body["timings"]["prompt_n"] == n_live
+
+
+# checkpoint appendix: magic(4) version(4) count(4), then per checkpoint
+# n_tokens(8) pos_min(4) pos_max(4) and three blobs (target, draft, speculative), each size(8) + data
+def parse_ckpt_appendix(data):
+    off = data.find(struct.pack(" 0
+    count = struct.unpack_from(" n_live
+
+
+# the draft blobs of the checkpoint appendix are not covered by the main payload checks,
+# so restoring into a server with another draft KV cache type must not abort
+@pytest.mark.parametrize("ctkd_restore", ["f16", "q8_0"])
+def test_slot_restore_checkpoints_draft_kv_type_change(swa_server, ctkd_restore):
+    server = swa_server
+    server.model_draft_hf_repo = "ggml-org/tinygemma3-GGUF:Q8_0"  # same file as the target, already in the HF cache
+    server.spec_type = "draft-simple"
+    server.ctkd = "f16"
+    server.start()
+
+    base = "The quick brown fox jumps over the lazy dog. " * 20
+
+    res = server.make_request("POST", "/completion", data={
+        "prompt": base + "The first ending of this story is a happy one.",
+        "id_slot": 1,
+        "cache_prompt": True,
+    })
+    assert res.status_code == 200
+
+    res = server.make_request("POST", "/completion", data={
+        "prompt": base + "But the second ending was different and sad.",
+        "id_slot": 1,
+        "cache_prompt": True,
+    })
+    assert res.status_code == 200
+    n_live = res.body["timings"]["prompt_n"]
+
+    res = server.make_request("POST", "/slots/1?action=erase")
+    assert res.status_code == 200
+
+    res = server.make_request("POST", "/completion", data={
+        "prompt": base + "The first ending of this story is a happy one.",
+        "id_slot": 1,
+        "cache_prompt": True,
+    })
+    assert res.status_code == 200
+
+    res = server.make_request("POST", "/slots/1?action=save", data={
+        "filename": "ckpt_draft_slot1.bin",
+    })
+    assert res.status_code == 200
+
+    server.stop()
+    server.ctkd = ctkd_restore
+    server.start()
+
+    res = server.make_request("POST", "/slots/1?action=restore", data={
+        "filename": "ckpt_draft_slot1.bin",
+    })
+    assert res.status_code == 200
+
+    res = server.make_request("POST", "/completion", data={
+        "prompt": base + "But the second ending was different and sad.",
+        "id_slot": 1,
+        "cache_prompt": True,
+    })
+    assert res.status_code == 200
+    assert res.body["timings"]["prompt_n"] == n_live
diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py
index 90c4ebac8..76e7b2bf4 100644
--- a/tools/server/tests/utils.py
+++ b/tools/server/tests/utils.py
@@ -65,6 +65,7 @@ class ServerProcess:
     model_url: str | None = None
     model_file: str | None = None
     model_draft: str | None = None
+    model_draft_hf_repo: str | None = None
     n_threads: int | None = None
     n_gpu_layer: int | None = None
     n_batch: int | None = None
@@ -80,6 +81,7 @@ class ServerProcess:
     n_slots: int | None = None
     ctk: str | None = None
     ctv: str | None = None
+    ctkd: str | None = None
     fa: str | None = None
     server_continuous_batching: bool | None = False
     server_embeddings: bool | None = False
@@ -171,6 +173,8 @@ class ServerProcess:
             server_args.extend(["--model-url", self.model_url])
         if self.model_draft:
             server_args.extend(["--model-draft", self.model_draft])
+        if self.model_draft_hf_repo:
+            server_args.extend(["--hf-repo-draft", self.model_draft_hf_repo])
         if self.model_hf_repo:
             server_args.extend(["--hf-repo", self.model_hf_repo])
         if self.model_hf_file:
@@ -221,6 +225,8 @@ class ServerProcess:
             server_args.extend(["-ctk", self.ctk])
         if self.ctv:
             server_args.extend(["-ctv", self.ctv])
+        if self.ctkd:
+            server_args.extend(["-ctkd", self.ctkd])
         if self.fa is not None:
             server_args.extend(["-fa", self.fa])
         if self.n_predict: