ui: configure allow-remote=root in .npmrc for npm 12+ (#32158) (#32196)

* ci: reduce qdc timeouts to avoid stalling pipelines and runners

* ci: fix flake8 warnings

.devops/nix/python-scripts.nix

@@ -22,10 +22,6 @@ let
     gitpython
     tabulate
 
-    # for examples/pydantic-models-to-grammar-examples.py
-    docstring-parser
-    pydantic
-
   ];
 
   llama-python-test-deps = with python3Packages; [

.github/workflows/release.yml

@@ -1239,6 +1239,7 @@ jobs:
         with:
           key: release-windows-2022-${{ matrix.arch }}-cuda-${{ matrix.cuda }}
           evict-old-files: 1d
+          max-size: "1G"
 
       - name: Build
         id: cmake_build

CMakeLists.txt

@@ -3,8 +3,8 @@ project("llama.cpp" C CXX)
 include(CheckIncludeFileCXX)
 
 ### llama.cpp version
-set(LLAMA_VERSION_MAJOR 0)
-set(LLAMA_VERSION_MINOR 9)
+set(LLAMA_VERSION_MAJOR 1)
+set(LLAMA_VERSION_MINOR 0)
 set(LLAMA_VERSION_PATCH 0)
 set(LLAMA_VERSION_BASE "${LLAMA_VERSION_MAJOR}.${LLAMA_VERSION_MINOR}.${LLAMA_VERSION_PATCH}")
 

README.md

@@ -7,7 +7,7 @@
 <b>LLM inference in C/C++</b>
 
 [![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](https://opensource.org/licenses/MIT)
-[![Release](https://img.shields.io/github/v/release/ggml-org/llama.cpp?filter=v*&color=brightgreen)](https://github.com/ggml-org/llama.cpp/releases?q=tag:v0)
+[![Release](https://img.shields.io/github/v/release/ggml-org/llama.cpp?filter=v*&color=brightgreen)](https://github.com/ggml-org/llama.cpp/releases?q=tag:v1)
 [![Nightly](https://img.shields.io/github/v/release/ggml-org/llama.cpp?label=nightly&filter=b*&color=orange)](https://github.com/ggml-org/llama.cpp/releases?q=b)
 [![Server](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/server.yml?label=Server)](https://github.com/ggml-org/llama.cpp/actions/workflows/server.yml)
 [![Docker](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/docker.yml?label=Docker)](https://github.com/ggml-org/llama.cpp/actions/workflows/docker.yml)

common/arg.cpp

@@ -2501,15 +2501,6 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.ppl_output_type = value;
         }
     ).set_examples({LLAMA_EXAMPLE_PERPLEXITY}));
-    add_opt(common_arg(
-        {"-dt", "--defrag-thold"}, "N",
-        string_format("KV cache defragmentation threshold (DEPRECATED)"),
-        [](common_params & params, const std::string & value) {
-            GGML_UNUSED(params);
-            GGML_UNUSED(value);
-            LOG_WRN("DEPRECATED: --defrag-thold is deprecated and no longer necessary to specify\n");
-        }
-    ).set_env("LLAMA_ARG_DEFRAG_THOLD"));
     if (ex == LLAMA_EXAMPLE_SERVER) {
         // this is to make sure this option appears in the server-specific section of the help message
         add_opt(common_arg(

common/common.h

@@ -1149,7 +1149,7 @@ struct common_batch {
     // attach a state embedding (e.g. the target hidden state for MTP) to the entry at idx, can only be set once per entry
     bool set_embd_state(int32_t idx, llama_embd state);
 
-    // the state of entry idx is row `row` of the h_nextn device buffer of the state source context, see llama_set_embd_state_src()
+    // the state of entry idx is row `row` of the h_nextn device buffer of the state source context, see llama_set_ctx_state_src()
     bool set_embd_state_row(int32_t idx, int32_t row);
 
     // add an embedding-only entry (no token id)

common/parsers/k2-horizon.cpp

@@ -48,6 +48,11 @@ common_chat_params common_chat_params_init_k2_horizon(const common_chat_template
     data.thinking_start_tag = THINK_START;
     data.thinking_end_tags  = { THINK_END };
 
+    // The models sometimes close medium/low reasoning with the generic </ifm|think>
+    if (think != "ifm|think") {
+        data.thinking_end_tags.push_back("</ifm|think>");
+    }
+
     data.preserved_tokens = data.thinking_end_tags;
     data.preserved_tokens.insert(data.preserved_tokens.end(), {
         THINK_START, SECTION_START, SECTION_END, CALL_START, CALL_END,

common/speculative.cpp

@@ -10,8 +10,6 @@
 #include "ngram-mod.h"
 #include "sampling.h"
 
-#include "../src/llama-ext.h" // staging API: llama_set_embeddings_nextn / llama_get_embeddings_nextn_ith (used by MTP)
-
 #include <algorithm>
 #include <cassert>
 #include <cmath>
@@ -1650,9 +1648,9 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
 
         // shared memory and chained heads keep the host rows
         if (!is_mem_shared && !chain_heads) {
-            dev_rows = llama_set_embeddings_nextn_dev(ctx_tgt, true) && llama_set_embd_state_src(ctx_dft, ctx_tgt);
+            dev_rows = llama_set_embeddings_nextn_dev(ctx_tgt, true) && llama_set_ctx_state_src(ctx_dft, ctx_tgt);
             if (!dev_rows) {
-                llama_set_embd_state_src(ctx_dft, nullptr);
+                llama_set_ctx_state_src(ctx_dft, nullptr);
                 llama_set_embeddings_nextn_dev(ctx_tgt, false);
             }
         }
@@ -1662,7 +1660,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
     ~common_speculative_impl_draft_mtp() override {
         auto * ctx_dft = this->params.ctx_dft;
         if (dev_rows) {
-            llama_set_embd_state_src(ctx_dft, nullptr);
+            llama_set_ctx_state_src(ctx_dft, nullptr);
             llama_set_embeddings_nextn_dev(this->params.ctx_tgt, false);
         }
         for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) backend_chains.size(); ++seq_id) {
@@ -3021,6 +3019,26 @@ common_speculative * common_speculative_init(common_params_speculative & params,
         return nullptr;
     }
 
+    // the n-gram drafters ignore p_min, and DSpark uses it as the cutoff of its confidence head
+    const bool uses_p_min = std::any_of(configs.begin(), configs.end(), [](const common_speculative_config & c) {
+        switch (c.type) {
+            case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK:
+            case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE:
+            case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K:
+            case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V:
+            case COMMON_SPECULATIVE_TYPE_NGRAM_MOD:
+            case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE:
+                return false;
+            default:
+                return true;
+        }
+    });
+
+    if (params.draft.probabilistic && params.draft.p_min > 0.0f && uses_p_min) {
+        LOG_WRN("%s: --spec-draft-p-min %.2f stops sampled drafts early, consider setting it to 0 with --spec-draft-sampling probabilistic\n",
+                __func__, (double) params.draft.p_min);
+    }
+
     common_speculative_ptr result(new common_speculative {
         /* .dparams     = */ common_speculative_draft_params_vec(n_seq),
         /* .impls       = */ std::move(impls),

docs/backend/snapdragon/developer.md

@@ -60,7 +60,7 @@ Implementation: [`htp-repack.cpp`](../../../ggml/src/ggml-hexagon/htp-repack.cpp
 ### Tile Dimensions and Grid Layout
 
 Each tile represents a submatrix of `N_rows = 32` rows and `K_cols` columns:
-- Standard types (`Q4_0`, `Q4_1`, `Q8_0`, `IQ4_NL`, `MXFP4`, and K-quants): `K_cols = 32` (`HTP_TENSOR_TILE_N_COLS = 32`).
+- Standard types (`Q4_0`, `Q4_1`, `Q5_0`, `Q5_1`, `Q8_0`, `IQ4_NL`, `MXFP4`, and K-quants): `K_cols = 32` (`HTP_TENSOR_TILE_N_COLS = 32`).
 - `Q2_0`: `K_cols = 64` (`HTP_TENSOR_TILE_K_Q2_0 = 64`).
 - `Q1_0`: `K_cols = 128` (`HTP_TENSOR_TILE_K_Q1_0 = 128`).
 
@@ -123,6 +123,17 @@ Block-32 formats use a tile size of 32 rows x 32 columns (`K_cols = 32`):
   - **Scales** (bytes 512..639, 128 bytes):
     32 pairs of FP16 values (`ggml_half`): `scale[2 * row + 0] = d` (scale) and `scale[2 * row + 1] = m` (min) for each row `row = 0..31`.
 
+- **Q5_0** (704 bytes per tile, `HTP_TENSOR_TILE_SIZE_Q5_0`):
+  - **Quants** (bytes 0..511): same 4-bit nibble transposition and storage layout as `Q4_0`, holding the low 4 bits.
+  - **5th-bit plane** (bytes 512..639): byte `32*c + row` (`c = 0..3`) holds a flag for `k = 2*(4i + c)` at bit `2i` and for `k = 2*(4i + c) + 1` at bit `2i + 1`, for `i = 0..3`. The flag is set when the 5th bit of the quant is clear.
+  - **Scales** (bytes 640..703, 64 bytes): 32 FP16 values, storing scale `d` for each row `row = 0..31`.
+  - **Dequantization**: `0xF0` is OR-ed into the nibble where the flag is set, giving the signed value `q - 16` (-16..15), scaled by `d`.
+
+- **Q5_1** (768 bytes per tile, `HTP_TENSOR_TILE_SIZE_Q5_1`, same layout as the Q5_K tile):
+  - **Quants** (bytes 0..511) and **scales** (bytes 512..639): same as `Q4_1`, holding the low 4 bits and `(d, m)` pairs.
+  - **5th-bit plane** (bytes 640..767): same as `Q5_K`; a bit is set when the 5th bit of the quant is set.
+  - Uses the `Q5_K` HVX/HMX kernels (Q8_1 activations with block sums).
+
 ### Q1_0 Format (1-Bit Quantization)
 
 `Q1_0` packs 128 columns per tile (`K_cols = 128`, `HTP_TENSOR_TILE_K_Q1_0 = 128`):

docs/build.md

@@ -608,7 +608,7 @@ You can test with:
 If the following info is output on screen, you are using `llama.cpp` with the CANN backend:
 ```bash
 llm_load_tensors:       CANN model buffer size = 13313.00 MiB
-llama_new_context_with_model:       CANN compute buffer size =  1260.81 MiB
+llama_init_from_model:       CANN compute buffer size =  1260.81 MiB
 ```
 
 For detailed info, such as model/device supports, CANN install, please refer to [llama.cpp for CANN](./backend/CANN.md).

docs/ops.md

@@ -40,7 +40,6 @@ Legend:
 |          CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           CUMSUM | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             DIAG | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                    DIAG_MASK_INF | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              DIV | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                     DSV4_HC_COMB | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                     DSV4_HC_POST | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
@@ -127,7 +126,7 @@ Legend:
 |                             TANH | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |               TIMESTEP_EMBEDDING | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                            TOP_K | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
-|                       TOP_K_MASK | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                       TOP_K_MASK | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
 |                              TRI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                            TRUNC | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          UPSCALE | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |

docs/ops/BLAS.csv

@@ -8741,9 +8741,6 @@
 "BLAS","ROUND","type=f32,ne=[1024,1024,1,1]","support","0","no","BLAS"
 "BLAS","TRUNC","type=f32,ne=[7,1,5,3]","support","0","no","BLAS"
 "BLAS","TRUNC","type=f32,ne=[1024,1024,1,1]","support","0","no","BLAS"
-"BLAS","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","0","no","BLAS"
-"BLAS","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","0","no","BLAS"
-"BLAS","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","BLAS"
 "BLAS","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","BLAS"
 "BLAS","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","BLAS"
 "BLAS","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","BLAS"

docs/ops/CANN.csv

@@ -8764,9 +8764,6 @@
 "CANN0","ROUND","type=f32,ne=[1024,1024,1,1]","support","0","no","CANN"
 "CANN0","TRUNC","type=f32,ne=[7,1,5,3]","support","0","no","CANN"
 "CANN0","TRUNC","type=f32,ne=[1024,1024,1,1]","support","0","no","CANN"
-"CANN0","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","1","yes","CANN"
-"CANN0","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","1","yes","CANN"
-"CANN0","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","1","yes","CANN"
 "CANN0","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CANN"
 "CANN0","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CANN"
 "CANN0","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CANN"

docs/ops/CPU.csv

@@ -12962,9 +12962,6 @@
 "CPU","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","CPU"
 "CPU","TRUNC","type=f32,ne=[7,1,5,3]","support","1","yes","CPU"
 "CPU","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","CPU"
-"CPU","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","1","yes","CPU"
-"CPU","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","1","yes","CPU"
-"CPU","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","1","yes","CPU"
 "CPU","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CPU"
 "CPU","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CPU"
 "CPU","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CPU"

docs/ops/CUDA.csv

@@ -8725,9 +8725,6 @@
 "CUDA0","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","CUDA"
 "CUDA0","TRUNC","type=f32,ne=[7,1,5,3]","support","1","yes","CUDA"
 "CUDA0","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","CUDA"
-"CUDA0","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","1","yes","CUDA"
-"CUDA0","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","1","yes","CUDA"
-"CUDA0","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","1","yes","CUDA"
 "CUDA0","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CUDA"
 "CUDA0","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CUDA"
 "CUDA0","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CUDA"

docs/ops/ET.csv

@@ -9637,9 +9637,6 @@
 "ET","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","ET"
 "ET","TRUNC","type=f32,ne=[7,1,5,3]","support","0","no","ET"
 "ET","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","ET"
-"ET","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","0","no","ET"
-"ET","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","0","no","ET"
-"ET","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","ET"
 "ET","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","ET"
 "ET","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","ET"
 "ET","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","ET"

docs/ops/Hexagon.csv

@@ -14564,9 +14564,6 @@
 "HTP0","ROUND","type=f32,ne=[1024,1024,1,1]","support","0","no","HTP"
 "HTP0","TRUNC","type=f32,ne=[7,1,5,3]","support","0","no","HTP"
 "HTP0","TRUNC","type=f32,ne=[1024,1024,1,1]","support","0","no","HTP"
-"HTP0","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","0","no","HTP"
-"HTP0","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","0","no","HTP"
-"HTP0","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","HTP"
 "HTP0","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","HTP"
 "HTP0","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","HTP"
 "HTP0","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","HTP"

docs/ops/Metal.csv

@@ -9628,9 +9628,6 @@
 "MTL0","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","MTL"
 "MTL0","TRUNC","type=f32,ne=[7,1,5,3]","support","1","yes","MTL"
 "MTL0","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","MTL"
-"MTL0","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","0","no","MTL"
-"MTL0","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","0","no","MTL"
-"MTL0","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","MTL"
 "MTL0","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","MTL"
 "MTL0","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","MTL"
 "MTL0","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","MTL"

docs/ops/OpenCL.csv

@@ -14887,9 +14887,6 @@
 "OpenCL","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","OpenCL"
 "OpenCL","TRUNC","type=f32,ne=[7,1,5,3]","support","1","yes","OpenCL"
 "OpenCL","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","OpenCL"
-"OpenCL","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","1","yes","OpenCL"
-"OpenCL","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","1","yes","OpenCL"
-"OpenCL","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","OpenCL"
 "OpenCL","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","OpenCL"
 "OpenCL","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","OpenCL"
 "OpenCL","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","OpenCL"

docs/ops/SYCL.csv

@@ -14534,9 +14534,6 @@
 "SYCL0","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","SYCL"
 "SYCL0","TRUNC","type=f32,ne=[7,1,5,3]","support","1","yes","SYCL"
 "SYCL0","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","SYCL"
-"SYCL0","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","1","yes","SYCL"
-"SYCL0","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","1","yes","SYCL"
-"SYCL0","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","1","yes","SYCL"
 "SYCL0","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","SYCL"
 "SYCL0","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","SYCL"
 "SYCL0","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","SYCL"

docs/ops/Vulkan.csv

@@ -12313,9 +12313,6 @@
 "Vulkan0","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","Vulkan"
 "Vulkan0","TRUNC","type=f32,ne=[7,1,5,3]","support","1","yes","Vulkan"
 "Vulkan0","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","Vulkan"
-"Vulkan0","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","1","yes","Vulkan"
-"Vulkan0","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","1","yes","Vulkan"
-"Vulkan0","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","1","yes","Vulkan"
 "Vulkan0","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","Vulkan"
 "Vulkan0","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","Vulkan"
 "Vulkan0","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","Vulkan"
@@ -19442,3 +19439,31 @@
 "Vulkan0","LIGHTNING_INDEXER","hsk=128,nh=64,kv=256,nb=512,ns=4,nm=1,type_K=q4_1","support","0","no","Vulkan"
 "Vulkan0","LIGHTNING_INDEXER","hsk=128,nh=64,kv=256,nb=512,ns=4,nm=1,type_K=q4_0","support","0","no","Vulkan"
 "Vulkan0","LIGHTNING_INDEXER","hsk=128,nh=64,kv=256,nb=512,ns=4,nm=1,type_K=iq4_nl","support","0","no","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=256,n_blk=4,n_pick=64,n_top=16,n_extra=3,score=1,nb=32,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=1001,n_blk=4,n_pick=70,n_top=16,n_extra=3,score=1,nb=7,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=256,n_blk=4,n_pick=64,n_top=64,n_extra=0,score=1,nb=33,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=512,n_blk=8,n_pick=16,n_top=32,n_extra=7,score=1,nb=5,nh=1,ns=4,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=4096,n_blk=4,n_pick=1024,n_top=512,n_extra=3,score=1,nb=1,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=4096,n_blk=4,n_pick=1024,n_top=512,n_extra=3,score=1,nb=64,nh=1,ns=2,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=256,n_blk=0,n_pick=256,n_top=16,n_extra=0,score=0,nb=32,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=1001,n_blk=0,n_pick=1001,n_top=64,n_extra=0,score=0,nb=7,nh=1,ns=3,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=4096,n_blk=0,n_pick=4096,n_top=2048,n_extra=0,score=0,nb=5,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=512,n_blk=0,n_pick=512,n_top=32,n_extra=5,score=1,nb=9,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=512,n_blk=16,n_pick=32,n_top=4,n_extra=0,score=1,nb=6,nh=4,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=1024,n_blk=8,n_pick=128,n_top=16,n_extra=8,score=0,nb=3,nh=2,ns=2,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=1001,n_blk=4,n_pick=70,n_top=16,n_extra=3,score=1,nb=7,nh=1,ns=1,real=0,bcast=1","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=512,n_blk=8,n_pick=16,n_top=32,n_extra=7,score=1,nb=5,nh=1,ns=4,real=0,bcast=1","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=256,n_blk=4,n_pick=64,n_top=16,n_extra=3,score=1,nb=32,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=1001,n_blk=4,n_pick=70,n_top=16,n_extra=3,score=1,nb=7,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=256,n_blk=4,n_pick=64,n_top=64,n_extra=0,score=1,nb=33,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=512,n_blk=8,n_pick=16,n_top=32,n_extra=7,score=1,nb=5,nh=1,ns=4,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=4096,n_blk=4,n_pick=1024,n_top=512,n_extra=3,score=1,nb=1,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=4096,n_blk=4,n_pick=1024,n_top=512,n_extra=3,score=1,nb=64,nh=1,ns=2,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=256,n_blk=0,n_pick=256,n_top=16,n_extra=0,score=0,nb=32,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=1001,n_blk=0,n_pick=1001,n_top=64,n_extra=0,score=0,nb=7,nh=1,ns=3,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=4096,n_blk=0,n_pick=4096,n_top=2048,n_extra=0,score=0,nb=5,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=512,n_blk=0,n_pick=512,n_top=32,n_extra=5,score=1,nb=9,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=512,n_blk=16,n_pick=32,n_top=4,n_extra=0,score=1,nb=6,nh=4,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=1024,n_blk=8,n_pick=128,n_top=16,n_extra=8,score=0,nb=3,nh=2,ns=2,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=1001,n_blk=4,n_pick=70,n_top=16,n_extra=3,score=1,nb=7,nh=1,ns=1,real=0,bcast=1","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=512,n_blk=8,n_pick=16,n_top=32,n_extra=7,score=1,nb=5,nh=1,ns=4,real=0,bcast=1","support","1","yes","Vulkan"

docs/ops/WebGPU.csv

@@ -10560,9 +10560,6 @@
 "WebGPU: WebGPU","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","WebGPU"
 "WebGPU: WebGPU","TRUNC","type=f32,ne=[7,1,5,3]","support","1","yes","WebGPU"
 "WebGPU: WebGPU","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","WebGPU"
-"WebGPU: WebGPU","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","0","no","WebGPU"
-"WebGPU: WebGPU","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","0","no","WebGPU"
-"WebGPU: WebGPU","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","WebGPU"
 "WebGPU: WebGPU","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","WebGPU"
 "WebGPU: WebGPU","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","WebGPU"
 "WebGPU: WebGPU","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","WebGPU"

docs/ops/ZenDNN.csv

@@ -9319,9 +9319,6 @@
 "ZenDNN","ROUND","type=f32,ne=[1024,1024,1,1]","support","0","no","ZenDNN"
 "ZenDNN","TRUNC","type=f32,ne=[7,1,5,3]","support","0","no","ZenDNN"
 "ZenDNN","TRUNC","type=f32,ne=[1024,1024,1,1]","support","0","no","ZenDNN"
-"ZenDNN","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","0","no","ZenDNN"
-"ZenDNN","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","0","no","ZenDNN"
-"ZenDNN","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","ZenDNN"
 "ZenDNN","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","ZenDNN"
 "ZenDNN","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","ZenDNN"
 "ZenDNN","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","ZenDNN"

docs/ops/zDNN.csv

@@ -14550,9 +14550,6 @@
 "zDNN","ROUND","type=f32,ne=[1024,1024,1,1]","support","0","no","zDNN"
 "zDNN","TRUNC","type=f32,ne=[7,1,5,3]","support","0","no","zDNN"
 "zDNN","TRUNC","type=f32,ne=[1024,1024,1,1]","support","0","no","zDNN"
-"zDNN","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","0","no","zDNN"
-"zDNN","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","0","no","zDNN"
-"zDNN","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","zDNN"
 "zDNN","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","zDNN"
 "zDNN","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","zDNN"
 "zDNN","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","zDNN"

docs/speculative.md

@@ -189,7 +189,7 @@ llama-server ... --spec-type ngram-mod --spec-ngram-mod-n-match 24 --spec-ngram-
 
 Applications:
 
-- Iterating over a block of text/code (e.g. in llama.vim)
+- Iterating over a block of text/code (e.g. in [llama.vim](https://github.com/ggml-org/llama.vim))
 - Reasoning models (when they have to repeat their thinking in the final answer)
 - Summarization
 

examples/eval-callback/README.md

@@ -20,12 +20,12 @@ Will print:
 ```shell
 llm_load_tensors: offloaded 33/33 layers to GPU
 ...
-llama_new_context_with_model: n_ctx      = 512
+llama_init_from_model: n_ctx      = 512
 ...
-llama_new_context_with_model:      CUDA0 compute buffer size =   105.00 MiB
-llama_new_context_with_model:  CUDA_Host compute buffer size =     6.01 MiB
-llama_new_context_with_model: graph nodes  = 1225
-llama_new_context_with_model: graph splits = 2
+llama_init_from_model:      CUDA0 compute buffer size =   105.00 MiB
+llama_init_from_model:  CUDA_Host compute buffer size =     6.01 MiB
+llama_init_from_model: graph nodes  = 1225
+llama_init_from_model: graph splits = 2
 ggml_debug:                 inp_embd = (f32)   GET_ROWS(token_embd.weight{2560, 51200, 1, 1}, inp_tokens{1, 1, 1, 1}}) = {2560, 1, 1, 1}
                                      [
                                       [

examples/json_schema_pydantic_example.py

@@ -1,82 +0,0 @@
-# Usage:
-#! ./llama-server -m some-model.gguf &
-#! pip install pydantic
-#! python json_schema_pydantic_example.py
-
-from pydantic import BaseModel, Field, TypeAdapter
-from annotated_types import MinLen
-from typing import Annotated, List, Optional
-import json, requests
-
-if True:
-
-    def create_completion(*, response_model=None, endpoint="http://localhost:9931/v1/chat/completions", messages, **kwargs):
-        '''
-        Creates a chat completion using an OpenAI-compatible endpoint w/ JSON schema support
-        (llama.cpp server, llama-cpp-python, Anyscale / Together...)
-
-        The response_model param takes a type (+ supports Pydantic) and behaves just as w/ Instructor (see below)
-        '''
-        response_format = None
-        type_adapter = None
-
-        if response_model:
-            type_adapter = TypeAdapter(response_model)
-            schema = type_adapter.json_schema()
-            messages = [{
-                "role": "system",
-                "content": f"You respond in JSON format with the following schema: {json.dumps(schema, indent=2)}"
-            }] + messages
-            response_format={"type": "json_object", "schema": schema}
-
-        data = requests.post(endpoint, headers={"Content-Type": "application/json"},
-                             json=dict(messages=messages, response_format=response_format, **kwargs)).json()
-        if 'error' in data:
-            raise Exception(data['error']['message'])
-
-        content = data["choices"][0]["message"]["content"]
-        return type_adapter.validate_json(content) if type_adapter else content
-
-else:
-
-    # This alternative branch uses Instructor + OpenAI client lib.
-    # Instructor support streamed iterable responses, retry & more.
-    # (see https://python.useinstructor.com/)
-    #! pip install instructor openai
-    import instructor, openai
-    client = instructor.patch(
-        openai.OpenAI(api_key="123", base_url="http://localhost:9931"),
-        mode=instructor.Mode.JSON_SCHEMA)
-    create_completion = client.chat.completions.create
-
-
-if __name__ == '__main__':
-
-    class QAPair(BaseModel):
-        class Config:
-            extra = 'forbid'  # triggers additionalProperties: false in the JSON schema
-        question: str
-        concise_answer: str
-        justification: str
-        stars: Annotated[int, Field(ge=1, le=5)]
-
-    class PyramidalSummary(BaseModel):
-        class Config:
-            extra = 'forbid'  # triggers additionalProperties: false in the JSON schema
-        title: str
-        summary: str
-        question_answers: Annotated[List[QAPair], MinLen(2)]
-        sub_sections: Optional[Annotated[List['PyramidalSummary'], MinLen(2)]]
-
-    print("# Summary\n", create_completion(
-        model="...",
-        response_model=PyramidalSummary,
-        messages=[{
-            "role": "user",
-            "content": f"""
-                You are a highly efficient corporate document summarizer.
-                Create a pyramidal summary of an imaginary internal document about our company processes
-                (starting high-level, going down to each sub sections).
-                Keep questions short, and answers even shorter (trivia / quizz style).
-            """
-        }]))

examples/llama.android/lib/src/main/cpp/ai_chat.cpp

@@ -99,7 +99,7 @@ static llama_context *init_context(llama_model *model, const int n_ctx = DEFAULT
     ctx_params.n_threads_batch = n_threads;
     auto *context = llama_init_from_model(g_model, ctx_params);
     if (context == nullptr) {
-        LOGe("%s: llama_new_context_with_model() returned null)", __func__);
+        LOGe("%s: llama_init_from_model() returned null)", __func__);
     }
     return context;
 }

examples/llama.vim

@@ -1,783 +0,0 @@
-" LLM-based text completion using llama.cpp
-"
-" requires:
-"
-"   - neovim or vim
-"   - curl
-"   - llama.cpp server instance
-"   - FIM-compatible model
-"
-" sample config:
-"
-"   - Tab       - accept the current suggestion
-"   - Shift+Tab - accept just the first line of the suggestion
-"   - Ctrl+F    - toggle FIM completion manually
-"
-" make symlink or copy this file to ~/.config/nvim/autoload/llama.vim
-"
-" start the llama.cpp server with a FIM-compatible model. for example:
-"
-"   $ llama-server -m {model.gguf} --port 8012 -ngl 99 -fa --ubatch-size 512 --batch-size 1024 --cache-reuse 256
-"
-"   --batch-size [512, model max context]
-"
-"     adjust the batch size to control how much of the provided local context will be used during the inference
-"     lower values will use smaller part of the context around the cursor, which will result in faster processing
-"
-"   --ubatch-size [64, 2048]
-"
-"     chunks the batch into smaller chunks for faster processing
-"     depends on the specific hardware. use llama-bench to profile and determine the best size
-"
-"   --cache-reuse (ge:llama_config.n_predict, 1024]
-"
-"     this should be either 0 (disabled) or strictly larger than g:llama_config.n_predict
-"     using non-zero value enables context reuse on the server side which dramatically improves the performance at
-"     large contexts. a value of 256 should be good for all cases
-"
-" run this once to initialise llama.vim:
-"
-"   :call llama#init()
-"
-" more info: https://github.com/ggml-org/llama.cpp/pull/9787
-"
-
-" colors (adjust to your liking)
-highlight llama_hl_hint guifg=#ff772f ctermfg=202
-highlight llama_hl_info guifg=#77ff2f ctermfg=119
-
-" general parameters:
-"
-"   endpoint:         llama.cpp server endpoint
-"   n_prefix:         number of lines before the cursor location to include in the local prefix
-"   n_suffix:         number of lines after  the cursor location to include in the local suffix
-"   n_predict:        max number of tokens to predict
-"   t_max_prompt_ms:  max allotted time for the prompt processing (TODO: not yet supported)
-"   t_max_predict_ms: max allotted time for the prediction
-"   show_info:        show extra info about the inference (0 - disabled, 1 - statusline, 2 - inline)
-"   auto_fim:         trigger FIM completion automatically on cursor movement
-"   max_line_suffix:  do not auto-trigger FIM completion if there are more than this number of characters to the right of the cursor
-"
-" ring buffer of chunks, accumulated with time upon:
-"
-"  - completion request
-"  - yank
-"  - entering a buffer
-"  - leaving a buffer
-"  - writing a file
-"
-" parameters for the ring-buffer with extra context:
-"
-"   ring_n_chunks:    max number of chunks to pass as extra context to the server (0 to disable)
-"   ring_chunk_size:  max size of the chunks (in number of lines)
-"                     note: adjust these numbers so that you don't overrun your context
-"                           at ring_n_chunks = 64 and ring_chunk_size = 64 you need ~32k context
-"   ring_scope:       the range around the cursor position (in number of lines) for gathering chunks after FIM
-"   ring_update_ms:   how often to process queued chunks in normal mode
-"
-let s:default_config = {
-    \ 'endpoint':         'http://127.0.0.1:8012/infill',
-    \ 'n_prefix':         256,
-    \ 'n_suffix':         64,
-    \ 'n_predict':        128,
-    \ 't_max_prompt_ms':  500,
-    \ 't_max_predict_ms': 3000,
-    \ 'show_info':        2,
-    \ 'auto_fim':         v:true,
-    \ 'max_line_suffix':  8,
-    \ 'ring_n_chunks':    64,
-    \ 'ring_chunk_size':  64,
-    \ 'ring_scope':       1024,
-    \ 'ring_update_ms':   1000,
-    \ }
-
-let g:llama_config = get(g:, 'llama_config', s:default_config)
-
-function! s:get_indent(str)
-    let l:count = 0
-    for i in range(len(a:str))
-        if a:str[i] == "\t"
-            let l:count += &tabstop - 1
-        else
-            break
-        endif
-    endfor
-    return l:count
-endfunction
-
-function! s:rand(i0, i1) abort
-    return a:i0 + rand() % (a:i1 - a:i0 + 1)
-endfunction
-
-function! llama#init()
-    if !executable('curl')
-        echohl WarningMsg
-        echo 'llama.vim requires the "curl" command to be available'
-        echohl None
-        return
-    endif
-
-    let s:pos_x = 0 " cursor position upon start of completion
-    let s:pos_y = 0
-
-    let s:line_cur = ''
-
-    let s:line_cur_prefix = ''
-    let s:line_cur_suffix = ''
-
-    let s:ring_chunks = [] " current set of chunks used as extra context
-    let s:ring_queued = [] " chunks that are queued to be sent for processing
-    let s:ring_n_evict = 0
-
-    let s:hint_shown = v:false
-    let s:pos_y_pick = -9999 " last y where we picked a chunk
-    let s:pos_dx = 0
-    let s:content = []
-    let s:can_accept = v:false
-
-    let s:timer_fim = -1
-    let s:t_fim_start = reltime() " used to measure total FIM time
-    let s:t_last_move = reltime() " last time the cursor moved
-
-    let s:current_job = v:null
-
-    let s:ghost_text_nvim = exists('*nvim_buf_get_mark')
-    let s:ghost_text_vim = has('textprop')
-
-    if s:ghost_text_vim
-        let s:hlgroup_hint = 'llama_hl_hint'
-        let s:hlgroup_info = 'llama_hl_info'
-
-        if empty(prop_type_get(s:hlgroup_hint))
-            call prop_type_add(s:hlgroup_hint, {'highlight': s:hlgroup_hint})
-        endif
-        if empty(prop_type_get(s:hlgroup_info))
-            call prop_type_add(s:hlgroup_info, {'highlight': s:hlgroup_info})
-        endif
-    endif
-
-    augroup llama
-        autocmd!
-        autocmd InsertEnter     * inoremap <expr> <silent> <C-F> llama#fim_inline(v:false)
-        autocmd InsertLeavePre  * call llama#fim_cancel()
-
-        autocmd CursorMoved     * call s:on_move()
-        autocmd CursorMovedI    * call s:on_move()
-        autocmd CompleteChanged * call llama#fim_cancel()
-
-        if g:llama_config.auto_fim
-            autocmd CursorMovedI * call llama#fim(v:true)
-        endif
-
-        " gather chunks upon yanking
-        autocmd TextYankPost    * if v:event.operator ==# 'y' | call s:pick_chunk(v:event.regcontents, v:false, v:true) | endif
-
-        " gather chunks upon entering/leaving a buffer
-        autocmd BufEnter        * call timer_start(100, {-> s:pick_chunk(getline(max([1, line('.') - g:llama_config.ring_chunk_size/2]), min([line('.') + g:llama_config.ring_chunk_size/2, line('$')])), v:true, v:true)})
-        autocmd BufLeave        * call                      s:pick_chunk(getline(max([1, line('.') - g:llama_config.ring_chunk_size/2]), min([line('.') + g:llama_config.ring_chunk_size/2, line('$')])), v:true, v:true)
-
-        " gather chunk upon saving the file
-        autocmd BufWritePost    * call s:pick_chunk(getline(max([1, line('.') - g:llama_config.ring_chunk_size/2]), min([line('.') + g:llama_config.ring_chunk_size/2, line('$')])), v:true, v:true)
-    augroup END
-
-    silent! call llama#fim_cancel()
-
-    " init background update of the ring buffer
-    if g:llama_config.ring_n_chunks > 0
-        call s:ring_update()
-    endif
-endfunction
-
-" compute how similar two chunks of text are
-" 0 - no similarity, 1 - high similarity
-" TODO: figure out something better
-function! s:chunk_sim(c0, c1)
-    let l:lines0 = len(a:c0)
-    let l:lines1 = len(a:c1)
-
-    let l:common = 0
-
-    for l:line0 in a:c0
-        for l:line1 in a:c1
-            if l:line0 == l:line1
-                let l:common += 1
-                break
-            endif
-        endfor
-    endfor
-
-    return 2.0 * l:common / (l:lines0 + l:lines1)
-endfunction
-
-" pick a random chunk of size g:llama_config.ring_chunk_size from the provided text and queue it for processing
-"
-" no_mod   - do not pick chunks from buffers with pending changes
-" do_evict - evict chunks that are very similar to the new one
-"
-function! s:pick_chunk(text, no_mod, do_evict)
-    " do not pick chunks from buffers with pending changes or buffers that are not files
-    if a:no_mod && (getbufvar(bufnr('%'), '&modified') || !buflisted(bufnr('%')) || !filereadable(expand('%')))
-        return
-    endif
-
-    " if the extra context option is disabled - do nothing
-    if g:llama_config.ring_n_chunks <= 0
-        return
-    endif
-
-    " don't pick very small chunks
-    if len(a:text) < 3
-        return
-    endif
-
-    if len(a:text) + 1 < g:llama_config.ring_chunk_size
-        let l:chunk = a:text
-    else
-        let l:l0 = s:rand(0, max([0, len(a:text) - g:llama_config.ring_chunk_size/2]))
-        let l:l1 = min([l:l0 + g:llama_config.ring_chunk_size/2, len(a:text)])
-
-        let l:chunk = a:text[l:l0:l:l1]
-    endif
-
-    let l:chunk_str = join(l:chunk, "\n") . "\n"
-
-    " check if this chunk is already added
-    let l:exist = v:false
-
-    for i in range(len(s:ring_chunks))
-        if s:ring_chunks[i].data == l:chunk
-            let l:exist = v:true
-            break
-        endif
-    endfor
-
-    for i in range(len(s:ring_queued))
-        if s:ring_queued[i].data == l:chunk
-            let l:exist = v:true
-            break
-        endif
-    endfor
-
-    if l:exist
-        return
-    endif
-
-    " evict queued chunks that are very similar to the new one
-    for i in range(len(s:ring_queued) - 1, 0, -1)
-        if s:chunk_sim(s:ring_queued[i].data, l:chunk) > 0.9
-            if a:do_evict
-                call remove(s:ring_queued, i)
-                let s:ring_n_evict += 1
-            else
-                return
-            endif
-        endif
-    endfor
-
-    " also from s:ring_chunks
-    for i in range(len(s:ring_chunks) - 1, 0, -1)
-        if s:chunk_sim(s:ring_chunks[i].data, l:chunk) > 0.9
-            if a:do_evict
-                call remove(s:ring_chunks, i)
-                let s:ring_n_evict += 1
-            else
-                return
-            endif
-        endif
-    endfor
-
-    " TODO: become parameter ?
-    if len(s:ring_queued) == 16
-        call remove(s:ring_queued, 0)
-    endif
-
-    call add(s:ring_queued, {'data': l:chunk, 'str': l:chunk_str, 'time': reltime(), 'filename': expand('%')})
-
-    "let &statusline = 'extra context: ' . len(s:ring_chunks) . ' / ' . len(s:ring_queued)
-endfunction
-
-" picks a queued chunk, sends it for processing and adds it to s:ring_chunks
-" called every g:llama_config.ring_update_ms
-function! s:ring_update()
-    call timer_start(g:llama_config.ring_update_ms, {-> s:ring_update()})
-
-    " update only if in normal mode or if the cursor hasn't moved for a while
-    if mode() !=# 'n' && reltimefloat(reltime(s:t_last_move)) < 3.0
-        return
-    endif
-
-    if len(s:ring_queued) == 0
-        return
-    endif
-
-    " move the first queued chunk to the ring buffer
-    if len(s:ring_chunks) == g:llama_config.ring_n_chunks
-        call remove(s:ring_chunks, 0)
-    endif
-
-    call add(s:ring_chunks, remove(s:ring_queued, 0))
-
-    "let &statusline = 'updated context: ' . len(s:ring_chunks) . ' / ' . len(s:ring_queued)
-
-    " send asynchronous job with the new extra context so that it is ready for the next FIM
-    let l:extra_context = []
-    for l:chunk in s:ring_chunks
-        call add(l:extra_context, {
-            \ 'text':     l:chunk.str,
-            \ 'time':     l:chunk.time,
-            \ 'filename': l:chunk.filename
-            \ })
-    endfor
-
-    " no samplers needed here
-    let l:request = json_encode({
-        \ 'input_prefix':     "",
-        \ 'input_suffix':     "",
-        \ 'input_extra':      l:extra_context,
-        \ 'prompt':           "",
-        \ 'n_predict':        1,
-        \ 'temperature':      0.0,
-        \ 'stream':           v:false,
-        \ 'samplers':         ["temperature"],
-        \ 'cache_prompt':     v:true,
-        \ 't_max_prompt_ms':  1,
-        \ 't_max_predict_ms': 1
-        \ })
-
-    let l:curl_command = [
-        \ "curl",
-        \ "--silent",
-        \ "--no-buffer",
-        \ "--request", "POST",
-        \ "--url", g:llama_config.endpoint,
-        \ "--header", "Content-Type: application/json",
-        \ "--data", l:request
-        \ ]
-
-    " no callbacks because we don't need to process the response
-    if s:ghost_text_nvim
-        call jobstart(l:curl_command, {})
-    elseif s:ghost_text_vim
-        call job_start(l:curl_command, {})
-    endif
-endfunction
-
-" necessary for 'inoremap <expr>'
-function! llama#fim_inline(is_auto) abort
-    call llama#fim(a:is_auto)
-    return ''
-endfunction
-
-" the main FIM call
-" takes local context around the cursor and sends it together with the extra context to the server for completion
-function! llama#fim(is_auto) abort
-    " we already have a suggestion for the current cursor position
-    if s:hint_shown && !a:is_auto
-        call llama#fim_cancel()
-        return
-    endif
-
-    call llama#fim_cancel()
-
-    " avoid sending repeated requests too fast
-    if reltimefloat(reltime(s:t_fim_start)) < 0.6
-        if s:timer_fim != -1
-            call timer_stop(s:timer_fim)
-            let s:timer_fim = -1
-        endif
-
-        let s:t_fim_start = reltime()
-        let s:timer_fim = timer_start(600, {-> llama#fim(v:true)})
-        return
-    endif
-
-    let s:t_fim_start = reltime()
-
-    let s:content = []
-    let s:can_accept = v:false
-
-    let s:pos_x = col('.') - 1
-    let s:pos_y = line('.')
-    let l:max_y = line('$')
-
-    let l:lines_prefix = getline(max([1, s:pos_y - g:llama_config.n_prefix]), s:pos_y - 1)
-    let l:lines_suffix = getline(s:pos_y + 1, min([l:max_y, s:pos_y + g:llama_config.n_suffix]))
-
-    let s:line_cur = getline('.')
-
-    let s:line_cur_prefix = strpart(s:line_cur, 0, s:pos_x)
-    let s:line_cur_suffix = strpart(s:line_cur, s:pos_x)
-
-    if a:is_auto && len(s:line_cur_suffix) > g:llama_config.max_line_suffix
-        return
-    endif
-
-    let l:prefix = ""
-        \ . join(l:lines_prefix, "\n")
-        \ . "\n"
-
-    let l:prompt = ""
-        \ . s:line_cur_prefix
-
-    let l:suffix = ""
-        \ . s:line_cur_suffix
-        \ . "\n"
-        \ . join(l:lines_suffix, "\n")
-        \ . "\n"
-
-    " prepare the extra context data
-    let l:extra_context = []
-    for l:chunk in s:ring_chunks
-        call add(l:extra_context, {
-            \ 'text':     l:chunk.str,
-            \ 'time':     l:chunk.time,
-            \ 'filename': l:chunk.filename
-            \ })
-    endfor
-
-    " the indentation of the current line
-    let l:indent = strlen(matchstr(s:line_cur_prefix, '^\s*'))
-
-    let l:request = json_encode({
-        \ 'input_prefix':     l:prefix,
-        \ 'input_suffix':     l:suffix,
-        \ 'input_extra':      l:extra_context,
-        \ 'prompt':           l:prompt,
-        \ 'n_predict':        g:llama_config.n_predict,
-        \ 'n_indent':         l:indent,
-        \ 'top_k':            40,
-        \ 'top_p':            0.99,
-        \ 'stream':           v:false,
-        \ 'samplers':         ["top_k", "top_p", "infill"],
-        \ 'cache_prompt':     v:true,
-        \ 't_max_prompt_ms':  g:llama_config.t_max_prompt_ms,
-        \ 't_max_predict_ms': g:llama_config.t_max_predict_ms
-        \ })
-
-    let l:curl_command = [
-        \ "curl",
-        \ "--silent",
-        \ "--no-buffer",
-        \ "--request", "POST",
-        \ "--url", g:llama_config.endpoint,
-        \ "--header", "Content-Type: application/json",
-        \ "--data", l:request
-        \ ]
-
-    if s:current_job != v:null
-        if s:ghost_text_nvim
-            call jobstop(s:current_job)
-        elseif s:ghost_text_vim
-            call job_stop(s:current_job)
-        endif
-    endif
-
-    " send the request asynchronously
-    if s:ghost_text_nvim
-        let s:current_job = jobstart(l:curl_command, {
-            \ 'on_stdout': function('s:fim_on_stdout', [s:pos_x, s:pos_y, a:is_auto]),
-            \ 'on_exit':   function('s:fim_on_exit'),
-            \ 'stdout_buffered': v:true
-            \ })
-    elseif s:ghost_text_vim
-        let s:current_job = job_start(l:curl_command, {
-            \ 'out_cb': function('s:fim_on_stdout', [s:pos_x, s:pos_y, a:is_auto]),
-            \ 'exit_cb':   function('s:fim_on_exit')
-            \ })
-    endif
-
-    " TODO: per-file location
-    let l:delta_y = abs(s:pos_y - s:pos_y_pick)
-
-    " gather some extra context nearby and process it in the background
-    " only gather chunks if the cursor has moved a lot
-    " TODO: something more clever? reranking?
-    if a:is_auto && l:delta_y > 32
-        " expand the prefix even further
-        call s:pick_chunk(getline(max([1,       s:pos_y - g:llama_config.ring_scope]), max([1,       s:pos_y - g:llama_config.n_prefix])), v:false, v:false)
-
-        " pick a suffix chunk
-        call s:pick_chunk(getline(min([l:max_y, s:pos_y + g:llama_config.n_suffix]),   min([l:max_y, s:pos_y + g:llama_config.n_suffix + g:llama_config.ring_chunk_size])), v:false, v:false)
-
-        let s:pos_y_pick = s:pos_y
-    endif
-endfunction
-
-" if first_line == v:true accept only the first line of the response
-function! llama#fim_accept(first_line)
-    " insert the suggestion at the cursor location
-    if s:can_accept && len(s:content) > 0
-        call setline(s:pos_y, s:line_cur[:(s:pos_x - 1)] . s:content[0])
-        if len(s:content) > 1
-            if !a:first_line
-                call append(s:pos_y, s:content[1:-1])
-            endif
-        endif
-
-        " move the cursor to the end of the accepted text
-        if !a:first_line && len(s:content) > 1
-            call cursor(s:pos_y + len(s:content) - 1, s:pos_x + s:pos_dx + 1)
-        else
-            call cursor(s:pos_y, s:pos_x + len(s:content[0]))
-        endif
-    endif
-
-    call llama#fim_cancel()
-endfunction
-
-function! llama#fim_cancel()
-    let s:hint_shown = v:false
-
-    " clear the virtual text
-    let l:bufnr = bufnr('%')
-
-    if s:ghost_text_nvim
-        let l:id_vt_fim = nvim_create_namespace('vt_fim')
-        call nvim_buf_clear_namespace(l:bufnr, l:id_vt_fim,  0, -1)
-    elseif s:ghost_text_vim
-        call prop_remove({'type': s:hlgroup_hint, 'all': v:true})
-        call prop_remove({'type': s:hlgroup_info, 'all': v:true})
-    endif
-
-    " remove the mappings
-    silent! iunmap <buffer> <Tab>
-    silent! iunmap <buffer> <S-Tab>
-    silent! iunmap <buffer> <Esc>
-endfunction
-
-function! s:on_move()
-    let s:t_last_move = reltime()
-
-    call llama#fim_cancel()
-endfunction
-
-" callback that processes the FIM result from the server and displays the suggestion
-function! s:fim_on_stdout(pos_x, pos_y, is_auto, job_id, data, event = v:null)
-    if s:ghost_text_nvim
-        let l:raw = join(a:data, "\n")
-    elseif s:ghost_text_vim
-        let l:raw = a:data
-    endif
-
-    if len(l:raw) == 0
-        return
-    endif
-
-    if a:pos_x != col('.') - 1 || a:pos_y != line('.')
-        return
-    endif
-
-    " show the suggestion only in insert mode
-    if mode() !=# 'i'
-        return
-    endif
-
-    let s:pos_x = a:pos_x
-    let s:pos_y = a:pos_y
-
-    let s:can_accept = v:true
-    let l:has_info   = v:false
-
-    if s:can_accept && v:shell_error
-        if !a:is_auto
-            call add(s:content, "<| curl error: is the server on? |>")
-        endif
-        let s:can_accept = v:false
-    endif
-
-    let l:n_prompt    = 0
-    let l:t_prompt_ms = 1.0
-    let l:s_prompt    = 0
-
-    let l:n_predict    = 0
-    let l:t_predict_ms = 1.0
-    let l:s_predict    = 0
-
-    " get the generated suggestion
-    if s:can_accept
-        let l:response = json_decode(l:raw)
-
-        for l:part in split(get(l:response, 'content', ''), "\n", 1)
-            call add(s:content, l:part)
-        endfor
-
-        " remove trailing new lines
-        while len(s:content) > 0 && s:content[-1] == ""
-            call remove(s:content, -1)
-        endwhile
-
-        let l:generation_settings = get(l:response, 'generation_settings', {})
-        let l:n_ctx = get(l:generation_settings, 'n_ctx', 0)
-
-        let l:n_cached  = get(l:response, 'tokens_cached', 0)
-        let l:truncated = get(l:response, 'truncated', v:false)
-
-        " if response.timings is available
-        if len(get(l:response, 'timings', {})) > 0
-            let l:has_info = v:true
-            let l:timings  = get(l:response, 'timings', {})
-
-            let l:n_prompt    = get(l:timings, 'prompt_n', 0)
-            let l:t_prompt_ms = get(l:timings, 'prompt_ms', 1)
-            let l:s_prompt    = get(l:timings, 'prompt_per_second', 0)
-
-            let l:n_predict    = get(l:timings, 'predicted_n', 0)
-            let l:t_predict_ms = get(l:timings, 'predicted_ms', 1)
-            let l:s_predict    = get(l:timings, 'predicted_per_second', 0)
-        endif
-    endif
-
-    if len(s:content) == 0
-        call add(s:content, "")
-        let s:can_accept = v:false
-    endif
-
-    if len(s:content) == 0
-        return
-    endif
-
-    " NOTE: the following is logic for discarding predictions that repeat existing text
-    "       the code is quite ugly and there is very likely a simpler and more canonical way to implement this
-    "
-    "       still, I wonder if there is some better way that avoids having to do these special hacks?
-    "       on one hand, the LLM 'sees' the contents of the file before we start editing, so it is normal that it would
-    "       start generating whatever we have given it via the extra context. but on the other hand, it's not very
-    "       helpful to re-generate the same code that is already there
-
-    " truncate the suggestion if the first line is empty
-    if len(s:content) == 1 && s:content[0] == ""
-        let s:content = [""]
-    endif
-
-    " ... and the next lines are repeated
-    if len(s:content) > 1 && s:content[0] == "" && s:content[1:] == getline(s:pos_y + 1, s:pos_y + len(s:content) - 1)
-        let s:content = [""]
-    endif
-
-    " truncate the suggestion if it repeats the suffix
-    if len(s:content) == 1 && s:content[0] == s:line_cur_suffix
-        let s:content = [""]
-    endif
-
-    " find the first non-empty line (strip whitespace)
-    let l:cmp_y = s:pos_y + 1
-    while l:cmp_y < line('$') && getline(l:cmp_y) =~? '^\s*$'
-        let l:cmp_y += 1
-    endwhile
-
-    if (s:line_cur_prefix . s:content[0]) == getline(l:cmp_y)
-        " truncate the suggestion if it repeats the next line
-        if len(s:content) == 1
-            let s:content = [""]
-        endif
-
-        " ... or if the second line of the suggestion is the prefix of line l:cmp_y + 1
-        if len(s:content) == 2 && s:content[-1] == getline(l:cmp_y + 1)[:len(s:content[-1]) - 1]
-            let s:content = [""]
-        endif
-
-        " ... or if the middle chunk of lines of the suggestion is the same as [l:cmp_y + 1, l:cmp_y + len(s:content) - 1)
-        if len(s:content) > 2 && join(s:content[1:-1], "\n") == join(getline(l:cmp_y + 1, l:cmp_y + len(s:content) - 1), "\n")
-            let s:content = [""]
-        endif
-    endif
-
-    " keep only lines that have the same or larger whitespace prefix as s:line_cur_prefix
-    "let l:indent = strlen(matchstr(s:line_cur_prefix, '^\s*'))
-    "for i in range(1, len(s:content) - 1)
-    "    if strlen(matchstr(s:content[i], '^\s*')) < l:indent
-    "        let s:content = s:content[:i - 1]
-    "        break
-    "    endif
-    "endfor
-
-    let s:pos_dx = len(s:content[-1])
-
-    let s:content[-1] .= s:line_cur_suffix
-
-    call llama#fim_cancel()
-
-    " display virtual text with the suggestion
-    let l:bufnr = bufnr('%')
-
-    if s:ghost_text_nvim
-        let l:id_vt_fim = nvim_create_namespace('vt_fim')
-    endif
-
-    " construct the info message
-    if g:llama_config.show_info > 0 && l:has_info
-        let l:prefix = '   '
-
-        if l:truncated
-            let l:info = printf("%s | WARNING: the context is full: %d / %d, increase the server context size or reduce g:llama_config.ring_n_chunks",
-                \ g:llama_config.show_info == 2 ? l:prefix : 'llama.vim',
-                \ l:n_cached, l:n_ctx
-                \ )
-        else
-            let l:info = printf("%s | c: %d / %d, r: %d / %d, e: %d, q: %d / 16 | p: %d (%.2f ms, %.2f t/s) | g: %d (%.2f ms, %.2f t/s) | t: %.2f ms",
-                \ g:llama_config.show_info == 2 ? l:prefix : 'llama.vim',
-                \ l:n_cached,  l:n_ctx, len(s:ring_chunks), g:llama_config.ring_n_chunks, s:ring_n_evict, len(s:ring_queued),
-                \ l:n_prompt,  l:t_prompt_ms,  l:s_prompt,
-                \ l:n_predict, l:t_predict_ms, l:s_predict,
-                \ 1000.0 * reltimefloat(reltime(s:t_fim_start))
-                \ )
-        endif
-
-        if g:llama_config.show_info == 1
-            " display the info in the statusline
-            let &statusline = l:info
-            let l:info = ''
-        endif
-    endif
-
-    " display the suggestion and append the info to the end of the first line
-    if s:ghost_text_nvim
-        call nvim_buf_set_extmark(l:bufnr, l:id_vt_fim, s:pos_y - 1, s:pos_x - 1, {
-            \ 'virt_text': [[s:content[0], 'llama_hl_hint'], [l:info, 'llama_hl_info']],
-            \ 'virt_text_win_col': virtcol('.') - 1
-            \ })
-
-        call nvim_buf_set_extmark(l:bufnr, l:id_vt_fim, s:pos_y - 1, 0, {
-            \ 'virt_lines': map(s:content[1:], {idx, val -> [[val, 'llama_hl_hint']]}),
-            \ 'virt_text_win_col': virtcol('.')
-            \ })
-    elseif s:ghost_text_vim
-        let l:new_suffix = s:content[0]
-        if !empty(l:new_suffix)
-            call prop_add(s:pos_y, s:pos_x + 1, {
-                        \ 'type': s:hlgroup_hint,
-                        \ 'text': l:new_suffix
-                        \ })
-        endif
-        for line in s:content[1:]
-            call prop_add(s:pos_y, 0, {
-                        \ 'type': s:hlgroup_hint,
-                        \ 'text': line,
-                        \ 'text_padding_left': s:get_indent(line),
-                        \ 'text_align': 'below'
-                        \ })
-        endfor
-        if !empty(l:info)
-            call prop_add(s:pos_y, 0, {
-                        \ 'type': s:hlgroup_info,
-                        \ 'text': l:info,
-                        \ 'text_padding_left': col('$'),
-                        \ 'text_wrap': 'truncate'
-                        \ })
-        endif
-    endif
-
-    " setup accept shortcuts
-    inoremap <buffer> <Tab>   <C-O>:call llama#fim_accept(v:false)<CR>
-    inoremap <buffer> <S-Tab> <C-O>:call llama#fim_accept(v:true)<CR>
-
-    let s:hint_shown = v:true
-endfunction
-
-function! s:fim_on_exit(job_id, exit_code, event = v:null)
-    if a:exit_code != 0
-        echom "Job failed with exit code: " . a:exit_code
-    endif
-
-    let s:current_job = v:null
-endfunction

examples/pydantic_models_to_grammar_examples.py

@@ -1,312 +0,0 @@
-#!/usr/bin/env python3
-
-"""Function calling example using pydantic models."""
-
-from __future__ import annotations
-
-import argparse
-import datetime
-import json
-import logging
-import textwrap
-import sys
-from enum import Enum
-from typing import Optional, Union
-
-import requests
-from pydantic import BaseModel, Field
-from pydantic_models_to_grammar import (add_run_method_to_dynamic_model, convert_dictionary_to_pydantic_model,
-                                        create_dynamic_model_from_function, generate_gbnf_grammar_and_documentation)
-
-
-def create_completion(host, prompt, gbnf_grammar):
-    """Calls the /completion API on llama-server.
-
-    See
-    https://github.com/ggml-org/llama.cpp/tree/HEAD/tools/server#api-endpoints
-    """
-    print(f"  Request:\n    Grammar:\n{textwrap.indent(gbnf_grammar, '      ')}\n    Prompt:\n{textwrap.indent(prompt.rstrip(), '      ')}")
-    headers = {"Content-Type": "application/json"}
-    data = {"prompt": prompt, "grammar": gbnf_grammar}
-    result = requests.post(f"http://{host}/completion", headers=headers, json=data).json()
-    assert data.get("error") is None, data
-    logging.info("Result: %s", result)
-    content = result["content"]
-    print(f"  Model: {result['model']}")
-    print(f"  Result:\n{textwrap.indent(json.dumps(json.loads(content), indent=2), '    ')}")
-    return content
-
-
-# A function for the agent to send a message to the user.
-class SendMessageToUser(BaseModel):
-    """Send a message to the User."""
-    chain_of_thought: str = Field(..., description="Your chain of thought while sending the message.")
-    message: str = Field(..., description="Message you want to send to the user.")
-
-    def run(self):
-        print(f"SendMessageToUser: {self.message}")
-
-
-def example_rce(host):
-    """Minimal test case where the LLM call an arbitrary python function."""
-    print("- example_rce")
-    tools = [SendMessageToUser]
-    gbnf_grammar, documentation = generate_gbnf_grammar_and_documentation(
-        pydantic_model_list=tools, outer_object_name="function",
-        outer_object_content="function_parameters", model_prefix="Function", fields_prefix="Parameters")
-    system_message = "You are an advanced AI, tasked to assist the user by calling functions in JSON format. The following are the available functions and their parameters and types:\n\n" + documentation
-    user_message = "What is 42 * 42?"
-    prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{user_message}<|im_end|>\n<|im_start|>assistant"
-    text = create_completion(host, prompt, gbnf_grammar)
-    json_data = json.loads(text)
-    tools_map = {tool.__name__:tool for tool in tools}
-    # This finds "SendMessageToUser":
-    tool = tools_map.get(json_data["function"])
-    if not tool:
-        print(f"Error: unknown tool {json_data['function']}")
-        return 1
-    tool(**json_data["function_parameters"]).run()
-    return 0
-
-
-# Enum for the calculator tool.
-class MathOperation(Enum):
-    ADD = "add"
-    SUBTRACT = "subtract"
-    MULTIPLY = "multiply"
-    DIVIDE = "divide"
-
-
-# Simple pydantic calculator tool for the agent that can add, subtract,
-# multiply, and divide. Docstring and description of fields will be used in
-# system prompt.
-class Calculator(BaseModel):
-    """Perform a math operation on two numbers."""
-    number_one: Union[int, float] = Field(..., description="First number.")
-    operation: MathOperation = Field(..., description="Math operation to perform.")
-    number_two: Union[int, float] = Field(..., description="Second number.")
-
-    def run(self):
-        if self.operation == MathOperation.ADD:
-            return self.number_one + self.number_two
-        elif self.operation == MathOperation.SUBTRACT:
-            return self.number_one - self.number_two
-        elif self.operation == MathOperation.MULTIPLY:
-            return self.number_one * self.number_two
-        elif self.operation == MathOperation.DIVIDE:
-            return self.number_one / self.number_two
-        else:
-            raise ValueError("Unknown operation.")
-
-
-def example_calculator(host):
-    """Have the LLM ask to get a calculation done.
-
-    Here the grammar gets generated by passing the available function models to
-    generate_gbnf_grammar_and_documentation function. This also generates a
-    documentation usable by the LLM.
-
-    pydantic_model_list is the list of pydantic models outer_object_name is an
-    optional name for an outer object around the actual model object. Like a
-    "function" object with "function_parameters" which contains the actual model
-    object. If None, no outer object will be generated outer_object_content is
-    the name of outer object content.
-
-    model_prefix is the optional prefix for models in the documentation. (Default="Output Model")
-    fields_prefix is the prefix for the model fields in the documentation. (Default="Output Fields")
-    """
-    print("- example_calculator")
-    tools = [SendMessageToUser, Calculator]
-    gbnf_grammar, documentation = generate_gbnf_grammar_and_documentation(
-        pydantic_model_list=tools, outer_object_name="function",
-        outer_object_content="function_parameters", model_prefix="Function", fields_prefix="Parameters")
-    system_message = "You are an advanced AI, tasked to assist the user by calling functions in JSON format. The following are the available functions and their parameters and types:\n\n" + documentation
-    user_message1 = "What is 42 * 42?"
-    prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{user_message1}<|im_end|>\n<|im_start|>assistant"
-    text = create_completion(host, prompt, gbnf_grammar)
-    json_data = json.loads(text)
-    expected = {
-        "function": "Calculator",
-        "function_parameters": {
-            "number_one": 42,
-            "operation": "multiply",
-            "number_two": 42
-        }
-    }
-    if json_data != expected:
-        print("  Result is not as expected!")
-    tools_map = {tool.__name__:tool for tool in tools}
-    # This finds "Calculator":
-    tool = tools_map.get(json_data["function"])
-    if not tool:
-        print(f"Error: unknown tool {json_data['function']}")
-        return 1
-    result = tool(**json_data["function_parameters"]).run()
-    print(f"  Call {json_data['function']} gave result {result}")
-    return 0
-
-
-class Category(Enum):
-    """The category of the book."""
-    Fiction = "Fiction"
-    NonFiction = "Non-Fiction"
-
-
-class Book(BaseModel):
-    """Represents an entry about a book."""
-    title: str = Field(..., description="Title of the book.")
-    author: str = Field(..., description="Author of the book.")
-    published_year: Optional[int] = Field(..., description="Publishing year of the book.")
-    keywords: list[str] = Field(..., description="A list of keywords.")
-    category: Category = Field(..., description="Category of the book.")
-    summary: str = Field(..., description="Summary of the book.")
-
-
-def example_struct(host):
-    """A example structured output based on pydantic models.
-
-    The LLM will create an entry for a Book database out of an unstructured
-    text. We need no additional parameters other than our list of pydantic
-    models.
-    """
-    print("- example_struct")
-    tools = [Book]
-    gbnf_grammar, documentation = generate_gbnf_grammar_and_documentation(pydantic_model_list=tools)
-    system_message = "You are an advanced AI, tasked to create a dataset entry in JSON for a Book. The following is the expected output model:\n\n" + documentation
-    text = """The Feynman Lectures on Physics is a physics textbook based on some lectures by Richard Feynman, a Nobel laureate who has sometimes been called "The Great Explainer". The lectures were presented before undergraduate students at the California Institute of Technology (Caltech), during 1961–1963. The book's co-authors are Feynman, Robert B. Leighton, and Matthew Sands."""
-    prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant"
-    text = create_completion(host, prompt, gbnf_grammar)
-    json_data = json.loads(text)
-    # In this case, there's no function nor function_parameters.
-    # Here the result will vary based on the LLM used.
-    keys = sorted(["title", "author", "published_year", "keywords", "category", "summary"])
-    if keys != sorted(json_data.keys()):
-        print(f"Unexpected result: {sorted(json_data.keys())}")
-        return 1
-    book = Book(**json_data)
-    print(f"  As a Book object: %s" % book)
-    return 0
-
-
-def get_current_datetime(output_format: Optional[str] = None):
-    """Get the current date and time in the given format.
-
-    Args:
-         output_format: formatting string for the date and time, defaults to '%Y-%m-%d %H:%M:%S'
-    """
-    return datetime.datetime.now().strftime(output_format or "%Y-%m-%d %H:%M:%S")
-
-
-# Example function to get the weather.
-def get_current_weather(location, unit):
-    """Get the current weather in a given location"""
-    if "London" in location:
-        return json.dumps({"location": "London", "temperature": "42", "unit": unit.value})
-    elif "New York" in location:
-        return json.dumps({"location": "New York", "temperature": "24", "unit": unit.value})
-    elif "North Pole" in location:
-        return json.dumps({"location": "North Pole", "temperature": "-42", "unit": unit.value})
-    return json.dumps({"location": location, "temperature": "unknown"})
-
-
-def example_concurrent(host):
-    """An example for parallel function calling with a Python function, a pydantic
-    function model and an OpenAI like function definition.
-    """
-    print("- example_concurrent")
-    # Function definition in OpenAI style.
-    current_weather_tool = {
-        "type": "function",
-        "function": {
-            "name": "get_current_weather",
-            "description": "Get the current weather in a given location",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "location": {
-                        "type": "string",
-                        "description": "The city and state, e.g. San Francisco, CA",
-                    },
-                    "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
-                },
-                "required": ["location"],
-            },
-        },
-    }
-    # Convert OpenAI function definition into pydantic model.
-    current_weather_tool_model = convert_dictionary_to_pydantic_model(current_weather_tool)
-    # Add the actual function to a pydantic model.
-    current_weather_tool_model = add_run_method_to_dynamic_model(current_weather_tool_model, get_current_weather)
-
-    # Convert normal Python function to a pydantic model.
-    current_datetime_model = create_dynamic_model_from_function(get_current_datetime)
-
-    tools = [SendMessageToUser, Calculator, current_datetime_model, current_weather_tool_model]
-    gbnf_grammar, documentation = generate_gbnf_grammar_and_documentation(
-        pydantic_model_list=tools, outer_object_name="function",
-        outer_object_content="params", model_prefix="Function", fields_prefix="Parameters", list_of_outputs=True)
-    system_message = "You are an advanced AI assistant. You are interacting with the user and with your environment by calling functions. You call functions by writing JSON objects, which represent specific function calls.\nBelow is a list of your available function calls:\n\n" + documentation
-    text = """Get the date and time, get the current weather in celsius in London and solve the following calculation: 42 * 42"""
-    prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant"
-    text = create_completion(host, prompt, gbnf_grammar)
-    json_data = json.loads(text)
-    expected = [
-      {
-        "function": "get_current_datetime",
-        "params": {
-          "output_format": "%Y-%m-%d %H:%M:%S"
-        }
-      },
-      {
-        "function": "get_current_weather",
-        "params": {
-          "location": "London",
-          "unit": "celsius"
-        }
-      },
-      {
-        "function": "Calculator",
-        "params": {
-          "number_one": 42,
-          "operation": "multiply",
-          "number_two": 42
-        }
-      }
-    ]
-    res = 0
-    if json_data != expected:
-        print("  Result is not as expected!")
-        print("  This can happen on highly quantized models")
-        res = 1
-    tools_map = {tool.__name__:tool for tool in tools}
-    for call in json_data:
-      tool = tools_map.get(call["function"])
-      if not tool:
-          print(f"Error: unknown tool {call['function']}")
-          return 1
-      result = tool(**call["params"]).run()
-      print(f"  Call {call['function']} returned {result}")
-    # Should output something like this:
-    #   Call get_current_datetime returned 2024-07-15 09:50:38
-    #   Call get_current_weather returned {"location": "London", "temperature": "42", "unit": "celsius"}
-    #   Call Calculator returned 1764
-    return res
-
-
-def main():
-    parser = argparse.ArgumentParser(description=sys.modules[__name__].__doc__)
-    parser.add_argument("--host", default="localhost:9931", help="llama.cpp server")
-    parser.add_argument("-v", "--verbose", action="store_true", help="enables logging")
-    args = parser.parse_args()
-    logging.basicConfig(level=logging.INFO if args.verbose else logging.ERROR)
-    ret = 0
-    # Comment out below to only run the example you want.
-    ret = ret or example_rce(args.host)
-    ret = ret or example_calculator(args.host)
-    ret = ret or example_struct(args.host)
-    ret = ret or example_concurrent(args.host)
-    return ret
-
-
-if __name__ == "__main__":
-    sys.exit(main())

examples/realtime_client.py

@@ -1,211 +0,0 @@
-#!/usr/bin/env python3
-"""Example client for the llama-server /v1/realtime websocket, with the openai SDK.
-
-Start a server with an audio model, e.g. LFM3-Audio:
-
-    llama-server -hf LiquidAI/LFM3-Audio-2B-GGUF
-
-Spoken conversation, the audio of each wav file is one user turn. It is sent in chunks as it would come from a
-microphone, the server prefills it as it arrives so the response starts right after the last chunk:
-
-    python examples/realtime_client.py --instructions "Respond with interleaved text and audio." question.wav
-
-Output modalities: left out, the model picks; "text", "audio" or "text,audio" (interleaved) force the response.
-Without --instructions, a forced mode comes with the system prompt of the LFM3-Audio or LFM3-Omni model card:
-
-    python examples/realtime_client.py --output-modalities text,audio --out answer.wav question.wav
-
-Tool calls: --tools gives the conversation a demo tool, the client runs its calls and sends the results back:
-
-    python examples/realtime_client.py --tools --text "Set the volume to 30."
-
-Text turns go in with --text, transcription sessions transcribe each file:
-
-    python examples/realtime_client.py --text "What is the capital of France?"
-    python examples/realtime_client.py --type transcription --language en speech.wav
-
-Live transcription models (Voxtral Realtime, LFM3-Audio live ASR) transcribe while the audio is sent, a transcription
-session opens with ?intent=transcription:
-
-    python examples/realtime_client.py --type transcription speech.wav
-
-Requirements: pip install "openai[realtime]" numpy
-"""
-
-import argparse
-import asyncio
-import base64
-import json
-import time
-import wave
-
-import numpy as np
-from openai import AsyncOpenAI
-
-
-# system prompts of the LiquidAI/LFM3-Audio-2B and LiquidAI/LFM3-Omni-2B model cards for a forced output mode
-MODE_INSTRUCTIONS = {
-    "audio": {"text,audio": "Respond with interleaved text and audio."},
-    "omni":  {"text": "Respond in text.", "text,audio": "Respond in interleaved text and audio."},
-}
-# and with tools: calls, or an answer in interleaved text and speech
-TOOLS_INSTRUCTIONS = "Respond with function calls, and interleaved text and audio.\n\nList of tools:"
-
-# the demo tool of --tools
-DEMO_TOOLS = [{
-    "type": "function",
-    "name": "set_volume",
-    "description": "Set the device speaker volume.",
-    "parameters": {
-        "type": "object",
-        "properties": {"volume": {"type": "integer", "description": "Volume from 0 to 100."}},
-        "required": ["volume"],
-    },
-}]
-
-
-def run_tool(name, arguments):
-    if name == "set_volume":
-        return {"volume": json.loads(arguments)["volume"]}
-    return {"error": f"unknown tool {name}"}
-
-
-def read_wav(path, rate):
-    # PCM16 wav -> mono int16 at the given rate
-    with wave.open(path, "rb") as f:
-        assert f.getsampwidth() == 2, "the example reads PCM16 wav files"
-        x = np.frombuffer(f.readframes(f.getnframes()), dtype=np.int16).astype(np.float32)
-        x = x.reshape(-1, f.getnchannels()).mean(axis=1)
-        sr = f.getframerate()
-    if sr != rate:
-        n = int(round(len(x) * rate / sr))
-        x = np.interp(np.arange(n) * sr / rate, np.arange(len(x)), x)
-    return np.clip(x, -32768, 32767).astype(np.int16)
-
-
-def write_wav(path, pcm, rate):
-    with wave.open(path, "wb") as f:
-        f.setnchannels(1)
-        f.setsampwidth(2)
-        f.setframerate(rate)
-        f.writeframes(pcm.tobytes())
-
-
-async def main(args):
-    # the SDK derives a wss:// URL from base_url, a plain http server needs ws://
-    client = AsyncOpenAI(base_url=args.url, api_key=args.api_key, websocket_base_url=args.url.replace("http", "ws", 1))
-    query = {"intent": "transcription"} if args.type == "transcription" else {}
-    async with client.realtime.connect(model=args.model, extra_query=query) as conn:
-        ev = await conn.recv()
-        if ev.type == "error":
-            print("error:", ev.error.message)
-            return
-        # the audio goes in at the rate of the model, PCM16 mono
-        session_info = ev.model_dump(warnings=False)["session"]
-        audio_fmt = session_info["audio"]
-        card = "omni" if "omni" in str(session_info.get("model", "")).lower() else "audio"
-        in_rate = audio_fmt["input"]["format"]["rate"]
-        out_rate = (audio_fmt.get("output") or {}).get("format", {}).get("rate", 24000)
-
-        session = {"type": args.type}
-        if args.type == "realtime":
-            instructions = args.instructions
-            if instructions is None:
-                instructions = TOOLS_INSTRUCTIONS if args.tools else MODE_INSTRUCTIONS[card].get(args.output_modalities)
-            if args.tools:
-                session["tools"] = DEMO_TOOLS
-            if instructions is not None:
-                session["instructions"] = instructions
-            if args.output_modalities:
-                session["output_modalities"] = args.output_modalities.split(",")
-        elif args.language:
-            session["audio"] = {"input": {"transcription": {"language": args.language}}}
-        await conn.session.update(session=session)  # ty: ignore[invalid-argument-type]
-        ev = await conn.recv()
-        if ev.type == "error":
-            print("error:", ev.error.message)
-            return
-
-        turns = [("audio", path) for path in args.wav] + [("text", text) for text in args.text]
-        for i, (kind, src) in enumerate(turns):
-            t_last = 0.0
-
-            # the input is sent while the events are read: live transcription arrives while the audio is sent
-            async def send_turn():
-                nonlocal t_last
-                if kind == "audio":
-                    pcm = read_wav(src, in_rate)
-                    step = in_rate * args.chunk_ms // 1000
-                    for j in range(0, len(pcm), step):
-                        chunk = pcm[j:j + step]
-                        await conn.input_audio_buffer.append(audio=base64.b64encode(chunk.tobytes()).decode())
-                        if args.realtime:
-                            await asyncio.sleep(len(chunk) / in_rate)
-                    if args.type == "transcription":
-                        # "final" ends the input, models that transcribe with a delay then complete the last item
-                        await conn.send({"type": "input_audio_buffer.commit", "final": i == len(args.wav) - 1})  # ty: ignore[invalid-argument-type]
-                else:
-                    await conn.conversation.item.create(item={
-                        "type": "message", "role": "user", "content": [{"type": "input_text", "text": src}]})
-                if args.type == "realtime":
-                    # the audio not committed yet ends the user turn
-                    await conn.response.create()
-                t_last = time.perf_counter()
-
-            print(f"--- turn {i} ({kind}: {src})")
-            sender = asyncio.create_task(send_turn())
-            t_first = None
-            audio = []
-            async for ev in conn:
-                if ev.type == "response.output_text.delta" or ev.type == "conversation.item.input_audio_transcription.delta":
-                    t_first = t_first or time.perf_counter()
-                    print(ev.delta, end="", flush=True)
-                elif ev.type == "response.output_audio.delta":
-                    t_first = t_first or time.perf_counter()
-                    audio.append(np.frombuffer(base64.b64decode(ev.delta), dtype=np.int16))
-                elif ev.type == "response.done":
-                    # the calls of the response are run, their results go back and the model answers them
-                    calls = [o for o in ev.model_dump()["response"]["output"] if o["type"] == "function_call"]
-                    for call in calls:
-                        result = json.dumps(run_tool(call["name"], call["arguments"]))
-                        print(f"\n[call {call['name']}({call['arguments']}) -> {result}]", flush=True)
-                        await conn.conversation.item.create(item={
-                            "type": "function_call_output", "call_id": call["call_id"], "output": result})
-                    if not calls:
-                        break
-                    await conn.response.create()
-                elif ev.type == "conversation.item.input_audio_transcription.completed":
-                    break
-                elif ev.type == "error":
-                    print("\nerror:", ev.error.message)
-                    return
-            await sender
-            print()
-            if t_first is not None and t_first > t_last:
-                print(f"first output {1000 * (t_first - t_last):.0f} ms after the end of the input")
-            elif t_first is not None:
-                print(f"output started {1000 * (t_last - t_first):.0f} ms before the end of the input, "
-                      f"done {1000 * (time.perf_counter() - t_last):.0f} ms after it")
-            if audio:
-                pcm_out = np.concatenate(audio)
-                path = f"{args.out.removesuffix('.wav')}_{i}.wav" if len(turns) > 1 else args.out
-                write_wav(path, pcm_out, out_rate)
-                print(f"audio: {len(pcm_out) / out_rate:.2f} s -> {path}")
-
-
-if __name__ == "__main__":
-    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-    p.add_argument("wav", nargs="*", help="PCM16 wav files, one user turn each")
-    p.add_argument("--text", action="append", default=[], help="a user turn as text, after the audio turns, repeatable")
-    p.add_argument("--type", choices=["realtime", "transcription"], default="realtime", help="session type")
-    p.add_argument("--instructions", help="system prompt of the conversation")
-    p.add_argument("--tools", action="store_true", help="give the conversation a demo tool (set_volume), its calls are run")
-    p.add_argument("--output-modalities", choices=["text", "audio", "text,audio"], help="default: the model picks")
-    p.add_argument("--language", help="transcription language, e.g. en")
-    p.add_argument("--chunk-ms", type=int, default=100, help="audio chunk size")
-    p.add_argument("--realtime", action=argparse.BooleanOptionalAction, default=True, help="send the audio in real time")
-    p.add_argument("--out", default="response.wav", help="where the audio of the responses goes")
-    p.add_argument("--url", default="http://127.0.0.1:9931/v1")
-    p.add_argument("--api-key", default="none")
-    p.add_argument("--model", default="default", help="model name, used by the router")
-    asyncio.run(main(p.parse_args()))

examples/reason-act.sh

@@ -1,16 +0,0 @@
-#!/usr/bin/env bash
-
-cd `dirname $0`
-cd ..
-
-# get -m model parameter otherwise defer to default
-if [ "$1" == "-m" ]; then
-  MODEL="-m $2 "
-fi
-
-./llama-cli $MODEL --color \
-    -f ./prompts/reason-act.txt \
-    -i --interactive-first \
-    --top_k 10000 --temp 0.2 --repeat_penalty 1 -t 7 -c 2048 \
-    -r "Question:" -r "Observation:" --in-prefix " " \
-    -n -1

examples/server-llama2-13B.sh

@@ -1,26 +0,0 @@
-#!/usr/bin/env bash
-
-set -e
-
-cd "$(dirname "$0")/.." || exit
-
-# Specify the model you want to use here:
-MODEL="${MODEL:-./models/llama-2-13b-chat.ggmlv3.q5_K_M.bin}"
-PROMPT_TEMPLATE=${PROMPT_TEMPLATE:-./prompts/chat-system.txt}
-
-# Adjust to the number of CPU cores you want to use.
-N_THREAD="${N_THREAD:-12}"
-
-# Note: you can also override the generation options by specifying them on the command line:
-GEN_OPTIONS="${GEN_OPTIONS:---ctx_size 4096 --batch-size 1024}"
-
-
-# shellcheck disable=SC2086 # Intended splitting of GEN_OPTIONS
-./llama-server $GEN_OPTIONS \
-  --model "$MODEL" \
-  --threads "$N_THREAD" \
-  --rope-freq-scale 1.0 \
-  "$@"
-
-# I used this to test the model with mps, but omitted it from the general purpose. If you want to use it, just specify it on the command line.
-# -ngl 1 \

examples/server_embd.py

@@ -1,35 +0,0 @@
-import asyncio
-import asyncio.threads
-import requests
-import numpy as np
-
-
-n = 8
-
-result = []
-
-async def requests_post_async(*args, **kwargs):
-    return await asyncio.threads.to_thread(requests.post, *args, **kwargs)
-
-async def main():
-    model_url = "http://127.0.0.1:6900"
-    responses: list[requests.Response] = await asyncio.gather(*[requests_post_async(
-        url= f"{model_url}/embedding",
-        json= {"content": "a "*1022}
-    ) for i in range(n)])
-
-    for response in responses:
-        embedding = response.json()["embedding"]
-        print(embedding[-8:])
-        result.append(embedding)
-
-asyncio.run(main())
-
-# compute cosine similarity
-
-for i in range(n-1):
-    for j in range(i+1, n):
-        embedding1 = np.array(result[i])
-        embedding2 = np.array(result[j])
-        similarity = np.dot(embedding1, embedding2) / (np.linalg.norm(embedding1) * np.linalg.norm(embedding2))
-        print(f"Similarity between {i} and {j}: {similarity:.2f}")

examples/test-cmake/CMakeLists.txt

@@ -8,7 +8,7 @@ option(LLAMA_TEST_USE_SUBDIR "Use add_subdirectory instead of find_package" OFF)
 if(LLAMA_TEST_USE_SUBDIR)
     add_subdirectory(../../ llama.cpp)
 else()
-    find_package(llama 0.1.0 REQUIRED)
+    find_package(llama CONFIG REQUIRED PATHS ${CMAKE_PREFIX_PATH} NO_DEFAULT_PATH)
 endif()
 
 add_executable(test-cmake test-cmake.cpp)

ggml/CMakeLists.txt

@@ -3,8 +3,8 @@ cmake_minimum_required(VERSION 3.14...3.28) # for add_link_options and implicit
 project("ggml" C CXX ASM)
 
 ### GGML Version
-set(GGML_VERSION_MAJOR 0)
-set(GGML_VERSION_MINOR 29)
+set(GGML_VERSION_MAJOR 1)
+set(GGML_VERSION_MINOR 0)
 set(GGML_VERSION_PATCH 0)
 set(GGML_VERSION_BASE "${GGML_VERSION_MAJOR}.${GGML_VERSION_MINOR}.${GGML_VERSION_PATCH}")
 

ggml/include/ggml-rpc.h

@@ -6,26 +6,31 @@
 extern "C" {
 #endif
 
-#define RPC_PROTO_MAJOR_VERSION    9
+#define RPC_PROTO_MAJOR_VERSION    10
 #define RPC_PROTO_MINOR_VERSION    0
 #define RPC_PROTO_PATCH_VERSION    0
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif
 
 #define GGML_RPC_MAX_SERVERS       16
 
 // backend API
 GGML_BACKEND_API ggml_backend_t ggml_backend_rpc_init(const char * endpoint, uint32_t device);
+
 GGML_BACKEND_API bool ggml_backend_is_rpc(ggml_backend_t backend);
 
 GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, uint32_t device);
 
 GGML_BACKEND_API void ggml_backend_rpc_get_device_memory(const char * endpoint, uint32_t device, size_t * free, size_t * total);
 
-GGML_BACKEND_API void ggml_backend_rpc_start_server(const char * endpoint, const char * cache_dir,
-                                                    size_t n_threads, size_t n_devices, ggml_backend_dev_t * devices);
+GGML_BACKEND_API void ggml_backend_rpc_start_server(
+        const char * endpoint,
+        const char * cache_dir,
+        size_t n_threads,
+        size_t n_devices,
+        ggml_backend_dev_t * devices);
 
 GGML_BACKEND_API ggml_backend_reg_t ggml_backend_rpc_reg(void);
 GGML_BACKEND_API ggml_backend_reg_t ggml_backend_rpc_add_server(const char * endpoint);

ggml/include/ggml.h

@@ -1,178 +1,5 @@
 #pragma once
 
-//
-// GGML Tensor Library
-//
-// This documentation is still a work in progress.
-// If you wish some specific topics to be covered, feel free to drop a comment:
-//
-//   https://github.com/ggml-org/whisper.cpp/issues/40
-//
-// ## Overview
-//
-// This library implements:
-//
-//  - a set of tensor operations
-//  - automatic differentiation
-//  - basic optimization algorithms
-//
-// The aim of this library is to provide a minimalistic approach for various machine learning tasks. This includes,
-// but is not limited to, the following:
-//
-//  - linear regression
-//  - support vector machines
-//  - neural networks
-//
-// The library allows the user to define a certain function using the available tensor operations. This function
-// definition is represented internally via a computation graph. Each tensor operation in the function definition
-// corresponds to a node in the graph. Having the computation graph defined, the user can choose to compute the
-// function's value and/or its gradient with respect to the input variables. Optionally, the function can be optimized
-// using one of the available optimization algorithms.
-//
-// For example, here we define the function: f(x) = a*x^2 + b
-//
-//   {
-//       struct ggml_init_params params = {
-//           .mem_size   = 16*1024*1024,
-//           .mem_buffer = NULL,
-//       };
-//
-//       // memory allocation happens here
-//       struct ggml_context * ctx = ggml_init(params);
-//
-//       struct ggml_tensor * x = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-//
-//       ggml_set_param(ctx, x); // x is an input variable
-//
-//       struct ggml_tensor * a  = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-//       struct ggml_tensor * b  = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-//       struct ggml_tensor * x2 = ggml_mul(ctx, x, x);
-//       struct ggml_tensor * f  = ggml_add(ctx, ggml_mul(ctx, a, x2), b);
-//
-//       ...
-//   }
-//
-// Notice that the function definition above does not involve any actual computation. The computation is performed only
-// when the user explicitly requests it. For example, to compute the function's value at x = 2.0:
-//
-//   {
-//       ...
-//
-//       struct ggml_cgraph * gf = ggml_new_graph(ctx);
-//       ggml_build_forward_expand(gf, f);
-//
-//       // set the input variable and parameter values
-//       ggml_set_f32(x, 2.0f);
-//       ggml_set_f32(a, 3.0f);
-//       ggml_set_f32(b, 4.0f);
-//
-//       ggml_graph_compute_with_ctx(ctx, &gf, n_threads);
-//
-//       printf("f = %f\n", ggml_get_f32_1d(f, 0));
-//
-//       ...
-//   }
-//
-// The actual computation is performed in the ggml_graph_compute() function.
-//
-// The ggml_new_tensor_...() functions create new tensors. They are allocated in the memory buffer provided to the
-// ggml_init() function. You have to be careful not to exceed the memory buffer size. Therefore, you have to know
-// in advance how much memory you need for your computation. Alternatively, you can allocate a large enough memory
-// and after defining the computation graph, call the ggml_used_mem() function to find out how much memory was
-// actually needed.
-//
-// The ggml_set_param() function marks a tensor as an input variable. This is used by the automatic
-// differentiation and optimization algorithms.
-//
-// The described approach allows to define the function graph once and then compute its forward or backward graphs
-// multiple times. All computations will use the same memory buffer allocated in the ggml_init() function. This way
-// the user can avoid the memory allocation overhead at runtime.
-//
-// The library supports multi-dimensional tensors - up to 4 dimensions. The FP16 and FP32 data types are first class
-// citizens, but in theory the library can be extended to support FP8 and integer data types.
-//
-// Each tensor operation produces a new tensor. Initially the library was envisioned to support only the use of unary
-// and binary operations. Most of the available operations fall into one of these two categories. With time, it became
-// clear that the library needs to support more complex operations. The way to support these operations is not clear
-// yet, but a few examples are demonstrated in the following operations:
-//
-//   - ggml_permute()
-//   - ggml_conv_1d_1s()
-//   - ggml_conv_1d_2s()
-//
-// For each tensor operator, the library implements a forward and backward computation function. The forward function
-// computes the output tensor value given the input tensor values. The backward function computes the adjoint of the
-// input tensors given the adjoint of the output tensor. For a detailed explanation of what this means, take a
-// calculus class, or watch the following video:
-//
-//   What is Automatic Differentiation?
-//   https://www.youtube.com/watch?v=wG_nF1awSSY
-//
-//
-// ## Tensor data (struct ggml_tensor)
-//
-// The tensors are stored in memory via the ggml_tensor struct. The structure provides information about the size of
-// the tensor, the data type, and the memory buffer where the tensor data is stored. Additionally, it contains
-// pointers to the "source" tensors - i.e. the tensors that were used to compute the current tensor. For example:
-//
-//   {
-//       struct ggml_tensor * c = ggml_add(ctx, a, b);
-//
-//       assert(c->src[0] == a);
-//       assert(c->src[1] == b);
-//   }
-//
-// The multi-dimensional tensors are stored in row-major order. The ggml_tensor struct contains fields for the
-// number of elements in each dimension ("ne") as well as the number of bytes ("nb", a.k.a. stride). This allows
-// to store tensors that are not contiguous in memory, which is useful for operations such as transposition and
-// permutation. All tensor operations have to take the stride into account and not assume that the tensor is
-// contiguous in memory.
-//
-// The data of the tensor is accessed via the "data" pointer. For example:
-//
-//   {
-//       const int nx = 2;
-//       const int ny = 3;
-//
-//       struct ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nx, ny);
-//
-//       for (int y = 0; y < ny; y++) {
-//           for (int x = 0; x < nx; x++) {
-//               *(float *) ((char *) a->data + y*a->nb[1] + x*a->nb[0]) = x + y;
-//           }
-//       }
-//
-//       ...
-//   }
-//
-// Alternatively, there are helper functions, such as ggml_get_f32_1d() and ggml_set_f32_1d() that can be used.
-//
-// ## The matrix multiplication operator (ggml_mul_mat)
-//
-// TODO
-//
-//
-// ## Multi-threading
-//
-// TODO
-//
-//
-// ## Overview of ggml.c
-//
-// TODO
-//
-//
-// ## SIMD optimizations
-//
-// TODO
-//
-//
-// ## Debugging ggml
-//
-// TODO
-//
-//
-
 #ifdef GGML_SHARED
 #    if defined(_WIN32) && !defined(__MINGW32__)
 #        ifdef GGML_BUILD
@@ -445,7 +272,6 @@ extern "C" {
     // the precision parameters are stored as ggml_tensor.op_params to the respective ops
     enum ggml_prec {
         GGML_PREC_UNDEFINED = 0,
-        GGML_PREC_DEFAULT   = 0,  // note: deprecated, use GGML_PREC_UNDEFINED
         GGML_PREC_F32       = 10,
         GGML_PREC_BF16      = 15,
         GGML_PREC_F16       = 20,
@@ -499,7 +325,6 @@ extern "C" {
         GGML_OP_DUP,
         GGML_OP_ADD,
         GGML_OP_ADD_ID,
-        GGML_OP_ADD1,
         GGML_OP_ACC,
         GGML_OP_SUB,
         GGML_OP_MUL,
@@ -541,8 +366,6 @@ extern "C" {
         GGML_OP_GET_ROWS_BACK,
         GGML_OP_SET_ROWS,
         GGML_OP_DIAG,
-        GGML_OP_DIAG_MASK_INF,
-        GGML_OP_DIAG_MASK_ZERO,
         GGML_OP_SOFT_MAX,
         GGML_OP_SOFT_MAX_BACK,
         GGML_OP_ROPE,
@@ -765,10 +588,6 @@ extern "C" {
     GGML_API size_t  ggml_type_size(enum ggml_type type);             // size in bytes for all elements in a block
     GGML_API size_t  ggml_row_size (enum ggml_type type, int64_t ne); // size in bytes for all elements in a row
 
-    GGML_DEPRECATED(
-    GGML_API double ggml_type_sizef(enum ggml_type type), // ggml_type_size()/ggml_blck_size() as float
-    "use ggml_row_size() instead");
-
     GGML_API const char * ggml_type_name(enum ggml_type type);
     GGML_API const char * ggml_op_name  (enum ggml_op   op);
     GGML_API const char * ggml_op_symbol(enum ggml_op   op);
@@ -937,18 +756,6 @@ extern "C" {
             struct ggml_tensor  * b,
             struct ggml_tensor  * ids);
 
-    GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_add1(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            struct ggml_tensor  * b),
-        "use ggml_add instead");
-
-    GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_add1_inplace(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            struct ggml_tensor  * b),
-        "use ggml_add_inplace instead");
-
     // dst = a
     // view(dst, nb1, nb2, nb3, offset) += b
     // return dst
@@ -1500,13 +1307,6 @@ extern "C" {
             struct ggml_tensor  * a,
             struct ggml_tensor  * b);
 
-    // change the precision of a matrix multiplication
-    // set to GGML_PREC_F32 for higher precision (useful for phi-2)
-    GGML_DEPRECATED(GGML_API void ggml_mul_mat_set_prec(
-            struct ggml_tensor * a,
-            enum ggml_prec       prec),
-        "use ggml_prec_set_acc() instead");
-
     // change the hint of a matrix multiplication
     GGML_API void ggml_mul_mat_set_hint(
             struct ggml_tensor * a,
@@ -1778,30 +1578,6 @@ extern "C" {
         struct ggml_context     * ctx,
         struct ggml_tensor      * a);
 
-    // set elements above the diagonal to -INF
-    GGML_API struct ggml_tensor * ggml_diag_mask_inf(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            int                   n_past);
-
-    // in-place, returns view(a)
-    GGML_API struct ggml_tensor * ggml_diag_mask_inf_inplace(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            int                   n_past);
-
-    // set elements above the diagonal to 0
-    GGML_API struct ggml_tensor * ggml_diag_mask_zero(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            int                   n_past);
-
-    // in-place, returns view(a)
-    GGML_API struct ggml_tensor * ggml_diag_mask_zero_inplace(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            int                   n_past);
-
     GGML_API struct ggml_tensor * ggml_clamp(
             struct ggml_context * ctx,
             struct ggml_tensor  * a,
@@ -2005,36 +1781,6 @@ extern "C" {
             float                 beta_fast,
             float                 beta_slow);
 
-    GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_rope_custom(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            struct ggml_tensor  * b,
-            int                   n_dims,
-            int                   mode,
-            int                   n_ctx_orig,
-            float                 freq_base,
-            float                 freq_scale,
-            float                 ext_factor,
-            float                 attn_factor,
-            float                 beta_fast,
-            float                 beta_slow),
-        "use ggml_rope_ext instead");
-
-    GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_rope_custom_inplace(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            struct ggml_tensor  * b,
-            int                   n_dims,
-            int                   mode,
-            int                   n_ctx_orig,
-            float                 freq_base,
-            float                 freq_scale,
-            float                 ext_factor,
-            float                 attn_factor,
-            float                 beta_fast,
-            float                 beta_slow),
-        "use ggml_rope_ext_inplace instead");
-
     // rotary position embedding backward, i.e compute dx from dy
     // a - dy
     GGML_API struct ggml_tensor * ggml_rope_ext_back(
@@ -2340,7 +2086,7 @@ extern "C" {
         GGML_SCALE_MODE_BILINEAR = 1,
         GGML_SCALE_MODE_BICUBIC  = 2,
 
-        GGML_SCALE_MODE_COUNT
+        GGML_SCALE_MODE_COUNT = 3
     };
 
     enum ggml_scale_flag {
@@ -2356,18 +2102,6 @@ extern "C" {
             int                   scale_factor,
             enum ggml_scale_mode  mode);
 
-    // interpolate
-    // interpolate scale to specified dimensions
-    GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_upscale_ext(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            int                   ne0,
-            int                   ne1,
-            int                   ne2,
-            int                   ne3,
-            enum ggml_scale_mode  mode),
-        "use ggml_interpolate instead");
-
     // Up- or downsamples the input to the specified size.
     // 2D scale modes (eg. bilinear) are applied to the first two dimensions.
     GGML_API struct ggml_tensor * ggml_interpolate(
@@ -2536,11 +2270,6 @@ extern "C" {
             float                 max_bias,
             float                 logit_softcap);
 
-    GGML_DEPRECATED(GGML_API void ggml_flash_attn_ext_set_prec(
-            struct ggml_tensor * a,
-            enum ggml_prec       prec),
-        "use ggml_prec_set_acc() instead");
-
     GGML_API enum ggml_prec ggml_flash_attn_ext_get_prec(
             const struct ggml_tensor * a);
 

ggml/src/ggml-alloc.c

@@ -23,11 +23,8 @@ bool ggml_op_can_inplace(enum ggml_op op) {
     switch (op) {
         case GGML_OP_FILL:
         case GGML_OP_SCALE:
-        case GGML_OP_DIAG_MASK_ZERO:
-        case GGML_OP_DIAG_MASK_INF:
         case GGML_OP_ADD:
         case GGML_OP_ADD_ID:
-        case GGML_OP_ADD1:
         case GGML_OP_SUB:
         case GGML_OP_MUL:
         case GGML_OP_DIV:

ggml/src/ggml-backend-meta.cpp

@@ -939,7 +939,6 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             case GGML_OP_ADD_ID: {
                 split_state = handle_bin_bcast(src_ss);
             } break;
-            case GGML_OP_ADD1:
             case GGML_OP_ACC: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ true);
             } break;
@@ -1020,9 +1019,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             case GGML_OP_SET_ROWS: {
                 split_state = handle_set_rows(src_ss);
             } break;
-            case GGML_OP_DIAG:
-            case GGML_OP_DIAG_MASK_INF:
-            case GGML_OP_DIAG_MASK_ZERO: {
+            case GGML_OP_DIAG: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ true);
             } break;
             case GGML_OP_SOFT_MAX:

ggml/src/ggml-cann/aclnn_ops.cpp

@@ -1347,31 +1347,6 @@ void ggml_cann_rms_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst) {
     GGML_CANN_CALL_ACLNN_OP(ctx, RmsNorm, acl_src.get(), acl_gamma.get(), eps, acl_dst.get(), acl_rstd.get());
 }
 
-// TODO: performace is low.
-void ggml_cann_diag_mask(ggml_backend_cann_context & ctx, ggml_tensor * dst, float value) {
-    ggml_tensor * src = dst->src[0];
-
-    acl_tensor_ptr acl_src = ggml_cann_create_tensor(src);
-    acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst);
-
-    const int n_past = ((int32_t *) dst->op_params)[0];
-
-    ggml_cann_pool_alloc one_tensor_allocator(ctx.pool(), ggml_nbytes(src));
-    void *               buffer = one_tensor_allocator.get();
-
-    acl_tensor_ptr mask_tensor = ggml_cann_create_tensor(buffer, ggml_cann_type_mapping(src->type),
-                                                         ggml_type_size(src->type), src->ne, src->nb, GGML_MAX_DIMS);
-
-    aclnn_fill_scalar(ctx, value, mask_tensor.get());
-
-    float          alphaValue = 1.0f;
-    acl_scalar_ptr alpha      = ggml_cann_create_scalar(&alphaValue, aclDataType::ACL_FLOAT);
-
-    GGML_CANN_CALL_ACLNN_OP(ctx, InplaceTriu, mask_tensor.get(), n_past + 1);
-    GGML_CANN_CALL_ACLNN_OP(ctx, Tril, acl_src.get(), n_past + 1, acl_dst.get());
-    GGML_CANN_CALL_ACLNN_OP(ctx, InplaceAdd, acl_dst.get(), mask_tensor.get(), alpha.get());
-}
-
 /**
  * @brief Permutes the dimensions of a tensor according to a specified order.
  *

ggml/src/ggml-cann/aclnn_ops.h

@@ -466,21 +466,6 @@ void ggml_cann_dup(ggml_backend_cann_context & ctx, ggml_tensor * dst);
  */
 void ggml_cann_rms_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst);
 
-/**
- * @brief   Applies a diagonal mask to the tensor with a specified value.
- *
- * @details This function creates a mask tensor filled with ones, then applies
- *          an upper triangular and lower triangular operation to it based on
- *          the number of past elements specified. Afterward, it adds the masked
- *          tensor to the destination tensor in-place.
- *
- * @param ctx The backend CANN context used for operations.
- * @param dst The destination tensor where the result will be stored. dst->op is
- *            `GGML_OP_DIAG_MASK`
- * @param value The value to use for masking.
- */
-void ggml_cann_diag_mask(ggml_backend_cann_context & ctx, ggml_tensor * dst, float value);
-
 /**
  * @brief   Performs an image-to-column transformation on the input tensor.
  *

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

@@ -1793,7 +1793,6 @@ static bool ggml_cann_compute_forward(ggml_backend_cann_context & ctx, struct gg
             ggml_cann_dup(ctx, dst);
             break;
         case GGML_OP_ADD:
-        case GGML_OP_ADD1:
             ggml_cann_binary_op<aclnn_add>(ctx, dst);
             break;
         case GGML_OP_SUB:
@@ -1958,9 +1957,6 @@ static bool ggml_cann_compute_forward(ggml_backend_cann_context & ctx, struct gg
         case GGML_OP_PERMUTE:
         case GGML_OP_TRANSPOSE:
             break;
-        case GGML_OP_DIAG_MASK_INF:
-            ggml_cann_diag_mask(ctx, dst, -INFINITY);
-            break;
         case GGML_OP_SOFT_MAX:
             ggml_cann_softmax(ctx, dst);
             break;
@@ -2605,6 +2601,9 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
             }
         case GGML_OP_SUM:
             return ggml_is_contiguous_rows(op->src[0]);
+        case GGML_OP_RMS_NORM:
+            // the cached gamma of ones does not track its type, so F32 and BF16 can not share it
+            return op->src[0]->type != GGML_TYPE_BF16;
         case GGML_OP_L2_NORM:
         case GGML_OP_CROSS_ENTROPY_LOSS:
         case GGML_OP_DUP:
@@ -2618,15 +2617,12 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
         case GGML_OP_TRANSPOSE:
         case GGML_OP_NORM:
         case GGML_OP_ADD:
-        case GGML_OP_ADD1:
         case GGML_OP_SUB:
         case GGML_OP_MUL:
         case GGML_OP_DIV:
-        case GGML_OP_RMS_NORM:
         case GGML_OP_SQR:
         case GGML_OP_SQRT:
         case GGML_OP_CLAMP:
-        case GGML_OP_DIAG_MASK_INF:
         case GGML_OP_SUM_ROWS:
         case GGML_OP_ARGSORT:
         case GGML_OP_ACC:

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

@@ -1782,10 +1782,6 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_add_id(params, tensor);
             } break;
-        case GGML_OP_ADD1:
-            {
-                ggml_compute_forward_add1(params, tensor);
-            } break;
         case GGML_OP_ACC:
             {
                 ggml_compute_forward_acc(params, tensor);
@@ -1926,14 +1922,6 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_diag(params, tensor);
             } break;
-        case GGML_OP_DIAG_MASK_INF:
-            {
-                ggml_compute_forward_diag_mask_inf(params, tensor);
-            } break;
-        case GGML_OP_DIAG_MASK_ZERO:
-            {
-                ggml_compute_forward_diag_mask_zero(params, tensor);
-            } break;
         case GGML_OP_SOFT_MAX:
             {
                 ggml_compute_forward_soft_max(params, tensor);
@@ -2285,7 +2273,6 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_CONT:
         case GGML_OP_ADD:
         case GGML_OP_ADD_ID:
-        case GGML_OP_ADD1:
         case GGML_OP_ACC:
         case GGML_OP_CUMSUM:
         case GGML_OP_TRI:
@@ -2407,8 +2394,6 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
             {
                 n_tasks = 1;
             } break;
-        case GGML_OP_DIAG_MASK_ZERO:
-        case GGML_OP_DIAG_MASK_INF:
         case GGML_OP_SOFT_MAX_BACK:
         case GGML_OP_ROPE:
         case GGML_OP_ROPE_BACK:
@@ -2889,7 +2874,6 @@ struct ggml_cplan ggml_graph_plan(
                     } break;
                 case GGML_OP_ADD:
                 case GGML_OP_ADD_ID:
-                case GGML_OP_ADD1:
                     {
                         if (ggml_is_quantized(node->src[0]->type) || node->src[0]->type == GGML_TYPE_F8_E4M3) {
                             cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks;

ggml/src/ggml-cpu/ops.cpp

@@ -4144,6 +4144,50 @@ static void ggml_compute_forward_rms_norm_f32(
     }
 }
 
+static void ggml_compute_forward_rms_norm_bf16(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+
+    const ggml_tensor * src0 = dst->src[0];
+
+    GGML_ASSERT(ggml_are_same_shape(src0, dst));
+    GGML_ASSERT(src0->nb[0] == sizeof(ggml_bf16_t));
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    GGML_TENSOR_UNARY_OP_LOCALS
+
+    float eps;
+    memcpy(&eps, dst->op_params, sizeof(float));
+    GGML_ASSERT(eps >= 0.0f);
+
+    for (int64_t i03 = 0; i03 < ne03; i03++) {
+        for (int64_t i02 = 0; i02 < ne02; i02++) {
+            for (int64_t i01 = ith; i01 < ne01; i01 += nth) {
+                const ggml_bf16_t * x = (ggml_bf16_t *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03);
+                ggml_bf16_t       * y = (ggml_bf16_t *) ((char *)  dst->data + i01*nb1  + i02*nb2  + i03*nb3);
+
+                ggml_float sum = 0.0;
+                for (int64_t i00 = 0; i00 < ne00; i00++) {
+                    const float v = GGML_BF16_TO_FP32(x[i00]);
+                    sum += (ggml_float)(v*v);
+                }
+
+                const float mean  = sum/ne00;
+                const float scale = 1.0f/sqrtf(mean + eps);
+
+                // if you hit this, likely you got an inf somewhere earlier
+                assert(scale > 0.0f);
+
+                for (int64_t i00 = 0; i00 < ne00; i00++) {
+                    y[i00] = GGML_FP32_TO_BF16(GGML_BF16_TO_FP32(x[i00])*scale);
+                }
+            }
+        }
+    }
+}
+
 void ggml_compute_forward_rms_norm(
         const ggml_compute_params * params,
         ggml_tensor * dst) {
@@ -4155,6 +4199,10 @@ void ggml_compute_forward_rms_norm(
             {
                 ggml_compute_forward_rms_norm_f32<GGML_RMS_NORM_FUSE_OP_NONE>(params, dst);
             } break;
+        case GGML_TYPE_BF16:
+            {
+                ggml_compute_forward_rms_norm_bf16(params, dst);
+            } break;
         default:
             {
                 GGML_ABORT("fatal error");
@@ -5704,94 +5752,6 @@ void ggml_compute_forward_diag(
     }
 }
 
-// ggml_compute_forward_diag_mask_inf
-
-static void ggml_compute_forward_diag_mask_f32(
-        const ggml_compute_params * params,
-        ggml_tensor * dst,
-        const float value) {
-
-    const ggml_tensor * src0 = dst->src[0];
-
-    const int ith = params->ith;
-    const int nth = params->nth;
-
-    const int  n_past  = ((int32_t *) dst->op_params)[0];
-    const bool inplace = src0->data == dst->data;
-
-    GGML_ASSERT(n_past >= 0);
-
-    if (!inplace) {
-        if (ith == 0) {
-            // memcpy needs to be synchronized across threads to avoid race conditions.
-            // => do it in INIT phase
-            GGML_ASSERT(ggml_nelements(dst) == ggml_nelements(src0));
-            GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src0));
-            memcpy(
-                ((char *)  dst->data),
-                ((char *) src0->data),
-                ggml_nbytes(dst));
-        }
-        ggml_barrier(params->threadpool);
-    }
-
-    // TODO: handle transposed/permuted matrices
-
-    const int n  = ggml_nrows(src0);
-    const int nc = src0->ne[0];
-    const int nr = src0->ne[1];
-    const int nz = n/nr;
-
-    GGML_ASSERT( dst->nb[0] == sizeof(float));
-    GGML_ASSERT(src0->nb[0] == sizeof(float));
-
-    for (int k = 0; k < nz; k++) {
-        for (int j = ith; j < nr; j += nth) {
-            for (int i = n_past; i < nc; i++) {
-                if (i > n_past + j) {
-                    *(float *)((char *) dst->data + k*dst->nb[2] + j*dst->nb[1] + i*dst->nb[0]) = value;
-                }
-            }
-        }
-    }
-}
-
-void ggml_compute_forward_diag_mask_inf(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-
-    const ggml_tensor * src0 = dst->src[0];
-
-    switch (src0->type) {
-        case GGML_TYPE_F32:
-            {
-                ggml_compute_forward_diag_mask_f32(params, dst, -INFINITY);
-            } break;
-        default:
-            {
-                GGML_ABORT("fatal error");
-            }
-    }
-}
-
-void ggml_compute_forward_diag_mask_zero(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-
-    const ggml_tensor * src0 = dst->src[0];
-
-    switch (src0->type) {
-        case GGML_TYPE_F32:
-            {
-                ggml_compute_forward_diag_mask_f32(params, dst, 0);
-            } break;
-        default:
-            {
-                GGML_ABORT("fatal error");
-            }
-    }
-}
-
 // ggml_compute_forward_soft_max
 
 static void ggml_compute_forward_soft_max_f32(

ggml/src/ggml-cpu/ops.h

@@ -56,8 +56,6 @@ void ggml_compute_forward_get_rows(const struct ggml_compute_params * params, st
 void ggml_compute_forward_get_rows_back(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_set_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_diag(const struct ggml_compute_params * params, struct ggml_tensor * dst);
-void ggml_compute_forward_diag_mask_inf(const struct ggml_compute_params * params, struct ggml_tensor * dst);
-void ggml_compute_forward_diag_mask_zero(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_soft_max(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_soft_max_ext_back(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_rope(const struct ggml_compute_params * params, struct ggml_tensor * dst);

ggml/src/ggml-cpu/spacemit/ime.cpp

@@ -995,7 +995,8 @@ class tensor_traits_common : public tensor_traits_base {
                         spacemit_kernels::rvv::forward_rms_norm_f32(params, op);
                         return true;
                     default:
-                        GGML_ABORT("fatal error");
+                        ggml_compute_forward_rms_norm(params, op);
+                        return true;
                 }
             case GGML_OP_ADD:
                 switch (op->src[0]->type) {

ggml/src/ggml-cuda/CMakeLists.txt

@@ -29,6 +29,7 @@ if (CUDAToolkit_FOUND)
         # 86     == RTX 3000, needs CUDA v11.1
         # 89     == RTX 4000, needs CUDA v11.8
         # 90     == Hopper H100/200, needs CUDA v11.8
+        # 103    == Blackwell GB300, needs CUDA v12.9
         # 120    == Blackwell, needs CUDA v12.8, FP4 tensor cores
         #
         # XX-virtual == compile CUDA code as PTX, do JIT compilation to binary code on first run
@@ -64,7 +65,7 @@ if (CUDAToolkit_FOUND)
                 list(APPEND CMAKE_CUDA_ARCHITECTURES 120a-real)
             endif()
             if (CUDAToolkit_VERSION VERSION_GREATER_EQUAL "12.9")
-                list(APPEND CMAKE_CUDA_ARCHITECTURES 121a-real)
+                list(APPEND CMAKE_CUDA_ARCHITECTURES 103-real 121a-real)
             endif()
         endif()
     endif()

ggml/src/ggml-cuda/diagmask.cu

@@ -1,40 +0,0 @@
-#include "diagmask.cuh"
-
-static __global__ void diag_mask_inf_f32(const float * x, float * dst, const int ncols, const int rows_per_channel, const int n_past) {
-    const int col = blockDim.y*blockIdx.y + threadIdx.y;
-    const int row = blockDim.x*blockIdx.x + threadIdx.x;
-
-    if (col >= ncols) {
-        return;
-    }
-
-    const int i = row*ncols + col;
-    //dst[i] = col > (n_past + row % rows_per_channel) ? -INFINITY : x[i];
-    //dst[i] = x[i] - (col > n_past + row % rows_per_channel) * INT_MAX; // equivalent within rounding error but slightly faster on GPU
-    dst[i] = x[i] - (col > n_past + row % rows_per_channel) * FLT_MAX;
-}
-
-static void diag_mask_inf_f32_cuda(const float * x, float * dst, const int ncols_x, const int nrows_x, const int rows_per_channel, const int n_past, cudaStream_t stream) {
-    const dim3 block_dims(1, CUDA_DIAG_MASK_INF_BLOCK_SIZE, 1);
-    const int block_num_x = (ncols_x + CUDA_DIAG_MASK_INF_BLOCK_SIZE - 1) / CUDA_DIAG_MASK_INF_BLOCK_SIZE;
-    const dim3 block_nums(nrows_x, block_num_x, 1);
-    diag_mask_inf_f32<<<block_nums, block_dims, 0, stream>>>(x, dst, ncols_x, rows_per_channel, n_past);
-}
-
-void ggml_cuda_op_diag_mask_inf(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
-    const ggml_tensor * src0 = dst->src[0];
-    const float * src0_d = (const float *)src0->data;
-    float * dst_d = (float *)dst->data;
-    cudaStream_t stream = ctx.stream();
-
-    GGML_ASSERT(src0->type == GGML_TYPE_F32);
-    GGML_ASSERT( dst->type == GGML_TYPE_F32);
-
-    const int64_t ne00 = src0->ne[0];
-    const int64_t ne01 = src0->ne[1];
-    const int nrows0 = ggml_nrows(src0);
-
-    const int n_past = ((int32_t *) dst->op_params)[0];
-
-    diag_mask_inf_f32_cuda(src0_d, dst_d, ne00, nrows0, ne01, n_past, stream);
-}

ggml/src/ggml-cuda/diagmask.cuh

@@ -1,5 +0,0 @@
-#include "common.cuh"
-
-#define CUDA_DIAG_MASK_INF_BLOCK_SIZE 32
-
-void ggml_cuda_op_diag_mask_inf(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

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

@@ -25,7 +25,6 @@
 #include "ggml-cuda/cpy.cuh"
 #include "ggml-cuda/cross-entropy-loss.cuh"
 #include "ggml-cuda/cumsum.cuh"
-#include "ggml-cuda/diagmask.cuh"
 #include "ggml-cuda/diag.cuh"
 #include "ggml-cuda/fattn.cuh"
 #include "ggml-cuda/fp8.cuh"
@@ -2742,7 +2741,6 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
             ggml_cuda_dup(ctx, dst);
             break;
         case GGML_OP_ADD:
-        case GGML_OP_ADD1: // TODO: more efficient implementation
             ggml_cuda_op_add(ctx, dst);
             break;
         case GGML_OP_ADD_ID:
@@ -2937,9 +2935,6 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_DIAG:
             ggml_cuda_op_diag(ctx, dst);
             break;
-        case GGML_OP_DIAG_MASK_INF:
-            ggml_cuda_op_diag_mask_inf(ctx, dst);
-            break;
         case GGML_OP_SOFT_MAX:
             ggml_cuda_op_soft_max(ctx, dst);
             break;
@@ -4170,22 +4165,25 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph *                cgraph,
             add = cgraph->nodes[node_idx+2];
         }
 
-        GGML_ASSERT(rms_norm->src[0]->type == GGML_TYPE_F32);
-        GGML_ASSERT(rms_norm->type == GGML_TYPE_F32);
-
-        //rms norm only supports F32
-        if (mul->src[0]->type != GGML_TYPE_F32 ||
-            mul->src[1]->type != GGML_TYPE_F32 ||
-            mul->type != GGML_TYPE_F32) {
+        // the mul operand is a norm weight, which model files store as F32;
+        // the add operand is an F32 bias or a residual with the type of x
+        const ggml_type type = rms_norm->type;
+        if ((type != GGML_TYPE_F32 && type != GGML_TYPE_BF16) || rms_norm->src[0]->type != type) {
             return false;
         }
 
-        if (add && (add->src[0]->type != GGML_TYPE_F32 ||
-            add->src[1]->type != GGML_TYPE_F32 ||
-            add->type != GGML_TYPE_F32) ) {
+        const ggml_tensor * mul_src = mul->src[0] == rms_norm ? mul->src[1] : mul->src[0];
+        if (mul_src->type != GGML_TYPE_F32 || mul->type != type) {
             return false;
         }
 
+        if (add) {
+            const ggml_tensor * add_src = add->src[0] == mul ? add->src[1] : add->src[0];
+            if ((add_src->type != GGML_TYPE_F32 && add_src->type != type) || add->type != type) {
+                return false;
+            }
+        }
+
         //if rms norm is the B operand, then we don't handle broadcast
         if (rms_norm == mul->src[1] && !ggml_are_same_shape(mul->src[0], rms_norm)) {
             return false;
@@ -4207,8 +4205,9 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph *                cgraph,
         const ggml_tensor * rms_norm = cgraph->nodes[node_idx];
         const ggml_tensor * scale    = cgraph->nodes[node_idx+1];
 
-        GGML_ASSERT(rms_norm->src[0]->type == GGML_TYPE_F32);
-        GGML_ASSERT(rms_norm->type == GGML_TYPE_F32);
+        if (rms_norm->src[0]->type != GGML_TYPE_F32 || rms_norm->type != GGML_TYPE_F32) {
+            return false;
+        }
 
         float bias;
         memcpy(&bias, (const float *) scale->op_params + 1, sizeof(float));
@@ -6529,9 +6528,11 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
         case GGML_OP_SILU_BACK:
             return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_NORM:
-        case GGML_OP_RMS_NORM:
         case GGML_OP_L2_NORM:
             return ggml_is_contiguous_rows(op->src[0]);
+        case GGML_OP_RMS_NORM:
+            return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_BF16) && op->type == op->src[0]->type &&
+                   ggml_is_contiguous_rows(op->src[0]);
         case GGML_OP_RMS_NORM_BACK:
             return ggml_is_contiguous(op->src[0]);
         case GGML_OP_NONE:
@@ -6540,7 +6541,6 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
         case GGML_OP_PERMUTE:
         case GGML_OP_TRANSPOSE:
         case GGML_OP_ADD_ID:
-        case GGML_OP_ADD1:
         case GGML_OP_SQR:
         case GGML_OP_SQRT:
         case GGML_OP_SIN:
@@ -6587,8 +6587,6 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
         }
         case GGML_OP_CONT:
             return true;
-        case GGML_OP_DIAG_MASK_INF:
-            return true;
         case GGML_OP_SOFT_MAX:
             return true;
         case GGML_OP_SOFT_MAX_BACK: {

ggml/src/ggml-cuda/mmq-config-gcn.cuh

@@ -1,31 +1,36 @@
 static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_gcn(ggml_type type, int J, bool fallback) {
     CASE(GGML_TYPE_Q1_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 512, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 512, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q1_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 512, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 512, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
 
     CASE(GGML_TYPE_Q2_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q2_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q2_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
 
     CASE(GGML_TYPE_Q4_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q4_0, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q4_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q4_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q4_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q4_0, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
@@ -37,8 +42,11 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
 
     CASE(GGML_TYPE_Q4_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q4_1, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q4_1, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 512, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q4_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q4_1, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q4_1, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
@@ -50,43 +58,45 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
     CASE(GGML_TYPE_Q5_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q5_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q5_0, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 512, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q5_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q5_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q5_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q5_0, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q5_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q5_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q5_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
 
     CASE(GGML_TYPE_Q5_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q5_1, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q5_1, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q5_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q5_1, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q5_1, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q5_1, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q5_1, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q5_1, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
 
     CASE(GGML_TYPE_Q8_0, 512, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q8_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q8_0, 512, 2,  64,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q8_0, 512, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q8_0, 512, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q8_0, 512, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q8_0, 512, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q8_0, 512, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q8_0, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q8_0, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q8_0, 512, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q8_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q8_0, 512, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q8_0, 512, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q8_0, 512, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
@@ -95,7 +105,10 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
 
     CASE(GGML_TYPE_Q2_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q2_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q2_K, 512, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q2_K, 512, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q2_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q2_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
@@ -108,18 +121,20 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
     CASE(GGML_TYPE_Q3_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q3_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q3_K, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 512, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q3_K, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q3_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q3_K, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 512, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q3_K, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
 
     CASE(GGML_TYPE_Q4_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q4_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_Q4_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q4_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q4_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q4_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
@@ -131,6 +146,7 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
 
     CASE(GGML_TYPE_Q5_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q5_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q5_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q5_K, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q5_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
@@ -143,7 +159,11 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
 
     CASE(GGML_TYPE_Q6_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q6_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q6_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  56, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q6_K, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q6_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q6_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
@@ -151,126 +171,151 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
     CASE(GGML_TYPE_Q6_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q6_K, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q6_K, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  56, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_Q6_K, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
 
 // ---------------------------------------------------------------------------------------------
 
     CASE(GGML_TYPE_IQ1_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    // CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ1_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    // CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
 
     CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    // CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    // CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
 
     CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    // CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    // CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
 
     CASE(GGML_TYPE_IQ2_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    // CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ2_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    // CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
 
     CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    // CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    // CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
 
     CASE(GGML_TYPE_IQ3_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    // CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ3_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    // CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
 
     CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
 
     CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
 
 // ---------------------------------------------------------------------------------------------
 
     CASE(GGML_TYPE_MXFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_MXFP4, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_MXFP4, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_MXFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_MXFP4, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_MXFP4, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_MXFP4, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_MXFP4, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_MXFP4, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
 
     CASE(GGML_TYPE_NVFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
-    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_NVFP4, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_NVFP4, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_NVFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
-    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 512, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_NVFP4, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_NVFP4, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_NVFP4, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);

ggml/src/ggml-cuda/norm.cu

@@ -1,4 +1,5 @@
 #include "norm.cuh"
+#include "convert.cuh"
 #include <cstdint>
 
 template <int block_size>
@@ -81,9 +82,35 @@ static __global__ void group_norm_f32(const float * x, float * dst, const int gr
     }
 }
 
-template <int block_size, int vpt, bool do_vec, bool do_multiply = false, bool do_add = false, bool do_scale = false>
-static __global__ void rms_norm_f32(const float * x,
-                                    float *       dst,
+// 4 consecutive values of x as float4, x + 4*i is aligned to 4 values
+template <typename T>
+static __device__ __forceinline__ float4 rms_norm_load4(const T * x, const int i) {
+    if constexpr (std::is_same_v<T, nv_bfloat16>) {
+        __align__(8) nv_bfloat162 v[2];
+        ggml_cuda_memcpy_1<sizeof(v)>(v, x + 4*i);
+        const float2 lo = ggml_cuda_cast<float2>(v[0]);
+        const float2 hi = ggml_cuda_cast<float2>(v[1]);
+        return make_float4(lo.x, lo.y, hi.x, hi.y);
+    } else {
+        return ((const float4 *) x)[i];
+    }
+}
+
+template <typename T>
+static __device__ __forceinline__ void rms_norm_store4(T * dst, const int i, const float4 r) {
+    if constexpr (std::is_same_v<T, nv_bfloat16>) {
+        const __align__(8) nv_bfloat16 v[4] = {
+            ggml_cuda_cast<nv_bfloat16>(r.x), ggml_cuda_cast<nv_bfloat16>(r.y), ggml_cuda_cast<nv_bfloat16>(r.z), ggml_cuda_cast<nv_bfloat16>(r.w) };
+        ggml_cuda_memcpy_1<sizeof(v)>(dst + 4*i, v);
+    } else {
+        ((float4 *) dst)[i] = r;
+    }
+}
+
+// x and dst are float or nv_bfloat16, mul is float, add is float or the type of x
+template <typename T, int block_size, int vpt, bool do_vec, bool do_multiply = false, bool do_add = false, bool do_scale = false, typename T_add = float>
+static __global__ void rms_norm_f32(const T *     x,
+                                    T *           dst,
                                     const int     ncols,
                                     const int     nchannels,
                                     const int     nsamples,
@@ -99,7 +126,7 @@ static __global__ void rms_norm_f32(const float * x,
                                     const uint3   mul_nrows_packed     = make_uint3(0, 0, 0),
                                     const uint3   mul_nchannels_packed = make_uint3(0, 0, 0),
                                     const uint3   mul_nsamples_packed  = make_uint3(0, 0, 0),
-                                    const float * add                  = nullptr,
+                                    const T_add * add                  = nullptr,
                                     const int64_t add_stride_row       = 0,
                                     const int64_t add_stride_channel   = 0,
                                     const int64_t add_stride_sample    = 0,
@@ -123,8 +150,8 @@ static __global__ void rms_norm_f32(const float * x,
     // 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 T * xc   = x   + sample*stride_sample + channel*stride_channel + row*stride_row;
+            T       * dstc = dst + ((sample*nchannels + channel)*nrows + row)*ncols;
 
             [[maybe_unused]] const float * mulc = nullptr;
             if constexpr (do_multiply) {
@@ -134,7 +161,7 @@ static __global__ void rms_norm_f32(const float * x,
                 mulc = mul + mul_sample * mul_stride_sample + mul_channel * mul_stride_channel + mul_row * mul_stride_row;
             }
 
-            [[maybe_unused]] const float * addc = nullptr;
+            [[maybe_unused]] const T_add * addc = nullptr;
             if constexpr (do_add) {
                 const int add_row     = fastmodulo(row, add_nrows_packed);
                 const int add_channel = fastmodulo(channel, add_nchannels_packed);
@@ -144,10 +171,7 @@ static __global__ void rms_norm_f32(const float * x,
 
             if constexpr (do_vec) {
                 const int      ncols4 = ncols / 4;
-                const float4 * x4     = (const float4 *) xc;
                 const float4 * mul4   = (const float4 *) mulc;
-                const float4 * add4   = (const float4 *) addc;
-                float4 *       dst4   = (float4 *) dstc;
 
                 float4 xs[vpt];
                 float  tmp = 0.0f;
@@ -156,7 +180,7 @@ static __global__ void rms_norm_f32(const float * x,
                 for (int k = 0; k < vpt; k++) {
                     const int i = tid + k*block_size;
                     if (i < ncols4) {
-                        xs[k] = x4[i];
+                        xs[k] = rms_norm_load4(xc, i);
                         tmp += xs[k].x*xs[k].x + xs[k].y*xs[k].y + xs[k].z*xs[k].z + xs[k].w*xs[k].w;
                     }
                 }
@@ -176,7 +200,7 @@ static __global__ void rms_norm_f32(const float * x,
                             r.w *= m.w;
                         }
                         if constexpr (do_add) {
-                            const float4 a = add4[i];
+                            const float4 a = rms_norm_load4(addc, i);
                             r.x += a.x;
                             r.y += a.y;
                             r.z += a.z;
@@ -188,14 +212,14 @@ static __global__ void rms_norm_f32(const float * x,
                             r.z *= scale_out;
                             r.w *= scale_out;
                         }
-                        dst4[i] = r;
+                        rms_norm_store4(dstc, i, r);
                     }
                 }
             } else {
                 float tmp = 0.0f; // partial sum for thread in warp
 
                 for (int col = tid; col < ncols; col += block_size) {
-                    const float xi = xc[col];
+                    const float xi = ggml_cuda_cast<float>(xc[col]);
                     tmp += xi * xi;
                 }
 
@@ -206,17 +230,18 @@ static __global__ void rms_norm_f32(const float * x,
                 const float scale = rsqrtf(mean + eps);
 
                 for (int col = tid; col < ncols; col += block_size) {
+                    const float xi = ggml_cuda_cast<float>(xc[col]);
                     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];
+                        dstc[col]         = ggml_cuda_cast<T>(scale * xi * mulc[mul_col] + ggml_cuda_cast<float>(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];
+                        dstc[col]         = ggml_cuda_cast<T>(scale * xi * mulc[mul_col]);
                     } else if constexpr (do_scale) {
-                        dstc[col] = scale_out * (scale * xc[col]);
+                        dstc[col] = ggml_cuda_cast<T>(scale_out * (scale * xi));
                     } else {
-                        dstc[col] = scale * xc[col];
+                        dstc[col] = ggml_cuda_cast<T>(scale * xi);
                     }
                 }
             }
@@ -387,7 +412,7 @@ static void group_norm_f32_cuda(
     do { \
         const dim3 block_dims(BLOCK_SIZE, 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<BLOCK_SIZE, VPT, VEC, do_multiply, do_add, do_scale>, launch_params, \
+        ggml_cuda_kernel_launch(rms_norm_f32<T, BLOCK_SIZE, VPT, VEC, do_multiply, do_add, do_scale, T_add>, launch_params, \
             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, \
@@ -395,10 +420,10 @@ static void group_norm_f32_cuda(
             add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, scale_out); \
     } while (0)
 
-template <bool do_multiply, bool do_add, bool do_scale>
+template <bool do_multiply, bool do_add, bool do_scale, typename T_add = float, typename T>
 static void rms_norm_f32_dispatch(
         const dim3     blocks_num,
-        const ggml_tensor * x_t,   float * dst,   const int ncols, const int nchannels, const int nsamples,
+        const ggml_tensor * x_t,   T * 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 ggml_tensor * mul_t, const int64_t mul_stride_row, const int64_t mul_stride_channel, const int64_t mul_stride_sample,
         const uint3    mul_ncols_packed, const uint3 mul_nrows_packed, const uint3 mul_nchannels_packed, const uint3 mul_nsamples_packed,
@@ -408,20 +433,20 @@ static void rms_norm_f32_dispatch(
     static_assert(!do_add || do_multiply, "fusing add is not supported without multiplying");
     static_assert(!do_scale || !do_multiply, "fusing scale is not supported with multiplying");
 
-    const float * x   = (const float *) x_t->data;
+    const T     * x   = (const T *) x_t->data;
     const float * mul = mul_t ? (const float *) mul_t->data : nullptr;
-    const float * add = add_t ? (const float *) add_t->data : nullptr;
+    const T_add * add = add_t ? (const T_add *) add_t->data : nullptr;
 
     // dst is written at row * ncols, i.e. always at a multiple of ncols, so only its base needs checking here
-    bool vec = ncols % 4 == 0 && ncols > 0 && ((uintptr_t) dst & 15) == 0;
-    vec = vec && ggml_cuda_is_aligned(x_t, 16);
+    bool vec = ncols % 4 == 0 && ncols > 0 && ((uintptr_t) dst % (4*sizeof(T))) == 0;
+    vec = vec && ggml_cuda_is_aligned(x_t, 4*sizeof(T));
     if constexpr (do_multiply) {
         vec = vec && mul_ncols_packed.z == (uint32_t) ncols
             && ggml_cuda_is_aligned(mul_t, 16);
     }
     if constexpr (do_add) {
         vec = vec && add_ncols_packed.z == (uint32_t) ncols
-            && ggml_cuda_is_aligned(add_t, 16);
+            && ggml_cuda_is_aligned(add_t, 4*sizeof(T_add));
     }
 
     const int n4 = ncols / 4;
@@ -449,9 +474,9 @@ static void rms_norm_f32_dispatch(
 }
 #undef RMS_NORM_LAUNCH
 
-template <bool do_scale = false>
+template <bool do_scale = false, typename T>
 static void rms_norm_f32_cuda(
-        const ggml_tensor * x_t, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples,
+        const ggml_tensor * x_t, T * 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, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX));
@@ -462,10 +487,11 @@ static void rms_norm_f32_cuda(
         stream, scale_out);
 }
 
+template <typename T>
 static void rms_norm_mul_f32_cuda(const ggml_tensor * x_t,
                                   const ggml_tensor * mul_t,
                                   const ggml_tensor * add_t,
-                                  float             * dst,
+                                  T                 * dst,
                                   const int      ncols,
                                   const int      nrows,
                                   const int      nchannels,
@@ -511,12 +537,32 @@ static void rms_norm_mul_f32_cuda(const ggml_tensor * x_t,
         const uint3 add_nrows_packed     = init_fastdiv_values(add_nrows);
         const uint3 add_nchannels_packed = init_fastdiv_values(add_nchannels);
         const uint3 add_nsamples_packed  = init_fastdiv_values(add_nsamples);
-        rms_norm_f32_dispatch</*do_multiply*/true, /*do_add*/true, /*do_scale*/false>(
-            blocks_num, x_t, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
-            mul_t, mul_stride_row, mul_stride_channel, mul_stride_sample,
-            mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
-            add_t, add_stride_row, add_stride_channel, add_stride_sample,
-            add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, stream);
+        // the add operand is F32 or has the type of x (e.g. a BF16 residual)
+        if (add_t->type == GGML_TYPE_F32) {
+            rms_norm_f32_dispatch</*do_multiply*/true, /*do_add*/true, /*do_scale*/false, float>(
+                blocks_num, x_t, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
+                mul_t, mul_stride_row, mul_stride_channel, mul_stride_sample,
+                mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
+                add_t, add_stride_row, add_stride_channel, add_stride_sample,
+                add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, stream);
+        } else {
+            rms_norm_f32_dispatch</*do_multiply*/true, /*do_add*/true, /*do_scale*/false, T>(
+                blocks_num, x_t, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
+                mul_t, mul_stride_row, mul_stride_channel, mul_stride_sample,
+                mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
+                add_t, add_stride_row, add_stride_channel, add_stride_sample,
+                add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, stream);
+        }
+    }
+}
+
+// calls launch with the data of dst as float * or nv_bfloat16 *, the output types of the RMS_NORM kernels
+template <typename F>
+static void rms_norm_launch_typed(ggml_tensor * dst, const F & launch) {
+    if (dst->type == GGML_TYPE_BF16) {
+        launch((nv_bfloat16 *) dst->data);
+    } else {
+        launch((float *) dst->data);
     }
 }
 
@@ -590,11 +636,10 @@ void ggml_cuda_op_group_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
 
 void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const ggml_tensor * src0 = dst->src[0];
-    float * dst_d = (float *) dst->data;
     cudaStream_t stream = ctx.stream();
 
-    GGML_ASSERT(src0->type == GGML_TYPE_F32);
-    GGML_ASSERT( dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_BF16);
+    GGML_ASSERT( dst->type == src0->type);
 
     GGML_TENSOR_UNARY_OP_LOCALS;
 
@@ -608,7 +653,9 @@ void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const int64_t s02 = nb02 / ts0;
     const int64_t s03 = nb03 / ts0;
 
-    rms_norm_f32_cuda(src0, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream);
+    rms_norm_launch_typed(dst, [&](auto * dst_d) {
+        rms_norm_f32_cuda(src0, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream);
+    });
 }
 
 void ggml_cuda_op_rms_norm_scale_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * scale_tensor) {
@@ -653,12 +700,12 @@ void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor *
         GGML_ASSERT(false);
     }
 
-    float * dst_d = (float *) mul_tensor->data;
     cudaStream_t stream = ctx.stream();
 
-    GGML_ASSERT(rms_norm_src->type == GGML_TYPE_F32);
-    GGML_ASSERT(dst->type == GGML_TYPE_F32);
-    GGML_ASSERT(mul_tensor->type == GGML_TYPE_F32);
+    GGML_ASSERT(rms_norm_src->type == GGML_TYPE_F32 || rms_norm_src->type == GGML_TYPE_BF16);
+    GGML_ASSERT(dst->type == rms_norm_src->type);
+    GGML_ASSERT(mul_tensor->type == rms_norm_src->type);
+    GGML_ASSERT(mul_src->type == GGML_TYPE_F32);
     GGML_ASSERT(eps >= 0.0f);
 
     const int64_t ne00 = rms_norm_src->ne[0];
@@ -683,14 +730,16 @@ void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor *
     const int mul_nchannels = mul_src->ne[2];
     const int mul_nsamples  = mul_src->ne[3];
 
-    rms_norm_mul_f32_cuda(rms_norm_src, mul_src, nullptr, dst_d,
-                          ne00, ne01, ne02, ne03,
-                          /*s00*/ s01, s02, s03,
-                          /*mul_s00*/ mul_s01, mul_s02, mul_s03,
-                          mul_ncols, mul_nrows, mul_nchannels, mul_nsamples,
-                          /*add_s00*/ 0, 0, 0,
-                          0, 0, 0, 0,
-                          eps, stream);
+    rms_norm_launch_typed(mul_tensor, [&](auto * dst_d) {
+        rms_norm_mul_f32_cuda(rms_norm_src, mul_src, nullptr, dst_d,
+                              ne00, ne01, ne02, ne03,
+                              /*s00*/ s01, s02, s03,
+                              /*mul_s00*/ mul_s01, mul_s02, mul_s03,
+                              mul_ncols, mul_nrows, mul_nchannels, mul_nsamples,
+                              /*add_s00*/ 0, 0, 0,
+                              0, 0, 0, 0,
+                              eps, stream);
+    });
 }
 
 void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx,
@@ -722,13 +771,14 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx,
         GGML_ASSERT(false);
     }
 
-    float *      dst_d  = (float *) add_tensor->data;
     cudaStream_t stream = ctx.stream();
 
-    GGML_ASSERT(rms_norm_src->type == GGML_TYPE_F32);
-    GGML_ASSERT(dst->type == GGML_TYPE_F32);
-    GGML_ASSERT(mul_tensor->type == GGML_TYPE_F32);
-    GGML_ASSERT(add_tensor->type == GGML_TYPE_F32);
+    GGML_ASSERT(rms_norm_src->type == GGML_TYPE_F32 || rms_norm_src->type == GGML_TYPE_BF16);
+    GGML_ASSERT(dst->type == rms_norm_src->type);
+    GGML_ASSERT(mul_tensor->type == rms_norm_src->type);
+    GGML_ASSERT(add_tensor->type == rms_norm_src->type);
+    GGML_ASSERT(mul_src->type == GGML_TYPE_F32);
+    GGML_ASSERT(add_src->type == GGML_TYPE_F32 || add_src->type == rms_norm_src->type);
     GGML_ASSERT(eps >= 0.0f);
 
     const int64_t ne00 = rms_norm_src->ne[0];
@@ -764,14 +814,16 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx,
     const int add_nchannels = add_src->ne[2];
     const int add_nsamples  = add_src->ne[3];
 
-    rms_norm_mul_f32_cuda(rms_norm_src, mul_src,add_src,dst_d,
-                          ne00,ne01, ne02, ne03,
-                          /*s00*/ s01, s02, s03,
-                          /*mul_s00*/ mul_s01, mul_s02, mul_s03,
-                          mul_ncols, mul_nrows, mul_nchannels, mul_nsamples,
-                          /*add_s00*/ add_s01, add_s02, add_s03,
-                          add_ncols, add_nrows, add_nchannels, add_nsamples,
-                          eps, stream);
+    rms_norm_launch_typed(add_tensor, [&](auto * dst_d) {
+        rms_norm_mul_f32_cuda(rms_norm_src, mul_src, add_src, dst_d,
+                              ne00, ne01, ne02, ne03,
+                              /*s00*/ s01, s02, s03,
+                              /*mul_s00*/ mul_s01, mul_s02, mul_s03,
+                              mul_ncols, mul_nrows, mul_nchannels, mul_nsamples,
+                              /*add_s00*/ add_s01, add_s02, add_s03,
+                              add_ncols, add_nrows, add_nchannels, add_nsamples,
+                              eps, stream);
+    });
 }
 
 void ggml_cuda_op_rms_norm_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {

ggml/src/ggml-cuda/pad_reflect_1d.cu

@@ -18,11 +18,12 @@ static __global__ __launch_bounds__(CUDA_PAD_REFLECT_1D_BLOCK_SIZE, 1) void
         const int64_t             nb2,
         const int64_t             nb3,
         const int                 p0,
-        const int                 p1) {
+        const int                 p1,
+        const uint32_t            block_offset) {
     const int64_t i3 = blockIdx.z;
     const int64_t i2 = blockIdx.y;
 
-    const uint2   div_mod_packed = fast_div_modulo(blockIdx.x, ne01);
+    const uint2   div_mod_packed = fast_div_modulo(blockIdx.x + block_offset, ne01);
     const int64_t tile1          = div_mod_packed.y;  // i1
     const int64_t tile0          = div_mod_packed.x;  // nth i0 tile
     const int64_t i1             = tile1;
@@ -77,15 +78,21 @@ void ggml_cuda_op_pad_reflect_1d(ggml_backend_cuda_context & ctx, ggml_tensor *
     // sanity: padded length matches
     GGML_ASSERT(ne0 == ne00 + p0 + p1);
 
-    constexpr int64_t bx     = CUDA_PAD_REFLECT_1D_BLOCK_SIZE;  // threads per block (x)
-    const int64_t     tiles0 = (ne0 + bx - 1) / bx;             // number of tiles along i0
-    // grid.x covers i1 and all tiles of i0: [ne01 * tiles0]
-    // grid.y covers i2: [ne02]
-    // grid.z covers i3: [ne03]
-    const dim3        grid_dims((unsigned) (ne01 * tiles0), (unsigned) ne02, (unsigned) ne03);
+    constexpr int64_t bx         = CUDA_PAD_REFLECT_1D_BLOCK_SIZE;  // threads per block (x)
+    constexpr int64_t max_blocks = UINT32_MAX / bx;                 // HIP refuses more than 2^32 threads per grid dimension
+    const int64_t     tiles0     = (ne0 + bx - 1) / bx;             // number of tiles along i0
+    const int64_t     nblocks    = ne01 * tiles0;                   // blocks covering i1 and all tiles of i0
+    GGML_ASSERT(nblocks <= UINT32_MAX);
     const dim3        block_dims((unsigned) bx, 1, 1);
 
-    pad_reflect_1d_kernel_f32<<<grid_dims, block_dims, 0, stream>>>(
-        src0->data, dst->data, ne0, ne00, ne01_packed, ne02, ne03, src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3],
-        dst->nb[0], dst->nb[1], dst->nb[2], dst->nb[3], p0, p1);
+    // grid.x covers a chunk of at most max_blocks of the nblocks
+    // grid.y covers i2: [ne02]
+    // grid.z covers i3: [ne03]
+    for (int64_t offset = 0; offset < nblocks; offset += max_blocks) {
+        const dim3 grid_dims((unsigned) std::min(nblocks - offset, max_blocks), (unsigned) ne02, (unsigned) ne03);
+
+        pad_reflect_1d_kernel_f32<<<grid_dims, block_dims, 0, stream>>>(
+            src0->data, dst->data, ne0, ne00, ne01_packed, ne02, ne03, src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3],
+            dst->nb[0], dst->nb[1], dst->nb[2], dst->nb[3], p0, p1, (uint32_t) offset);
+    }
 }

ggml/src/ggml-hexagon/CMakeLists.txt

@@ -33,10 +33,11 @@ add_library(htp_iface OBJECT
 set_target_properties(htp_iface PROPERTIES POSITION_INDEPENDENT_CODE ON)
 target_include_directories(htp_iface PUBLIC
     ${HEXAGON_SDK_ROOT}/incs
-    ${HEXAGON_SDK_ROOT}/incs/stddef
     ${HEXAGON_SDK_ROOT}/utils/examples
     ${CMAKE_CURRENT_SOURCE_DIR}/htp
     ${CMAKE_CURRENT_BINARY_DIR})
+target_include_directories(htp_iface PRIVATE
+    ${HEXAGON_SDK_ROOT}/incs/stddef)
 
 build_idl(htp/htp_iface.idl htp_iface)
 
@@ -65,8 +66,14 @@ ggml_add_backend_library(${TARGET_NAME}
     libdl.h
     ../../include/ggml-hexagon.h)
 
+if (WIN32)
+    # some clang-cl drivers do not predefine ARM64 macros for winnt.h
+    target_compile_definitions(${TARGET_NAME} PRIVATE _M_ARM64=1 _ARM64_=1)
+endif()
+
 target_link_libraries(${TARGET_NAME} PRIVATE htp_iface)
 target_include_directories(${TARGET_NAME} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/htp ${CMAKE_CURRENT_BINARY_DIR})
+target_include_directories(${TARGET_NAME} SYSTEM PRIVATE ${HEXAGON_SDK_ROOT}/incs/stddef)
 
 # Build HTP skels
 set(HTP_SKELS)

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

@@ -316,6 +316,7 @@ struct ggml_hexagon_session {
     ggml_hexagon_opbatch* op_batch;
     ggml_hexagon_opqueue* op_queue;
 
+    std::unordered_set<ggml_hexagon_shared_buffer *>                     buffers;
     std::unordered_map<int, std::unique_ptr<ggml_hexagon_shared_buffer>> cloned_buffers;
     std::unordered_set<ggml_hexagon_session *>                           virt_peers;
     std::unordered_set<ggml_hexagon_session *>                           phys_peers;
@@ -337,6 +338,7 @@ struct ggml_hexagon_session {
     ggml_hexagon_mdev_group                mdev;
     ggml_backend_dev_t                     dev       = nullptr;
     ggml_backend_hexagon_device_context *  dev_ctx   = nullptr;
+    ggml_hexagon_shared_buffer *           shm_buf   = nullptr;
     ggml_hexagon_fence_buffer *            fence_buf = nullptr;
 
     ggml_hexagon_session(const ggml_hexagon_device_config & config, ggml_backend_dev_t dev = nullptr, uint32_t mdev_idx = 0, uint32_t mdev_count = 0) noexcept(false);
@@ -373,8 +375,8 @@ struct ggml_hexagon_session {
     void flush_peers();
     void flush_pending(bool all = true);
 
-    ggml_hexagon_shared_buffer * mmap_tensor(const ggml_tensor * t);
-    bool clone_buffer(const ggml_hexagon_shared_buffer*);
+    void mmap_buffers();
+    void clone_buffer(const ggml_hexagon_shared_buffer*);
     void release_buffer(const ggml_hexagon_shared_buffer*);
     void unclone_buffer(const ggml_hexagon_shared_buffer*);
 
@@ -434,7 +436,7 @@ struct ggml_hexagon_rpcmem_block {
     int       fd   = -1;
     size_t    size = 0;
 
-    std::unordered_set<ggml_hexagon_session *> mapped_clones;
+    std::unordered_set<ggml_hexagon_session *> clones;
 
     ggml_hexagon_rpcmem_block(size_t size) {
         base = (uint8_t *) rpcmem_alloc2(RPCMEM_HEAP_ID_SYSTEM, RPCMEM_DEFAULT_FLAGS, size);
@@ -468,16 +470,24 @@ struct ggml_hexagon_shared_buffer {
     size_t       size()   const { return mem ? mem->size : 0;  }
     int          fd()     const { return mem ? mem->fd   : -1; }
 
-    void mmap(bool extended = false) {
+    void mmap() {
         if (!this->mem)   return;
         if (this->mapped) return;
 
-        GGML_ASSERT(!this->pinned || !extended);
+        GGML_ASSERT(!this->pinned || !this->extended);
 
-        this->extended = extended;
-        fastrpc_map_flags flags = this->pinned ? FASTRPC_MAP_FD : (extended ? FASTRPC_MAP_FD_DELAYED_EXTENDED : FASTRPC_MAP_FD_DELAYED);
+        fastrpc_map_flags flags = this->pinned ? FASTRPC_MAP_FD : (this->extended ? FASTRPC_MAP_FD_DELAYED_EXTENDED : FASTRPC_MAP_FD_DELAYED);
 
         int err = fastrpc_mmap(sess->domain_id, fd(), (void *) base(), 0, size(), flags);
+        if (err != 0) {
+            if (this->extended) {
+                flags = FASTRPC_MAP_FD_EXTENDED;
+            }
+            HEX_VERBOSE("ggml-hex: %s mmap failed (err 0x%08x), retrying in 100ms...\n", sess->c_name(), (unsigned) err);
+            std::this_thread::sleep_for(std::chrono::milliseconds(100));
+            err = fastrpc_mmap(sess->domain_id, fd(), (void *) base(), 0, size(), flags);
+        }
+
         if (err != 0) {
             GGML_LOG_ERROR("ggml-hex: %s buffer mapping failed : domain_id %d size %zu fd %d error 0x%08x\n", sess->c_name(),
                     sess->domain_id, size(), fd(), (unsigned) err);
@@ -550,10 +560,10 @@ struct ggml_hexagon_shared_buffer {
 
     // Clone constructor for cross-session mapping
     ggml_hexagon_shared_buffer(ggml_hexagon_session * sess, const ggml_hexagon_shared_buffer & other) {
-        this->sess   = sess;
-        this->mem    = other.mem;
-        this->mapped = false;
-        this->pinned = other.pinned;
+        this->sess     = sess;
+        this->mem      = other.mem;
+        this->mapped   = false;
+        this->pinned   = false;
         this->extended = other.extended;
     }
 
@@ -563,19 +573,30 @@ struct ggml_hexagon_shared_buffer {
 };
 
 struct ggml_hexagon_fence_buffer : public ggml_hexagon_shared_buffer {
+    uint8_t *             data       = nullptr;
     uint32_t              slot_count = 0;
     uint32_t              slot_head  = 0;
     std::vector<uint32_t> free_slots;
     ggml_backend_buffer   backend_buffer{};
 
-    ggml_hexagon_fence_buffer(ggml_hexagon_session * sess, ggml_backend_buffer_type_t buft, size_t size)
-        : ggml_hexagon_shared_buffer(sess, size, false /* pinned */),
-          slot_count(size / HTP_FENCE_SLOT_SIZE),
+    static size_t size_of(size_t n_slots) {
+        return n_slots * HTP_FENCE_SLOT_SIZE;
+    }
+
+    ggml_hexagon_fence_buffer(ggml_hexagon_shared_buffer * shm_buf, ggml_backend_buffer_type_t buft, size_t offset, size_t n_slots)
+        : ggml_hexagon_shared_buffer(shm_buf->sess, *shm_buf),
+          data(shm_buf->base() + offset),
+          slot_count(n_slots),
           slot_head(0) {
+        GGML_ASSERT((uintptr_t) data % HEX_L2_LINE_SIZE == 0);
+        this->mapped           = shm_buf->mapped;
         backend_buffer.buft    = buft;
         backend_buffer.context = static_cast<ggml_hexagon_shared_buffer *>(this);
-        backend_buffer.size    = size;
-        mmap(false);
+        backend_buffer.size    = size_of(n_slots);
+    }
+
+    ~ggml_hexagon_fence_buffer() {
+        this->mapped = false;
     }
 
     uint8_t * alloc_slot(uint32_t n_slots = 1) {
@@ -583,11 +604,11 @@ struct ggml_hexagon_fence_buffer : public ggml_hexagon_shared_buffer {
         if (n_slots == 1 && !free_slots.empty()) {
             uint32_t slot = free_slots.back();
             free_slots.pop_back();
-            ptr = base() + (size_t) slot * HTP_FENCE_SLOT_SIZE;
+            ptr = data + (size_t) slot * HTP_FENCE_SLOT_SIZE;
         } else if (slot_head + n_slots <= slot_count) {
             uint32_t slot = slot_head;
             slot_head += n_slots;
-            ptr = base() + (size_t) slot * HTP_FENCE_SLOT_SIZE;
+            ptr = data + (size_t) slot * HTP_FENCE_SLOT_SIZE;
         }
         if (ptr) {
             memset(ptr, 0, (size_t) n_slots * HTP_FENCE_SLOT_SIZE);
@@ -597,7 +618,7 @@ struct ggml_hexagon_fence_buffer : public ggml_hexagon_shared_buffer {
 
     void free_slot(void * ptr, uint32_t n_slots = 1) {
         if (!ptr) return;
-        uint32_t slot = ((uint8_t *) ptr - base()) / HTP_FENCE_SLOT_SIZE;
+        uint32_t slot = ((uint8_t *) ptr - data) / HTP_FENCE_SLOT_SIZE;
         for (uint32_t i = 0; i < n_slots; i++) {
             free_slots.push_back(slot + i);
         }
@@ -636,6 +657,7 @@ inline ggml_tensor ggml_hexagon_session::make_fence_tensor(const volatile void *
 
 static void ggml_backend_hexagon_buffer_free_buffer(ggml_backend_buffer_t buffer) {
     auto sbuf = static_cast<ggml_hexagon_shared_buffer *>(buffer->context);
+    sbuf->sess->buffers.erase(sbuf);
     sbuf->sess->unclone_buffer(sbuf);
     delete sbuf;
 }
@@ -975,31 +997,44 @@ 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<ggml_backend_hexagon_alloc_buffer_n_plan_item>;
-
-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;
 };
 
+struct ggml_backend_hexagon_alloc_buffer_n_plan_item {
+    size_t size;
+    std::vector<ggml_backend_hexagon_alloc_unit> units;
+};
+
+using ggml_backend_hexagon_alloc_buffer_n_plan_t = std::vector<ggml_backend_hexagon_alloc_buffer_n_plan_item>;
+
+static bool ggml_hexagon_tensor_layer(const struct ggml_tensor * t, int & kind, int & layer) {
+    if (strncmp(t->name, "blk.", 4) == 0 && isdigit((unsigned char) t->name[4])) {
+        char * end = NULL;
+        long l = strtol(t->name + 4, &end, 10);
+        if (end && *end == '.') {
+            kind  = 0;
+            layer = (int) l;
+            return true;
+        }
+    }
+    if (strncmp(t->name, "cache_", 6) == 0) {
+        const char * p = strstr(t->name, "_l");
+        if (p && isdigit((unsigned char) p[2])) {
+            char * end = NULL;
+            long l = strtol(p + 2, &end, 10);
+            if (end && *end == '\0') {
+                kind  = 1;
+                layer = (int) l;
+                return true;
+            }
+        }
+    }
+    return false;
+}
+
 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;
@@ -1012,61 +1047,72 @@ static ggml_backend_hexagon_alloc_buffer_n_plan_t ggml_backend_hexagon_alloc_buf
     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);
+        if (t->data != NULL || t->view_src != NULL) {
+            i++;
+            continue;
         }
 
-        const char * layer_suffix = ggml_hexagon_kv_layer_suffix(t);
+        size_t unit_size = GGML_PAD(ggml_backend_buft_get_alloc_size(buft, t), alignment);
+        int unit_first = i;
+        int unit_last  = i + 1;
 
-        while (unit_last < n_tensors) {
+        int kind = -1, layer = -1;
+        const bool has_layer = ggml_hexagon_tensor_layer(t, kind, layer);
+
+        while (has_layer && unit_last < n_tensors) {
             struct ggml_tensor * next = tensors[unit_last];
-
-            if (next->view_src != NULL) {
+            if (next->data != NULL || 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;
-                }
+            int next_kind = -1, next_layer = -1;
+            if (!ggml_hexagon_tensor_layer(next, next_kind, next_layer) ||
+                next_kind != kind || next_layer != layer) {
+                break;
             }
 
-            break;
+            size_t next_size = GGML_PAD(ggml_backend_buft_get_alloc_size(buft, next), alignment);
+            if (unit_size + next_size > max_size) {
+                break;
+            }
+
+            unit_size += next_size;
+            unit_last++;
         }
 
         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;
+        bool placed = false;
+        for (auto & item : plan) {
+            if (item.size + unit.size <= max_size) {
+                item.size += unit.size;
+                item.units.push_back(unit);
+                placed = true;
+                break;
+            }
+        }
+        if (!placed) {
+            plan.push_back({ unit.size, { unit } });
         }
-
-        cur_buf_size += unit.size;
     }
 
-    if (cur_buf_size > 0) {
-        plan.push_back({ cur_buf_size, cur_buf_first, n_tensors });
+    std::stable_sort(plan.begin(), plan.end(), [](const auto & a, const auto & b) {
+        return a.size > b.size;
+    });
+
+    if (!plan.empty()) {
+        for (const auto & unit : units) {
+            if (unit.size == 0) {
+                plan[0].units.push_back(unit);
+            }
+        }
     }
 
     return plan;
@@ -1092,28 +1138,19 @@ static ggml_backend_buffer_t ggml_backend_hexagon_buffer_type_alloc_buffer_n(
         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) {
+        for (const auto & unit : item.units) {
+            for (int j = unit.first; j < unit.last; j++) {
+                struct ggml_tensor * t = tensors[j];
+                if (t->data == NULL && 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) {
+                break;
+            }
         }
         if (t_failed != NULL) {
             GGML_LOG_ERROR("%s: failed to initialize tensor %s\n", __func__, t_failed->name);
@@ -1131,6 +1168,19 @@ static ggml_backend_buffer_t ggml_backend_hexagon_buffer_type_alloc_buffer_n(
         return NULL;
     }
 
+    for (int j = 0; j < n_tensors; j++) {
+        struct ggml_tensor * t = tensors[j];
+        if (t->view_src != NULL && t->buffer == NULL) {
+            if (ggml_backend_view_init(t) != GGML_STATUS_SUCCESS) {
+                GGML_LOG_ERROR("%s: failed to initialize tensor %s\n", __func__, t->name);
+                for (ggml_backend_buffer_t b : buffers) {
+                    ggml_backend_buffer_free(b);
+                }
+                return NULL;
+            }
+        }
+    }
+
     if (buffers.size() == 1) {
         return buffers[0];
     }
@@ -1200,7 +1250,6 @@ struct ggml_hexagon_opbatch {
     ggml_hexagon_session*            sess;
 
     std::vector<htp_opnode>          ops;       // htp_opnode of ops
-
     std::vector<htp_buf_desc>        h_bufs;    // htp buffer descriptors
     std::vector<htp_tensor>          h_tens;    // htp tensor descriptors
     std::vector<htp_op_desc>         h_ops;     // htp op descriptors
@@ -1529,45 +1578,45 @@ struct ggml_hexagon_registry {
 struct ggml_hexagon_opqueue {
     // Shared buffer for storing batches
     ggml_hexagon_shared_buffer *shm_buf;
-    size_t                      shm_blk_size;
+    uint8_t *                   data;
+    size_t                      slot_size;
     size_t                      depth;
 
     using opvec = std::vector<htp_opnode>;
 
     std::vector<opvec>          op_cache;       // per batch op cache
 
-    ggml_hexagon_opqueue(ggml_hexagon_session *sess, size_t batch_size, size_t depth) : depth(depth) {
+    static size_t size_of(size_t batch_size, size_t n_threads, size_t depth = 1) {
         size_t n_bufs    = HTP_OP_MAX_BUFS;
         size_t n_ops     = batch_size;
         size_t n_tensors = std::min<size_t>(n_ops * 4 + HTP_OP_MAX_INPUTS + HTP_OP_MAX_OUTPUTS, HTP_OP_MAX_TENSORS);
 
         size_t tr_size = 0;
         if (opt_profile == 3) {
-            tr_size = sess->n_threads * opt_optrace * sizeof(htp_trace_desc);
+            tr_size = n_threads * opt_optrace * sizeof(htp_trace_desc);
         }
 
-        shm_blk_size = sizeof(htp_buf_desc)  * n_bufs    +
-                       sizeof(htp_tensor)    * n_tensors +
-                       sizeof(htp_op_desc)   * n_ops     +
-                       sizeof(htp_prof_desc) * n_ops     +
-                       tr_size;
-
-        shm_buf = new ggml_hexagon_shared_buffer(sess, shm_blk_size * depth, true /* pinned */);
+        size_t slot_size = sizeof(htp_buf_desc)  * n_bufs    +
+                           sizeof(htp_tensor)    * n_tensors +
+                           sizeof(htp_op_desc)   * n_ops     +
+                           sizeof(htp_prof_desc) * n_ops     +
+                           tr_size;
+        return slot_size * depth;
+    }
 
+    ggml_hexagon_opqueue(ggml_hexagon_shared_buffer *shm_buf, size_t offset, size_t batch_size, size_t depth)
+        : shm_buf(shm_buf),
+          data(shm_buf->base() + offset),
+          slot_size(size_of(batch_size, shm_buf->sess->n_threads)),
+          depth(depth) {
         op_cache.resize(depth);
 
         if (opt_verbose) {
-            GGML_LOG_INFO("ggml-hex: %s allocated opqueue : batch-size %zu depth %zu shm-size %zu shm-block-size %zu\n",
-                    sess->c_name(), batch_size, depth, shm_buf->size(), shm_blk_size);
+            GGML_LOG_INFO("ggml-hex: %s allocated opqueue : batch-size %zu depth %zu shm-size %zu slot-size %zu\n",
+                    shm_buf->c_name(), batch_size, depth, size_of(batch_size, shm_buf->sess->n_threads, depth), slot_size);
         }
     }
 
-    ~ggml_hexagon_opqueue() {
-        delete shm_buf;
-    }
-
-    size_t shm_size() const { return shm_buf ? shm_buf->size() : 0; }
-
     // push new batch
     bool push(htp_opbatch_req& req, dspqueue_buffer& dbuf, const ggml_hexagon_opbatch* op_batch, uint64_t seq) {
         static_assert(sizeof(htp_opbatch_req) % 8 == 0, "sizeof(htp_opbatch_req) must be multiple of 8");
@@ -1603,13 +1652,13 @@ struct ggml_hexagon_opqueue {
             req.n_traces = 0;
         }
 
-        dbuf.ptr      = shm_buf->base() + ((size_t) slot * shm_blk_size);
+        dbuf.ptr      = data + ((size_t) slot * slot_size);
         dbuf.fd       = shm_buf->fd();
         dbuf.flags    = DSPQUEUE_BUFFER_FLAG_FLUSH_SENDER | DSPQUEUE_BUFFER_FLAG_INVALIDATE_RECIPIENT;
         dbuf.offset   = (uint8_t*) dbuf.ptr - (uint8_t*) shm_buf->base();
         dbuf.size     = b_size + t_size + o_size + p_size + tr_size;
 
-        GGML_ASSERT(dbuf.size <= shm_blk_size);
+        GGML_ASSERT(dbuf.size <= slot_size);
 
         uint8_t * m_ptr = (uint8_t*) dbuf.ptr;
         uint8_t * b_ptr = m_ptr; m_ptr += b_size;
@@ -1657,7 +1706,7 @@ struct ggml_hexagon_opqueue {
         }
 
         const size_t m_size = b_size + t_size + o_size + p_size + tr_size;
-        GGML_ASSERT(m_size <= shm_blk_size);
+        GGML_ASSERT(m_size <= slot_size);
 
         HEX_VERBOSE("ggml-hex: %s opqueue-pop batch #%llu : n-bufs %u n-tensors %u n-ops %u : m-size %zu b-size %zu t-size %zu o-size %zu\n",
                 shm_buf->sess->c_name(), (unsigned long long) rsp.seq, rsp.n_bufs, rsp.n_tensors, rsp.n_ops,
@@ -1775,6 +1824,7 @@ void ggml_hexagon_session::flush_batch(size_t min_ops) {
     if (op_batch->n_ops < min_ops) { return; }
 
     op_batch->sort_buffers();
+    this->mmap_buffers();
 
     htp_opbatch_req req {};
     dspqueue_buffer dbuf{};
@@ -1819,10 +1869,17 @@ void ggml_hexagon_session::flush_batch(size_t min_ops) {
 }
 
 void ggml_hexagon_session::enqueue_op(const htp_opnode & node) {
-    auto clone_tensor_buffer = [this](const ggml_tensor * t) {
-        auto sbuf = this->mmap_tensor(t);
-        if (!sbuf) return;
-        if (sbuf->sess != this) {
+    auto register_tensor = [&](const ggml_tensor * t) {
+        if (!t) return;
+        auto * sbuf = static_cast<ggml_hexagon_shared_buffer *>(t->buffer->context);
+        if (ggml_backend_buffer_get_usage(t->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
+            sbuf->extended = opt_dma64;
+        }
+        if (sbuf->sess == this) {
+            if (!sbuf->mapped) {
+                this->buffers.insert(sbuf);
+            }
+        } else {
             this->clone_buffer(sbuf);
         }
         for (auto & sub : this->mdev.sessions) {
@@ -1830,11 +1887,11 @@ void ggml_hexagon_session::enqueue_op(const htp_opnode & node) {
         }
     };
 
-    for (auto t : node.get_inputs()) {
-        clone_tensor_buffer(t);
+    for (const auto * src : node.get_inputs()) {
+        register_tensor(src);
     }
-    for (auto t : node.get_outputs()) {
-        clone_tensor_buffer(t);
+    for (const auto * dst : node.get_outputs()) {
+        register_tensor(dst);
     }
 
     const bool fits_tensors = op_batch->fit_tensors(node);
@@ -1945,46 +2002,61 @@ void ggml_hexagon_session::enqueue_allreduce(
     this->enqueue_op(ar_node);
 }
 
-ggml_hexagon_shared_buffer * ggml_hexagon_session::mmap_tensor(const ggml_tensor * t) {
-    if (!t) return nullptr;
+void ggml_hexagon_session::mmap_buffers() {
+    auto collect_session = [](ggml_hexagon_session * sess, std::unordered_set<ggml_hexagon_shared_buffer *> & unmapped) {
+        for (auto * sbuf : sess->buffers) {
+            if (!sbuf->mapped) {
+                unmapped.insert(sbuf);
+            }
+        }
+        for (auto & [fd, clone] : sess->cloned_buffers) {
+            if (!clone->mapped) {
+                unmapped.insert(clone.get());
+            }
+        }
+    };
 
-    auto sbuf = static_cast<ggml_hexagon_shared_buffer *>(t->buffer->context);
-    if (!sbuf->mapped) {
-        const bool is_weight = ggml_backend_buffer_get_usage(t->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS;
-        const bool extended  = opt_dma64 && is_weight;
-        sbuf->mmap(extended);
+    std::unordered_set<ggml_hexagon_shared_buffer *> unmapped;
+    collect_session(this, unmapped);
+    for (auto & sub : this->mdev.sessions) {
+        collect_session(sub.get(), unmapped);
+    }
+
+    if (!unmapped.empty()) {
+        std::vector<ggml_hexagon_shared_buffer *> sorted(unmapped.begin(), unmapped.end());
+        std::sort(sorted.begin(), sorted.end(), [](const auto * a, const auto * b) {
+            if (a->size() != b->size()) {
+                return a->size() > b->size();
+            }
+            if (a->fd() != b->fd()) {
+                return a->fd() < b->fd();
+            }
+            return a->sess->domain_id < b->sess->domain_id;
+        });
+        for (auto * sbuf : sorted) {
+            sbuf->mmap();
+        }
     }
-    return sbuf;
 }
 
-bool ggml_hexagon_session::clone_buffer(const ggml_hexagon_shared_buffer * sbuf)
+void ggml_hexagon_session::clone_buffer(const ggml_hexagon_shared_buffer * sbuf)
 {
     GGML_ASSERT(sbuf && sbuf->mem);
-    if (sbuf->sess == this) return true;
+    if (sbuf->sess == this) return;
 
     auto mem = sbuf->mem;
     int   fd = mem->fd;
 
     GGML_ASSERT(fd >= 0);
 
-    if (this->cloned_buffers.find(fd) != this->cloned_buffers.end()) return true;
-
-    GGML_ASSERT(sbuf->mapped);
+    if (this->cloned_buffers.find(fd) != this->cloned_buffers.end()) return;
 
     HEX_VERBOSE("ggml-hex: %s clone-buffer: %s base %p size %zu fd %d\n", this->name.c_str(),
                 sbuf->c_name(), sbuf->base(), sbuf->size(), fd);
 
     auto clone = std::make_unique<ggml_hexagon_shared_buffer>(this, *sbuf);
-    try {
-        clone->mmap(sbuf->extended);
-    } catch (const std::exception & exc) {
-        GGML_LOG_ERROR("ggml-hex: %s lazy mapping of buffer context failed: %s\n", this->c_name(), exc.what());
-        return false;
-    }
-
     this->cloned_buffers[fd] = std::move(clone);
-    mem->mapped_clones.insert(this);
-    return true;
+    mem->clones.insert(this);
 }
 
 void ggml_hexagon_session::release_buffer(const ggml_hexagon_shared_buffer * sbuf) {
@@ -2000,14 +2072,14 @@ void ggml_hexagon_session::release_buffer(const ggml_hexagon_shared_buffer * sbu
         auto clone = std::move(it->second);
         this->cloned_buffers.erase(it);
     }
-    mem->mapped_clones.erase(this);
+    mem->clones.erase(this);
 }
 
 void ggml_hexagon_session::unclone_buffer(const ggml_hexagon_shared_buffer * sbuf) {
     GGML_ASSERT(sbuf && sbuf->mem);
 
     auto mem = sbuf->mem;
-    std::vector<ggml_hexagon_session *> sessions(mem->mapped_clones.begin(), mem->mapped_clones.end());
+    std::vector<ggml_hexagon_session *> sessions(mem->clones.begin(), mem->clones.end());
 
     for (auto * sess : sessions) {
         sess->release_buffer(sbuf);
@@ -2276,9 +2348,14 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
     }
 
     // Allocate buffers and state for op batching
-    this->op_queue = new ggml_hexagon_opqueue(this, opt_opbatch, opt_opqueue);
+    const size_t n_fences     = 512;
+    const size_t opq_size     = ggml_hexagon_opqueue::size_of(opt_opbatch, this->n_threads, opt_opqueue);
+    const size_t fence_offset = GGML_PAD(opq_size, HEX_L2_LINE_SIZE);
+    const size_t fence_size   = ggml_hexagon_fence_buffer::size_of(n_fences);
 
-    this->fence_buf = new ggml_hexagon_fence_buffer(this, &dev_ctx->fence_buffer_type, 64 * 1024);
+    this->shm_buf   = new ggml_hexagon_shared_buffer(this, fence_offset + fence_size, true /* pinned */);
+    this->op_queue  = new ggml_hexagon_opqueue(this->shm_buf, 0, opt_opbatch, opt_opqueue);
+    this->fence_buf = new ggml_hexagon_fence_buffer(this->shm_buf, &dev_ctx->fence_buffer_type, fence_offset, n_fences);
     if (this->mdev.count > 1) {
         this->mdev_fence_slot = this->alloc_fence(this->mdev.count);
     }
@@ -2287,7 +2364,7 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
         opt_vmem = ggml_hexagon_measure_max_vmem(this);
         GGML_LOG_INFO("ggml-hex: %s measured max vmem %zu\n", this->c_name(), opt_vmem);
     }
-    const size_t shm_size = this->op_queue->shm_size();
+    const size_t shm_size = this->shm_buf->size();
     this->max_vmem = (opt_vmem > shm_size) ? (opt_vmem - shm_size) : opt_vmem;
 
     this->op_batch = new ggml_hexagon_opbatch(this, opt_opbatch, this->max_vmem);
@@ -2328,8 +2405,6 @@ void ggml_hexagon_session::release() noexcept(true) {
         }
     }
 
-    delete this->op_batch;
-    delete this->op_queue;
     for (auto & it : this->cpy_fence_slots) {
         free_fence((void *) it.second, 1);
     }
@@ -2340,6 +2415,14 @@ void ggml_hexagon_session::release() noexcept(true) {
         delete this->fence_buf;
         this->fence_buf = nullptr;
     }
+
+    delete this->op_batch;
+    delete this->op_queue;
+    delete this->shm_buf;
+    this->op_batch = nullptr;
+    this->op_queue = nullptr;
+    this->shm_buf  = nullptr;
+
     while (!this->cloned_buffers.empty()) {
         release_buffer(this->cloned_buffers.begin()->second.get());
     }
@@ -2734,19 +2817,12 @@ static bool ggml_hexagon_cpy_tensor_async_phys(ggml_backend_t backend_src, ggml_
     auto sess_src = static_cast<ggml_hexagon_session *>(backend_src->context);
     auto sess_dst = static_cast<ggml_hexagon_session *>(backend_dst->context);
 
-    sess_src->mmap_tensor(src);
-    auto sbuf_dst = sess_dst->mmap_tensor(dst);
-
-    if (!sess_src->clone_buffer(sbuf_dst)) { return false; }
-
     const uint64_t src_key = ggml_hexagon_session_key(sess_src);
     auto & fence_slot = sess_dst->cpy_fence_slots[src_key];
     if (!fence_slot) {
         fence_slot = (volatile uint32_t *) sess_dst->alloc_fence(1);
     }
 
-    if (!sess_src->clone_buffer(sess_dst->fence_buf)) { return false; }
-
     uint32_t fence_seq = sess_dst->next_fence_seq();
 
     HEX_VERBOSE("ggml-hex: %s cpy-tensor-async %s -> %s size %zu : seq 0x%x\n",
@@ -2766,11 +2842,6 @@ static bool ggml_hexagon_cpy_tensor_async_virt(ggml_backend_t backend_src, ggml_
     auto sess_src = static_cast<ggml_hexagon_session *>(backend_src->context);
     auto sess_dst = static_cast<ggml_hexagon_session *>(backend_dst->context);
 
-    auto sbuf_src = sess_src->mmap_tensor(src);
-    sess_dst->mmap_tensor(dst);
-
-    if (!sess_dst->clone_buffer(sbuf_src)) { return false; }
-
     HEX_VERBOSE("ggml-hex: %s cpy-tensor-async %s -> %s size %zu\n",
                 sess_dst->name.c_str(), src->name, dst->name, ggml_nbytes(src));
 
@@ -2893,7 +2964,6 @@ static void ggml_backend_hexagon_event_wait(ggml_backend_t backend, ggml_backend
         return;
     }
 
-    sess->clone_buffer(hex_event->fence_sess->fence_buf);
     sess->add_peer(hex_event->sess);
     sess->enqueue_fence(&hex_event->fence_tensor, hex_event->seq, /* wait = */ true);
 }
@@ -3428,6 +3498,10 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
                   "please update hexagon_type to match ggml_type");
     static_assert((unsigned int) HTP_TYPE_Q4_1 == (unsigned int) GGML_TYPE_Q4_1,
                   "please update hexagon_type to match ggml_type");
+    static_assert((unsigned int) HTP_TYPE_Q5_0 == (unsigned int) GGML_TYPE_Q5_0,
+                  "please update hexagon_type to match ggml_type");
+    static_assert((unsigned int) HTP_TYPE_Q5_1 == (unsigned int) GGML_TYPE_Q5_1,
+                  "please update hexagon_type to match ggml_type");
     static_assert((unsigned int) HTP_TYPE_Q8_0 == (unsigned int) GGML_TYPE_Q8_0,
                   "please update hexagon_type to match ggml_type");
     static_assert((unsigned int) HTP_TYPE_MXFP4 == (unsigned int) GGML_TYPE_MXFP4,

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

@@ -953,6 +953,8 @@ static bool htp_matmul_check_types(
             case GGML_TYPE_Q2_0:
             case GGML_TYPE_Q4_0:
             case GGML_TYPE_Q4_1:
+            case GGML_TYPE_Q5_0:
+            case GGML_TYPE_Q5_1:
             case GGML_TYPE_Q8_0:
             case GGML_TYPE_IQ4_NL:
             case GGML_TYPE_MXFP4:
@@ -995,6 +997,8 @@ static bool htp_matmul_check_types(
         case GGML_TYPE_Q2_0:
         case GGML_TYPE_Q4_0:
         case GGML_TYPE_Q4_1:
+        case GGML_TYPE_Q5_0:
+        case GGML_TYPE_Q5_1:
         case GGML_TYPE_Q8_0:
         case GGML_TYPE_IQ4_NL:
         case GGML_TYPE_MXFP4:

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

@@ -15,7 +15,7 @@
 #    include <arm_neon.h>
 #endif
 
-// --- 4-bit block-32 nibble helpers (shared by Q4_0, IQ4_NL, Q4_1, MXFP4) ---
+// --- 4-bit block-32 nibble helpers (shared by Q4_0, IQ4_NL, Q4_1, Q5_0, Q5_1, MXFP4) ---
 
 static inline void unpack_nibbles_32(uint8_t * dst, const uint8_t * qs) {
     for (unsigned int i = 0; i < 16; ++i) {
@@ -48,6 +48,28 @@ static inline void transpose_tile_to_nibbles(uint8_t tile_quants[32][32], const
     }
 }
 
+// 5th-bit plane of one row (qh bit k = quant k), same bit layout as the Q5_K plane
+static inline void pack_q5_plane(uint8_t * plane, int row, uint32_t qh) {
+    for (int c = 0; c < 4; c++) {
+        uint32_t v = 0;
+        for (int i = 0; i < 4; i++) {
+            v |= ((qh >> (8 * i + 2 * c)) & 3) << (2 * i);
+        }
+        plane[c * 32 + row] = (uint8_t) v;
+    }
+}
+
+static inline uint32_t unpack_q5_plane(const uint8_t * plane, int row) {
+    uint32_t qh = 0;
+    for (int c = 0; c < 4; c++) {
+        const uint32_t v = plane[c * 32 + row];
+        for (int i = 0; i < 4; i++) {
+            qh |= ((v >> (2 * i)) & 3) << (8 * i + 2 * c);
+        }
+    }
+    return qh;
+}
+
 #if defined(__ARM_NEON) || defined(__aarch64__)
 static void repack_neon_transpose8x8(const uint8x8_t rows[8], uint8x8_t columns[8]) {
     const uint8x8x2_t t0 = vtrn_u8(rows[0], rows[1]);
@@ -433,7 +455,7 @@ static void unpack_from_tiled_grid_32(void * data, const ggml_tensor * t, size_t
     }
 }
 
-// --- Block-32 Quantized Types (Q4_0, IQ4_NL, Q4_1, Q8_0, MXFP4) ---
+// --- Block-32 Quantized Types (Q4_0, IQ4_NL, Q4_1, Q5_0, Q5_1, Q8_0, MXFP4) ---
 
 static void repack_q4_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
     repack_to_tiled_grid_32<block_q4_0, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_0,
@@ -546,6 +568,141 @@ static void repack_tiled_q4_1(void * data, const ggml_tensor * t, size_t offset,
         });
 }
 
+static void repack_q5_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
+    repack_to_tiled_grid_32<block_q5_0, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q5_0,
+        [](uint8_t * tile_dst, const block_q5_0 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
+            uint8_t *   plane_dst = tile_dst + 512;
+            ggml_half * scale_dst = (ggml_half *) (tile_dst + 640);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_4bit_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    const block_q5_0 & src_block = src_slice[(ct * 32 + row) * row_stride + kt];
+                    uint32_t qh;
+                    memcpy(&qh, src_block.qh, sizeof(qh));
+                    pack_q5_plane(plane_dst, row, ~qh);
+                    scale_dst[row] = src_block.d;
+                }
+                return;
+            }
+#endif
+            uint8_t tile_quants[32][32];
+            for (int row = 0; row < 32; row++) {
+                const int64_t r = ct * 32 + row;
+                if (r < ne1 && kt < ne0 / 32) {
+                    const block_q5_0 & b = src_slice[r * row_stride + kt];
+                    unpack_nibbles_32(tile_quants[row], b.qs);
+                    uint32_t qh;
+                    memcpy(&qh, b.qh, sizeof(qh));
+                    pack_q5_plane(plane_dst, row, ~qh);
+                    scale_dst[row] = b.d;
+                } else {
+                    memset(tile_quants[row], 0, 32);
+                    pack_q5_plane(plane_dst, row, 0);
+                    scale_dst[row] = 0;
+                }
+            }
+            transpose_nibbles_to_tile(tile_dst, tile_quants);
+        });
+}
+
+static void repack_tiled_q5_0(void * data, const ggml_tensor * t, size_t offset, size_t size) {
+    unpack_from_tiled_grid_32<block_q5_0, 1>(data, t, offset, size, HTP_TENSOR_TILE_SIZE_Q5_0,
+        [](block_q5_0 * dst_slice, const uint8_t * tile_src, int64_t ne0, int64_t start_row, int64_t end_row, int64_t row_stride, int ct, int kt) {
+            uint8_t tile_quants[32][32];
+            transpose_tile_to_nibbles(tile_quants, tile_src);
+            const uint8_t *   plane_src = tile_src + 512;
+            const ggml_half * scale_src = (const ggml_half *) (tile_src + 640);
+            for (int row = 0; row < 32; row++) {
+                const int64_t r = ct * 32 + row;
+                if (r >= start_row && r < end_row && kt < ne0 / 32) {
+                    block_q5_0 & b = dst_slice[(r - start_row) * row_stride + kt];
+                    pack_nibbles_32(b.qs, tile_quants[row]);
+                    const uint32_t qh = ~unpack_q5_plane(plane_src, row);
+                    memcpy(b.qh, &qh, sizeof(qh));
+                    b.d = scale_src[row];
+                }
+            }
+        });
+}
+
+// Q5_1 uses the Q5_K tile layout
+static void repack_q5_1_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
+    repack_to_tiled_grid_32<block_q5_1, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q5_1,
+        [](uint8_t * tile_dst, const block_q5_1 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
+            ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
+            uint8_t *   plane_dst = tile_dst + 640;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_4bit_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    const block_q5_1 & src_block = src_slice[(ct * 32 + row) * row_stride + kt];
+                    uint32_t qh;
+                    memcpy(&qh, src_block.qh, sizeof(qh));
+                    pack_q5_plane(plane_dst, row, qh);
+                    scale_dst[2 * row + 0] = src_block.d;
+                    scale_dst[2 * row + 1] = src_block.m;
+                }
+                return;
+            }
+#endif
+            uint8_t tile_quants[32][32];
+            for (int row = 0; row < 32; row++) {
+                const int64_t r = ct * 32 + row;
+                if (r < ne1 && kt < ne0 / 32) {
+                    const block_q5_1 & b = src_slice[r * row_stride + kt];
+                    unpack_nibbles_32(tile_quants[row], b.qs);
+                    uint32_t qh;
+                    memcpy(&qh, b.qh, sizeof(qh));
+                    pack_q5_plane(plane_dst, row, qh);
+                    scale_dst[2 * row + 0] = b.d;
+                    scale_dst[2 * row + 1] = b.m;
+                } else {
+                    memset(tile_quants[row], 0, 32);
+                    pack_q5_plane(plane_dst, row, 0);
+                    scale_dst[2 * row + 0] = 0;
+                    scale_dst[2 * row + 1] = 0;
+                }
+            }
+            transpose_nibbles_to_tile(tile_dst, tile_quants);
+        });
+}
+
+static void repack_tiled_q5_1(void * data, const ggml_tensor * t, size_t offset, size_t size) {
+    unpack_from_tiled_grid_32<block_q5_1, 1>(data, t, offset, size, HTP_TENSOR_TILE_SIZE_Q5_1,
+        [](block_q5_1 * dst_slice, const uint8_t * tile_src, int64_t ne0, int64_t start_row, int64_t end_row, int64_t row_stride, int ct, int kt) {
+            uint8_t tile_quants[32][32];
+            transpose_tile_to_nibbles(tile_quants, tile_src);
+            const ggml_half * scale_src = (const ggml_half *) (tile_src + 512);
+            const uint8_t *   plane_src = tile_src + 640;
+            for (int row = 0; row < 32; row++) {
+                const int64_t r = ct * 32 + row;
+                if (r >= start_row && r < end_row && kt < ne0 / 32) {
+                    block_q5_1 & b = dst_slice[(r - start_row) * row_stride + kt];
+                    pack_nibbles_32(b.qs, tile_quants[row]);
+                    const uint32_t qh = unpack_q5_plane(plane_src, row);
+                    memcpy(b.qh, &qh, sizeof(qh));
+                    b.d = scale_src[2 * row + 0];
+                    b.m = scale_src[2 * row + 1];
+                }
+            }
+        });
+}
+
 static void repack_mxfp4_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
     repack_to_tiled_grid_32<block_mxfp4, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_MXFP4,
         [](uint8_t * tile_dst, const block_mxfp4 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
@@ -1538,6 +1695,12 @@ void htp_repack_tensor_tiled(ggml_tensor * tensor, const void * data, size_t siz
         case GGML_TYPE_Q4_1:
             repack_q4_1_tiled(tensor, data, 0, size);
             break;
+        case GGML_TYPE_Q5_0:
+            repack_q5_0_tiled(tensor, data, 0, size);
+            break;
+        case GGML_TYPE_Q5_1:
+            repack_q5_1_tiled(tensor, data, 0, size);
+            break;
         case GGML_TYPE_Q8_0:
             repack_q8_0_tiled(tensor, data, 0, size);
             break;
@@ -1593,6 +1756,12 @@ void htp_unpack_tensor_tiled(void * data, const ggml_tensor * tensor, size_t off
         case GGML_TYPE_Q4_1:
             repack_tiled_q4_1(data, tensor, offset, size);
             break;
+        case GGML_TYPE_Q5_0:
+            repack_tiled_q5_0(data, tensor, offset, size);
+            break;
+        case GGML_TYPE_Q5_1:
+            repack_tiled_q5_1(data, tensor, offset, size);
+            break;
         case GGML_TYPE_Q8_0:
             repack_tiled_q8_0(data, tensor, offset, size);
             break;

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

@@ -10,6 +10,7 @@
 static inline bool htp_is_tiled_weight_type(enum ggml_type type) {
     return type == GGML_TYPE_Q1_0  || type == GGML_TYPE_Q2_0   ||
            type == GGML_TYPE_Q4_0  || type == GGML_TYPE_Q4_1   ||
+           type == GGML_TYPE_Q5_0  || type == GGML_TYPE_Q5_1   ||
            type == GGML_TYPE_MXFP4 || type == GGML_TYPE_IQ4_NL ||
            type == GGML_TYPE_Q8_0  || type == GGML_TYPE_Q6_K   ||
            type == GGML_TYPE_Q2_K  || type == GGML_TYPE_Q3_K   ||

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

@@ -647,6 +647,106 @@ static void dequantize_tiled_weight_to_fp16_task_q5_k(
     }
 }
 
+static void dequantize_tiled_weight_to_fp16_task_q5_0(
+        const tiled_dequantize_state_t *state,
+        uint32_t start_tile, uint32_t end_tile) {
+
+    const HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F);
+
+    for (uint32_t t = start_tile; t < end_tile; t++) {
+        const uint8_t * tile_src = state->src + t * state->aligned_tile_size;
+        __fp16 * dst_ptr = state->dst + t * HTP_MM_HMX_TILE_N_ELMS;
+
+        HVX_Vector v_sc = hvx_vmem(tile_src + 640);
+        HVX_Vector v_scale_duplicated = Q6_V_lo_W(Q6_W_vshuff_VVR(v_sc, v_sc, -2));
+
+        // Load all 4 groups in parallel
+        HVX_Vector vq0 = hvx_vmem(tile_src + 0 * 128);
+        HVX_Vector vq1 = hvx_vmem(tile_src + 1 * 128);
+        HVX_Vector vq2 = hvx_vmem(tile_src + 2 * 128);
+        HVX_Vector vq3 = hvx_vmem(tile_src + 3 * 128);
+
+        // Nibble extraction
+        HVX_Vector v_lo0 = Q6_V_vand_VV(vq0, mask_h4);
+        HVX_Vector v_hi0 = Q6_Vub_vlsr_VubR(vq0, 4);
+        HVX_Vector v_lo1 = Q6_V_vand_VV(vq1, mask_h4);
+        HVX_Vector v_hi1 = Q6_Vub_vlsr_VubR(vq1, 4);
+        HVX_Vector v_lo2 = Q6_V_vand_VV(vq2, mask_h4);
+        HVX_Vector v_hi2 = Q6_Vub_vlsr_VubR(vq2, 4);
+        HVX_Vector v_lo3 = Q6_V_vand_VV(vq3, mask_h4);
+        HVX_Vector v_hi3 = Q6_Vub_vlsr_VubR(vq3, 4);
+
+        // Q5_0: OR 0xF0 where the plane flags a clear 5th bit (q - 16)
+        HVX_Vector v_plane = hvx_vmem(tile_src + 512);
+        v_lo0 = hvx_q5_0_or_neg(v_lo0, v_plane, 0);
+        v_hi0 = hvx_q5_0_or_neg(v_hi0, v_plane, 1);
+        v_lo1 = hvx_q5_0_or_neg(v_lo1, v_plane, 2);
+        v_hi1 = hvx_q5_0_or_neg(v_hi1, v_plane, 3);
+        v_lo2 = hvx_q5_0_or_neg(v_lo2, v_plane, 4);
+        v_hi2 = hvx_q5_0_or_neg(v_hi2, v_plane, 5);
+        v_lo3 = hvx_q5_0_or_neg(v_lo3, v_plane, 6);
+        v_hi3 = hvx_q5_0_or_neg(v_hi3, v_plane, 7);
+
+        // Shuffling
+        HVX_VectorPair vp_shuf0 = Q6_W_vshuff_VVR(v_hi0, v_lo0, -1);
+        HVX_VectorPair vp_shuf1 = Q6_W_vshuff_VVR(v_hi1, v_lo1, -1);
+        HVX_VectorPair vp_shuf2 = Q6_W_vshuff_VVR(v_hi2, v_lo2, -1);
+        HVX_VectorPair vp_shuf3 = Q6_W_vshuff_VVR(v_hi3, v_lo3, -1);
+
+        // Unpack to 16-bit
+        HVX_VectorPair vp_int16_lo0 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf0));
+        HVX_VectorPair vp_int16_hi0 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf0));
+        HVX_VectorPair vp_int16_lo1 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf1));
+        HVX_VectorPair vp_int16_hi1 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf1));
+        HVX_VectorPair vp_int16_lo2 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf2));
+        HVX_VectorPair vp_int16_hi2 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf2));
+        HVX_VectorPair vp_int16_lo3 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf3));
+        HVX_VectorPair vp_int16_hi3 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf3));
+
+        // Convert and scale multiplication
+        HVX_Vector v_grp0_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo0)), v_scale_duplicated));
+        HVX_Vector v_grp0_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo0)), v_scale_duplicated));
+        HVX_Vector v_grp0_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi0)), v_scale_duplicated));
+        HVX_Vector v_grp0_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi0)), v_scale_duplicated));
+
+        HVX_Vector v_grp1_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo1)), v_scale_duplicated));
+        HVX_Vector v_grp1_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo1)), v_scale_duplicated));
+        HVX_Vector v_grp1_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi1)), v_scale_duplicated));
+        HVX_Vector v_grp1_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi1)), v_scale_duplicated));
+
+        HVX_Vector v_grp2_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo2)), v_scale_duplicated));
+        HVX_Vector v_grp2_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo2)), v_scale_duplicated));
+        HVX_Vector v_grp2_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi2)), v_scale_duplicated));
+        HVX_Vector v_grp2_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi2)), v_scale_duplicated));
+
+        HVX_Vector v_grp3_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo3)), v_scale_duplicated));
+        HVX_Vector v_grp3_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo3)), v_scale_duplicated));
+        HVX_Vector v_grp3_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi3)), v_scale_duplicated));
+        HVX_Vector v_grp3_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi3)), v_scale_duplicated));
+
+        // Parallel Stores
+        hvx_vmem(dst_ptr +  0 * 64) = v_grp0_0;
+        hvx_vmem(dst_ptr +  1 * 64) = v_grp0_1;
+        hvx_vmem(dst_ptr +  2 * 64) = v_grp0_2;
+        hvx_vmem(dst_ptr +  3 * 64) = v_grp0_3;
+
+        hvx_vmem(dst_ptr +  4 * 64) = v_grp1_0;
+        hvx_vmem(dst_ptr +  5 * 64) = v_grp1_1;
+        hvx_vmem(dst_ptr +  6 * 64) = v_grp1_2;
+        hvx_vmem(dst_ptr +  7 * 64) = v_grp1_3;
+
+        hvx_vmem(dst_ptr +  8 * 64) = v_grp2_0;
+        hvx_vmem(dst_ptr +  9 * 64) = v_grp2_1;
+        hvx_vmem(dst_ptr + 10 * 64) = v_grp2_2;
+        hvx_vmem(dst_ptr + 11 * 64) = v_grp2_3;
+
+        hvx_vmem(dst_ptr + 12 * 64) = v_grp3_0;
+        hvx_vmem(dst_ptr + 13 * 64) = v_grp3_1;
+        hvx_vmem(dst_ptr + 14 * 64) = v_grp3_2;
+        hvx_vmem(dst_ptr + 15 * 64) = v_grp3_3;
+    }
+}
+
 // Q6_K stores 6-bit weights and one fp16 scale per 16 k, see HTP_TENSOR_TILE_SIZE_Q6_K.
 // A k-group holds 4 k per row, the HMX tile holds 2, so each group is dealt into two tiles.
 static void dequantize_tiled_weight_to_fp16_task_q6_k(

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

@@ -82,6 +82,7 @@ struct htp_ops_context {
 struct htp_context {
     dspqueue_t             dsp_queue;
 
+    struct htp_mmap        shm;
     struct htp_mmap        mmap[HTP_MAX_MMAPS];
     dma_queue_t            dma[HTP_MAX_NTHREADS];
     struct htp_thread_trace trace[HTP_MAX_NTHREADS];

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

@@ -21,6 +21,8 @@ enum htp_data_type {
     HTP_TYPE_F16    = 1,
     HTP_TYPE_Q4_0   = 2,
     HTP_TYPE_Q4_1   = 3,
+    HTP_TYPE_Q5_0   = 6,
+    HTP_TYPE_Q5_1   = 7,
     HTP_TYPE_Q8_0   = 8,
     HTP_TYPE_Q2_K   = 10,
     HTP_TYPE_Q3_K   = 11,

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

@@ -17,6 +17,8 @@
 #define HTP_TENSOR_TILE_SIZE_F16    2048
 #define HTP_TENSOR_TILE_SIZE_Q4_0   576
 #define HTP_TENSOR_TILE_SIZE_Q4_1   640
+#define HTP_TENSOR_TILE_SIZE_Q5_0   704
+#define HTP_TENSOR_TILE_SIZE_Q5_1   768
 #define HTP_TENSOR_TILE_SIZE_Q8_0   1088
 #define HTP_TENSOR_TILE_SIZE_IQ4_NL 576
 #define HTP_TENSOR_TILE_SIZE_MXFP4  544
@@ -30,7 +32,7 @@
 #define HTP_TENSOR_TILE_K_Q2_0      64
 #define HTP_TENSOR_TILE_SIZE_Q2_0   576
 
-// Tiled weight tensor layouts (Q4_0, Q8_0, IQ4_NL, MXFP4, Q1_0, Q2_0, K-quants):
+// Tiled weight tensor layouts (Q4_0, Q5_0, Q5_1, Q8_0, IQ4_NL, MXFP4, Q1_0, Q2_0, K-quants):
 // See docs/backend/snapdragon/developer.md ("Tiled Weight Formats").
 
 struct htp_tensor_mdev_range {
@@ -143,6 +145,8 @@ static inline uint32_t htp_tensor_get_tile_size(int type) {
         case HTP_TYPE_Q4_0:   return HTP_TENSOR_TILE_SIZE_Q4_0;
         case HTP_TYPE_IQ4_NL: return HTP_TENSOR_TILE_SIZE_IQ4_NL;
         case HTP_TYPE_Q4_1:   return HTP_TENSOR_TILE_SIZE_Q4_1;
+        case HTP_TYPE_Q5_0:   return HTP_TENSOR_TILE_SIZE_Q5_0;
+        case HTP_TYPE_Q5_1:   return HTP_TENSOR_TILE_SIZE_Q5_1;
         case HTP_TYPE_Q4_K:   return HTP_TENSOR_TILE_SIZE_Q4_K;
         case HTP_TYPE_Q8_0:   return HTP_TENSOR_TILE_SIZE_Q8_0;
         case HTP_TYPE_Q5_K:   return HTP_TENSOR_TILE_SIZE_Q5_K;
@@ -182,6 +186,8 @@ static inline uint32_t htp_tensor_get_row_size(int type, uint32_t ne00) {
         case HTP_TYPE_Q4_0:
         case HTP_TYPE_IQ4_NL: return (ne00 / 32) * 18;
         case HTP_TYPE_Q4_1:   return (ne00 / 32) * 20;
+        case HTP_TYPE_Q5_0:   return (ne00 / 32) * 22;
+        case HTP_TYPE_Q5_1:   return (ne00 / 32) * 24;
         case HTP_TYPE_MXFP4:  return (ne00 / 32) * 17;
         case HTP_TYPE_Q2_K:   return (ne00 / 256) * 84;
         case HTP_TYPE_Q2_0:   return (ne00 / 64) * 18;

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

@@ -639,6 +639,74 @@ static inline HVX_VectorPair accum_5bit_32x2(
     return Q6_W_vcombine_VV(v_sum1, v_sum0);
 }
 
+// OR 0xF0 into every lane of v whose flag j is set in the plane, i.e. 4-bit value | 0xF0 = q - 16 where the 5th bit is clear (Q5_0)
+static inline HVX_Vector hvx_q5_0_or_neg(HVX_Vector v, HVX_Vector v_plane, int j) {
+    HVX_VectorPred q = Q6_Q_vand_VR(v_plane, 0x01010101u << j);
+    return Q6_V_vandor_VQR(v, q, 0xF0F0F0F0);
+}
+
+// Q5_0: signed bytes (q - 16) straight from the nibbles and the plane
+static inline HVX_VectorPair unpack_and_interleave_q5_0_x2(HVX_Vector v_src, HVX_Vector v_plane, int i, HVX_Vector mask_h4) {
+    HVX_Vector v_lo = hvx_q5_0_or_neg(Q6_V_vand_VV(v_src, mask_h4), v_plane, 2 * i);
+    HVX_Vector v_hi = hvx_q5_0_or_neg(Q6_Vub_vlsr_VubR(v_src, 4), v_plane, 2 * i + 1);
+    HVX_VectorPair v01_pair = Q6_W_vshuff_VVR(v_hi, v_lo, -1);
+    HVX_Vector v01_lo = Q6_V_lo_W(v01_pair);
+    HVX_Vector v01_hi = Q6_V_hi_W(v01_pair);
+
+    HVX_Vector v23_lo = Q6_V_valign_VVR(v01_hi, v01_lo, 64);
+    HVX_Vector v_W0 = Q6_V_lo_W(Q6_W_vshuff_VVR(v23_lo, v01_lo, -2));
+
+    HVX_Vector v67_lo = Q6_V_valign_VVR(v01_lo, v01_hi, 64);
+    HVX_Vector v_W1 = Q6_V_lo_W(Q6_W_vshuff_VVR(v67_lo, v01_hi, -2));
+
+    return Q6_W_vcombine_VV(v_W1, v_W0);
+}
+
+static inline HVX_Vector accum_q5_0_32x1(
+    const HVX_Vector * restrict vptr,
+    const HVX_Vector * restrict v_act
+) {
+    HVX_Vector v_sum0 = Q6_V_vzero();
+    HVX_Vector v_sum1 = Q6_V_vzero();
+    HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F);
+    HVX_Vector v_plane = vptr[4];
+
+    #pragma unroll
+    for (int i = 0; i < 4; i++) {
+        HVX_VectorPair v_W_pair = unpack_and_interleave_q5_0_x2(vptr[i], v_plane, i, mask_h4);
+        v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, Q6_V_lo_W(v_W_pair), v_act[i * 2 + 0]);
+        v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, Q6_V_hi_W(v_W_pair), v_act[i * 2 + 1]);
+    }
+
+    return Q6_Vw_vadd_VwVw(v_sum0, v_sum1);
+}
+
+static inline HVX_VectorPair accum_q5_0_32x2(
+    const HVX_Vector * restrict vptr,
+    const HVX_Vector * restrict v_act0,
+    const HVX_Vector * restrict v_act1
+) {
+    HVX_Vector v_sum0 = Q6_V_vzero();
+    HVX_Vector v_sum1 = Q6_V_vzero();
+    HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F);
+    HVX_Vector v_plane = vptr[4];
+
+    #pragma unroll
+    for (int i = 0; i < 4; i++) {
+        HVX_VectorPair v_W_pair = unpack_and_interleave_q5_0_x2(vptr[i], v_plane, i, mask_h4);
+        HVX_Vector v_W0 = Q6_V_lo_W(v_W_pair);
+        HVX_Vector v_W1 = Q6_V_hi_W(v_W_pair);
+
+        v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, v_W0, v_act0[i * 2 + 0]);
+        v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, v_W1, v_act0[i * 2 + 1]);
+
+        v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, v_W0, v_act1[i * 2 + 0]);
+        v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, v_W1, v_act1[i * 2 + 1]);
+    }
+
+    return Q6_W_vcombine_VV(v_sum1, v_sum0);
+}
+
 // Q6_K weights are stored unsigned (0..63), see HTP_TENSOR_TILE_SIZE_Q6_K. Unpack k-group g of a tile to signed bytes (q - 32)
 static inline HVX_Vector unpack_q6_k_group(const HVX_Vector * restrict vptr, int g, HVX_Vector mask_0f, HVX_Vector mask_03, HVX_Vector i32) {
     HVX_Vector v_lo = (g & 1) ? Q6_Vub_vlsr_VubR(vptr[g >> 1], 4) : Q6_V_vand_VV(vptr[g >> 1], mask_0f);
@@ -1251,6 +1319,82 @@ static void tiled_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float
     }
 }
 
+static void tiled_vec_dot_q5_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
+    const uint8_t * restrict tile_ptr = vx;
+    const uint8_t * restrict y_q = vy;
+
+    HVX_Vector v_sum_float = Q6_V_vzero();
+
+    uint32_t n_k_tiles = n / 32;
+    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
+        const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 768);
+        const HVX_Vector * restrict v_act = (const HVX_Vector *) (y_q + kt * 1152);
+
+        HVX_Vector v_sum = accum_q5_0_32x1(vptr, v_act);
+        HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum);
+
+        HVX_Vector v_scale_w = vptr[5];
+        HVX_Vector v_scale_a = v_act[8];
+        HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a);
+        HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb);
+
+        v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
+    }
+
+    if (sz) {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+    } else {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    }
+}
+
+static void tiled_vec_dot_q5_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
+    const uint8_t * restrict tile_ptr = vx;
+    const uint8_t * restrict y0_q = vy0;
+    const uint8_t * restrict y1_q = vy1;
+
+    HVX_Vector v_sum_float_c0 = Q6_V_vzero();
+    HVX_Vector v_sum_float_c1 = Q6_V_vzero();
+
+    uint32_t n_k_tiles = n / 32;
+    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
+        const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 768);
+        const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1152);
+        const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1152);
+
+        HVX_VectorPair v_sums = accum_q5_0_32x2(vptr, v_act0, v_act1);
+        HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums);
+        HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums);
+
+        HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0);
+        HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1);
+
+        HVX_Vector v_scale_w = vptr[5];
+        HVX_Vector v_scale_a_c0 = v_act0[8];
+        HVX_Vector v_scale_a_c1 = v_act1[8];
+
+        HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a_c0);
+        HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a_c1);
+
+        HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0);
+        HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1);
+
+        v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0);
+        v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
+    }
+
+    if (sz0) {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+    } else {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+    }
+    if (sz1) {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+    } else {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    }
+}
+
 static void tiled_vec_dot_q5_k_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y_q = vy;

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

@@ -1089,7 +1089,12 @@ static int execute_op(struct htp_ops_context * octx) {
 }
 
 static inline bool reuse_buf(struct htp_context *ctx, uint32_t *m_reuse, struct htp_buf_desc *b) {
-    b->base = NULL;
+    if (b->fd == ctx->shm.fd) {
+        b->base = ctx->shm.base;
+        return true;
+    }
+
+    b->base = 0;
 
     for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
         struct htp_mmap *m = ctx->mmap + i;
@@ -1203,7 +1208,7 @@ static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uin
 
         for (uint32_t i=0; i < n_bufs; i++) {
             struct htp_buf_desc *b = bufs + i;
-            if (!(b->flags & HTP_BUF_EXTENDED)) {
+            if (!(b->flags & HTP_BUF_EXTENDED) && b->fd != ctx->shm.fd) {
                 b->base = 0;
             }
             if (!mmap_buf(ctx, b)) {
@@ -1316,6 +1321,9 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
             n_bufs, n_tens, n_ops, req->n_traces, dbuf->size, b_size, t_size, o_size);
 
     // Setup descriptor pointers
+    ctx->shm.fd   = dbuf->fd;
+    ctx->shm.base = (uintptr_t) dbuf->ptr - dbuf->offset;
+
     uint8_t * m_ptr = dbuf->ptr;
     struct htp_buf_desc* bufs = (struct htp_buf_desc*)  m_ptr; m_ptr += b_size;
     struct htp_tensor*   tens = (struct htp_tensor*)    m_ptr; m_ptr += t_size;

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

@@ -1664,6 +1664,12 @@ static int hvx_mm_init_tiled_vec_dot(struct htp_mm_context * mmctx, enum htp_dat
             mmctx->vec_dot_32x1 = tiled_vec_dot_q8_0_32x1;
             mmctx->vec_dot_32x2 = tiled_vec_dot_q8_0_32x2;
             return 0;
+        case HTP_TYPE_Q5_0:
+            mmctx->type         = "q5_0_tiled-f32";
+            mmctx->vec_dot_32x1 = tiled_vec_dot_q5_0_32x1;
+            mmctx->vec_dot_32x2 = tiled_vec_dot_q5_0_32x2;
+            return 0;
+        case HTP_TYPE_Q5_1:
         case HTP_TYPE_Q5_K:
             mmctx->type         = "q5_k_tiled-f32";
             mmctx->vec_dot_32x1 = tiled_vec_dot_q5_k_32x1;
@@ -2316,6 +2322,8 @@ static inline tiled_dequant_task_fn_t hmx_get_dequant_worker_fn(int weight_type,
         case HTP_TYPE_Q4_K:   return dequantize_tiled_weight_to_fp16_task_q4_1;
         case HTP_TYPE_MXFP4:  return dequantize_tiled_weight_to_fp16_task_mxfp4;
         case HTP_TYPE_Q8_0:   return dequantize_tiled_weight_to_fp16_task_q8_0;
+        case HTP_TYPE_Q5_0:   return dequantize_tiled_weight_to_fp16_task_q5_0;
+        case HTP_TYPE_Q5_1:
         case HTP_TYPE_Q5_K:   return dequantize_tiled_weight_to_fp16_task_q5_k;
         case HTP_TYPE_Q6_K:   return dequantize_tiled_weight_to_fp16_task_q6_k;
         case HTP_TYPE_Q3_K:   return dequantize_tiled_weight_to_fp16_task_q3_k;

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

@@ -185,7 +185,7 @@ next_nc:
 // (Q2_K: per-16 k sums, q8_1_s16)
 static inline bool htp_mm_weight_has_offset(int weight_type) {
     return weight_type == HTP_TYPE_Q4_1 || weight_type == HTP_TYPE_Q4_K || weight_type == HTP_TYPE_Q5_K ||
-           weight_type == HTP_TYPE_Q2_K || weight_type == HTP_TYPE_Q2_0;
+           weight_type == HTP_TYPE_Q2_K || weight_type == HTP_TYPE_Q2_0 || weight_type == HTP_TYPE_Q5_1;
 }
 
 // --- Activation/Row Size Helpers ---

ggml/src/ggml-metal/ggml-metal-device.m

@@ -1708,8 +1708,9 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
         case GGML_OP_ARGMAX:
             return has_simdgroup_reduction;
         case GGML_OP_NORM:
-        case GGML_OP_RMS_NORM:
             return has_simdgroup_reduction && (ggml_is_contiguous_rows(op->src[0]));
+        case GGML_OP_RMS_NORM:
+            return has_simdgroup_reduction && op->src[0]->type == GGML_TYPE_F32 && (ggml_is_contiguous_rows(op->src[0]));
         case GGML_OP_ROPE:
         case GGML_OP_ROPE_BACK:
             return ggml_rope_get_truncate(op);

ggml/src/ggml-opencl/CMakeLists.txt

@@ -76,7 +76,6 @@ set(GGML_OPENCL_KERNELS
     clamp
     cpy
     cvt
-    diag_mask_inf
     diag
     div
     dsv4_hc

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

@@ -798,7 +798,6 @@ struct ggml_backend_opencl_context {
     cl_program program_moe_add_id_glu;
     cl_program program_clamp;
     cl_program program_cvt;
-    cl_program program_diag_mask_inf;
     cl_program program_gelu;
     cl_program program_gemv_noshuffle_general;
     cl_program program_gemv_noshuffle;
@@ -886,7 +885,6 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_rms_norm_mul_add = nullptr;   // fused rms_norm(x)*w + b (residual)
     cl_kernel kernel_l2_norm_f32;
     cl_kernel kernel_group_norm, kernel_group_norm_mul_add;
-    cl_kernel kernel_diag_mask_inf, kernel_diag_mask_inf_8;
     cl_kernel kernel_diag_f32;
     cl_kernel kernel_soft_max, kernel_soft_max_4;
     cl_kernel kernel_soft_max_f16, kernel_soft_max_4_f16;
@@ -2291,23 +2289,6 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
         GGML_LOG_CONT(".");
     }
 
-    // diag_mask_inf
-    {
-#ifdef GGML_OPENCL_EMBED_KERNELS
-        const std::string kernel_src {
-            #include "diag_mask_inf.cl.h"
-        };
-#else
-        const std::string kernel_src = read_file("diag_mask_inf.cl");
-#endif
-        backend_ctx->program_diag_mask_inf =
-            build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
-
-        CL_CHECK((backend_ctx->kernel_diag_mask_inf_8 = clCreateKernel(backend_ctx->program_diag_mask_inf, "kernel_diag_mask_inf_8", &err), err));
-        CL_CHECK((backend_ctx->kernel_diag_mask_inf   = clCreateKernel(backend_ctx->program_diag_mask_inf, "kernel_diag_mask_inf", &err), err));
-        GGML_LOG_CONT(".");
-    }
-
     // diag
     {
 #ifdef GGML_OPENCL_EMBED_KERNELS
@@ -11985,7 +11966,7 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
         case GGML_OP_NORM:
             return true;
         case GGML_OP_RMS_NORM:
-            return op->ne[0] % 4 == 0 && ggml_is_contiguous_rows(op->src[0]);
+            return op->src[0]->type == GGML_TYPE_F32 && op->ne[0] % 4 == 0 && ggml_is_contiguous_rows(op->src[0]);
         case GGML_OP_L2_NORM:
             return ggml_is_contiguous_rows(op->src[0]);
         case GGML_OP_REPEAT:
@@ -12214,8 +12195,6 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
             return true;
         case GGML_OP_DIAG:
             return true;
-        case GGML_OP_DIAG_MASK_INF:
-            return op->ne[3] == 1;
         case GGML_OP_ROPE: {
             const int mode = ((const int32_t *) op->op_params)[2];
             if (!ggml_rope_get_truncate(op)) {
@@ -37948,68 +37927,6 @@ static void ggml_cl_set(ggml_backend_t backend, const ggml_tensor * src0, const
     backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
 }
 
-static void ggml_cl_diag_mask_inf(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
-    GGML_ASSERT(src0);
-    GGML_ASSERT(src0->extra);
-    GGML_ASSERT(dst);
-    GGML_ASSERT(dst->extra);
-
-    UNUSED(src1);
-
-    int n_past = ((int32_t *)(dst->op_params))[0];
-
-    const int  ne00 = src0 ? src0->ne[0] : 0;
-    const int  ne01 = src0 ? src0->ne[1] : 0;
-    const int  ne02 = src0 ? src0->ne[2] : 0;
-
-    ggml_backend_opencl_context *backend_ctx = (ggml_backend_opencl_context *)backend->context;
-
-    ggml_tensor_extra_cl * extra0 = (ggml_tensor_extra_cl *)src0->extra;
-    ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *)dst->extra;
-
-    cl_ulong offset0 = extra0->offset + src0->view_offs;
-    cl_ulong offsetd = extrad->offset + dst->view_offs;
-
-    cl_kernel kernel;
-
-    if (ne00%8 == 0) {
-        kernel = backend_ctx->kernel_diag_mask_inf_8;
-
-        CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem),   &extra0->data_device));
-        CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset0));
-        CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem),   &extrad->data_device));
-        CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offsetd));
-        CL_CHECK(clSetKernelArg(kernel, 4, sizeof(int),      &ne00));
-        CL_CHECK(clSetKernelArg(kernel, 5, sizeof(int),      &ne01));
-        CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int),      &n_past));
-
-        size_t global_work_size[] = {(size_t)ne00*ne01*ne02/8, 1, 1};
-        size_t local_work_size[] = {64, 1, 1};
-
-        backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
-    } else {
-        kernel = backend_ctx->kernel_diag_mask_inf;
-
-        CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem),   &extra0->data_device));
-        CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset0));
-        CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem),   &extrad->data_device));
-        CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offsetd));
-        CL_CHECK(clSetKernelArg(kernel, 4, sizeof(int),      &ne00));
-        CL_CHECK(clSetKernelArg(kernel, 5, sizeof(int),      &ne01));
-        CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int),      &n_past));
-
-        size_t global_work_size[] = {(size_t)ne00, (size_t)ne01, (size_t)ne02};
-        size_t local_work_size[] = {64, 1, 1};
-
-        size_t * local_work_size_ptr = local_work_size;
-        if (ne00 % 64 != 0 && !backend_ctx->non_uniform_workgroups) {
-            local_work_size_ptr = nullptr;  // Let driver choose the work-group sizes.
-        }
-
-        backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size_ptr, dst);
-    }
-}
-
 static void ggml_cl_diag(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
     GGML_ASSERT(src0);
     GGML_ASSERT(src0->extra);
@@ -39935,12 +39852,6 @@ bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor
             }
             func = ggml_cl_diag;
             break;
-        case GGML_OP_DIAG_MASK_INF:
-            if (!any_on_device) {
-                return false;
-            }
-            func = ggml_cl_diag_mask_inf;
-            break;
         case GGML_OP_SOFT_MAX:
             if (!any_on_device) {
                 return false;

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

@@ -1,58 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-
-//------------------------------------------------------------------------------
-// diag_mask_inf kernels
-//------------------------------------------------------------------------------
-kernel void kernel_diag_mask_inf(
-        global float * src0,
-        ulong offset0,
-        global float * dst,
-        ulong offsetd,
-        int ne00,
-        int ne01,
-        int n_past
-) {
-    src0 = (global float*)((global char*)src0 + offset0);
-    dst = (global float*)((global char*)dst + offsetd);
-
-    int i02 = get_global_id(2);
-    int i01 = get_global_id(1);
-    int i00 = get_global_id(0);
-
-    if (i00 > n_past + i01) {
-        dst[i02*ne01*ne00 + i01*ne00 + i00] = -INFINITY;
-    } else {
-        dst[i02*ne01*ne00 + i01*ne00 + i00] = src0[i02*ne01*ne00 + i01*ne00 + i00];
-    }
-}
-
-kernel void kernel_diag_mask_inf_8(
-        global float4 * src0,
-        ulong offset0,
-        global float4 * dst,
-        ulong offsetd,
-        int ne00,
-        int ne01,
-        int n_past
-) {
-    src0 = (global float4*)((global char*)src0 + offset0);
-    dst = (global float4*)((global char*)dst + offsetd);
-
-    int i = 2*get_global_id(0);
-
-    dst[i+0] = src0[i+0];
-    dst[i+1] = src0[i+1];
-    int i4 = 4*i;
-    int i02 = i4/(ne00*ne01); i4 -= i02*ne00*ne01;
-    int i01 = i4/(ne00);      i4 -= i01*ne00;
-    int i00 = i4;
-    for (int k = 3; k >= 0; --k) {
-        if (i00 + 4 + k <= n_past + i01) {
-            break;
-        }
-        (&dst[i+1])[k] = -INFINITY;
-        if (i00 + k > n_past + i01) {
-            (&dst[i])[k] = -INFINITY;
-        }
-    }
-}

ggml/src/ggml-openvino/ggml-openvino.cpp

@@ -1455,6 +1455,12 @@ static ggml_openvino_op_support is_op_supported_case(const ggml_tensor * op) {
         }
         break;
     }
+    case GGML_OP_RMS_NORM: {
+        if (op->type == GGML_TYPE_BF16) {
+            return {false, "RMS_NORM with BF16 type is not supported"};
+        }
+        break;
+    }
     case GGML_OP_ADD_ID: {
         // Keep support aligned with the CPU backend implementation, which only handles f32 inputs/output and i32 ids.
         if (op->type != GGML_TYPE_F32 || op->src[0]->type != GGML_TYPE_F32 || op->src[1]->type != GGML_TYPE_F32 ||

ggml/src/ggml-openvino/openvino/op_table.cpp

@@ -33,7 +33,6 @@ std::unordered_map<std::string, CreatorFunction> get_supported_ops() {
     using namespace ov::op;
     return {
         {"GGML_OP_ADD",             op::translate_add                              },
-        {"GGML_OP_ADD1",            op::translate_1to1_match_2_inputs<v1::Add>     },
         {"GGML_OP_ADD_ID",          op::translate_add_id                           },
         {"GGML_OP_CONCAT",          op::translate_concat                           },
         {"GGML_OP_CONT",            op::translate_cont                             },

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

@@ -2517,23 +2517,6 @@ k_argsort_f32_i32(const float *x, int *dst, const int ncols, int ncols_pad,
     }
 }
 
-static void diag_mask_inf_f32(const float * x, float * dst, const int ncols, const int rows_per_channel, const int n_past,
-                              const sycl::nd_item<3> &item_ct1) {
-    const int col = item_ct1.get_local_range(1) * item_ct1.get_group(1) +
-                    item_ct1.get_local_id(1);
-    const int row = item_ct1.get_local_range(2) * item_ct1.get_group(2) +
-                    item_ct1.get_local_id(2);
-
-    if (col >= ncols) {
-        return;
-    }
-
-    const int i = row*ncols + col;
-    //dst[i] = col > (n_past + row % rows_per_channel) ? -INFINITY : x[i];
-    //dst[i] = x[i] - (col > n_past + row % rows_per_channel) * INT_MAX; // equivalent within rounding error but slightly faster on GPU
-    dst[i] = x[i] - (col > n_past + row % rows_per_channel) * FLT_MAX;
-}
-
 static void scale_f32(const float * x, float * dst, const float scale, const float bias, const int k,
                       const sycl::nd_item<3> &item_ct1) {
     const int i = item_ct1.get_local_range(2) * item_ct1.get_group(2) +
@@ -3139,20 +3122,6 @@ static void argmax_f32_i32_sycl(const float *x, int *dst, const int ncols,
             });
     });
 }
-static void diag_mask_inf_f32_sycl(const float *x, float *dst,
-                                   const int ncols_x, const int nrows_x,
-                                   const int rows_per_channel, const int n_past,
-                                   queue_ptr stream) {
-    const sycl::range<3> block_dims(1, SYCL_DIAG_MASK_INF_BLOCK_SIZE, 1);
-    const int block_num_x = (ncols_x + SYCL_DIAG_MASK_INF_BLOCK_SIZE - 1) / SYCL_DIAG_MASK_INF_BLOCK_SIZE;
-    const sycl::range<3> block_nums(1, block_num_x, nrows_x);
-    stream->parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
-                         [=](sycl::nd_item<3> item_ct1) {
-                             diag_mask_inf_f32(x, dst, ncols_x,
-                                               rows_per_channel, n_past,
-                                               item_ct1);
-                         });
-}
 
 static dpct::err0 ggml_sycl_cpy_tensor_2d(void *dst,
                                           const struct ggml_tensor *src,
@@ -3536,25 +3505,6 @@ inline void ggml_sycl_op_argmax(ggml_backend_sycl_context & ctx, ggml_tensor * d
     argmax_f32_i32_sycl(src0_dd, dst_dd, ncols, nrows, main_stream);
 }
 
-inline void ggml_sycl_op_diag_mask_inf(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    GGML_ASSERT(dst->src[0]->type == GGML_TYPE_F32);
-    GGML_ASSERT( dst->type == GGML_TYPE_F32);
-    dpct::queue_ptr main_stream = ctx.stream();
-    SYCL_CHECK(ggml_sycl_set_device(ctx.device));
-    const float * src0_dd = static_cast<const float *>(dst->src[0]->data);
-    float *       dst_dd  = static_cast<float *>(dst->data);
-
-    const int64_t ne00 = dst->src[0]->ne[0];
-    const int64_t ne01 = dst->src[0]->ne[1];
-    const int nrows0 = ggml_nrows(dst->src[0]);
-    // the kernel indexes in 32-bit
-    GGML_ASSERT(ne00 * nrows0 <= INT32_MAX);
-
-    const int n_past = ((int32_t *) dst->op_params)[0];
-
-    diag_mask_inf_f32_sycl(src0_dd, dst_dd, ne00, nrows0, ne01, n_past, main_stream);
-}
-
 static void tri_f32_sycl(
     const float * src,
     float * dst,
@@ -6045,11 +5995,6 @@ static void ggml_sycl_scale(ggml_backend_sycl_context & ctx, ggml_tensor * dst)
     ggml_sycl_op_scale(ctx, dst);
 }
 
-static void ggml_sycl_diag_mask_inf(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/1);
-    ggml_sycl_op_diag_mask_inf(ctx, dst);
-}
-
 static void ggml_sycl_pool2d(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/1);
     ggml_sycl_op_pool2d(ctx, dst);
@@ -6193,7 +6138,6 @@ static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct gg
             ggml_sycl_dup(ctx, dst);
             break;
         case GGML_OP_ADD:
-        case GGML_OP_ADD1: // TODO: more efficient implementation
             ggml_sycl_add(ctx, dst);
             break;
         case GGML_OP_ADD_ID:
@@ -6396,9 +6340,6 @@ static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct gg
         case GGML_OP_TRI:
             ggml_sycl_op_tri(ctx, dst);
             break;
-        case GGML_OP_DIAG_MASK_INF:
-            ggml_sycl_diag_mask_inf(ctx, dst);
-            break;
         case GGML_OP_SOFT_MAX:
             ggml_sycl_op_soft_max(ctx, dst);
             break;
@@ -7639,7 +7580,6 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
         case GGML_OP_PERMUTE:
         case GGML_OP_TRANSPOSE:
         case GGML_OP_ADD:
-        case GGML_OP_ADD1:
         case GGML_OP_ADD_ID:
         case GGML_OP_SUB:
         case GGML_OP_COUNT_EQUAL:
@@ -7658,8 +7598,9 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
         case GGML_OP_NORM:
         case GGML_OP_L2_NORM:
         case GGML_OP_GROUP_NORM:
-        case GGML_OP_RMS_NORM:
             return true;
+        case GGML_OP_RMS_NORM:
+            return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
         case GGML_OP_RMS_NORM_BACK:
             return ggml_is_contiguous(op->src[0]);
         case GGML_OP_SCALE:
@@ -7674,8 +7615,6 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
                        op->type == GGML_TYPE_F32 &&
                        ggml_is_contiguous(src0);
             }
-        case GGML_OP_DIAG_MASK_INF:
-            return true;
         case GGML_OP_SOFT_MAX:
             return true;
         case GGML_OP_SOFT_MAX_BACK: {

ggml/src/ggml-sycl/presets.hpp

@@ -35,7 +35,6 @@
 #define SYCL_SCALE_BLOCK_SIZE 256
 #define SYCL_ROPE_BLOCK_SIZE 256
 #define SYCL_ALIBI_BLOCK_SIZE 32
-#define SYCL_DIAG_MASK_INF_BLOCK_SIZE 32
 #define SYCL_QUANTIZE_BLOCK_SIZE 256
 #define SYCL_DEQUANTIZE_BLOCK_SIZE 256
 #define SYCL_GET_ROWS_BLOCK_SIZE 256

ggml/src/ggml-vulkan/ggml-vulkan-common.h

@@ -109,6 +109,7 @@ void ggml_vk_quantize_q8_1(ggml_backend_vk_context * ctx, vk_context& subctx, co
 void ggml_vk_dsv4_hc_comb(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * mixes, const ggml_tensor * scale, const ggml_tensor * base, ggml_tensor * dst);
 void ggml_vk_dsv4_hc_pre(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * x, const ggml_tensor * weights, ggml_tensor * dst);
 void ggml_vk_dsv4_hc_post(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * x, const ggml_tensor * residual, const ggml_tensor * post, const ggml_tensor * comb, ggml_tensor * dst, const ggml_tensor * gate_scale_in = nullptr);
+void ggml_vk_top_k_mask(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst);
 void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx);
 bool ggml_vk_use_mul_mat_vec_id(const struct ggml_cgraph * cgraph, int node_idx);
 void ggml_vk_mul_mat_id(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx);
@@ -172,7 +173,6 @@ void ggml_vk_l2_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const gg
 void ggml_vk_unary(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst);
 void ggml_vk_xielu(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst);
 void ggml_vk_glu(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst);
-void ggml_vk_diag_mask_inf(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst);
 void ggml_vk_soft_max(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst);
 void ggml_vk_soft_max_back(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst);
 void ggml_vk_topk_moe(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_cgraph * cgraph, int node_idx);

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

@@ -1035,8 +1035,6 @@ void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph
         } else if (tensor->op == GGML_OP_SCALE) {
             const float * params = (const float *)tensor->op_params;
             tensor_clone = ggml_scale_bias(ggml_ctx, src_clone[0], params[0], params[1]);
-        } else if (tensor->op == GGML_OP_ADD1) {
-            tensor_clone = ggml_add1(ggml_ctx, src_clone[0], src_clone[1]);
         } else if (tensor->op == GGML_OP_ARANGE) {
             const float start = ggml_get_op_params_f32(tensor, 0);
             const float stop = ggml_get_op_params_f32(tensor, 1);
@@ -1101,8 +1099,6 @@ void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph
             }
         } else if (tensor->op == GGML_OP_SOFT_MAX_BACK) {
             tensor_clone = ggml_soft_max_ext_back(ggml_ctx, src_clone[0], src_clone[1], ((float *)tensor->op_params)[0], ((float *)tensor->op_params)[1]);
-        } else if (tensor->op == GGML_OP_DIAG_MASK_INF) {
-            tensor_clone = ggml_diag_mask_inf(ggml_ctx, src_clone[0], tensor->op_params[0]);
         } else if (tensor->op == GGML_OP_ROPE || tensor->op == GGML_OP_ROPE_BACK) {
             const int n_dims      = ((int32_t *) tensor->op_params)[1];
             const int mode        = ((int32_t *) tensor->op_params)[2];

ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h

@@ -219,6 +219,29 @@ struct vk_op_dsv4_hc_post_push_constants {
     float    gate_scale_out;
 };
 
+struct vk_op_top_k_mask_push_constants {
+    uint32_t n_kv;
+    uint32_t n_top;
+    uint32_t n_pick;
+    uint32_t n_blk;
+    uint32_t n_extra;
+    uint32_t has_score;
+    uint32_t ninf;
+
+    uint32_t ne1; uint32_t ne2;
+
+    uint32_t m1; uint32_t m2; uint32_t m3;
+    uint32_t nbm1; uint32_t nbm2; uint32_t nbm3;
+    uint32_t nbd1; uint32_t nbd2; uint32_t nbd3;
+
+    uint32_t m_offset;
+    uint32_t k_offset;
+    uint32_t s_offset;
+    uint32_t b_offset;
+    uint32_t e_offset;
+    uint32_t d_offset;
+};
+
 struct vk_op_count_experts_push_constants {
     uint32_t ne00;
     uint32_t ne01;
@@ -462,12 +485,6 @@ struct vk_op_add_id_push_constants {
     uint32_t s21;
 };
 
-struct vk_op_diag_mask_push_constants {
-    uint32_t ncols;
-    uint32_t rows_per_channel;
-    int32_t n_past;
-};
-
 struct vk_op_rope_push_constants {
     uint32_t rope_mode;
     uint32_t nrows;
@@ -978,6 +995,18 @@ template <> inline void init_pushconst_tensor_offsets(ggml_backend_vk_context *
     p.d_offset = get_misalign_bytes(ctx, dst)  / ggml_type_size(dst->type);
 }
 
+// the fifth source, extra, comes from dst
+template <> inline void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_top_k_mask_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) {
+    const ggml_tensor * extra = dst->src[4];
+
+    p.m_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type);
+    p.k_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type);
+    p.s_offset = src2  ? get_misalign_bytes(ctx, src2)  / ggml_type_size(src2->type)  : 0;
+    p.b_offset = src3  ? get_misalign_bytes(ctx, src3)  / ggml_type_size(src3->type)  : 0;
+    p.e_offset = extra ? get_misalign_bytes(ctx, extra) / ggml_type_size(extra->type) : 0;
+    p.d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type);
+}
+
 template <typename T> size_t push_constant_size(const T &t) {
     static_assert(std::is_class<T>::value, "T must be a struct/class");
     GGML_UNUSED(t);

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

@@ -943,7 +943,6 @@ struct vk_device_struct {
 
     vk_pipeline pipeline_leaky_relu[2];
     vk_pipeline pipeline_silu_back_f32;
-    vk_pipeline pipeline_diag_mask_inf_f32;
     vk_pipeline pipeline_soft_max_f32, pipeline_soft_max_f32_f16;
     vk_pipeline pipeline_soft_max_f32_wg512, pipeline_soft_max_f32_f16_wg512;
     vk_pipeline pipeline_soft_max_back_f32;
@@ -978,6 +977,8 @@ struct vk_device_struct {
     vk_pipeline pipeline_dsv4_hc_pre_gated_f32;
     vk_pipeline pipeline_dsv4_hc_post_f32;
     vk_pipeline pipeline_dsv4_hc_post_nocomb_f32;
+    vk_pipeline pipeline_top_k_mask_f16;
+    vk_pipeline pipeline_top_k_mask_f32;
     std::map<vk_solve_tri_pipeline_state, vk_pipeline> pipeline_solve_tri_f32;
     vk_pipeline pipeline_im2col_f32, pipeline_im2col_f32_f16;
     vk_pipeline pipeline_im2col_3d_f32, pipeline_im2col_3d_f32_f16;
@@ -1280,6 +1281,10 @@ struct ggml_backend_vk_context {
 
     vk_device device;
 
+    // Workaround for drivers moving VRAM contents off the GPU when idle
+    std::promise<void> keepalive_stop;
+    std::thread keepalive_thread;
+
     size_t semaphore_idx, binary_semaphore_idx, event_idx;
     ggml_vk_garbage_collector gc;
     size_t prealloc_size_x, prealloc_size_y, prealloc_size_split_k, prealloc_size_add_rms_partials, prealloc_size_add_rms_partials_offset;

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

@@ -3676,10 +3676,6 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     CREATE_UNARY_MUL(softplus, 3)
 #undef CREATE_UNARY_MUL
 
-    ggml_vk_create_pipeline(device, device->pipeline_add1_f16_f16, "add1_f16_f16", add1_f16_f16_len, add1_f16_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1);
-    ggml_vk_create_pipeline(device, device->pipeline_add1_f16_f32, "add1_f16_f32", add1_f16_f32_len, add1_f16_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1);
-    ggml_vk_create_pipeline(device, device->pipeline_add1_f32_f32, "add1_f32_f32", add1_f32_f32_len, add1_f32_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1);
-
     ggml_vk_create_pipeline(device, device->pipeline_arange_f32, "arange_f32", arange_f32_len, arange_f32_data, "main", 1, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
 
     ggml_vk_create_pipeline(device, device->pipeline_fill_f32, "fill_f32", fill_f32_len, fill_f32_data, "main", 1, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
@@ -3700,8 +3696,6 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
 
     ggml_vk_create_pipeline(device, device->pipeline_silu_back_f32, "silu_back_f32", silu_back_f32_len, silu_back_f32_data, "main", 3, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
 
-    ggml_vk_create_pipeline(device, device->pipeline_diag_mask_inf_f32, "diag_mask_inf_f32", diag_mask_inf_f32_len, diag_mask_inf_f32_data, "main", 2, sizeof(vk_op_diag_mask_push_constants), {1, 512, 1}, {}, 1, true);
-
     ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32, "soft_max_f32", soft_max_f32_len, soft_max_f32_data, "main", 4, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
     ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_wg512, "soft_max_f32_wg512", soft_max_f32_len, soft_max_f32_data, "main", 4, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 512 }, 1);
     ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_f16, "soft_max_f32_f16", soft_max_f32_f16_len, soft_max_f32_f16_data, "main", 4, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
@@ -3829,6 +3823,10 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_post_f32,       "dsv4_hc_post_f32",       dsv4_hc_post_f32_len, dsv4_hc_post_f32_data, "main", 5, sizeof(vk_op_dsv4_hc_post_push_constants), {256, 1, 1}, { 256, 1 }, 1);
     ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_post_nocomb_f32,"dsv4_hc_post_nocomb_f32",dsv4_hc_post_f32_len, dsv4_hc_post_f32_data, "main", 5, sizeof(vk_op_dsv4_hc_post_push_constants), {256, 1, 1}, { 256, 0 }, 1);
 
+    const uint32_t top_k_mask_wg = std::min({256u, device->properties.limits.maxComputeWorkGroupInvocations, device->properties.limits.maxComputeWorkGroupSize[0]});
+    ggml_vk_create_pipeline(device, device->pipeline_top_k_mask_f16, "top_k_mask_f16", top_k_mask_f16_len, top_k_mask_f16_data, "main", 6, sizeof(vk_op_top_k_mask_push_constants), {1, 1, 1}, {top_k_mask_wg}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_top_k_mask_f32, "top_k_mask_f32", top_k_mask_f32_len, top_k_mask_f32_data, "main", 6, sizeof(vk_op_top_k_mask_push_constants), {1, 1, 1}, {top_k_mask_wg}, 1);
+
     for (auto &s : device->pipeline_solve_tri_f32) {
         const vk_solve_tri_pipeline_state &state = s.first;
 
@@ -4243,6 +4241,8 @@ static bool ggml_vk_has_rebar(const vk::PhysicalDevice & physical_device) {
     return false;
 }
 
+static void ggml_vk_device_start_keepalive(vk_device& device);
+
 vk_device ggml_vk_get_device(size_t idx) {
     VK_LOG_DEBUG("ggml_vk_get_device(" << idx << ")");
 
@@ -7629,6 +7629,46 @@ void ggml_vk_dsv4_hc_post(ggml_backend_vk_context * ctx, vk_context& subctx, con
     ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { x_buf, r_buf, p_buf, c_buf, d_buf }, pc, { n_embd, n_tokens, 1 });
 }
 
+void ggml_vk_top_k_mask(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
+    const ggml_tensor * mask   = dst->src[0];
+    const ggml_tensor * top_k  = dst->src[1];
+    const ggml_tensor * score  = dst->src[2];
+    const ggml_tensor * blocks = dst->src[3];
+    const ggml_tensor * extra  = dst->src[4];
+    VK_LOG_DEBUG("ggml_vk_top_k_mask(" << mask << ", " << top_k << ", " << score << ", " << blocks << ", " << extra << ", " << dst << ")");
+
+    vk_pipeline pipeline = dst->type == GGML_TYPE_F16 ? ctx->device->pipeline_top_k_mask_f16 : ctx->device->pipeline_top_k_mask_f32;
+    GGML_ASSERT(pipeline != nullptr);
+
+    ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+
+    const vk_subbuffer m_buf = ggml_vk_tensor_subbuffer(ctx, mask,  true);
+    const vk_subbuffer k_buf = ggml_vk_tensor_subbuffer(ctx, top_k, true);
+    const vk_subbuffer s_buf = score  ? ggml_vk_tensor_subbuffer(ctx, score,  true) : k_buf;
+    const vk_subbuffer b_buf = blocks ? ggml_vk_tensor_subbuffer(ctx, blocks, true) : k_buf;
+    const vk_subbuffer e_buf = extra  ? ggml_vk_tensor_subbuffer(ctx, extra,  true) : k_buf;
+    const vk_subbuffer d_buf = ggml_vk_tensor_subbuffer(ctx, dst,   true);
+
+    vk_op_top_k_mask_push_constants pc = {
+        (uint32_t)mask->ne[0],
+        (uint32_t)top_k->ne[0],
+        (uint32_t)(blocks ? blocks->ne[1] : score ? score->ne[0] : mask->ne[0]),
+        blocks ? (uint32_t)blocks->ne[0] : 0,
+        extra  ? (uint32_t)extra->ne[0]  : 0,
+        score  ? 1u : 0u,
+        dst->type == GGML_TYPE_F16 ? 0xFC00u : 0xFF800000u,
+        (uint32_t)dst->ne[1], (uint32_t)dst->ne[2],
+        (uint32_t)mask->ne[1], (uint32_t)mask->ne[2], (uint32_t)mask->ne[3],
+        ggml_vk_nb_elem(mask, 1), ggml_vk_nb_elem(mask, 2), ggml_vk_nb_elem(mask, 3),
+        ggml_vk_nb_elem(dst,  1), ggml_vk_nb_elem(dst,  2), ggml_vk_nb_elem(dst,  3),
+        0, 0, 0, 0, 0, 0,
+    };
+    init_pushconst_tensor_offsets(ctx, pc, mask, top_k, score, blocks, dst);
+
+    ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { m_buf, k_buf, s_buf, b_buf, e_buf, d_buf }, pc,
+        { (uint32_t)dst->ne[1], (uint32_t)dst->ne[2], (uint32_t)dst->ne[3] });
+}
+
 void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx) {
     ggml_tensor * dst = cgraph->nodes[node_idx];
     ggml_tensor * src0 = dst->src[0];
@@ -9477,11 +9517,6 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
                 break;
         }
         return nullptr;
-    case GGML_OP_DIAG_MASK_INF:
-        if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
-            return ctx->device->pipeline_diag_mask_inf_f32;
-        }
-        return nullptr;
     case GGML_OP_SOFT_MAX:
         GGML_ASSERT(!src1 || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);
         GGML_ASSERT(!src2 || src2->type == GGML_TYPE_F32);
@@ -9860,17 +9895,6 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
             return pipeline;
         }
         return nullptr;
-    case GGML_OP_ADD1:
-        if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) {
-            return ctx->device->pipeline_add1_f16_f16;
-        }
-        if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) {
-            return ctx->device->pipeline_add1_f16_f32;
-        }
-        if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
-            return ctx->device->pipeline_add1_f32_f32;
-        }
-        return nullptr;
     case GGML_OP_ARANGE:
         if (dst->type == GGML_TYPE_F32) {
             return ctx->device->pipeline_arange_f32;
@@ -10009,9 +10033,6 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
             elements = { groups_x, groups_y, uint32_t(CEIL_DIV(uint64_t(total_workgroups), groups_xy)) };
             GGML_ASSERT(elements[2] <= ctx->device->properties.limits.maxComputeWorkGroupCount[2]);
         } break;
-    case GGML_OP_DIAG_MASK_INF:
-        elements = { (uint32_t)ggml_nrows(src0), (uint32_t)ne00, 1 };
-        break;
     case GGML_OP_ROPE:
     case GGML_OP_ROPE_BACK:
         {
@@ -10145,7 +10166,6 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
     case GGML_OP_SUB:
     case GGML_OP_DIV:
     case GGML_OP_MUL:
-    case GGML_OP_ADD1:
     case GGML_OP_OUT_PROD:
     case GGML_OP_ARANGE:
     case GGML_OP_FILL:
@@ -11040,21 +11060,6 @@ void ggml_vk_sqrt(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_
     ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_SQRT, vk_op_unary_push_constants_init(src0, dst));
 }
 
-void ggml_vk_add1(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
-    const uint32_t src0_type_size = ggml_type_size(src0->type);
-    const uint32_t src1_type_size = ggml_type_size(src1->type);
-    const uint32_t dst_type_size = ggml_type_size(dst->type);
-
-    ggml_vk_op_f32<vk_op_binary_push_constants>(ctx, subctx, src0, src1, nullptr, nullptr, dst, GGML_OP_ADD1, {
-        (uint32_t)ggml_nelements(src0),
-        (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size,
-        (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size,
-        (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] /  dst_type_size, (uint32_t) dst->nb[1] /  dst_type_size, (uint32_t) dst->nb[2] /  dst_type_size, (uint32_t) dst->nb[3] /  dst_type_size,
-        0,
-        0.0f, 0.0f, 0,
-    });
-}
-
 void ggml_vk_arange(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
     VK_LOG_DEBUG("ggml_vk_arange(dst=" << dst << ", ne=" << ggml_nelements(dst) << ")");
 
@@ -11569,11 +11574,6 @@ void ggml_vk_glu(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_t
         });
 }
 
-void ggml_vk_diag_mask_inf(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
-    int32_t * op_params = (int32_t *)dst->op_params;
-    ggml_vk_op_f32<vk_op_diag_mask_push_constants>(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_DIAG_MASK_INF, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0] });
-}
-
 void ggml_vk_soft_max(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) {
     float * op_params = (float *)dst->op_params;
 
@@ -12945,10 +12945,6 @@ bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, in
     case GGML_OP_UPSCALE:
         ggml_vk_upscale(ctx, compute_ctx, src0, node);
 
-        break;
-    case GGML_OP_ADD1:
-        ggml_vk_add1(ctx, compute_ctx, src0, src1, node);
-
         break;
     case GGML_OP_ARANGE:
         ggml_vk_arange(ctx, compute_ctx, node);
@@ -13099,10 +13095,6 @@ bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, in
         default:
             return false;
         }
-        break;
-    case GGML_OP_DIAG_MASK_INF:
-        ggml_vk_diag_mask_inf(ctx, compute_ctx, src0, node);
-
         break;
     case GGML_OP_SOFT_MAX:
         if (ctx->fused_topk_moe_mode != TOPK_MOE_COUNT) {
@@ -13263,6 +13255,11 @@ bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, in
 
         break;
 
+    case GGML_OP_TOP_K_MASK:
+        ggml_vk_top_k_mask(ctx, compute_ctx, node);
+
+        break;
+
     case GGML_OP_GATED_DELTA_NET:
         ggml_vk_gated_delta_net(ctx, compute_ctx, node);
 
@@ -13407,6 +13404,10 @@ void ggml_vk_graph_cleanup(ggml_backend_vk_context * ctx) {
 
 void ggml_vk_cleanup(ggml_backend_vk_context * ctx) {
     VK_LOG_DEBUG("ggml_vk_cleanup(" << ctx->name << ")");
+    if (ctx->keepalive_thread.joinable()) {
+        ctx->keepalive_stop.set_value();
+        ctx->keepalive_thread.join();
+    }
     // discard any unsubmitted command buffers
     ctx->compute_ctx.reset();
     // wait for any pending command buffers to finish
@@ -15773,11 +15774,34 @@ static ggml_guid_t ggml_backend_vk_guid() {
     return &guid;
 }
 
+static void ggml_vk_start_keepalive(ggml_backend_vk_context * ctx) {
+    const char * env = getenv("GGML_VK_KEEPALIVE");
+    const int s = env ? atoi(env) : 0;
+    if (s <= 0) {
+        return;
+    }
+
+    GGML_LOG_INFO("ggml_vulkan: GGML_VK_KEEPALIVE=%ds on %s\n", s, ctx->device->name.c_str());
+
+    vk_buffer buf = ggml_vk_create_buffer_check(ctx->device, 4, vk::MemoryPropertyFlagBits::eDeviceLocal);
+    ctx->keepalive_thread = std::thread([buf, stop = ctx->keepalive_stop.get_future(), s]() mutable {
+        while (stop.wait_for(std::chrono::seconds(s)) == std::future_status::timeout) {
+            try {
+                ggml_vk_buffer_memset(buf, 0, 0, 4);
+            } catch (const std::exception & e) {
+                GGML_LOG_WARN("ggml_vulkan: keepalive stopped: %s\n", e.what());
+                return;
+            }
+        }
+    });
+}
+
 ggml_backend_t ggml_backend_vk_init(size_t dev_num) {
     VK_LOG_DEBUG("ggml_backend_vk_init(" << dev_num << ")");
 
     ggml_backend_vk_context * ctx = new ggml_backend_vk_context;
     ggml_vk_init(ctx, dev_num);
+    ggml_vk_start_keepalive(ctx);
 
     ggml_backend_t vk_backend = new ggml_backend {
         /* .guid    = */ ggml_backend_vk_guid(),
@@ -16289,11 +16313,11 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
         case GGML_OP_VIEW:
         case GGML_OP_PERMUTE:
         case GGML_OP_TRANSPOSE:
-        case GGML_OP_RMS_NORM:
             return true;
         case GGML_OP_GROUP_NORM:
             return ggml_is_contiguous(op->src[0]);
         case GGML_OP_NORM:
+        case GGML_OP_RMS_NORM:
         case GGML_OP_L2_NORM:
             return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
         case GGML_OP_ADD:
@@ -16369,10 +16393,6 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
         case GGML_OP_CONCAT: {
             return ggml_vk_concat_supported(op->src[0], op->src[1], op);
         }
-        case GGML_OP_ADD1:
-            return (op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32)
-                || (op->src[0]->type == GGML_TYPE_F16 && op->src[1]->type == GGML_TYPE_F32)
-                || (op->src[0]->type == GGML_TYPE_F16 && op->src[1]->type == GGML_TYPE_F16);
         case GGML_OP_ARANGE:
             return op->type == GGML_TYPE_F32;
         case GGML_OP_FILL:
@@ -16383,8 +16403,6 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
         case GGML_OP_PAD_REFLECT_1D:
         case GGML_OP_ROLL:
             return op->src[0]->type == GGML_TYPE_F32;
-        case GGML_OP_DIAG_MASK_INF:
-            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_SOFT_MAX:
             return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32
                 && (!op->src[1] || (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16));
@@ -16519,6 +16537,30 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
                 }
                 return true;
             }
+        case GGML_OP_TOP_K_MASK:
+            {
+                const ggml_tensor * mask = op->src[0];
+                if ((op->type != GGML_TYPE_F16 && op->type != GGML_TYPE_F32) || mask->type != op->type) {
+                    return false;
+                }
+
+                // one workgroup per row
+                const auto & max_wg = device->properties.limits.maxComputeWorkGroupCount;
+                if (op->ne[1] > max_wg[0] || op->ne[2] > max_wg[1] || op->ne[3] > max_wg[2]) {
+                    return false;
+                }
+
+                // the shader expects contiguous rows in mask and dst, and contiguous tensors in the other inputs
+                if (!ggml_is_contiguous_rows(op) || !ggml_is_contiguous_rows(mask)) {
+                    return false;
+                }
+                for (int i = 1; i <= 4; ++i) {
+                    if (op->src[i] && !ggml_is_contiguous(op->src[i])) {
+                        return false;
+                    }
+                }
+                return true;
+            }
         case GGML_OP_GATED_DELTA_NET:
             {
                 const uint32_t S_v = op->src[2]->ne[0];

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

@@ -1,34 +0,0 @@
-#version 450
-
-#extension GL_EXT_shader_16bit_storage : require
-#extension GL_EXT_control_flow_attributes : enable
-
-layout (push_constant) uniform parameter
-{
-    uint ncols;
-    uint rows_per_channel;
-    uint n_past;
-} p;
-
-#include "types.glsl"
-
-layout(local_size_x = 1, local_size_y = 512, local_size_z = 1) in;
-
-layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
-layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
-
-void main() {
-    const uint col = gl_GlobalInvocationID.y;
-    const uint row = gl_GlobalInvocationID.x;
-
-    if (col >= p.ncols) {
-        return;
-    }
-
-    const uint i = row*p.ncols + col;
-    if (col > p.n_past + row % p.rows_per_channel) {
-        data_d[i] = D_TYPE(uintBitsToFloat(0xFF800000));
-    } else {
-        data_d[i] = D_TYPE(data_a[i]);
-    }
-}

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

@@ -0,0 +1,95 @@
+#version 450
+
+#extension GL_EXT_shader_16bit_storage : require
+#extension GL_KHR_memory_scope_semantics : enable
+
+// One workgroup per row: fill the row with -inf, then copy the mask cells of the live picks and the extra cells.
+// The op only moves bits, so f16 needs no float16 arithmetic and the score check does not depend on float controls.
+
+layout(constant_id = 0) const uint BLOCK_SIZE = 256;
+
+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+layout(push_constant) uniform parameter
+{
+    uint n_kv;
+    uint n_top;
+    uint n_pick;
+    uint n_blk;   // 0: a pick is a cell
+    uint n_extra;
+    uint has_score;
+    uint ninf;    // bits of -inf: a 16-bit literal would need shaderInt16
+
+    uint ne1; uint ne2;
+
+    uint m1; uint m2; uint m3;
+    uint nbm1; uint nbm2; uint nbm3;
+    uint nbd1; uint nbd2; uint nbd3;
+
+    uint m_offset;
+    uint k_offset;
+    uint s_offset;
+    uint b_offset;
+    uint e_offset;
+    uint d_offset;
+};
+
+layout(binding = 0, std430) readonly buffer M { MASK_TYPE data_m[]; };
+layout(binding = 1, std430) readonly buffer K { int data_k[]; };
+layout(binding = 2, std430) readonly buffer S { uint data_s[]; };
+layout(binding = 3, std430) readonly buffer B { int data_b[]; };
+layout(binding = 4, std430) readonly buffer E { int data_e[]; };
+// coherent: the fill and the copy of a cell can come from different invocations
+layout(binding = 5, std430) coherent writeonly buffer D { MASK_TYPE data_d[]; };
+
+void keep(const int c, const uint m_row, const uint d_row) {
+    if (c >= 0 && uint(c) < n_kv) {
+        data_d[d_row + uint(c)] = data_m[m_row + uint(c)];
+    }
+}
+
+// score > -inf
+bool is_live(const uint s) {
+    return s != 0xFF800000u && (s & 0x7FFFFFFFu) <= 0x7F800000u;
+}
+
+void main() {
+    const uint tid = gl_LocalInvocationIndex;
+
+    const uint i1 = gl_WorkGroupID.x;
+    const uint i2 = gl_WorkGroupID.y;
+    const uint i3 = gl_WorkGroupID.z;
+
+    const uint row   = i1 + ne1*(i2 + ne2*i3);
+    const uint m_row = m_offset + (i1 % m1)*nbm1 + (i2 % m2)*nbm2 + (i3 % m3)*nbm3;
+    const uint d_row = d_offset + i1*nbd1 + i2*nbd2 + i3*nbd3;
+
+    for (uint i = tid; i < n_kv; i += BLOCK_SIZE) {
+        data_d[d_row + i] = MASK_TYPE(ninf);
+    }
+
+    controlBarrier(gl_ScopeWorkgroup, gl_ScopeWorkgroup, gl_StorageSemanticsBuffer, gl_SemanticsAcquireRelease);
+
+    // one invocation per pick or extra cell; cells kept twice get the same value
+    for (uint i = tid; i < n_top + n_extra; i += BLOCK_SIZE) {
+        if (i >= n_top) {
+            keep(data_e[e_offset + row*n_extra + (i - n_top)], m_row, d_row);
+            continue;
+        }
+
+        const int p = data_k[k_offset + row*n_top + i];
+        if (p < 0 || uint(p) >= n_pick) {
+            continue;
+        }
+        if (has_score != 0 && !is_live(data_s[s_offset + row*n_pick + uint(p)])) {
+            continue;
+        }
+        if (n_blk == 0) {
+            keep(p, m_row, d_row);
+            continue;
+        }
+        for (uint k = 0; k < n_blk; ++k) {
+            keep(data_b[b_offset + uint(p)*n_blk + k], m_row, d_row);
+        }
+    }
+}

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

@@ -1075,8 +1075,6 @@ void process_shaders() {
 
     string_to_spv("silu_back_f32",  "silu_back.comp",   {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}});
 
-    string_to_spv("diag_mask_inf_f32", "diag_mask_inf.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
-
     string_to_spv("soft_max_f32", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
     string_to_spv("soft_max_f32_f16", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}}));
     string_to_spv("soft_max_back_f32", "soft_max_back.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
@@ -1122,6 +1120,8 @@ void process_shaders() {
     string_to_spv("dsv4_hc_comb_f32", "dsv4_hc_comb.comp", {});
     string_to_spv("dsv4_hc_pre_f32",  "dsv4_hc_pre.comp",  {});
     string_to_spv("dsv4_hc_post_f32", "dsv4_hc_post.comp", {});
+    string_to_spv("top_k_mask_f16", "top_k_mask.comp", {{"MASK_TYPE", "uint16_t"}});
+    string_to_spv("top_k_mask_f32", "top_k_mask.comp", {{"MASK_TYPE", "uint"}});
     string_to_spv("cumsum_f32", "cumsum.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
     string_to_spv("cumsum_multipass1_f32", "cumsum_multipass1.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
     string_to_spv("cumsum_multipass2_f32", "cumsum_multipass2.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));

ggml/src/ggml.c

@@ -714,13 +714,13 @@ static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = {
         .from_float_ref           = (ggml_from_float_t) quantize_row_q4_1_ref,
     },
     [4] = { // GGML_TYPE_Q4_2
-        .type_name                = "DEPRECATED",
+        .type_name                = "REMOVED",
         .blck_size                = 0,
         .type_size                = 0,
         .is_quantized             = false,
     },
     [5] = { // GGML_TYPE_Q4_3
-        .type_name                = "DEPRECATED",
+        .type_name                = "REMOVED",
         .blck_size                = 0,
         .type_size                = 0,
         .is_quantized             = false,
@@ -1009,7 +1009,6 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "DUP",
     "ADD",
     "ADD_ID",
-    "ADD1",
     "ACC",
     "SUB",
     "MUL",
@@ -1051,8 +1050,6 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "GET_ROWS_BACK",
     "SET_ROWS",
     "DIAG",
-    "DIAG_MASK_INF",
-    "DIAG_MASK_ZERO",
     "SOFT_MAX",
     "SOFT_MAX_BACK",
     "ROPE",
@@ -1117,7 +1114,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "GLU",
 };
 
-static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102");
+static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99");
 
 static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "none",
@@ -1125,7 +1122,6 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "x",
     "x+y",
     "x[i]+y",
-    "x+y",
     "view(x,nb,offset)+=y->x",
     "x-y",
     "x*y",
@@ -1167,8 +1163,6 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "get_rows_back(x)",
     "set_rows(x)",
     "diag(x)",
-    "diag_mask_inf(x)",
-    "diag_mask_zero(x)",
     "soft_max(x)",
     "soft_max_back(x)",
     "rope(x)",
@@ -1233,7 +1227,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "glu(x)",
 };
 
-static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102");
+static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99");
 
 static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
 
@@ -1360,12 +1354,6 @@ size_t ggml_row_size(enum ggml_type type, int64_t ne) {
     return ggml_type_size(type)*ne/ggml_blck_size(type);
 }
 
-double ggml_type_sizef(enum ggml_type type) {
-    assert(type >= 0);
-    assert(type < GGML_TYPE_COUNT);
-    return ((double)(type_traits[type].type_size))/type_traits[type].blck_size;
-}
-
 const char * ggml_type_name(enum ggml_type type) {
     assert(type >= 0);
     assert(type < GGML_TYPE_COUNT);
@@ -2172,39 +2160,6 @@ struct ggml_tensor * ggml_add_id(
     return result;
 }
 
-// ggml_add1
-
-static struct ggml_tensor * ggml_add1_impl(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        struct ggml_tensor  * b,
-        bool                  inplace) {
-    GGML_ASSERT(ggml_is_scalar(b));
-    GGML_ASSERT(ggml_is_padded_1d(a));
-
-    struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a);
-
-    result->op     = GGML_OP_ADD1;
-    result->src[0] = a;
-    result->src[1] = b;
-
-    return result;
-}
-
-struct ggml_tensor * ggml_add1(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        struct ggml_tensor  * b) {
-    return ggml_add1_impl(ctx, a, b, false);
-}
-
-struct ggml_tensor * ggml_add1_inplace(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        struct ggml_tensor  * b) {
-    return ggml_add1_impl(ctx, a, b, true);
-}
-
 // ggml_acc
 
 static struct ggml_tensor * ggml_acc_impl(
@@ -3419,16 +3374,6 @@ struct ggml_tensor * ggml_mul_mat(
     return result;
 }
 
-void ggml_mul_mat_set_prec(
-        struct ggml_tensor * a,
-        enum ggml_prec       prec) {
-    GGML_ASSERT(a->op == GGML_OP_MUL_MAT);
-
-    const int32_t prec_i32 = (int32_t) prec;
-
-    ggml_set_op_params_i32(a, 0, prec_i32);
-}
-
 void ggml_mul_mat_set_hint(
         struct ggml_tensor * a,
         enum ggml_op_hint    hint) {
@@ -4102,70 +4047,6 @@ struct ggml_tensor * ggml_diag(
     return result;
 }
 
-// ggml_diag_mask_inf
-
-static struct ggml_tensor * ggml_diag_mask_inf_impl(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        int                   n_past,
-        bool                  inplace) {
-    struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a);
-
-    int32_t params[] = { n_past };
-    ggml_set_op_params(result, params, sizeof(params));
-
-    result->op     = GGML_OP_DIAG_MASK_INF;
-    result->src[0] = a;
-
-    return result;
-}
-
-struct ggml_tensor * ggml_diag_mask_inf(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        int                   n_past) {
-    return ggml_diag_mask_inf_impl(ctx, a, n_past, false);
-}
-
-struct ggml_tensor * ggml_diag_mask_inf_inplace(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        int                   n_past) {
-    return ggml_diag_mask_inf_impl(ctx, a, n_past, true);
-}
-
-// ggml_diag_mask_zero
-
-static struct ggml_tensor * ggml_diag_mask_zero_impl(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        int                   n_past,
-        bool                  inplace) {
-    struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a);
-
-    int32_t params[] = { n_past };
-    ggml_set_op_params(result, params, sizeof(params));
-
-    result->op     = GGML_OP_DIAG_MASK_ZERO;
-    result->src[0] = a;
-
-    return result;
-}
-
-struct ggml_tensor * ggml_diag_mask_zero(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        int                   n_past) {
-    return ggml_diag_mask_zero_impl(ctx, a, n_past, false);
-}
-
-struct ggml_tensor * ggml_diag_mask_zero_inplace(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        int                   n_past) {
-    return ggml_diag_mask_zero_impl(ctx, a, n_past, true);
-}
-
 // ggml_clamp
 
 static struct ggml_tensor * ggml_clamp_impl(
@@ -4487,44 +4368,6 @@ struct ggml_tensor * ggml_rope_ext_inplace(
     );
 }
 
-struct ggml_tensor * ggml_rope_custom(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        struct ggml_tensor  * b,
-        int                   n_dims,
-        int                   mode,
-        int                   n_ctx_orig,
-        float                 freq_base,
-        float                 freq_scale,
-        float                 ext_factor,
-        float                 attn_factor,
-        float                 beta_fast,
-        float                 beta_slow) {
-    return ggml_rope_impl(
-        ctx, a, b, NULL, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale,
-        ext_factor, attn_factor, beta_fast, beta_slow, false
-    );
-}
-
-struct ggml_tensor * ggml_rope_custom_inplace(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        struct ggml_tensor  * b,
-        int                   n_dims,
-        int                   mode,
-        int                   n_ctx_orig,
-        float                 freq_base,
-        float                 freq_scale,
-        float                 ext_factor,
-        float                 attn_factor,
-        float                 beta_fast,
-        float                 beta_slow) {
-    return ggml_rope_impl(
-        ctx, a, b, NULL, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale,
-        ext_factor, attn_factor, beta_fast, beta_slow, true
-    );
-}
-
 // Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get
 // `corr_dim(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))`
 static float ggml_rope_yarn_corr_dim(int n_dims, int n_ctx_orig, float n_rot, float base) {
@@ -5264,17 +5107,6 @@ struct ggml_tensor * ggml_upscale(
     return ggml_interpolate_impl(ctx, a, a->ne[0] * scale_factor, a->ne[1] * scale_factor, a->ne[2], a->ne[3], mode);
 }
 
-struct ggml_tensor * ggml_upscale_ext(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        int                   ne0,
-        int                   ne1,
-        int                   ne2,
-        int                   ne3,
-        enum ggml_scale_mode  mode) {
-    return ggml_interpolate_impl(ctx, a, ne0, ne1, ne2, ne3, mode);
-}
-
 struct ggml_tensor * ggml_interpolate(
         struct ggml_context * ctx,
         struct ggml_tensor  * a,
@@ -5654,16 +5486,6 @@ struct ggml_tensor * ggml_flash_attn_ext_rows(
 }
 
 
-void ggml_flash_attn_ext_set_prec(
-        struct ggml_tensor * a,
-        enum ggml_prec       prec) {
-    GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT);
-
-    const int32_t prec_i32 = (int32_t) prec;
-
-    ggml_set_op_params_i32(a, 3, prec_i32); // scale is on first pos, max_bias on second
-}
-
 enum ggml_prec ggml_flash_attn_ext_get_prec(
         const struct ggml_tensor * a) {
     GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT);
@@ -6873,22 +6695,6 @@ static void ggml_acc_or_set(
     ggml_build_forward_expand(cgraph, cgraph->grads[isrc]);
 }
 
-static void ggml_add1_or_set(
-        struct ggml_context * ctx,
-        struct ggml_cgraph  * cgraph,
-        size_t                isrc,
-        struct ggml_tensor  * tensor) {
-    struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc];
-    GGML_ASSERT(src);
-    if (cgraph->grads[isrc]) {
-        cgraph->grads[isrc] = ggml_add1_impl(ctx, cgraph->grads[isrc], tensor, cgraph->grad_accs[isrc]);
-    } else {
-        cgraph->grads[isrc] = ggml_repeat(ctx, tensor, src);
-    }
-    ggml_format_name(cgraph->grads[isrc], "grad for %s", src->name);
-    ggml_build_forward_expand(cgraph, cgraph->grads[isrc]);
-}
-
 static void ggml_sub_or_set(
         struct ggml_context * ctx,
         struct ggml_cgraph  * cgraph,
@@ -6943,14 +6749,6 @@ static void ggml_compute_backward(
                 ggml_add_or_set(ctx, cgraph, isrc1, tmp);
             }
         } break;
-        case GGML_OP_ADD1: {
-            if (src0_needs_grads) {
-                ggml_add_or_set(ctx, cgraph, isrc0, grad);
-            }
-            if (src1_needs_grads) {
-                ggml_add_or_set(ctx, cgraph, isrc1, ggml_mean(ctx, grad)); // TODO: should probably be sum instead of mean
-            }
-        } break;
         case GGML_OP_ACC: {
             if (src0_needs_grads) {
                 ggml_add_or_set(ctx, cgraph, isrc0, grad);
@@ -7023,7 +6821,7 @@ static void ggml_compute_backward(
         } break;
         case GGML_OP_SUM: {
             if (src0_needs_grads) {
-                ggml_add1_or_set(ctx, cgraph, isrc0, grad);
+                ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat(ctx, grad, src0));
             }
         } break;
         case GGML_OP_SUM_ROWS: {
@@ -7033,7 +6831,7 @@ static void ggml_compute_backward(
         } break;
         case GGML_OP_MEAN: {
             if (src0_needs_grads) {
-                ggml_add1_or_set(ctx, cgraph, isrc0, ggml_scale_impl(ctx, grad, 1.0f/src0->ne[0], 0.0, false));
+                ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat(ctx, ggml_scale_impl(ctx, grad, 1.0f/src0->ne[0], 0.0, false), src0));
             }
         } break;
         case GGML_OP_REPEAT: {
@@ -7223,20 +7021,6 @@ static void ggml_compute_backward(
                 // noop
             }
         } break;
-        case GGML_OP_DIAG_MASK_INF: {
-            if (src0_needs_grads) {
-                /* ggml_diag_mask_inf_impl() shouldn't be here */
-                /* ref:  https://github.com/ggml-org/llama.cpp/pull/4203#discussion_r1412377992 */
-                const int n_past = ((const int32_t *) tensor->op_params)[0];
-                ggml_add_or_set(ctx, cgraph, isrc0, ggml_diag_mask_zero_impl(ctx, grad, n_past, false));
-            }
-        } break;
-        case GGML_OP_DIAG_MASK_ZERO: {
-            if (src0_needs_grads) {
-                const int n_past = ((const int32_t *) tensor->op_params)[0];
-                ggml_add_or_set(ctx, cgraph, isrc0, ggml_diag_mask_zero_impl(ctx, grad, n_past, false));
-            }
-        } break;
         case GGML_OP_TRI: {
             if (src0_needs_grads) {
                 ggml_add_or_set(ctx, cgraph, isrc0, ggml_tri(ctx, grad, (enum ggml_tri_type) ggml_get_op_params_i32(tensor, 0)));

include/llama.h

@@ -270,7 +270,7 @@ extern "C" {
         llama_pos    *  pos;
         int32_t      *  n_seq_id;
         llama_seq_id ** seq_id;
-        int8_t       *  logits;   // TODO: rename this to "output"
+        int8_t       *  logits;
     } llama_batch;
 
     enum llama_model_kv_override_type {
@@ -389,7 +389,6 @@ extern "C" {
         float    yarn_beta_fast;   // YaRN low correction dim
         float    yarn_beta_slow;   // YaRN high correction dim
         uint32_t yarn_orig_ctx;    // YaRN original context size
-        float    defrag_thold;     // [DEPRECATED] defragment the KV cache if holes/size > thold, <= 0 disabled (default)
 
         ggml_backend_sched_eval_callback cb_eval;
         void * cb_eval_user_data;
@@ -516,11 +515,6 @@ extern "C" {
                                    void * set_tensor_data_ud, // userdata for function
               struct llama_model_params   params);
 
-    DEPRECATED(LLAMA_API struct llama_model * llama_load_model_from_file(
-                             const char * path_model,
-              struct llama_model_params   params),
-            "use llama_model_load_from_file instead");
-
     // Load a model from a file
     // If the file is split into multiple parts, the file name must follow this pattern: <name>-%05d-of-%05d.gguf
     // If the split file name does not follow this pattern, use llama_model_load_from_splits
@@ -546,20 +540,12 @@ extern "C" {
             const struct llama_model * model,
                         const char * path_model);
 
-    DEPRECATED(LLAMA_API void llama_free_model(struct llama_model * model),
-            "use llama_model_free instead");
-
     LLAMA_API void llama_model_free(struct llama_model * model);
 
     LLAMA_API struct llama_context * llama_init_from_model(
                      struct llama_model * model,
             struct llama_context_params   params);
 
-    DEPRECATED(LLAMA_API struct llama_context * llama_new_context_with_model(
-                     struct llama_model * model,
-            struct llama_context_params   params),
-            "use llama_init_from_model instead");
-
     // Frees all allocated memory
     LLAMA_API void llama_free(struct llama_context * ctx);
 
@@ -584,13 +570,6 @@ extern "C" {
     LLAMA_API uint32_t llama_n_seq_max  (const struct llama_context * ctx);
     LLAMA_API uint32_t llama_n_rs_seq   (const struct llama_context * ctx);
 
-    DEPRECATED(LLAMA_API int32_t llama_n_ctx_train(const struct llama_model * model), "use llama_model_n_ctx_train instead");
-    DEPRECATED(LLAMA_API int32_t llama_n_embd     (const struct llama_model * model), "use llama_model_n_embd instead");
-    DEPRECATED(LLAMA_API int32_t llama_n_layer    (const struct llama_model * model), "use llama_model_n_layer instead");
-    DEPRECATED(LLAMA_API int32_t llama_n_head     (const struct llama_model * model), "use llama_model_n_head instead");
-
-    DEPRECATED(LLAMA_API int32_t llama_n_vocab    (const struct llama_vocab * vocab), "use llama_vocab_n_tokens instead");
-
     LLAMA_API const struct llama_model * llama_get_model   (const struct llama_context * ctx);
     LLAMA_API           llama_memory_t   llama_get_memory  (const struct llama_context * ctx);
     LLAMA_API  enum llama_pooling_type   llama_pooling_type(const struct llama_context * ctx); // TODO: rename to llama_get_pooling_type
@@ -684,6 +663,19 @@ extern "C" {
             const char * fname_out,
             const llama_model_quantize_params * params);
 
+    //
+    // NextN / speculative draft model queries
+    //
+
+    // DFlash2 selector: number of candidate tokens the draft keeps per block position from its logits
+    // (the packed lattice also holds the top_k x top_k pairwise transition scores); 0 if the model has no selector head
+    LLAMA_API int32_t llama_model_dflash_selector_top_k(const struct llama_model * model);
+
+    // returns pointer to the target-model layer indices
+    LLAMA_API const int32_t * llama_model_target_layer_ids  (const struct llama_model * model);
+    // returns the number of extracted layers from target model
+    LLAMA_API uint32_t        llama_model_target_layer_ids_n(const struct llama_model * model);
+
     //
     // Adapters
     //
@@ -830,8 +822,6 @@ extern "C" {
     // (logits, embedding and memory)
     // Only use when saving the state, not when restoring it, otherwise the size may be too small.
     LLAMA_API size_t llama_state_get_size(struct llama_context * ctx);
-    LLAMA_API DEPRECATED(size_t llama_get_state_size(struct llama_context * ctx),
-        "use llama_state_get_size instead");
 
     // Copies the state to the specified destination address.
     // Destination needs to have allocated enough memory.
@@ -840,10 +830,6 @@ extern "C" {
             struct llama_context * ctx,
                          uint8_t * dst,
                           size_t   size);
-    LLAMA_API DEPRECATED(size_t llama_copy_state_data(
-            struct llama_context * ctx,
-                         uint8_t * dst),
-        "use llama_state_get_data instead");
 
     // Set the state reading from the specified address
     // Returns the number of bytes read
@@ -851,10 +837,6 @@ extern "C" {
             struct llama_context * ctx,
                    const uint8_t * src,
                           size_t   size);
-    LLAMA_API DEPRECATED(size_t llama_set_state_data(
-            struct llama_context * ctx,
-                   const uint8_t * src),
-        "use llama_state_set_data instead");
 
     // Save/load session file
     LLAMA_API bool llama_state_load_file(
@@ -863,25 +845,12 @@ extern "C" {
                      llama_token * tokens_out,
                           size_t   n_token_capacity,
                           size_t * n_token_count_out);
-    LLAMA_API DEPRECATED(bool llama_load_session_file(
-            struct llama_context * ctx,
-                      const char * path_session,
-                     llama_token * tokens_out,
-                          size_t   n_token_capacity,
-                          size_t * n_token_count_out),
-        "use llama_state_load_file instead");
 
     LLAMA_API bool llama_state_save_file(
             struct llama_context * ctx,
                       const char * path_session,
                const llama_token * tokens,
                           size_t   n_token_count);
-    LLAMA_API DEPRECATED(bool llama_save_session_file(
-            struct llama_context * ctx,
-                      const char * path_session,
-               const llama_token * tokens,
-                          size_t   n_token_count),
-        "use llama_state_save_file instead");
 
     // Get the exact size needed to copy the state of a single sequence
     LLAMA_API size_t llama_state_seq_get_size(
@@ -1065,6 +1034,15 @@ extern "C" {
                                                int32_t   idx,
                                      struct llama_embd   embd);
 
+    // Add an entry whose embedding is row `row` of the layer input device buffer of ctx_other.
+    // All entries of the batch must be of this kind. Positions are set with llama_batch_ext_set_pos(), n_pos_per_embd of them.
+    // Returns the batch index (>= 0), or the llama_batch_ext_add() error codes.
+    LLAMA_API int32_t llama_batch_ext_add_embd_row(struct llama_batch_ext * batch, llama_seq_id seq_id, int32_t row);
+
+    // The state of entry idx is row `row` of the h_nextn device buffer of the context set with llama_set_ctx_state_src().
+    // All entries of a batch have a state row, or none.
+    LLAMA_API bool llama_batch_ext_set_embd_state_row(struct llama_batch_ext * batch, int32_t idx, int32_t row);
+
     // Set if output embeddings should be available for the token at index idx in the batch
     // Note: for now, this is equivalent to setting the output logits
     LLAMA_API bool llama_batch_ext_set_output_embd(
@@ -1118,14 +1096,6 @@ extern "C" {
     // Returns whether the context is currently using causal attention
     LLAMA_API bool llama_get_causal_attn(const struct llama_context * ctx);
 
-    // Set whether the model is in warmup mode or not
-    // If true, all model tensors are activated during llama_decode() to load and cache their weights.
-    //
-    // note: using this can cause extra graph reallocations because it changes the graph topology with MoE models,
-    //       so it is generally not recommended to use in practice. will be removed in the future
-    DEPRECATED(LLAMA_API void llama_set_warmup(struct llama_context * ctx, bool warmup),
-            "user code should do warmup runs manually [TAG_LLAMA_GRAPH_NO_WARMUP]");
-
     // Set abort callback
     LLAMA_API void llama_set_abort_callback(struct llama_context * ctx, ggml_abort_callback abort_callback, void * abort_callback_data);
 
@@ -1170,6 +1140,57 @@ extern "C" {
     // otherwise: float[n_embd] (1-dimensional)
     LLAMA_API float * llama_get_embeddings_seq(struct llama_context * ctx, llama_seq_id seq_id);
 
+    //
+    // NextN / speculative decoding API
+    //
+
+    // Set whether the context outputs nextn embeddings or not
+    // If masked == true,  output the embeddings only for the tokens with batch.logits != 0
+    // If masked == false, output the embeddings for all tokens in the batch regardless of batch.logits
+    LLAMA_API void llama_set_embeddings_nextn(struct llama_context * ctx, bool value, bool masked);
+
+    // Select which appended NextN block the DECODER_MTP graph runs (offset past
+    // the trunk: il = n_layer() + offset). Used by the speculative NextN driver to
+    // chain multiple trained NextN heads. Default 0 (first head).
+    LLAMA_API void llama_set_nextn_layer_offset(struct llama_context * ctx, int32_t offset);
+
+    // Get all nextn embeddings, mirroring llama_get_embeddings()
+    LLAMA_API float * llama_get_embeddings_nextn(struct llama_context * ctx);
+
+    // Get the nextn embeddings for the ith token, mirroring llama_get_embeddings_ith()
+    LLAMA_API float * llama_get_embeddings_nextn_ith(struct llama_context * ctx, int32_t i);
+
+    // Set whether the context outputs the input embeddings of a specific layer
+    LLAMA_API void llama_set_embeddings_layer_inp(struct llama_context * ctx, uint32_t lid, bool value);
+
+    // Get the input embeddings of a specific layer, mirroring llama_get_embeddings()
+    LLAMA_API float * llama_get_embeddings_layer_inp(struct llama_context * ctx, uint32_t lid);
+
+    // Write the enabled layer inputs to a device buffer [n_layer_inp*n_embd, n_batch] instead of host memory.
+    // Row r holds the token at batch index r, layers in ascending layer id order.
+    LLAMA_API void llama_set_embeddings_layer_inp_dev(struct llama_context * ctx, bool value);
+
+    // Whether decode accepts entries added with llama_batch_ext_add_embd_row()
+    // (needs a DFlash/DSpark decoder context whose ctx_other has a layer input device buffer).
+    LLAMA_API bool llama_supports_embd_row(const struct llama_context * ctx);
+
+    // Write the unmasked h_nextn rows (see llama_set_embeddings_nextn) of each decode to a device buffer
+    // [n_embd_out, n_batch + n_seq_max] instead of host memory: row r < n_batch is the token at batch index r of the last decode,
+    // row n_batch + s is free for the caller, e.g. for seq s (see llama_set_embeddings_nextn_rows()).
+    // Returns false if the buffer cannot be used, the rows are then extracted to host memory.
+    LLAMA_API bool llama_set_embeddings_nextn_dev(struct llama_context * ctx, bool value);
+
+    // Copy rows of the h_nextn device buffer to or from host memory, ordered after the last decode and its readers
+    LLAMA_API bool llama_get_embeddings_nextn_rows(struct llama_context * ctx, int64_t row0, int64_t n_rows, float * out);
+    LLAMA_API bool llama_set_embeddings_nextn_rows(struct llama_context * ctx, int64_t row0, int64_t n_rows, const float * data);
+
+    // Read the state rows of batches (see llama_batch_ext_set_embd_state_row()) from the h_nextn device buffer of ctx_src.
+    // Returns whether decode accepts batches with state rows.
+    LLAMA_API bool llama_set_ctx_state_src(struct llama_context * ctx, struct llama_context * ctx_src);
+
+    // The paired context (llama_context_params::ctx_other), e.g. the target context of a draft context; may be nullptr
+    LLAMA_API struct llama_context * llama_get_ctx_other(struct llama_context * ctx);
+
     //
     // backend sampling API [EXPERIMENTAL]
     // note: use only if the llama_context was created with at least one llama_sampler_seq_config
@@ -1238,31 +1259,6 @@ extern "C" {
     LLAMA_API llama_token llama_vocab_fim_rep(const struct llama_vocab * vocab);
     LLAMA_API llama_token llama_vocab_fim_sep(const struct llama_vocab * vocab);
 
-    DEPRECATED(LLAMA_API const char * llama_token_get_text(const struct llama_vocab * vocab, llama_token token), "use llama_vocab_get_text instead");
-    DEPRECATED(LLAMA_API float llama_token_get_score(const struct llama_vocab * vocab, llama_token token), "use llama_vocab_get_score instead");
-    DEPRECATED(LLAMA_API enum llama_token_attr llama_token_get_attr(const struct llama_vocab * vocab, llama_token token), "use llama_vocab_get_attr instead");
-    DEPRECATED(LLAMA_API bool llama_token_is_eog(const struct llama_vocab * vocab, llama_token token), "use llama_vocab_is_eog instead");
-    DEPRECATED(LLAMA_API bool llama_token_is_control(const struct llama_vocab * vocab, llama_token token), "use llama_vocab_is_control instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_bos(const struct llama_vocab * vocab), "use llama_vocab_bos instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_eos(const struct llama_vocab * vocab), "use llama_vocab_eos instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_eot(const struct llama_vocab * vocab), "use llama_vocab_eot instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_cls(const struct llama_vocab * vocab), "use llama_vocab_cls instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_sep(const struct llama_vocab * vocab), "use llama_vocab_sep instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_nl (const struct llama_vocab * vocab), "use llama_vocab_nl instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_pad(const struct llama_vocab * vocab), "use llama_vocab_pad instead");
-    DEPRECATED(LLAMA_API bool llama_add_bos_token(const struct llama_vocab * vocab), "use llama_vocab_get_add_bos instead");
-    DEPRECATED(LLAMA_API bool llama_add_eos_token(const struct llama_vocab * vocab), "use llama_vocab_get_add_eos instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_fim_pre(const struct llama_vocab * vocab), "use llama_vocab_fim_pre instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_fim_suf(const struct llama_vocab * vocab), "use llama_vocab_fim_suf instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_fim_mid(const struct llama_vocab * vocab), "use llama_vocab_fim_mid instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_fim_pad(const struct llama_vocab * vocab), "use llama_vocab_fim_pad instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_fim_rep(const struct llama_vocab * vocab), "use llama_vocab_fim_rep instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_fim_sep(const struct llama_vocab * vocab), "use llama_vocab_fim_sep instead");
-
-    // CLS is equivalent to BOS
-    DEPRECATED(LLAMA_API llama_token llama_vocab_cls(const struct llama_vocab * vocab), // classification
-            "use llama_vocab_bos instead");
-
     //
     // Tokenization
     //
@@ -1534,17 +1530,6 @@ extern "C" {
                           const char * grammar_str,
                           const char * grammar_root);
 
-    DEPRECATED(LLAMA_API struct llama_sampler * llama_sampler_init_grammar_lazy(
-            const struct llama_vocab * vocab,
-                          const char * grammar_str,
-                          const char * grammar_root,
-                         const char ** trigger_words,
-                                size_t num_trigger_words,
-                   const llama_token * trigger_tokens,
-                                size_t num_trigger_tokens),
-        "use llama_sampler_init_grammar_lazy_patterns instead");
-
-
     /// @details Lazy grammar sampler, introduced in https://github.com/ggml-org/llama.cpp/pull/9639
     /// @param trigger_patterns A list of patterns that will trigger the grammar sampler. Pattern will be matched from the start of the generation output, and grammar sampler will be fed content starting from its first match group.
     /// @param trigger_tokens A list of tokens that will trigger the grammar sampler. Grammar sampler will be fed content starting from the trigger token included.

pyproject.toml

@@ -6,7 +6,7 @@ version = "0.0.0"
 dynamic = ["classifiers"]
 readme = "README.md"
 authors = [{name = "GGML", email = "ggml@ggml.ai"}]
-requires-python = '>=3.10,<3.15'
+requires-python = '>=3.10,<3.16'
 dependencies = [
     'numpy (>=2.2.6,<3.0.0); python_version < "3.14"',
     'numpy (>=2.3.0,<3.0.0); python_version >= "3.14"',
@@ -40,9 +40,12 @@ packages = [
 
 [tool.poetry.dependencies]
 torch = [
-    { version = "==2.11.0", source = "pypi", markers = "sys_platform == 'darwin'" },
-    { version = "==2.11.0+cpu", source = "pytorch", markers = "sys_platform == 'linux'" },
-    { version = "==2.11.0", source = "pypi", markers = "sys_platform == 'win32'" },
+    { version = "==2.11.0", source = "pypi", markers = "sys_platform == 'darwin' and python_version < '3.15'" },
+    { version = "==2.11.0+cpu", source = "pytorch", markers = "sys_platform == 'linux' and python_version < '3.15'" },
+    { version = "==2.11.0", source = "pypi", markers = "sys_platform == 'win32' and python_version < '3.15'" },
+    { version = "==2.15.0", source = "pypi", markers = "sys_platform == 'darwin' and python_version >= '3.15'" },
+    { version = "==2.15.0+cpu", source = "pytorch", markers = "sys_platform == 'linux' and python_version >= '3.15'" },
+    { version = "==2.15.0", source = "pypi", markers = "sys_platform == 'win32' and python_version >= '3.15'" },
 ]
 
 [tool.poetry.group.dev.dependencies]

requirements/requirements-all.txt

@@ -4,7 +4,6 @@
 
 -r ./requirements-compare-llama-bench.txt
 -r ./requirements-server-bench.txt
--r ./requirements-pydantic.txt
 -r ./requirements-test-tokenizer-random.txt
 
 -r ./requirements-convert_hf_to_gguf.txt

requirements/requirements-compare-llama-bench.txt

@@ -1,3 +1,4 @@
 tabulate~=0.9.0
 GitPython~=3.1.43
-matplotlib~=3.10.0
+matplotlib~=3.10.0; python_version < "3.11"
+matplotlib~=3.11.2; python_version >= "3.11"

requirements/requirements-convert_hf_to_gguf.txt

@@ -2,4 +2,5 @@
 --extra-index-url https://download.pytorch.org/whl/cpu
 
 ## Embedding Gemma requires PyTorch 2.6.0 or later, bumped to 2.11.0 for compatibility
-torch==2.11.0
+torch==2.11.0; python_version < "3.15"
+torch==2.15.0; python_version >= "3.15"

requirements/requirements-convert_legacy_llama.txt

@@ -1,5 +1,5 @@
 numpy~=2.2.6; python_version < "3.14"
-numpy~=2.4.6; python_version >= "3.14"
+numpy~=2.5.4; python_version >= "3.14"
 sentencepiece>=0.1.98,<0.3.0
 
 transformers==4.57.6

requirements/requirements-gguf_editor_gui.txt

@@ -1,4 +1,4 @@
 numpy~=2.2.6; python_version < "3.14"
-numpy~=2.4.6; python_version >= "3.14"
+numpy~=2.5.4; python_version >= "3.14"
 PySide6~=6.9.0
 gguf>=0.17.0

requirements/requirements-pydantic.txt

@@ -1,3 +0,0 @@
-docstring_parser~=0.15
-pydantic~=2.11.7
-requests~=2.32.3

requirements/requirements-server-bench.txt

@@ -1,6 +1,7 @@
 datasets~=4.8.0
-matplotlib~=3.10.0
+matplotlib~=3.10.0; python_version < "3.11"
+matplotlib~=3.11.2; python_version >= "3.11"
 numpy~=2.2.6; python_version < "3.14"
-numpy~=2.4.6; python_version >= "3.14"
+numpy~=2.5.4; python_version >= "3.14"
 requests~=2.32.3
 tqdm~=4.67.1

requirements/requirements-test-tokenizer-random.txt

@@ -1 +1 @@
-cffi~=1.16.0
+cffi~=2.1.1

requirements/requirements-tool_bench.txt

@@ -1,9 +1,11 @@
 pytest~=8.3.3
-matplotlib~=3.10.0
+matplotlib~=3.10.0; python_version < "3.11"
+matplotlib~=3.11.2; python_version >= "3.11"
 numpy~=2.2.6; python_version < "3.14"
-numpy~=2.4.6; python_version >= "3.14"
+numpy~=2.5.4; python_version >= "3.14"
 openai~=2.14.0
-pandas~=2.3.3
+pandas~=2.3.3; python_version < "3.11"
+pandas~=3.0.6; python_version >= "3.11"
 prometheus-client~=0.20.0
 requests~=2.32.3
 wget~=3.2

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

@@ -83,6 +83,8 @@ def main():
         '_ZL18flash_attn_ext_f16ILi576ELi512ELi2ELi16ELb0ELb1ELb0EEvPKcS1_S1_S1_S1_PKiS3_PfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
         '_ZL18flash_attn_ext_f16ILi576ELi512ELi8ELi4ELb0ELb1ELb0EEvPKcS1_S1_S1_S1_PKiS3_PfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
         '_ZL18flash_attn_ext_vecILi256ELi1EL9ggml_type2ELS0_2ELb0EEvPKcS2_S2_S2_S2_PKiS4_PfP15HIP_vector_typeIfLj2EEffffjfiS6_IjLj3EEiiiiiiiiiiiliiliiiiil',
+        '_ZL18flash_attn_ext_vecILi256ELi1EL9ggml_type43ELS0_43ELb0EEvPKcS2_S2_S2_S2_PKiS4_PfP15HIP_vector_typeIfLj2EEffffjfiS6_IjLj3EEiiiiiiiiiiiliiliiiiil',
+        '_ZL18flash_attn_ext_vecILi256ELi1EL9ggml_type43ELS0_43ELb1EEvPKcS2_S2_S2_S2_PKiS4_PfP15HIP_vector_typeIfLj2EEffffjfiS6_IjLj3EEiiiiiiiiiiiliiliiiiil',
         '_ZL18flash_attn_ext_f16ILi320ELi256ELi2ELi32ELb0ELb0ELb0EEvPKcS1_S1_S1_S1_PKiS3_PfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
         '_ZL18flash_attn_ext_f16ILi576ELi512ELi1ELi16ELb0ELb1ELb1EEvPKcS1_S1_S1_S1_PKiS3_PfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
     }

scripts/snapdragon/qdc/run_qdc_jobs.py

@@ -68,8 +68,9 @@ log = logging.getLogger(__name__)
 
 POLL_INTERVAL        = 30
 JOB_TIMEOUT          = 3600
+RUNNER_TIMEOUT       = 300
 LOG_UPLOAD_TIMEOUT   = 600
-CAPACITY_TIMEOUT     = 1800
+CAPACITY_TIMEOUT     = 300
 CAPACITY_POLL        = 60
 MAX_CONCURRENT_JOBS  = 5
 DEFAULT_RETRIES      = 0
@@ -82,6 +83,10 @@ class DeviceUnavailableError(Exception):
     """Raised when the QDC device resource is not available (retryable)."""
 
 
+class RunnerUnavailableError(Exception):
+    """Raised when QDC runner or capacity is not available within timeout (non-retryable)."""
+
+
 _SCRIPTS_DIR = Path(__file__).parent
 _TESTS_DIR = _SCRIPTS_DIR / "tests"
 
@@ -253,8 +258,9 @@ DEVICE_PLATFORM: dict[str, Platform] = {
 # =============================================================================
 
 
-def wait_for_job(client, job_id: str, timeout: int) -> str:
+def wait_for_job(client, job_id: str, timeout: int, runner_timeout: int = RUNNER_TIMEOUT) -> str:
     elapsed = 0
+    submitted_elapsed = 0
     last_state = None
     consecutive_errors = 0
     max_consecutive_errors = 5
@@ -285,6 +291,26 @@ def wait_for_job(client, job_id: str, timeout: int) -> str:
         if raw != last_state:
             log.info("Job %s: %s", job_id, raw)
             last_state = raw
+
+        is_submitted = (status == JobState.SUBMITTED) or (
+            isinstance(raw, str) and raw.lower() == "submitted"
+        )
+        if is_submitted:
+            submitted_elapsed += POLL_INTERVAL
+            if submitted_elapsed >= runner_timeout:
+                try:
+                    qdc_api.abort_job(client, job_id)
+                    log.warning(
+                        "Aborted job %s: runner did not pick up job within %ds",
+                        job_id,
+                        runner_timeout,
+                    )
+                except Exception as e:
+                    log.warning("Failed to abort job %s: %s", job_id, e)
+                raise RunnerUnavailableError(
+                    f"Runner did not pick up job {job_id} within {runner_timeout}s"
+                )
+
         time.sleep(POLL_INTERVAL)
         elapsed += POLL_INTERVAL
     # Abort to free the QDC concurrency slot instead of leaking it.
@@ -312,10 +338,10 @@ def wait_for_log_upload(client, job_id: str) -> None:
     log.warning("Timed out waiting for log upload after %ds", LOG_UPLOAD_TIMEOUT)
 
 
-def wait_for_capacity(client, max_jobs: int = MAX_CONCURRENT_JOBS) -> None:
+def wait_for_capacity(client, timeout: int = CAPACITY_TIMEOUT, max_jobs: int = MAX_CONCURRENT_JOBS) -> None:
     """Block until the user's active (non-terminal) QDC job count is below max_jobs."""
     elapsed = 0
-    while elapsed < CAPACITY_TIMEOUT:
+    while elapsed < timeout:
         jobs_page = qdc_api.get_jobs_list(client, page_number=0, page_size=50)
         if jobs_page is None:
             log.warning(
@@ -338,7 +364,7 @@ def wait_for_capacity(client, max_jobs: int = MAX_CONCURRENT_JOBS) -> None:
         time.sleep(CAPACITY_POLL)
         elapsed += CAPACITY_POLL
     raise TimeoutError(
-        f"Capacity wait timed out after {CAPACITY_TIMEOUT}s"
+        f"Capacity wait timed out after {timeout}s"
     )
 
 
@@ -445,7 +471,11 @@ def fetch_logs_and_parse_tests(
     )
 
 
-def write_summary(result: JobResult, title: str = "QDC Test Results") -> None:
+def write_summary(
+    result: JobResult,
+    title: str = "QDC Test Results",
+    note: str | None = None,
+) -> None:
     summary_path = os.environ.get("GITHUB_STEP_SUMMARY")
     if not summary_path:
         return
@@ -454,8 +484,12 @@ def write_summary(result: JobResult, title: str = "QDC Test Results") -> None:
 
     lines = [
         f"## {title}\n",
-        f"Overall: {icon} {'PASSED' if result.passed else 'FAILED'}\n",
     ]
+    if note:
+        lines.append(f"{note}\n")
+    else:
+        lines.append(f"Overall: {icon} {'PASSED' if result.passed else 'FAILED'}\n")
+
     reportable = {n: ok for n, ok in result.tests.items() if "test_install" not in n}
     if reportable:
         lines += ["| Test | Result |", "| ---- | ------ |"]
@@ -464,7 +498,7 @@ def write_summary(result: JobResult, title: str = "QDC Test Results") -> None:
         passed_n = sum(1 for v in reportable.values() if v)
         failed_n = sum(1 for v in reportable.values() if not v)
         lines += ["", f"**{passed_n} passed, {failed_n} failed**"]
-    else:
+    elif not note:
         lines.append("_No per-test data available._")
 
     failed_names = [n for n, ok in reportable.items() if not ok]
@@ -519,6 +553,8 @@ def parse_args() -> argparse.Namespace:
                    help="Test suite to run (default: bench)")
     p.add_argument("--job-timeout", type=int, default=JOB_TIMEOUT, metavar="SECONDS",
                    help=f"Max seconds to wait for job completion (default: {JOB_TIMEOUT})")
+    p.add_argument("--runner-timeout", type=int, default=RUNNER_TIMEOUT, metavar="SECONDS",
+                   help=f"Max seconds to wait for runner to pick up job (default: {RUNNER_TIMEOUT})")
     p.add_argument("--retries", type=int, default=DEFAULT_RETRIES, metavar="N",
                    help="Number of retries when device is unavailable (default: 0)")
     p.add_argument("--retry-delay", type=int, default=RETRY_DELAY, metavar="SECONDS",
@@ -537,9 +573,11 @@ def _submit_and_run_job(client, args, spec, target_id, artifact_id) -> JobResult
     test failure).
     """
     try:
-        wait_for_capacity(client)
+        wait_for_capacity(client, timeout=args.runner_timeout)
     except TimeoutError:
-        raise DeviceUnavailableError("Capacity wait timed out — device busy")
+        raise RunnerUnavailableError(
+            f"Capacity wait timed out after {args.runner_timeout}s - device busy"
+        )
 
     job_name = spec.job_name_fmt.format(base="llama.cpp Hexagon tests")
 
@@ -559,11 +597,16 @@ def _submit_and_run_job(client, args, spec, target_id, artifact_id) -> JobResult
         job_parameters=[JobSubmissionParameter.WIFIENABLED],
     )
     if job_id is None:
-        raise DeviceUnavailableError("Job submission failed — device may be unavailable")
+        raise DeviceUnavailableError("Job submission failed - device may be unavailable")
     log.info("Job submitted: %s  (device=%s)", job_id, args.device)
 
     try:
-        job_status = wait_for_job(client, job_id, timeout=args.job_timeout)
+        job_status = wait_for_job(
+            client,
+            job_id,
+            timeout=args.job_timeout,
+            runner_timeout=args.runner_timeout,
+        )
     except TimeoutError as e:
         raise DeviceUnavailableError(str(e))
     log.info("Job %s finished: %s", job_id, job_status)
@@ -637,6 +680,14 @@ def main() -> int:
         try:
             result = _submit_and_run_job(client, args, spec, target_id, artifact_id)
             break
+        except RunnerUnavailableError as e:
+            log.warning("QDC runner unavailable for %s after %ds - job skipped: %s", args.device, args.runner_timeout, e)
+            write_summary(
+                JobResult(passed=True, tests={}),
+                title=f"QDC Tests ({args.device})",
+                note=f"SKIPPED: Runner did not pick up job within {args.runner_timeout}s. Skipped to avoid stalling CI.",
+            )
+            return 0
         except DeviceUnavailableError as e:
             if attempt < max_attempts:
                 log.warning(

scripts/sync-ggml.last

@@ -1 +1 @@
-57e2c2ca2a7cae756b1c4dfafa0cdfef9051d234
+417a024b4fb6bc3116d41af7c067166a039bdf12

src/llama-context.cpp

@@ -207,7 +207,6 @@ llama_context::llama_context(
     cparams.embeddings_nextn_masked = false;
     cparams.offload_kqv             = params.offload_kqv;
     cparams.no_perf                 = params.no_perf;
-    cparams.warmup                  = false;
 
     // +1: id n_layer() taps the output of the last layer ("input" of the head)
     cparams.embeddings_layer_inp.resize(hparams.n_layer() + 1, false);
@@ -1652,11 +1651,11 @@ bool llama_context::set_h_nextn_rows(int64_t row0, int64_t n_rows, const float *
     return true;
 }
 
-bool llama_context::set_state_src(llama_context * src) {
-    if (cparams.ctx_state_src != src) {
+bool llama_context::set_ctx_state_src(llama_context * ctx_src) {
+    if (cparams.ctx_state_src != ctx_src) {
         sched_need_reserve = true;
     }
-    cparams.ctx_state_src = src;
+    cparams.ctx_state_src = ctx_src;
     return supports_state_row();
 }
 
@@ -1699,14 +1698,6 @@ bool llama_context::get_causal_attn() const {
     return cparams.causal_attn;
 }
 
-void llama_context::set_warmup(bool value) {
-    LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value);
-
-    // no scheduler reserve needed because warmups are usually with small batches
-
-    cparams.warmup = value;
-}
-
 bool llama_context::set_sampler(llama_seq_id seq_id, llama_sampler * sampler) {
     if (!sampler && sampling.samplers.count(seq_id) == 0) {
         return true;
@@ -4444,7 +4435,6 @@ llama_context_params llama_context_default_params() {
         /*.yarn_beta_fast              =*/ -1.0f,
         /*.yarn_beta_slow              =*/ -1.0f,
         /*.yarn_orig_ctx               =*/ 0,
-        /*.defrag_thold                =*/ -1.0f,
         /*.cb_eval                     =*/ nullptr,
         /*.cb_eval_user_data           =*/ nullptr,
         /*.type_k                      =*/ GGML_TYPE_F16,
@@ -4613,13 +4603,6 @@ llama_context * llama_init_from_model(
     return nullptr;
 }
 
-// deprecated
-llama_context * llama_new_context_with_model(
-                 llama_model * model,
-        llama_context_params   params) {
-    return llama_init_from_model(model, params);
-}
-
 void llama_free(llama_context * ctx) {
     delete ctx;
 }
@@ -4695,10 +4678,6 @@ bool llama_get_causal_attn(const llama_context * ctx) {
     return ctx->get_causal_attn();
 }
 
-void llama_set_warmup(llama_context * ctx, bool warmup) {
-    ctx->set_warmup(warmup);
-}
-
 void llama_synchronize(llama_context * ctx) {
     ctx->synchronize();
 }
@@ -4757,8 +4736,8 @@ bool llama_set_embeddings_nextn_rows(llama_context * ctx, int64_t row0, int64_t
     return ctx->set_h_nextn_rows(row0, n_rows, data);
 }
 
-bool llama_set_embd_state_src(llama_context * ctx, llama_context * src) {
-    return ctx->set_state_src(src);
+bool llama_set_ctx_state_src(llama_context * ctx, llama_context * ctx_src) {
+    return ctx->set_ctx_state_src(ctx_src);
 }
 
 void llama_set_embeddings_layer_inp_dev(llama_context * ctx, bool value) {
@@ -4993,31 +4972,6 @@ bool llama_memory_can_shift(llama_memory_t mem) {
 
 // llama state API
 
-// deprecated
-size_t llama_get_state_size(llama_context * ctx) {
-    return llama_state_get_size(ctx);
-}
-
-// deprecated
-size_t llama_copy_state_data(llama_context * ctx, uint8_t * dst) {
-    return llama_state_get_data(ctx, dst, -1);
-}
-
-// deprecated
-size_t llama_set_state_data(llama_context * ctx, const uint8_t * src) {
-    return llama_state_set_data(ctx, src, -1);
-}
-
-// deprecated
-bool llama_load_session_file(llama_context * ctx, const char * path_session, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) {
-    return llama_state_load_file(ctx, path_session, tokens_out, n_token_capacity, n_token_count_out);
-}
-
-// deprecated
-bool llama_save_session_file(llama_context * ctx, const char * path_session, const llama_token * tokens, size_t n_token_count) {
-    return llama_state_save_file(ctx, path_session, tokens, n_token_count);
-}
-
 // Returns the *actual* size of the state.
 // Intended to be used when saving to state to a buffer.
 size_t llama_state_get_size(llama_context * ctx) {