sycl: use TILE for quantized KV decode on BMG (#26689)

Route quantized KV decode to TILE on Xe2 (BMG) only, keep VEC on other archs until validated there.

common/speculative.cpp

@@ -935,6 +935,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
     // dspark speculators
     bool sample_from_anchor = true;
 
+    // block-internal attention
+    bool causal_attn = false;
+
     const int32_t * target_layer_ids   = nullptr; // model_dft's extract layer indices
     uint32_t        target_layer_ids_n = 0;
 
@@ -972,12 +975,25 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
             if (llama_model_meta_val_str(model_dft, "dflash.sample_from_anchor", buf, sizeof(buf)) >= 0) {
                 sample_from_anchor = std::strcmp(buf, "true") == 0;
             }
+            if (llama_model_meta_val_str(model_dft, "dflash.attention.causal", buf, sizeof(buf)) >= 0) {
+                causal_attn = std::strcmp(buf, "true") == 0;
+            }
         }
 
         selector_top_k = llama_model_dflash_selector_top_k(model_dft);
         is_dflash2     = selector_top_k > 0;
         mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft));
 
+        if (is_dspark && this->params.p_min > 0.0f) {
+            char buf[16] = {};
+            const bool has_conf =
+                llama_model_meta_val_str(model_dft, "dflash.has_confidence_head", buf, sizeof(buf)) < 0 ||
+                std::strcmp(buf, "true") == 0;
+            if (!has_conf) {
+                throw std::runtime_error("DSpark draft has no confidence head: please set --spec-draft-p-min 0");
+            }
+        }
+
         LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str());
         LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min);
         LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u, sample_from_anchor=%s\n", __func__,
@@ -1036,7 +1052,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
 
         // DFlash2 reads its selector lattice from h_nextn and never consumes raw logits.
         llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2);
-        llama_set_causal_attn(ctx_dft, false); // DFlash needs non-causal attention
+        llama_set_causal_attn(ctx_dft, causal_attn); // DFlash needs non-causal attention unless the model says otherwise
     }
 
     ~common_speculative_impl_draft_dflash() override {

conversion/qwen.py

@@ -709,14 +709,20 @@ class DFlashModel(Qwen3Model):
             extract_layer_ids = [i + 1 for i in target_layer_ids]
             self.gguf_writer.add_target_layers(extract_layer_ids)
 
-        use_sliding_window = self.hparams.get("use_sliding_window", False)
-        sliding_window = self.hparams.get("sliding_window")
+        use_sliding_window = self.hparams.get("use_sliding_window", False) or dflash_config.get("use_swa", False)
+        sliding_window = dflash_config.get("swa_window_size") or self.hparams.get("sliding_window")
         layer_types = self.hparams.get("layer_types")
         if use_sliding_window and sliding_window and layer_types:
             is_swa = [lt == "sliding_attention" for lt in layer_types]
             self.gguf_writer.add_sliding_window(sliding_window)
             self.gguf_writer.add_sliding_window_pattern(is_swa)
 
+        causal = self.hparams.get("is_causal")
+        if causal is None:
+            causal = dflash_config.get("causal")
+        if causal is not None:
+            self.gguf_writer.add_causal_attention(bool(causal))
+
         # M-RoPE target: the draft ropes on the temporal dim only, so write
         # degenerate sections [n_rot/2, 0, 0, 0]
         if self._target_uses_mrope():
@@ -737,6 +743,8 @@ class DFlashModel(Qwen3Model):
         name, gen = item
         if not name.startswith("model."):
             name = "model." + name
+        if "sink" in name and not name.endswith(".weight"):
+            name += ".weight"
         return super().filter_tensors((name, gen))
 
     _ROPE_PERMUTE_SUFFIXES = (
@@ -815,6 +823,10 @@ class DSparkModel(DFlashModel):
         super().set_gguf_parameters()
         self.gguf_writer.add_sample_from_anchor(self._sample_from_anchor)
 
+        # confidence head is optional: vanilla-markov exports ship without it
+        has_conf = any("confidence_head.proj" in name for name in self.model_tensors)
+        self.gguf_writer.add_has_confidence_head(has_conf)
+
     @classmethod
     def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
         if item[0] == "t2d":  # not used at runtime
@@ -833,7 +845,7 @@ class DSparkModel(DFlashModel):
             self._d2t = data_torch
             return
 
-        if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith(("embed_tokens.weight", "lm_head.weight")):
+        if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith("lm_head.weight"):
             return
 
         # interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd

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

@@ -66,6 +66,7 @@ fa_vec_cfg_t fa_vec_baseline_cfg(int dk, int dv) {
 // One row per kept bucket, plus per-(dtype,dk,dv) ne11-collapsed domain defaults
 // (ne11_b = FA_VEC_NE11_DEFAULT, ne01_b = domain). To retune or add a device, re-run the
 // sweep and paste its output. See ggml-metal-tuning.h for the row/lookup semantics.
+// ref: https://github.com/ggml-org/llama.cpp/pull/27824
 constexpr fa_vec_entry_t fa_vec_tuned_table[] = {
     { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 2, 4 } },
     { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } },
@@ -449,6 +450,159 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = {
     { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
 
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 1, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 3, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 3, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 3, 3 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 3, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 128, 128, 2, 4 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 192, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 256, 256, 2, 3 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 320, 256, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 320, 256, 3, 0 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 2, 0 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 3, 0 }, { 4, 1 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 1, 3 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 2, 1 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 2, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 3, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 3, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 2, 0 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 2, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 2, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 3, 1 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 3, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 1, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 1, 4 }, { 1, 2 } },
+
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 64, 64, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 192, 128, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 256, 256, 3, 0 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 320, 256, 3, 0 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 512, 512, 3, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 1, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 1, 3 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 1, 4 }, { 1, 2 } },
+
     { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } },
     { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } },
     { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } },
@@ -640,7 +794,215 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = {
     { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
 
-        { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 1, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 3, 0 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 192, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 3, 2 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 3, 3 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, 1, 4 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 1, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 2, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 2, 3 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 2 }, { 4, 1 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 4 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 2, 0 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 2, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 1, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 3, 4 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 320, 256, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 512, 512, 3, 0 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 512, 512, 2, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 512, 512, 3, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
+
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 1, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 3, 4 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 128, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 128, 3, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 256, 256, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 320, 256, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 320, 256, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 512, 512, 2, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 512, 512, 2, 3 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 512, 512, 3, 3 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 3 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 2, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 2, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, 1, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, 2, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, 2, 3 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
+
+    { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, -1, 1 }, { 2, 4 } },
     { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, 1, 2 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, 1, 4 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, 2, 2 }, { 4, 4 } },

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

@@ -1,3 +1,4 @@
+#include <array>
 #include <cstdint>
 #include <cstdio>
 #include <cstring>
@@ -150,7 +151,8 @@ struct sdpa_partition {
 
 // Build + compile the contiguous-input GQA SDPA graph (MatMul->Divide->Add->SoftMax->MatMul), f32 out.
 // Mirrors the hardware-verified scratch/onednn_sdpa_probe.cpp build_gqa (partitions=1, sdp_primitive_kernel_t).
-static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int seq, int d) {
+static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int seq, int d,
+                                 const std::array<int64_t, 5> & k_str, const std::array<int64_t, 5> & v_str) try {
     using ltype = logical_tensor::layout_type;
     using dt    = logical_tensor::data_type;
     using ldims = logical_tensor::dims;
@@ -158,11 +160,12 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int
     const int   rep = H / Hkv;
     const ldims q_sz = {1, Hkv, rep, q, d}, kv_sz = {1, Hkv, 1, seq, d}, s_sz = {1, Hkv, rep, q, seq},
                 sc = {1, 1, 1, 1, 1}, msk = {1, 1, 1, q, seq}, o_sz = {1, Hkv, rep, q, d};
+    const ldims k_st(k_str.begin(), k_str.end()), v_st(v_str.begin(), v_str.end());
     int64_t        id = 0;
     sdpa_partition E;
 
     auto query  = logical_tensor(id++, t,  q_sz, ltype::strided);
-    auto key    = logical_tensor(id++, t,  kv_sz, ltype::strided);
+    auto key    = logical_tensor(id++, t,  kv_sz, k_st);
     auto score  = logical_tensor(id++, fi, s_sz, ltype::strided);
     auto bmm1   = op(id++, op::kind::MatMul, "bmm1");
     bmm1.set_attr<bool>(op::attr::transpose_b, true);          // key is [.., seq, d]
@@ -184,7 +187,7 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int
     smax.set_attr<std::string>(op::attr::mode, "inf_as_zero");
     smax.add_inputs({masked}); smax.add_outputs({probs});
 
-    auto value  = logical_tensor(id++, t,  kv_sz, ltype::strided);
+    auto value  = logical_tensor(id++, t,  kv_sz, v_st);
     // f16 output is REQUIRED to hit sdp_primitive_kernel_t (the systolic micro-kernel); an f32 output
     // falls to larger_partition_kernel_t which materializes N^2 (confirmed: scratch/onednn_sdpa_kernel_probe.cpp).
     // converted to the f32 ggml dst in the permute below.
@@ -198,6 +201,7 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int
 
     auto parts = g.get_partitions();
     if (parts.size() != 1 || !parts[0].is_supported()) {
+        GGML_LOG_WARN("%s: oneDNN did not fuse the SDPA graph; falling back to TILE kernel\n", __func__);
         return E;   // ok stays false -> caller falls back to TILE
     }
     E.ins      = parts[0].get_input_ports();
@@ -209,6 +213,12 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int
     E.ok       = true;
     return E;
 }
+catch (const std::exception & e) {
+    // compile() can reject a stride set the partitioner never inspects; memoise the failure so the
+    // fallback costs one build rather than one per call.
+    GGML_LOG_WARN("%s: oneDNN SDPA partition build failed (%s); falling back to TILE kernel\n", __func__, e.what());
+    return {};
+}
 
 void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tensor * dst) try {
     const ggml_tensor * Q    = dst->src[0];
@@ -234,13 +244,34 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
     ggml_sycl_pool_alloc<sycl::half> Qf(ctx.pool(), (size_t) H * q * d);
     cont_to_f16_sycl<float>((const char *) Q->data, Qf.get(), d, q, H, mb, Q->nb[1], Q->nb[2], Q->nb[3], stream);
 
-    // K/V: use pool-alloc for both F16 and dequant paths.
+    // K/V: bind the f16 cache in place. llama.cpp permutes it to [token][head][dim], so its head
+    // plane is strided rather than dense, which is what an explicit stride vector expresses.
+    // Quantized and f32 KV still stage a dense copy -- the layout the k_str/v_str defaults describe.
     sycl::half * K_ptr = nullptr;
     sycl::half * V_ptr = nullptr;
+    std::array<int64_t, 5> k_str{ Hkv * seq * d, seq * d, seq * d, d, 1 };
+    std::array<int64_t, 5> v_str = k_str;
     std::optional<ggml_sycl_pool_alloc<sycl::half>> Kf_pool;
     std::optional<ggml_sycl_pool_alloc<sycl::half>> Vf_pool;
 
-    if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16) {
+    auto bindable = [](const ggml_tensor * t) {
+        return t->nb[0] == sizeof(sycl::half) && t->nb[1] % sizeof(sycl::half) == 0 &&
+               t->nb[2] % sizeof(sycl::half) == 0 && t->nb[3] % sizeof(sycl::half) == 0;
+    };
+    auto elem_strides = [](const ggml_tensor * t) {
+        const int64_t s1 = (int64_t) (t->nb[1] / t->nb[0]);
+        const int64_t s2 = (int64_t) (t->nb[2] / t->nb[0]);
+        const int64_t s3 = (int64_t) (t->nb[3] / t->nb[0]);
+        // dims are {mb=1, Hkv, rep=1, seq, d}; the size-1 dims at 0 and 2 never advance an address.
+        return std::array<int64_t, 5>{ s3, s2, s2, s1, 1 };
+    };
+
+    if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && bindable(K) && bindable(V)) {
+        K_ptr = (sycl::half *) K->data;
+        V_ptr = (sycl::half *) V->data;
+        k_str = elem_strides(K);
+        v_str = elem_strides(V);
+    } else if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16) {
         Kf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d);
         Vf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d);
         cont_to_f16_sycl<sycl::half>((const char *) K->data, Kf_pool->get(), d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream);
@@ -341,19 +372,24 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
 
     ggml_sycl_pool_alloc<sycl::half> outf(ctx.pool(), (size_t) H * q * d);   // f16 contiguous SDPA out [mb,H,q,d]
 
-    // compile once per (device, shape), reuse across layers/calls.
+    // compile once per (device, shape, KV strides), reuse across layers/calls. Stride 2 always
+    // repeats stride 1 and stride 4 is always 1, so the key covers every entry that can differ.
     static std::unordered_map<std::string, sdpa_partition> cache;
-    char keyb[96];
-    snprintf(keyb, sizeof(keyb), "%d:%lld:%lld:%lld:%lld:%lld", ggml_sycl_get_device(),
-             (long long) H, (long long) Hkv, (long long) q, (long long) seq, (long long) d);
+    char keyb[256];
+    snprintf(keyb, sizeof(keyb), "%d:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld", ggml_sycl_get_device(),
+             (long long) H, (long long) Hkv, (long long) q, (long long) seq, (long long) d,
+             (long long) k_str[0], (long long) k_str[1], (long long) k_str[3],
+             (long long) v_str[0], (long long) v_str[1], (long long) v_str[3]);
     auto it = cache.find(keyb);
     if (it == cache.end()) {
-        it = cache.emplace(keyb, build_sdpa(eng, (int) H, (int) Hkv, (int) q, (int) seq, (int) d)).first;
+        it = cache.emplace(keyb, build_sdpa(eng, (int) H, (int) Hkv, (int) q, (int) seq, (int) d, k_str, v_str)).first;
     }
     sdpa_partition & E = it->second;
-    // _supported() is authoritative: if it accepted this op the partition must build.
-    // A failure here is a gap in _supported() -- surface it, don't mask it with a fallback.
-    GGML_ASSERT(E.ok && "oneDNN SDPA partition failed to build for a _supported() shape");
+    if (!E.ok) {
+        // oneDNN can decline a shape or a stride set that _supported() never sees; build_sdpa warns per key.
+        ggml_sycl_flash_attn_ext_tile(ctx, dst);
+        return;
+    }
 
     auto id2ptr = [&](size_t r) -> void * {
         if (r == E.id_q)     return Qf.get();

ggml/src/ggml-sycl/fattn.cpp

@@ -104,7 +104,6 @@ enum best_fattn_kernel {
 
 
 static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const ggml_tensor * dst) {
-    GGML_UNUSED(device);
 #ifndef SYCL_FLASH_ATTN
     GGML_UNUSED(dst);
     return BEST_FATTN_KERNEL_NONE;
@@ -263,6 +262,11 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const
             }
         } else {
             if (Q->ne[1] <= 2) {
+                // TILE is faster for quantized KV decode on Xe2 (BMG); keep VEC on untested archs
+                const gpu_arch arch = ggml_sycl_info().devices[device].hw_info.arch;
+                if (arch == gpu_arch::intel_gpu_bmg_g21 || arch == gpu_arch::intel_gpu_bmg_g31) {
+                    return BEST_FATTN_KERNEL_TILE;
+                }
                 return BEST_FATTN_KERNEL_VEC;
             }
         }

gguf-py/gguf/constants.py

@@ -167,6 +167,7 @@ class Keys:
         SELECTOR_RANK                     = "{arch}.selector_rank"
         SELECTOR_TOP_K                    = "{arch}.selector_top_k"
         SAMPLE_FROM_ANCHOR                = "{arch}.sample_from_anchor"
+        HAS_CONFIDENCE_HEAD               = "{arch}.has_confidence_head"
         NORM_BEFORE_RESIDUAL              = "{arch}.norm_before_residual"
         NORM_BEFORE_FC                    = "{arch}.norm_before_fc"
 

gguf-py/gguf/gguf_writer.py

@@ -1008,6 +1008,9 @@ class GGUFWriter:
     def add_sample_from_anchor(self, value: bool) -> None:
         self.add_bool(Keys.LLM.SAMPLE_FROM_ANCHOR.format(arch=self.arch), value)
 
+    def add_has_confidence_head(self, value: bool) -> None:
+        self.add_bool(Keys.LLM.HAS_CONFIDENCE_HEAD.format(arch=self.arch), value)
+
     def add_target_layers(self, value: Sequence[int]) -> None:
         self.add_array(Keys.LLM.TARGET_LAYERS.format(arch=self.arch), value)
 

src/llama-context.cpp

@@ -661,11 +661,19 @@ void llama_context::sched_reserve() {
 
     // reserve again with pp graph to avoid ggml-alloc reallocations during inference
     {
-        // TODO: not sure if the following graph would be worst case for multi-stream KV caches:
-        //
-        // auto * gf = graph_reserve(n_tokens, 1, n_tokens, mctx.get());
-        //
-        auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.hparams.no_alloc);
+        // TODO: the worst case graph is not always reached for `n_seqs > 1`
+        //       need to implement a more robust mechanism that tries a few different inputs and analyzes the results
+        ggml_cgraph * gf = nullptr;
+        switch (model.arch) {
+            case LLM_ARCH_MINIMAX_01:
+                // the `inp_diag_decay` tensor size scales with `n_seq_tokens^2` which
+                // makes `n_seqs == 1` use more memory for the compute graph compared to `n_seqs > 1`
+                gf = graph_reserve(n_tokens, 1,      n_outputs_pp, mctx.get(), model.hparams.no_alloc);
+                break;
+            default:
+                gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.hparams.no_alloc);
+        };
+
         if (!gf) {
             throw std::runtime_error("failed to allocate compute pp buffers");
         }
@@ -2892,13 +2900,83 @@ public:
         for (auto & [buft, mbuf] : mbufs_new) {
             const auto & mbuf_cur = mbufs.at(buft);
 
-            if (!mbuf_cur.buf || mbuf_cur.n_tensors != mbuf.n_tensors || mbuf_cur.total_size != mbuf.total_size) {
+            if (!mbuf_cur.buf || mbuf_cur.total_size != mbuf.total_size) {
                 GGML_ABORT("%s: memory buffer mismatch\n", __func__);
             }
 
-            for (size_t i = 0; i < mbuf_cur.org.size(); ++i) {
-                ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]);
+            if (mbuf_cur.n_tensors == mbuf.n_tensors) {
+                // same chunking: copy 1:1 by index
+                for (size_t i = 0; i < mbuf_cur.org.size(); ++i) {
+                    GGML_ASSERT(ggml_nbytes(mbuf_cur.cpy[i]) == ggml_nbytes(mbuf.org[i]));
+                    ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]);
+                }
+                continue;
             }
+
+            // different chunking: copy the write-side data (mbuf_cur.cpy) into the read-side targets (mbuf.org)
+            // with a byte cursor. Write and read enumerate the same logical data in the same order but may chunk
+            // it differently, so copy across tensor boundaries rather than 1:1 by index.
+            const size_t total = mbuf_cur.total_size;
+
+            ggml_init_params params_scratch = {
+                /*.mem_size   =*/ 2*(mbuf_cur.cpy.size() + mbuf.org.size())*ggml_tensor_overhead(),
+                /*.mem_buffer =*/ NULL,
+                /*.no_alloc   =*/ true,
+            };
+            ggml_context * ctx_scratch = ggml_init(params_scratch);
+
+            size_t src_pos  = 0;
+            size_t dst_pos  = 0;
+            size_t src_j    = 0;
+            size_t dst_i    = 0;
+            size_t src_base = 0;
+            size_t dst_base = 0;
+
+            while (src_pos < total) {
+                const auto & src_t = mbuf_cur.cpy[src_j];
+                const auto & dst_t = mbuf.org[dst_i];
+
+                const size_t src_size = ggml_nbytes(src_t);
+                const size_t dst_size = ggml_nbytes(dst_t);
+
+                const size_t src_off  = src_pos - src_base;
+                const size_t dst_off  = dst_pos - dst_base;
+
+                const size_t n_copy = std::min(src_size - src_off, dst_size - dst_off);
+
+                const size_t   el   = ggml_element_size(src_t);
+                const int64_t n_el = (int64_t) (n_copy / el);
+
+                auto * src_v = ggml_view_1d(ctx_scratch, src_t, n_el, src_off);
+                ggml_backend_view_init(src_v);
+                auto * dst_v = ggml_view_1d(ctx_scratch, dst_t, n_el, dst_off);
+                ggml_backend_view_init(dst_v);
+
+                ggml_backend_tensor_copy(src_v, dst_v);
+
+                src_pos += n_copy;
+                dst_pos += n_copy;
+
+                if (src_pos - src_base == src_size) {
+                    src_base = src_pos;
+                    ++src_j;
+                }
+                if (dst_pos - dst_base == dst_size) {
+                    dst_base = dst_pos;
+                    ++dst_i;
+                }
+            }
+
+            GGML_ASSERT(src_pos == total && dst_pos == total);
+            // any tensors left unvisited hold no data
+            for (size_t i = src_j; i < mbuf_cur.cpy.size(); ++i) {
+                GGML_ASSERT(ggml_nbytes(mbuf_cur.cpy[i]) == 0);
+            }
+            for (size_t i = dst_i; i < mbuf.org.size(); ++i) {
+                GGML_ASSERT(ggml_nbytes(mbuf.org[i]) == 0);
+            }
+
+            ggml_free(ctx_scratch);
         }
 
         GGML_ASSERT(buf_size == 0);

src/llama-model.h

@@ -672,6 +672,7 @@ struct llama_model {
     // dspark
     struct ggml_tensor * dspark_markov_w1   = nullptr;
     struct ggml_tensor * dspark_markov_w2   = nullptr;
+    struct ggml_tensor * dspark_markov_w2_s = nullptr;
     struct ggml_tensor * dspark_conf_proj   = nullptr;
     struct ggml_tensor * dspark_conf_proj_b = nullptr;
 

src/models/dflash.cpp

@@ -115,10 +115,11 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
     if (markov_meta) {
         const int64_t dspark_markov_rank = markov_meta->ne[0];
 
-        dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0);
-        dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0);
+        dspark_markov_w1   = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0);
+        dspark_markov_w2   = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0);
+        dspark_markov_w2_s = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "scale"),  { 1 }, TENSOR_NOT_REQUIRED);
 
-        dspark_conf_proj   = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, 0);
+        dspark_conf_proj   = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, TENSOR_NOT_REQUIRED);
         dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"),   { 1 },             TENSOR_NOT_REQUIRED);
 
         LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank);
@@ -219,6 +220,9 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
         layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0);
         layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0);
 
+        // optional per-head attention sinks (e.g. Nemotron DSpark)
+        layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), { n_head }, TENSOR_NOT_REQUIRED);
+
         layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), { n_embd }, 0);
         layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0);
         layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0);
@@ -290,7 +294,10 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &
 
     ggml_tensor * w1 = model.dspark_markov_w1;
     ggml_tensor * w2 = model.dspark_markov_w2;
-    GGML_ASSERT(w1 && w2 && model.dspark_conf_proj && "DSpark markov/confidence weights not loaded");
+    GGML_ASSERT(w1 && w2 && "DSpark markov weights not loaded");
+
+    // confidence head is optional
+    const bool has_conf = model.dspark_conf_proj != nullptr;
 
     ggml_tensor * base = res->t_logits; // [n_vocab, n_tokens]
     const int64_t n_vocab = base->ne[0];
@@ -321,23 +328,22 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &
     ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, token_stride, 0);
     prev = ggml_cont_1d(ctx0, prev, n_blocks);
 
-    // confidence head input: predicts per-position acceptance
-    ggml_tensor * conf_inp = res->t_embd; // [n_embd, n_tok]
-
     ggml_tensor * cat      = nullptr;
     ggml_tensor * cat_conf = nullptr;
 
     if (!sample_from_anchor) {
         // bonus anchor slot: pass the logits through unbiased, pad the (unread) confidence column
-        cat      = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0));
-        cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0)));
+        cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0));
+        if (has_conf) {
+            cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0)));
+        }
     }
 
     // TODO: the in-graph chain is greedy (argmax); sampling params affect only the final
     //       token pick, not the Markov conditioning path
     for (int64_t i = i_draft_beg; i < block_drafts; ++i) {
-        ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev);   // [R, n_blocks]
-        ggml_tensor * bias    = ggml_mul_mat(ctx0, w2, w1_prev); // [n_vocab_draft, n_blocks]
+        ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev);                          // [R, n_blocks]
+        ggml_tensor * bias    = g.build_lora_mm(w2, w1_prev, model.dspark_markov_w2_s); // [n_vocab_draft, n_blocks]
         if (model.d2t) {
             // reduced draft vocab: scatter the bias to the target rows (base is -inf on the others)
             const int64_t n_draft_vocab = bias->ne[0];
@@ -354,17 +360,21 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &
 
         cat = cat ? ggml_concat(ctx0, cat, col, 1) : col;
 
-        // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b)  -- [1, n_blocks]
-        ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks,
-                                                (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]);
-        ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0);
-        ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat);
-        if (model.dspark_conf_proj_b) {
-            conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b);
-        }
-        conf = ggml_sigmoid(ctx0, conf);
+        if (has_conf) {
+            // confidence head input: predicts per-position acceptance
+            ggml_tensor * conf_inp   = res->t_embd; // [n_embd, n_tok]
+            // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b)  -- [1, n_blocks]
+            ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks,
+                                                    (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]);
+            ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0);
+            ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat);
+            if (model.dspark_conf_proj_b) {
+                conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b);
+            }
+            conf = ggml_sigmoid(ctx0, conf);
 
-        cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf;
+            cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf;
+        }
 
         if (i + 1 < block_drafts) {
             prev = ggml_argmax(ctx0, col);
@@ -376,7 +386,7 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &
     out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, block_drafts, n_blocks]
     out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok);
 
-    {
+    if (has_conf) {
         ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, block_drafts);
         conf = ggml_cont(ctx0, ggml_permute(ctx0, conf, 0, 2, 1, 3));
         conf = ggml_reshape_2d(ctx0, conf, 1, n_tok);
@@ -707,8 +717,8 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
 
         // cache-aware, non-causal attention
         ggml_tensor * cur = use_iswa
-            ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il)
-            : build_attn(inp_attn,      layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+            ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il)
+            : build_attn(inp_attn,      layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il);
 
         if (attn_dynamic) {
             cur = build_dflash2_conv(*this, cur, attn_dynamic, layer.dflash_attn_conv_base, 1);

src/models/minimax-01.cpp

@@ -181,7 +181,7 @@ public:
         return res;
     }
 
-    const llama_hparams & hparams;
+    const llama_hparams hparams;
 
     ggml_tensor * inp_slopes     = nullptr; // F32 [n_head]
     ggml_tensor * inp_q_decay    = nullptr; // F32 [1, n_head, n_batch]

tests/CMakeLists.txt

@@ -149,6 +149,7 @@ if (LLAMA_LLGUIDANCE)
 endif ()
 
 llama_build(test-recurrent-state-rollback.cpp)
+llama_build(test-save-load-state.cpp)
 
 if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
     # these tests are disabled on Windows because they use internal functions not exported with LLAMA_API (when building with shared libraries)
@@ -237,6 +238,14 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
     set_tests_properties(test-recurrent-state-rollback-dsv4 PROPERTIES
         FIXTURES_REQUIRED generate-models
     )
+
+    # Test state save/load functionality across all architectures, using the generated dummy models
+    llama_test(
+        test-save-load-state
+        LABEL main
+        ARGS --models "${MODEL_DIR}"
+    )
+    set_tests_properties(test-save-load-state PROPERTIES FIXTURES_REQUIRED generate-models)
 endif()
 
 llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp)
@@ -299,10 +308,6 @@ llama_build_and_test(test-backend-sampler.cpp   LABEL "model")
 llama_build_and_test(test-state-restore-fragmented.cpp LABEL "model" ARGS -m "${MODEL_DEST}")
 set_tests_properties(test-state-restore-fragmented PROPERTIES FIXTURES_REQUIRED test-download-model)
 
-# Test state save/load functionality
-llama_build_and_test(test-save-load-state.cpp LABEL "model" ARGS -m "${MODEL_DEST}")
-set_tests_properties(test-save-load-state PROPERTIES FIXTURES_REQUIRED test-download-model)
-
 if (APPLE)
     llama_build(test-rset-release.cpp)
 endif()

tests/test-llama-archs.cpp

@@ -65,7 +65,7 @@ static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) {
 }
 
 static void usage(char ** argv) {
-    printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-v/--verbose]\n", argv[0]);
+    printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v/--verbose] [-h/--help]\n", argv[0]);
 }
 
 static std::vector<llama_token> get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed){
@@ -82,7 +82,7 @@ static std::vector<llama_token> get_tokens(const uint32_t n_tokens, const uint32
 static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
     gguf_context_ptr ret(gguf_init_empty());
     llama_model_saver ms(arch, ret.get());
-    const uint32_t n_ctx = 128;
+    const uint32_t n_ctx = 256;
 
     uint32_t n_vocab = 128;
     uint32_t n_embd  = 256;
@@ -256,10 +256,12 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
         ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector<uint32_t>(n_layer, 4));
     }
 
-    ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT,   arch == LLM_ARCH_MINIMAX_M3 || arch == LLM_ARCH_DEEPSEEK4 ? n_head : uint32_t(1));
+    // minimax-m3 keeps one indexer head per GQA head; the rest use a fixed 64 to match the fused
+    ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT,   arch == LLM_ARCH_MINIMAX_M3 ? n_head : uint32_t(64));
     // qwen4exp ropes indexer keys with the main rotary width, so its head can't be < n_rot
     ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH,
-              arch == LLM_ARCH_QWEN4EXP ? n_embd_head : uint32_t(64));
+              arch == LLM_ARCH_QWEN4EXP ? n_embd_head : uint32_t(128));
+
     ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K,        uint32_t(8));
     ms.add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE,   uint32_t(4));
     ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1));
@@ -762,6 +764,10 @@ int main(int argc, char ** argv) {
     std::string out;
 
     for (int i = 1; i < argc; i++) {
+        if (strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0) {
+            usage(argv);
+            return 0;
+        }
         if (strcmp(argv[i], "-a") == 0 || strcmp(argv[i], "--arch") == 0) {
             if (i + 1 < argc) {
                 const std::string arch_name = argv[++i];

tests/test-save-load-state.cpp

@@ -3,8 +3,12 @@
 #include "log.h"
 #include "llama-cpp.h"
 
+#include <algorithm>
 #include <clocale>
+#include <cstring>
+#include <filesystem>
 #include <random>
+#include <string>
 #include <vector>
 
 struct llama_batch_ptr {
@@ -53,7 +57,9 @@ static llama_tokens generate_tokens(llama_context * ctx, llama_sampler * smpl, i
 // - decode the last token
 // - generate n_predict tokens
 static llama_tokens test_baseline(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) {
-    auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))};
+    auto params_ctx = common_context_params_to_llama(params);
+    params_ctx.n_seq_max = 2;
+    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
 
     auto sparams = llama_sampler_chain_default_params();
     auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
@@ -161,7 +167,9 @@ static bool test_seq_rm_isolated(
 // - replay the last prompt token
 // - generate n_predict tokens and compare against expected result
 static bool test_state_load(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const llama_tokens & expected_result) {
-    auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))};
+    auto params_ctx = common_context_params_to_llama(params);
+    params_ctx.n_seq_max = 2;
+    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
 
     auto sparams = llama_sampler_chain_default_params();
     auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
@@ -347,38 +355,18 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p
 }
 
 
-int main(int argc, char ** argv) {
-    std::setlocale(LC_NUMERIC, "C");
-
-    common_params params;
-    params.prompt = "";
-    params.n_batch = 100;
-    params.out_file = "dump_state.bin";
-    params.sampling.seed = 1234;
-
-    common_init();
-
-    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) {
-        return 1;
-    }
-
-    if (params.n_parallel == 1) {
-        LOG_TRC("%s: n_parallel == 1, enabling unified kv cache\n", __func__);
-        params.kv_unified = true;
-    }
-
-    if (params.n_predict < 0) {
-        params.n_predict = 16;
-    }
-
-    ggml_backend_load_all();
+// Run the full save/load test suite (tests 1-5) for a single model.
+// Returns true if all tests pass, false otherwise.
+static bool run_save_load_tests_for_model(const std::string & model_path, const struct common_params & base_params) {
+    struct common_params params = base_params;
+    params.model.path = model_path;
 
     auto llama_init = common_init_from_params(params, true);
     auto * model = llama_init->model();
 
     if (model == nullptr) {
-        LOG_ERR("%s: failed to init\n", __func__);
-        return 1;
+        LOG_ERR("%s: failed to init model '%s'\n", __func__, model_path.c_str());
+        return false;
     }
 
     GGML_ASSERT(llama_init->context() == nullptr);
@@ -411,30 +399,127 @@ int main(int argc, char ** argv) {
     // Test 1: baseline (saves state to disk)
     auto result_baseline = test_baseline(model, params, tokens);
     if (result_baseline.empty()) {
-        return 1;
+        return false;
     }
 
     // Test 2: sequence removal isolation
     if (!test_seq_rm_isolated(model, params, tokens)) {
-        return 1;
+        return false;
     }
 
     // Test 3: state load
     if (!test_state_load(model, params, tokens, result_baseline)) {
-        return 1;
+        return false;
     }
 
     // Test 4: seq copy (host)
     if (!test_seq_cp_host(model, params, tokens, result_baseline)) {
-        return 1;
+        return false;
     }
 
     // Test 5: seq copy (device)
     if (!test_seq_cp_device(model, params, tokens, result_baseline)) {
-        return 1;
+        return false;
     }
 
     LOG("\nAll tests passed.\n");
 
-    return 0;
+    return true;
+}
+
+
+int main(int argc, char ** argv) {
+    std::setlocale(LC_NUMERIC, "C");
+
+    common_params params;
+    params.prompt = "";
+    params.n_batch = 100;
+    params.out_file = "dump_state.bin";
+    params.sampling.seed = 1234;
+
+    common_init();
+
+    // extract our own --models DIR option before handing the rest to the common arg parser
+    std::string models_dir;
+    std::vector<char *> filtered_argv;
+    filtered_argv.push_back(argv[0]);
+    for (int i = 1; i < argc; i++) {
+        if (strcmp(argv[i], "--models") == 0) {
+            if (i + 1 >= argc) {
+                LOG_ERR("%s: --models requires a directory argument\n", __func__);
+                return 1;
+            }
+            models_dir = argv[i + 1];
+            i++;
+        } else {
+            filtered_argv.push_back(argv[i]);
+        }
+    }
+    filtered_argv.push_back(nullptr);
+    const int fargc = (int)filtered_argv.size() - 1;
+
+    // in --models mode there is no single model; set a placeholder so the common parser's
+    // "--model is required" check passes (each model is set individually inside the loop)
+    if (!models_dir.empty()) {
+        params.model.path = models_dir;
+    }
+
+    if (!common_params_parse(fargc, filtered_argv.data(), params, LLAMA_EXAMPLE_COMMON)) {
+        return 1;
+    }
+
+    if (params.n_parallel == 1) {
+        LOG_TRC("%s: n_parallel == 1, enabling unified kv cache\n", __func__);
+        params.kv_unified = true;
+    }
+
+    if (params.n_predict < 0) {
+        params.n_predict = 16;
+    }
+
+    ggml_backend_load_all();
+
+    if (!models_dir.empty()) {
+        // run the suite over every dummy model in the directory
+        if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) {
+            LOG_ERR("%s: models directory '%s' does not exist\n", __func__, models_dir.c_str());
+            return 1;
+        }
+
+        std::vector<std::string> models;
+        for (const auto & entry : std::filesystem::directory_iterator(models_dir)) {
+            if (entry.is_regular_file() && entry.path().extension() == ".gguf") {
+                models.push_back(entry.path().string());
+            }
+        }
+        std::sort(models.begin(), models.end());
+
+        if (models.empty()) {
+            LOG_ERR("%s: no .gguf models found in '%s'\n", __func__, models_dir.c_str());
+            return 1;
+        }
+
+        LOG_INF("%s: running save/load tests over %zu models in '%s'\n", __func__, models.size(), models_dir.c_str());
+
+        size_t n_pass = 0;
+        size_t n_fail = 0;
+        for (const auto & model_path : models) {
+            LOG("\n================================================================\n");
+            LOG_INF("%s: model %s\n", __func__, model_path.c_str());
+
+            if (run_save_load_tests_for_model(model_path, params)) {
+                n_pass++;
+            } else {
+                n_fail++;
+            }
+        }
+
+        LOG("\n================================================================\n");
+        LOG_INF("%s: summary: %zu passed, %zu failed (of %zu)\n", __func__, n_pass, n_fail, models.size());
+
+        return n_fail == 0 ? 0 : 1;
+    }
+
+    // single-model mode
+    return run_save_load_tests_for_model(params.model.path, params) ? 0 : 1;
 }