tensor-split meta backend fixes (#26502)

* backend: propagate buffer usage in meta backend

* ggml-meta: make sure to call init_tensor for all new tensors

* meta: remove explicit check for meta backend in ggml_backend_meta_get_split_state

I can't seem to reproduce the original failure in the latest code.

.github/workflows/release.yml

@@ -1579,14 +1579,14 @@ jobs:
       - windows
       - windows-cpu
       - windows-cuda
-      #- windows-sycl
+      - windows-sycl
       - windows-rocm
       - windows-openvino
       #- ubuntu-22-rocm
       - ubuntu-cpu
       - ubuntu-vulkan
       - ubuntu-24-openvino
-      #- ubuntu-24-sycl
+      - ubuntu-24-sycl
       - android-arm64
       - macos-cpu
       - ios-xcode
@@ -1665,6 +1665,7 @@ jobs:
           tar -czvf release/llama-${{ steps.tag.outputs.name }}-ui.tar.gz --transform "s,^\.,llama-${{ steps.tag.outputs.name }}," -C ./ui-dist .
 
       - name: Attest release artifacts
+        id: attest
         uses: actions/attest@v4
         with:
           subject-path: 'release/*'
@@ -1696,6 +1697,9 @@ jobs:
             **Website:**
             - <https://llama.app>
 
+            **Attestations:**
+            - <${{ steps.attest.outputs.attestation-url }}>
+
             **macOS/iOS:**
             - [macOS Apple Silicon (arm64)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-macos-arm64.tar.gz)
             - macOS Apple Silicon (arm64, KleidiAI enabled) [DISABLED](https://github.com/ggml-org/llama.cpp/pull/23780)

common/arg.cpp

@@ -4658,6 +4658,12 @@ void common_params_add_preset_options(std::vector<common_arg> & args) {
         [](common_params &, int) { /* unused */ }
     ).set_env(COMMON_ARG_PRESET_STOP_TIMEOUT).set_preset_only());
 
+    args.push_back(common_arg(
+        {"dedup-cache-models"}, "0|1",
+        "in server router mode, hide a cached model from the model list when this preset resolves to the same model file",
+        [](common_params &, const std::string &) { /* unused */ }
+    ).set_env(COMMON_ARG_PRESET_DEDUP_CACHE_MODELS).set_preset_only());
+
     // args.push_back(common_arg(
     //     {"pin"},
     //     "in server router mode, do not unload this model if models_max is exceeded",

common/arg.h

@@ -11,8 +11,9 @@
 #include <memory>
 
 // pseudo-env variable to identify preset-only arguments
-#define COMMON_ARG_PRESET_LOAD_ON_STARTUP "__PRESET_LOAD_ON_STARTUP"
-#define COMMON_ARG_PRESET_STOP_TIMEOUT    "__PRESET_STOP_TIMEOUT"
+#define COMMON_ARG_PRESET_LOAD_ON_STARTUP    "__PRESET_LOAD_ON_STARTUP"
+#define COMMON_ARG_PRESET_STOP_TIMEOUT       "__PRESET_STOP_TIMEOUT"
+#define COMMON_ARG_PRESET_DEDUP_CACHE_MODELS "__PRESET_DEDUP_CACHE_MODELS"
 
 //
 // CLI argument parsing

common/common.cpp

@@ -1750,18 +1750,6 @@ struct ggml_threadpool_params ggml_threadpool_params_from_cpu_params(const commo
     return tpp;
 }
 
-namespace {
-
-bool can_share_threadpool(const ggml_threadpool_params & tpp1, const ggml_threadpool_params & tpp2) {
-    // n_threads does not matter -> we'll use what's larger
-    ggml_threadpool_params tpp_comparison = tpp1;
-    tpp_comparison.n_threads = tpp2.n_threads;
-
-    return ggml_threadpool_params_match(&tpp_comparison, &tpp2);
-}
-
-}  // namespace
-
 common_threadpools::~common_threadpools() {
     if (!free_fn) {
         return;
@@ -1790,9 +1778,9 @@ void common_threadpools::init(llama_context * ctx, const common_params & params)
     struct ggml_threadpool_params tpp =
             ggml_threadpool_params_from_cpu_params(params.cpuparams);
 
-    if (can_share_threadpool(tpp, tpp_batch)) {
-        tpp.n_threads = std::max(tpp.n_threads, tpp_batch.n_threads);
-    } else {
+    // each pool needs to match the respective n_threads exactly
+    // see: https://github.com/ggml-org/llama.cpp/pull/27138#issuecomment-5332307332
+    if (!ggml_threadpool_params_match(&tpp, &tpp_batch)) {
         threadpool_batch = ggml_threadpool_new_fn(&tpp_batch);
         if (!threadpool_batch) {
             COM_WRN("batch threadpool create failed : n_threads %d\n", tpp_batch.n_threads);

common/download.cpp

@@ -989,6 +989,26 @@ std::vector<common_cached_model_info> common_list_cached_models() {
     return result;
 }
 
+std::string common_download_resolve_path(const std::string & hf_repo_with_tag, const std::string & hf_file) {
+    auto [repo, tag] = common_download_split_repo_tag(hf_repo_with_tag);
+
+    auto files = hf_cache::get_cached_files(repo);
+    if (files.empty()) {
+        return "";
+    }
+
+    if (!hf_file.empty()) {
+        for (const auto & f : files) {
+            if (f.path == hf_file) {
+                return f.local_path;
+            }
+        }
+        return "";
+    }
+
+    return find_best_model(files, tag).local_path;
+}
+
 bool common_download_remove(const std::string & hf_repo_with_tag) {
     namespace fs = std::filesystem;
 

common/download.h

@@ -85,6 +85,10 @@ std::vector<std::string> common_download_get_all_parts(const std::string & url);
 // returns list of cached models
 std::vector<common_cached_model_info> common_list_cached_models();
 
+// resolve the local cached file path for a HF repo without network access (hf_file, if given, must match exactly)
+// returns an empty string if the model is not present in the cache
+std::string common_download_resolve_path(const std::string & hf_repo_with_tag, const std::string & hf_file = "");
+
 // download single file from url to local path
 // returns status code or -1 on error
 // skip_etag: if true, don't read/write .etag files (for HF cache where filename is the hash)

conversion/__init__.py

@@ -109,6 +109,8 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "GraniteSwitchForCausalLM": "granite",
     "GraniteSpeechForConditionalGeneration": "granite",
     "GraniteSpeechPlusForConditionalGeneration": "granite",
+    "GraniteSWAForCausalLM": "granite",
+    "GraniteMoeSWAForCausalLM": "granite",
     "Grok1ForCausalLM": "grok",
     "GrokForCausalLM": "grok",
     "GroveMoeForCausalLM": "grovemoe",

conversion/granite.py

@@ -74,6 +74,108 @@ class GraniteModel(LlamaModel):
         return super().filter_tensors(item)
 
 
+@ModelBase.register("GraniteSWAForCausalLM")
+class GraniteSWAModel(GraniteModel):
+    """Conversion for IBM's GraniteSWAForCausalLM (interleaved sliding window attention)"""
+    model_arch = gguf.MODEL_ARCH.GRANITE_SWA
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+
+        if name.endswith("sinks"):
+            name += ".weight"
+
+        return super().filter_tensors((name, gen))
+
+    def set_gguf_parameters(self):
+        """GraniteSWA uses Granite parameters plus sliding window configuration."""
+        super().set_gguf_parameters()
+
+        # Add sliding_window from config
+        sliding_window = self.hparams.get("sliding_window", 128)
+        self.gguf_writer.add_sliding_window(sliding_window)
+        logger.info("gguf: (granite_swa) sliding_window = %s", sliding_window)
+
+        # Derive sliding_window_pattern from layer_types
+        if layer_types := self.hparams.get("layer_types"):
+            is_swa = [t == "sliding_attention" for t in layer_types]
+            self.gguf_writer.add_sliding_window_pattern(is_swa)
+            logger.info("gguf: (granite_swa) sliding_window_pattern = %d SWA layers / %d total",
+                        sum(is_swa), len(is_swa))
+        else:
+            # Fall back to period-based pattern: i % 4 != 0
+            # This matches the transformers default pattern
+            n_layers = self.block_count
+            is_swa = [i % 4 != 0 for i in range(n_layers)]
+            self.gguf_writer.add_sliding_window_pattern(is_swa)
+            logger.info("gguf: (granite_swa) sliding_window_pattern (inferred) = %d SWA layers / %d total",
+                        sum(is_swa), n_layers)
+
+        # Add rope_pattern from no_rope_layers
+        if no_rope_layers := self.hparams.get("no_rope_layers"):
+            # Convert 1/0 to bool (1 = use RoPE, 0 = NoPE)
+            rope_pattern = [bool(x) for x in no_rope_layers]
+            self.gguf_writer.add_rope_pattern(rope_pattern)
+            logger.info("gguf: (granite_swa) rope_pattern = %d RoPE layers / %d total",
+                        sum(rope_pattern), len(rope_pattern))
+
+
+@ModelBase.register("GraniteMoeSWAForCausalLM")
+class GraniteMoeSWAModel(GraniteSWAModel):
+    """Conversion for IBM's GraniteMoeSWAForCausalLM (unified dense + MoE with iSWA)"""
+    model_arch = gguf.MODEL_ARCH.GRANITE_SWA
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        if shared_intermediate_size := self.hparams.get("shared_intermediate_size"):
+            self.gguf_writer.add_expert_shared_feed_forward_length(shared_intermediate_size)
+            logger.info("gguf: (granitemoewa) shared_intermediate_size = %s", shared_intermediate_size)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        """Split merged MoE tensors (gate+up) following standard MoE pattern."""
+
+        # Handle expert FFN tensors (merged gate+up) - swash format: experts.gate_up_proj
+        # Kept fused since inference (build_moe_ffn) supports a single gate_up_exps
+        # tensor for the routed experts.
+        if name.endswith("block_sparse_moe.experts.gate_up_proj"):
+            ffn_dim = self.hparams["intermediate_size"]
+            assert data_torch.shape[-2] == 2 * ffn_dim, f"Merged FFN tensor size must be 2 * intermediate_size, got {data_torch.shape[-2]}"
+            yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_UP_EXP, bid), bid)
+            return
+
+        # Handle expert FFN down projection - swash format: experts.down_proj
+        if name.endswith("block_sparse_moe.experts.down_proj"):
+            yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_EXP, bid), bid)
+            return
+
+        # Handle expert FFN tensors (merged gate+up) - standard granite format: input_linear.weight
+        # Kept fused since inference (build_moe_ffn) supports a single gate_up_exps
+        # tensor for the routed experts.
+        if name.endswith("block_sparse_moe.input_linear.weight"):
+            ffn_dim = self.hparams["intermediate_size"]
+            assert data_torch.shape[-2] == 2 * ffn_dim, "Merged FFN tensor size must be 2 * intermediate_size"
+            yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_UP_EXP, bid), bid)
+            return
+
+        # Handle shared expert FFN tensors (if present) - kept fused since
+        # inference (build_ffn) supports a single ffn_up_shexp tensor with
+        # LLM_FFN_SWIGLU for the shared expert.
+        if name.endswith("shared_mlp.input_linear.weight"):
+            ffn_dim = self.hparams.get("shared_intermediate_size", self.hparams["intermediate_size"])
+            assert data_torch.shape[-2] == 2 * ffn_dim, "Merged FFN tensor size must be 2 * shared_intermediate_size"
+            yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_SHEXP, bid), bid)
+            return
+
+        # Handle shared expert output (if present)
+        if name.endswith("shared_mlp.output_linear.weight"):
+            yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, bid), bid)
+            return
+
+        # Pass through to parent for all other tensors (including sinks)
+        yield from super().modify_tensors(data_torch, name, bid)
+
+
 @ModelBase.register("GraniteMoeForCausalLM", "GraniteMoeSharedForCausalLM")
 @ModelBase.example("ibm-granite/granite-3.1-3b-a800m-instruct")
 class GraniteMoeModel(GraniteModel):

ggml/include/ggml.h

@@ -1981,6 +1981,14 @@ extern "C" {
             float                 beta_fast,
             float                 beta_slow);
 
+    // set the offset dims for RoPE
+    // a must be GGML_OP_ROPE or GGML_OP_ROPE_BACK
+    // vision RoPE is not supported
+    // example: (marking: x = rotated, 0 = unrotated)
+    //     n_embd = 10, n_dims = 4, offset = 2 --> [00xxxx0000]
+    GGML_API struct ggml_tensor * ggml_rope_set_offset(
+            struct ggml_tensor  * a,
+            int                   n_offs);
 
     // clamp
     // in-place, returns view(a)

ggml/src/ggml-backend-impl.h

@@ -83,6 +83,7 @@ extern "C" {
     GGML_API ggml_backend_buffer_t ggml_backend_multi_buffer_alloc_buffer(ggml_backend_buffer_t * buffers, size_t n_buffers);
     GGML_API bool                  ggml_backend_buffer_is_multi_buffer(ggml_backend_buffer_t buffer);
     GGML_API void                  ggml_backend_multi_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage);
+    GGML_API void                  ggml_backend_meta_buffer_set_usage (ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage);
 
     //
     // Backend (meta)

ggml/src/ggml-backend-meta.cpp

@@ -1118,7 +1118,6 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
 }
 
 static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync) {
-    GGML_ASSERT(ggml_backend_buffer_is_meta(tensor->buffer));
     ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context;
     return ggml_backend_meta_get_split_state(buf_ctx->get_simple_tensor_container(tensor), tensor, assume_sync);
 }
@@ -1178,7 +1177,15 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m
         t_ij->flags = tensor->flags;
         memcpy(t_ij->op_params, tensor->op_params, sizeof(tensor->op_params));
         ggml_set_name(t_ij, tensor->name);
+
         t_ij->buffer = simple_buf;
+        if (simple_buf) {
+            // the backend that owns the buffer will set .extra
+            ggml_backend_buffer_init_tensor(simple_buf, t_ij);
+        } else {
+            t_ij->extra = tensor->extra;
+        }
+
         t_ij->view_src = tensor->view_src;
         t_ij->view_offs = tensor->view_offs;
         if (t_ij->view_src != nullptr && ggml_backend_buffer_is_meta(t_ij->view_src->buffer)) {
@@ -1209,7 +1216,6 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m
             t_ij->data = (char *) ggml_backend_buffer_get_base(simple_buf)
                 + size_t(tensor->data) - size_t(ggml_backend_buffer_get_base(tensor->buffer));
         }
-        t_ij->extra = tensor->extra;
         for (int i = 0; i < GGML_MAX_SRC; i++) {
             t_ij->src[i] = tensor->src[i];
             if (tensor->src[i] == tensor) {
@@ -1502,6 +1508,16 @@ bool ggml_backend_buffer_is_meta(ggml_backend_buffer_t buf) {
     return buf != nullptr && buf->iface.free_buffer == ggml_backend_meta_buffer_iface.free_buffer;
 }
 
+void ggml_backend_meta_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage) {
+    GGML_ASSERT(ggml_backend_buffer_is_meta(buffer));
+    ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context;
+    for (size_t i = 0; i < buf_ctx->bufs.size(); i++) {
+        if (buf_ctx->bufs[i]) {
+            ggml_backend_buffer_set_usage(buf_ctx->bufs[i].get(), usage);
+        }
+    }
+}
+
 static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
     const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft);
 

ggml/src/ggml-backend.cpp

@@ -182,6 +182,8 @@ void ggml_backend_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backe
     // FIXME: add a generic callback to the buffer interface
     if (ggml_backend_buffer_is_multi_buffer(buffer)) {
         ggml_backend_multi_buffer_set_usage(buffer, usage);
+    } else if (ggml_backend_buffer_is_meta(buffer)) {
+        ggml_backend_meta_buffer_set_usage(buffer, usage);
     }
 }
 

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

@@ -2534,6 +2534,9 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
             }
         case GGML_OP_ROPE:
             {
+                if (((const int32_t *) op->op_params)[15] != 0) {
+                    return false; // FIXME: support ggml_rope_set_offset
+                }
                 if (op->src[0]->ne[0] > 896) {
                     return false;
                 }

ggml/src/ggml-cpu/ops.cpp

@@ -5979,6 +5979,8 @@ static void ggml_compute_forward_rope_flt(
     memcpy(&beta_slow,   (int32_t *) dst->op_params + 10, sizeof(float));
     memcpy(&sections,    (int32_t *) dst->op_params + 11, sizeof(int)*4);
 
+    const int n_offs = ((int32_t *) dst->op_params)[15];
+
     GGML_TENSOR_UNARY_OP_LOCALS
 
     //printf("ne0: %d, ne1: %d, ne2: %d, ne3: %d\n", ne0, ne1, ne2, ne3);
@@ -5995,6 +5997,10 @@ static void ggml_compute_forward_rope_flt(
     GGML_ASSERT(n_dims <= ne0);
     GGML_ASSERT(n_dims % 2 == 0);
 
+    GGML_ASSERT(n_offs >= 0);
+    GGML_ASSERT(n_offs % 2 == 0);
+    GGML_ASSERT(n_offs + n_dims <= ne0);
+
     // rows per thread
     const int dr = (nr + nth - 1)/nth;
 
@@ -6020,6 +6026,7 @@ static void ggml_compute_forward_rope_flt(
 
     if (is_vision) {
         GGML_ASSERT(n_dims == ne0/2);
+        GGML_ASSERT(n_offs == 0);
     }
 
     const float * freq_factors = NULL;
@@ -6068,12 +6075,12 @@ static void ggml_compute_forward_rope_flt(
 
                 switch (mode) {
                     case GGML_ROPE_TYPE_NORMAL:
-                        rotate_pairs<T>(n_dims, 1, cache, src, dst_data, 1);
+                        rotate_pairs<T>(n_dims, 1, cache, src + n_offs, dst_data + n_offs, 1);
                         break;
                     case GGML_ROPE_TYPE_NEOX:
                     case GGML_ROPE_TYPE_MROPE:
                     case GGML_ROPE_TYPE_IMROPE:
-                        rotate_pairs<T>(n_dims, n_dims/2, cache, src, dst_data);
+                        rotate_pairs<T>(n_dims, n_dims/2, cache, src + n_offs, dst_data + n_offs);
                         break;
                     case GGML_ROPE_TYPE_VISION:
                         rotate_pairs<T>(ne0, n_dims, cache, src, dst_data);
@@ -6084,7 +6091,11 @@ static void ggml_compute_forward_rope_flt(
 
                 if (!is_vision) {
                     // fill the remain channels with data from src tensor
-                    for (int64_t i0 = n_dims; i0 < ne0; i0 += 2) {
+                    for (int64_t i0 = 0; i0 < ne0; i0 += 2) {
+                        if (i0 == n_offs) {
+                            i0 += n_dims - 2; // skip the rotated channels
+                            continue;
+                        }
                         const T * const src = (T *)((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
                         T * dst_data  = (T *)((char *)  dst->data + i3*nb3  + i2*nb2  + i1*nb1  + i0*nb0);
 

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

@@ -29,13 +29,15 @@ extern "C" {
 // FP16 to FP32 conversion
 
 // 16-bit float
-// on Arm, we use __fp16
+// on Arm, we use __fp16, which requires the IEEE fp16 format: implied on
+// AArch64, selected by -mfp16-format=ieee on 32 bit Arm, where the compiler
+// may otherwise reject the type
 // on x86, we use uint16_t
 //
 // for old CUDA compilers (<= 11), we use uint16_t: ref https://github.com/ggml-org/llama.cpp/pull/10616
 // for     MUSA compilers        , we use uint16_t: ref https://github.com/ggml-org/llama.cpp/pull/11843
 //
-#if defined(__ARM_NEON) && !(defined(__CUDACC__) && __CUDACC_VER_MAJOR__ <= 11) && !defined(__MUSACC__)
+#if defined(__ARM_NEON) && defined(__ARM_FP16_FORMAT_IEEE) && !(defined(__CUDACC__) && __CUDACC_VER_MAJOR__ <= 11) && !defined(__MUSACC__)
     #define GGML_CPU_COMPUTE_FP16_TO_FP32(x) neon_compute_fp16_to_fp32(x)
     #define GGML_CPU_COMPUTE_FP32_TO_FP16(x) neon_compute_fp32_to_fp16(x)
 
@@ -326,7 +328,7 @@ inline static float ggml_lookup_fp16_to_fp32(ggml_fp16_t f) {
     #define GGML_F16_VEC_REDUCE         GGML_F32Cx4_REDUCE
 #endif
 
-#elif defined(__ARM_NEON) && defined(__ARM_FEATURE_FMA)
+#elif defined(__ARM_NEON) && defined(__ARM_FEATURE_FMA) && defined(__ARM_FP16_FORMAT_IEEE)
 
 #define GGML_SIMD
 

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

@@ -2723,6 +2723,12 @@ static bool ggml_cuda_should_fuse_rms_norm_mul_rope(const ggml_tensor * rms_norm
         return false;
     }
 
+    // ggml_rope_set_offset is not yet supported in the fused kernel
+    const int n_offs = ((const int32_t *) rope->op_params)[15];
+    if (n_offs != 0) {
+        return false;
+    }
+
     return true;
 }
 

ggml/src/ggml-cuda/rope.cu

@@ -53,6 +53,7 @@ static __global__ void rope_norm(const T *            x,
                                  const int            s2,
                                  const int            s3,
                                  const int            n_dims,
+                                 const int            n_offs,
                                  const int32_t *      pos,
                                  const float          freq_scale,
                                  const float          ext_factor,
@@ -61,7 +62,8 @@ static __global__ void rope_norm(const T *            x,
                                  const float          theta_scale,
                                  const float *        freq_factors,
                                  const int64_t *      row_indices,
-                                 const int            set_rows_stride) {
+                                 const int            set_rows_stride,
+                                 const bool           inplace) {
     const int i0 = 2*(blockDim.y*blockIdx.y + threadIdx.y);
 
     if (i0 >= ne00) {
@@ -92,19 +94,24 @@ static __global__ void rope_norm(const T *            x,
             ggml_cuda_memcpy_1<4>(dst + idst, &v);
         }
     };
-    if (i0 >= n_dims) {
+    if (i0 < n_offs || i0 >= n_offs + n_dims) {
+        if (inplace) {
+            return;
+        }
         store_coaelsced(x[ix + 0], x[ix + 1]);
         return;
     }
 
-    const float theta_base = pos[i2]*powf(theta_scale, i0/2.0f);
+    const int iw = i0 - n_offs; // relative idx
 
-    const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f;
+    const float theta_base = pos[i2]*powf(theta_scale, iw/2.0f);
+
+    const float freq_factor = has_ff ? freq_factors[iw/2] : 1.0f;
 
     float cos_theta;
     float sin_theta;
 
-    rope_yarn<forward>(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta);
+    rope_yarn<forward>(theta_base/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor, cos_theta, sin_theta);
 
     const float x0 = x[ix + 0];
     const float x1 = x[ix + 1];
@@ -125,6 +132,7 @@ static __global__ void rope_neox(const T *            x,
                                  const int            s2,
                                  const int            s3,
                                  const int            n_dims,
+                                 const int            n_offs,
                                  const int32_t *      pos,
                                  const float          freq_scale,
                                  const float          ext_factor,
@@ -133,7 +141,8 @@ static __global__ void rope_neox(const T *            x,
                                  const float          theta_scale,
                                  const float *        freq_factors,
                                  const int64_t *      row_indices,
-                                 const int            set_rows_stride) {
+                                 const int            set_rows_stride,
+                                 const bool           inplace) {
     ggml_cuda_pdl_lc();
     const int i0 = 2*(blockDim.y*blockIdx.y + threadIdx.y);
 
@@ -158,27 +167,33 @@ static __global__ void rope_neox(const T *            x,
         idst += row_indices[i2] * set_rows_stride;
     }
 
-    if (i0 >= n_dims) {
+    if (i0 < n_offs || i0 >= n_offs + n_dims) {
+        if (inplace) {
+            return;
+        }
         dst[idst + i0 / 2 + 0] = ggml_cuda_cast<D>(x[ix + i0 / 2 + 0]);
         dst[idst + i0 / 2 + 1] = ggml_cuda_cast<D>(x[ix + i0 / 2 + 1]);
 
         return;
     }
 
-    const float theta_base = pos[i2]*powf(theta_scale, i0/2.0f);
+    const int iw = i0 - n_offs; // relative idx
 
-    const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f;
+    const float theta_base = pos[i2]*powf(theta_scale, iw/2.0f);
+
+    const float freq_factor = has_ff ? freq_factors[iw/2] : 1.0f;
 
     float cos_theta;
     float sin_theta;
 
-    rope_yarn<forward>(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta);
+    rope_yarn<forward>(theta_base/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor, cos_theta, sin_theta);
 
-    const float x0 = x[ix + 0];
-    const float x1 = x[ix + n_dims/2];
+    // idst/ix point at channel i0/2; the first channel of the rotated pair is n_offs + iw/2 = i0/2 + n_offs/2
+    const float x0 = x[ix + n_offs/2 + 0];
+    const float x1 = x[ix + n_offs/2 + n_dims/2];
 
-    dst[idst + 0]          = ggml_cuda_cast<D>(x0 * cos_theta - x1 * sin_theta);
-    dst[idst + n_dims / 2] = ggml_cuda_cast<D>(x0 * sin_theta + x1 * cos_theta);
+    dst[idst + n_offs/2 + 0]          = ggml_cuda_cast<D>(x0 * cos_theta - x1 * sin_theta);
+    dst[idst + n_offs/2 + n_dims / 2] = ggml_cuda_cast<D>(x0 * sin_theta + x1 * cos_theta);
 }
 
 template <bool forward, bool has_ff, typename T>
@@ -194,6 +209,7 @@ static __global__ void rope_multi(const T *            x,
                                   const int            s2,
                                   const int            s3,
                                   const int            n_dims,
+                                  const int            n_offs,
                                   const int32_t *      pos,
                                   const float          freq_scale,
                                   const float          ext_factor,
@@ -202,7 +218,8 @@ static __global__ void rope_multi(const T *            x,
                                   const float          theta_scale,
                                   const float *        freq_factors,
                                   const mrope_sections sections,
-                                  const bool           is_imrope) {
+                                  const bool           is_imrope,
+                                  const bool           inplace) {
     const int i0 = 2 * (blockDim.y * blockIdx.y + threadIdx.y);
 
     if (i0 >= ne00) {
@@ -219,52 +236,58 @@ static __global__ void rope_multi(const T *            x,
     const int ix   = i0 / 2 + i1 * s01 + i2 * s02 + i3 * s03;
 
     ggml_cuda_pdl_sync();
-    if (i0 >= n_dims) {
+    if (i0 < n_offs || i0 >= n_offs + n_dims) {
+        if (inplace) {
+            return;
+        }
         dst[idst + i0/2 + 0] = x[ix + i0/2 + 0];
         dst[idst + i0/2 + 1] = x[ix + i0/2 + 1];
 
         return;
     }
 
+    const int iw = i0 - n_offs; // relative idx
+
     const int sect_dims = sections.v[0] + sections.v[1] + sections.v[2] + sections.v[3];
     const int sec_w = sections.v[1] + sections.v[0];
-    const int sector = (i0 / 2) % sect_dims;
+    const int sector = (iw / 2) % sect_dims;
 
     float theta_base = 0.0;
     if (is_imrope) {
         if (sector % 3 == 1 && sector < 3 * sections.v[1]) {         // h
-            theta_base = pos[i2 + ne02 * 1] * powf(theta_scale, i0 / 2.0f);
+            theta_base = pos[i2 + ne02 * 1] * powf(theta_scale, iw / 2.0f);
         } else if (sector % 3 == 2 && sector < 3 * sections.v[2]) {  // w
-            theta_base = pos[i2 + ne02 * 2] * powf(theta_scale, i0 / 2.0f);
+            theta_base = pos[i2 + ne02 * 2] * powf(theta_scale, iw / 2.0f);
         } else if (sector % 3 == 0 && sector < 3 * sections.v[0]) {  // t
-            theta_base = pos[i2] * powf(theta_scale, i0 / 2.0f);
+            theta_base = pos[i2] * powf(theta_scale, iw / 2.0f);
         } else {
-            theta_base = pos[i2 + ne02 * 3] * powf(theta_scale, i0 / 2.0f);
+            theta_base = pos[i2 + ne02 * 3] * powf(theta_scale, iw / 2.0f);
         }
     } else {
         if (sector < sections.v[0]) {
-            theta_base = pos[i2] * powf(theta_scale, i0 / 2.0f);
+            theta_base = pos[i2] * powf(theta_scale, iw / 2.0f);
         } else if (sector >= sections.v[0] && sector < sec_w) {
-            theta_base = pos[i2 + ne02 * 1] * powf(theta_scale, i0 / 2.0f);
+            theta_base = pos[i2 + ne02 * 1] * powf(theta_scale, iw / 2.0f);
         } else if (sector >= sec_w && sector < sec_w + sections.v[2]) {
-            theta_base = pos[i2 + ne02 * 2] * powf(theta_scale, i0 / 2.0f);
+            theta_base = pos[i2 + ne02 * 2] * powf(theta_scale, iw / 2.0f);
         } else if (sector >= sec_w + sections.v[2]) {
-            theta_base = pos[i2 + ne02 * 3] * powf(theta_scale, i0 / 2.0f);
+            theta_base = pos[i2 + ne02 * 3] * powf(theta_scale, iw / 2.0f);
         }
     }
 
-    const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f;
+    const float freq_factor = has_ff ? freq_factors[iw/2] : 1.0f;
 
     float cos_theta;
     float sin_theta;
 
-    rope_yarn<forward>(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta);
+    rope_yarn<forward>(theta_base/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor, cos_theta, sin_theta);
 
-    const float x0 = x[ix + 0];
-    const float x1 = x[ix + n_dims/2];
+    // idst/ix point at channel i0/2; the first channel of the rotated pair is n_offs + iw/2 = i0/2 + n_offs/2
+    const float x0 = x[ix + n_offs/2 + 0];
+    const float x1 = x[ix + n_offs/2 + n_dims/2];
 
-    dst[idst + 0]        = x0*cos_theta - x1*sin_theta;
-    dst[idst + n_dims/2] = x0*sin_theta + x1*cos_theta;
+    dst[idst + n_offs/2 + 0]        = x0*cos_theta - x1*sin_theta;
+    dst[idst + n_offs/2 + n_dims/2] = x0*sin_theta + x1*cos_theta;
 }
 
 template <bool forward, bool has_ff, typename T>
@@ -344,6 +367,7 @@ static void rope_norm_cuda(const T *            x,
                            const int            s2,
                            const int            s3,
                            const int            n_dims,
+                           const int            n_offs,
                            const int            nr,
                            const int32_t *      pos,
                            const float          freq_scale,
@@ -354,6 +378,7 @@ static void rope_norm_cuda(const T *            x,
                            const float *        freq_factors,
                            const int64_t *      row_indices,
                            const int            set_rows_stride,
+                           const bool           inplace,
                            cudaStream_t         stream) {
     GGML_ASSERT(ne00 % 2 == 0);
     const dim3 block_dims(1, CUDA_ROPE_BLOCK_SIZE, 1);
@@ -364,12 +389,12 @@ static void rope_norm_cuda(const T *            x,
 
     if (freq_factors == nullptr) {
         rope_norm<forward, false><<<block_nums, block_dims, 0, stream>>>(
-            x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor,
-            attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride);
+            x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor,
+            attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, inplace);
     } else {
         rope_norm<forward, true><<<block_nums, block_dims, 0, stream>>>(
-            x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor,
-            attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride);
+            x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor,
+            attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, inplace);
     }
 }
 
@@ -386,6 +411,7 @@ static void rope_neox_cuda(const T *            x,
                            const int            s2,
                            const int            s3,
                            const int            n_dims,
+                           const int            n_offs,
                            const int            nr,
                            const int32_t *      pos,
                            const float          freq_scale,
@@ -396,6 +422,7 @@ static void rope_neox_cuda(const T *            x,
                            const float *        freq_factors,
                            const int64_t *      row_indices,
                            const int            set_rows_stride,
+                           const bool           inplace,
                            cudaStream_t         stream) {
     GGML_ASSERT(ne00 % 2 == 0);
     const dim3 block_dims(1, CUDA_ROPE_BLOCK_SIZE, 1);
@@ -407,12 +434,12 @@ static void rope_neox_cuda(const T *            x,
 
     if (freq_factors == nullptr) {
         ggml_cuda_kernel_launch(rope_neox<forward, false, T, D>, launch_params,
-            x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor,
-            attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride);
+            x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor,
+            attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, inplace);
     } else {
         ggml_cuda_kernel_launch(rope_neox<forward, true, T, D>, launch_params,
-            x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor,
-            attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride);
+            x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor,
+            attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, inplace);
     }
 }
 
@@ -429,6 +456,7 @@ static void rope_multi_cuda(const T *            x,
                             const int            s2,
                             const int            s3,
                             const int            n_dims,
+                            const int            n_offs,
                             const int            nr,
                             const int32_t *      pos,
                             const float          freq_scale,
@@ -439,6 +467,7 @@ static void rope_multi_cuda(const T *            x,
                             const float *        freq_factors,
                             const mrope_sections sections,
                             const bool           is_imrope,
+                            const bool           inplace,
                             cudaStream_t         stream) {
     GGML_ASSERT(ne00 % 2 == 0);
     const dim3 block_dims(1, CUDA_ROPE_BLOCK_SIZE, 1);
@@ -450,13 +479,13 @@ static void rope_multi_cuda(const T *            x,
     if (freq_factors == nullptr) {
         const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
         ggml_cuda_kernel_launch(rope_multi<forward, false, T>, launch_params,
-            x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor,
-            attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope);
+            x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor,
+            attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, inplace);
     } else {
         const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
         ggml_cuda_kernel_launch(rope_multi<forward, true, T>, launch_params,
-            x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor,
-            attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope);
+            x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor,
+            attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, inplace);
     }
 }
 
@@ -552,8 +581,12 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx,
     const int mode       = ((int32_t *) dst->op_params)[2];
     //const int n_ctx      = ((int32_t *) dst->op_params)[3];
     const int n_ctx_orig = ((int32_t *) dst->op_params)[4];
+    const int n_offs     = ((int32_t *) dst->op_params)[15];
     mrope_sections sections;
 
+    // when dst aliases src0, the channels outside the rotated window already hold the correct data
+    const bool inplace = dst_d == src0->data;
+
     // RoPE alteration for extended context
     float freq_base;
     float freq_scale;
@@ -581,6 +614,7 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx,
 
     if (is_vision) {
         GGML_ASSERT(n_dims == ne00/2);
+        GGML_ASSERT(n_offs == 0); // offset not supported for vision, as the rotated pairs span the whole row
     }
 
     const int32_t * pos = (const int32_t *) src1_d;
@@ -597,31 +631,31 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx,
     if (is_neox) {
         if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) {
             rope_neox_cuda<forward, float, float>((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02,
-                                                  s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base,
+                                                  s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
                                                   ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
-                                                  set_rows_stride, stream);
+                                                  set_rows_stride, inplace, stream);
         } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) {
             rope_neox_cuda<forward, float, half>((const float *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02,
-                                                 s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base,
+                                                 s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
                                                  ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
-                                                 set_rows_stride, stream);
+                                                 set_rows_stride, inplace, stream);
         } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) {
             rope_neox_cuda<forward, half, half>((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02,
-                                                s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base,
+                                                s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
                                                 ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
-                                                set_rows_stride, stream);
+                                                set_rows_stride, inplace, stream);
         } else {
             GGML_ABORT("fatal error");
         }
     } else if (is_mrope && !is_vision) {
         if (src0->type == GGML_TYPE_F32) {
             rope_multi_cuda<forward>((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02, s03, s1,
-                                     s2, s3, n_dims, nr, pos, freq_scale, freq_base, ext_factor, attn_factor,
-                                     corr_dims, freq_factors, sections, is_imrope, stream);
+                                     s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor,
+                                     corr_dims, freq_factors, sections, is_imrope, inplace, stream);
         } else if (src0->type == GGML_TYPE_F16) {
             rope_multi_cuda<forward>((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, s03, s1,
-                                     s2, s3, n_dims, nr, pos, freq_scale, freq_base, ext_factor, attn_factor,
-                                     corr_dims, freq_factors, sections, is_imrope, stream);
+                                     s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor,
+                                     corr_dims, freq_factors, sections, is_imrope, inplace, stream);
         } else {
             GGML_ABORT("fatal error");
         }
@@ -640,19 +674,19 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx,
     } else {
         if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) {
             rope_norm_cuda<forward, float, float>((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02,
-                                                  s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base,
+                                                  s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
                                                   ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
-                                                  set_rows_stride, stream);
+                                                  set_rows_stride, inplace, stream);
         } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) {
             rope_norm_cuda<forward, float, half>((const float *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02,
-                                                 s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base,
+                                                 s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
                                                  ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
-                                                 set_rows_stride, stream);
+                                                 set_rows_stride, inplace, stream);
         } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) {
             rope_norm_cuda<forward, half, half>((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02,
-                                                s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base,
+                                                s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
                                                 ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
-                                                set_rows_stride, stream);
+                                                set_rows_stride, inplace, stream);
         } else {
             GGML_ABORT("fatal error");
         }

ggml/src/ggml-et/ggml-et.cpp

@@ -1061,9 +1061,11 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm
                 const bool zero_view_offset = op->src[0]->view_src == nullptr || op->src[0]->view_offs == 0;
                 const bool has_sections = ggml_get_op_params_i32(op, 11) > 0 || ggml_get_op_params_i32(op, 12) > 0 ||
                                           ggml_get_op_params_i32(op, 13) > 0;
+                // FIXME: support ggml_rope_set_offset
+                const bool zero_rot_offset  = ggml_get_op_params_i32(op, 15) == 0;
 
                 supported =
-                    zero_view_offset && ndims <= 512 &&
+                    zero_view_offset && zero_rot_offset && ndims <= 512 &&
                     (is_normal || (is_neox && ndims % 16 == 0) || (is_imrope && ndims % 16 == 0 && has_sections));
             } else {
                 supported = false;

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

@@ -3180,6 +3180,10 @@ static bool ggml_hexagon_supported_argsort(const struct ggml_hexagon_session * s
 static bool ggml_hexagon_supported_rope(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
     const int32_t * op_params = &op->op_params[0];
 
+    if (op_params[15] != 0) {
+        return false; // FIXME: support ggml_rope_set_offset
+    }
+
     int mode = op_params[2];
 
     // n_dims == ne0/2, so the rotation spans the full row

ggml/src/ggml-hexagon/htp/flash-attn-ops.c

@@ -132,8 +132,8 @@ struct hmx_fa_context {
     __fp16 *     vtcm_v_tiles[2];      // V tiles (column-major, double-buffered)
     __fp16 *     vtcm_s_tiles[2];      // S = QK^T [g_br, Bc] (double-buffered)
     __fp16 *     vtcm_p_tiles[2];      // P = softmax(S) [g_br, Bc]
-    __fp16 *     vtcm_d_tiles;         // Diagonal rescale [g_br, g_br]
-    __fp16 *     vtcm_d_inv_l;         // Diagonal rescale (1/l) [g_br, g_br]
+    __fp16 *     vtcm_d_tiles[2];      // Diagonal rescale, g_br/32 packed diagonal tiles (double-buffered)
+    __fp16 *     vtcm_d_inv_l;         // Diagonal rescale (1/l), same packed layout
     HVX_Vector * vtcm_m_vec;           // Row max [g_br]
     HVX_Vector * vtcm_l_vec;           // Row sum [g_br]
     HVX_Vector * vtcm_s_rowmax;        // Softmax intermediate [g_br]
@@ -782,13 +782,14 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
             }
         }
 
-        // Initialize vtcm_d_tiles and vtcm_d_inv_l to 0
+        // Zero the whole rescale region: vtcm_d_tiles[0], the optional vtcm_d_tiles[1]
+        // and vtcm_d_inv_l are equal-sized and allocated back to back, so one run covers
+        // them all.  The scatter only ever writes the diagonal, ignore the rest.
         const size_t d_bytes_per_t = hex_align_up(d_tile_bytes / n, 128);
         const size_t d_start       = i * d_bytes_per_t;
         const size_t d_end         = hex_smin(d_start + d_bytes_per_t, d_tile_bytes);
         if (d_start < d_tile_bytes) {
-            hvx_splat_u8_a((char *) factx->vtcm_d_tiles + d_start, 0, d_end - d_start);
-            hvx_splat_u8_a((char *) factx->vtcm_d_inv_l + d_start, 0, d_end - d_start);
+            hvx_splat_u8_a((char *) factx->vtcm_d_tiles[0] + d_start, 0, d_end - d_start);
         }
     }
 
@@ -1432,17 +1433,19 @@ static inline void fa_softmax_impl(
         const HVX_VectorPred q_32_mask = Q6_Q_vsetq_R(32 * sizeof(__fp16));
         HVX_Vector           v_exp_m_diff = exp_m_diff_f16;
 
+        __fp16 * const d_tiles_out = factx->vtcm_d_tiles[args->buf_idx];
+
         size_t t0 = r_vec_idx * 2;
         if (t0 < args->n_row_tiles) {
             const HVX_Vector v_content = v_exp_m_diff;
-            __fp16 *         out_base  = factx->vtcm_d_tiles + t0 * (args->n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
+            __fp16 *         out_base  = d_tiles_out + t0 * HMX_FP16_TILE_N_ELMS;
             Q6_vscatter_QRMVhV(q_32_mask, (size_t) out_base, HMX_FP16_TILE_SIZE - 1, v_offsets, v_content);
         }
 
         size_t t1 = r_vec_idx * 2 + 1;
         if (t1 < args->n_row_tiles) {
             const HVX_Vector v_content = Q6_V_vror_VR(v_exp_m_diff, 64);
-            __fp16 *         out_base  = factx->vtcm_d_tiles + t1 * (args->n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
+            __fp16 *         out_base  = d_tiles_out + t1 * HMX_FP16_TILE_N_ELMS;
             Q6_vscatter_QRMVhV(q_32_mask, (size_t) out_base, HMX_FP16_TILE_SIZE - 1, v_offsets, v_content);
         }
     }
@@ -1506,7 +1509,7 @@ static __attribute__((noinline)) void fa_build_d_diag_inv_l(struct hmx_fa_contex
             v_content = Q6_V_vror_VR(v_content, 64);
         }
 
-        __fp16 * out_base = factx->vtcm_d_inv_l + i * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
+        __fp16 * out_base = factx->vtcm_d_inv_l + i * HMX_FP16_TILE_N_ELMS;
         Q6_vscatter_QRMVhV(q_32_mask, (size_t) out_base, HMX_FP16_TILE_SIZE - 1, v_offsets, v_content);
     }
 }
@@ -1615,7 +1618,7 @@ static void hmx_fa_o_update_worker(void * data) {
     const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS;
     const size_t v_stride = n_tiles_per_bc * HMX_FP16_TILE_N_ELMS;
     for (size_t r = 0; r < n_row_tiles; ++r) {
-        const __fp16 * d_diag     = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
+        const __fp16 * d_diag     = d_tiles + r * HMX_FP16_TILE_N_ELMS;
         const __fp16 * p_tile_in  = p_tiles + (r * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
         const __fp16 * o_rc       = o_prev + r * HMX_FP16_TILE_N_ELMS;
         const __fp16 * v_tile_in  = v_tiles;
@@ -1654,7 +1657,7 @@ static void hmx_fa_o_norm_worker(void * data) {
     asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)job->hmx_scales));
     const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS;
     for (size_t r = 0; r < n_row_tiles; ++r) {
-        const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
+        const __fp16 * d_diag = d_tiles + r * HMX_FP16_TILE_N_ELMS;
         const __fp16 * o_rc = o_prev + r * HMX_FP16_TILE_N_ELMS;
         __fp16 *       o_out = o_curr + r * DV_tiles * HMX_FP16_TILE_N_ELMS;
 
@@ -1882,7 +1885,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     factx.vtcm_s_tiles[1]     = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_s_tiles[1], pipeline);
     factx.vtcm_p_tiles[0]     = VTCM_LAYOUT_PTR(__fp16, base, L.off_p_tiles[0]);
     factx.vtcm_p_tiles[1]     = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_p_tiles[1], pipeline);
-    factx.vtcm_d_tiles        = VTCM_LAYOUT_PTR(__fp16, base, L.off_d_tiles);
+    factx.vtcm_d_tiles[0]     = VTCM_LAYOUT_PTR(__fp16, base, L.off_d_tiles[0]);
+    factx.vtcm_d_tiles[1]     = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_d_tiles[1], pipeline);
     factx.vtcm_d_inv_l        = VTCM_LAYOUT_PTR(__fp16, base, L.off_d_inv_l);
     factx.vtcm_m_vec          = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_m_vec);
     factx.vtcm_l_vec          = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_l_vec);
@@ -2039,7 +2043,30 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                             }
                         }
 
-                        // ---- 3. Pop and run K-prep for next block & push next QK-dot ----
+                        // ---- 3. Start HMX O update for block kv_blk - 1 (reads P[1 - buf_idx], V[1 - buf_idx], D) ----
+                        // O update relys on the previous block's P and V tiles.
+                        // O update MUST be pushed before the next block's QK-dot: hmx_queue_pop() retires the
+                        // oldest descriptor, so push order alone decides which pop waits for which job.
+                        // If OU went in after QK(i+1), the pop below would retire QK(i+1) and leave
+                        // OU(i-1) in flight into the next iteration, where V-prep overwrites V[prev_buf].
+                        if (kv_blk > 0) {
+                            const size_t prev_buf        = 1 - buf_idx;
+                            ou_job[prev_buf].o_curr      = o_tile_curr;
+                            ou_job[prev_buf].o_prev      = o_tile_prev;
+                            ou_job[prev_buf].p_tiles     = factx.vtcm_p_tiles[prev_buf];
+                            ou_job[prev_buf].v_tiles     = factx.vtcm_v_tiles[prev_buf];
+                            ou_job[prev_buf].d_tiles     = factx.vtcm_d_tiles[prev_buf];
+                            ou_job[prev_buf].hmx_scales  = factx.vtcm_hmx_scales_id;
+                            ou_job[prev_buf].n_row_tiles = n_row_tiles;
+                            ou_job[prev_buf].n_col_tiles =
+                                hmx_ceil_div(hex_smin(Bc, nek1 - (kv_blk - 1) * Bc), HMX_FP16_TILE_N_COLS);
+                            ou_job[prev_buf].n_row_tiles_g_br = n_row_tiles_g_br;
+                            ou_job[prev_buf].n_tiles_per_bc   = n_tiles_per_bc;
+                            ou_job[prev_buf].DV               = DV;
+                            hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job[prev_buf]));
+                        }
+
+                        // ---- 4. Pop and run K-prep for next block & push next QK-dot ----
                         if (kv_blk + 1 < factx.n_kv_blocks) {
                             const uint32_t next_start = (kv_blk + 1) * Bc;
                             const uint32_t next_rows  = hex_smin(Bc, nek1 - next_start);
@@ -2059,10 +2086,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                             hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job[next_buf]));
                         }
 
-                        // ---- 4. Wait for current block's QK-dot to finish ----
+                        // ---- 5. Wait for current block's QK-dot to finish ----
                         hmx_queue_pop(hmx_q);
 
-                        // ---- 5. Phase 2: softmax + build_D ----
+                        // ---- 6. Phase 2: softmax + build_D ----
                         fa_softmax_args_t sargs;
                         memset(&sargs, 0, sizeof(sargs));
                         sargs.factx                = &factx;
@@ -2085,23 +2112,6 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                         sargs.mask_vtcm_row_stride = factx.mask_buf_row_stride;
                         sargs.slopes               = factx.vtcm_slopes;
 
-                        // Start HMX O update for block kv_blk - 1 (reads P[1 - buf_idx], V[1 - buf_idx])
-                        if (kv_blk > 0) {
-                            const size_t prev_buf = 1 - buf_idx;
-                            ou_job[prev_buf].o_curr           = o_tile_curr;
-                            ou_job[prev_buf].o_prev           = o_tile_prev;
-                            ou_job[prev_buf].p_tiles          = factx.vtcm_p_tiles[prev_buf];
-                            ou_job[prev_buf].v_tiles          = factx.vtcm_v_tiles[prev_buf];
-                            ou_job[prev_buf].d_tiles          = factx.vtcm_d_tiles;
-                            ou_job[prev_buf].hmx_scales       = factx.vtcm_hmx_scales_id;
-                            ou_job[prev_buf].n_row_tiles      = n_row_tiles;
-                            ou_job[prev_buf].n_col_tiles      = hmx_ceil_div(hex_smin(Bc, nek1 - (kv_blk - 1) * Bc), HMX_FP16_TILE_N_COLS);
-                            ou_job[prev_buf].n_row_tiles_g_br = n_row_tiles_g_br;
-                            ou_job[prev_buf].n_tiles_per_bc   = n_tiles_per_bc;
-                            ou_job[prev_buf].DV               = DV;
-                            hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job[prev_buf]));
-                        }
-
                         // Run Softmax on HVX (blocking call)
                         fa_phase_softmax_and_build_d(&factx, &sargs, n_row_tiles, n_row_tiles_g_br);
 
@@ -2128,7 +2138,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                         ou_job[0].o_prev           = o_tile_prev;
                         ou_job[0].p_tiles          = factx.vtcm_p_tiles[1 - buf_idx];
                         ou_job[0].v_tiles          = factx.vtcm_v_tiles[1 - buf_idx];
-                        ou_job[0].d_tiles          = factx.vtcm_d_tiles;
+                        ou_job[0].d_tiles          = factx.vtcm_d_tiles[1 - buf_idx];
                         ou_job[0].hmx_scales       = factx.vtcm_hmx_scales_id;
                         ou_job[0].n_row_tiles      = n_row_tiles;
                         ou_job[0].n_col_tiles      = last_cols;
@@ -2232,7 +2242,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                             ou_job.o_prev           = o_tile_prev;
                             ou_job.p_tiles          = factx.vtcm_p_tiles[0];
                             ou_job.v_tiles          = factx.vtcm_v_tiles[0];
-                            ou_job.d_tiles          = factx.vtcm_d_tiles;
+                            ou_job.d_tiles          = factx.vtcm_d_tiles[0];
                             ou_job.hmx_scales       = factx.vtcm_hmx_scales_id;
                             ou_job.n_row_tiles      = n_row_tiles;
                             ou_job.n_col_tiles      = n_col_tiles;

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

@@ -109,7 +109,7 @@ struct hmx_fa_vtcm_layout {
     size_t off_v_tiles[2];
     size_t off_s_tiles[2];
     size_t off_p_tiles[2];
-    size_t off_d_tiles;
+    size_t off_d_tiles[2];
     size_t off_d_inv_l;
     size_t off_m_vec;
     size_t off_l_vec;
@@ -125,7 +125,7 @@ struct hmx_fa_vtcm_layout {
     size_t q_tile_bytes;
     size_t o_tile_bytes;
     size_t s_tile_bytes;       // S and P tiles (same size)
-    size_t d_tile_bytes;
+    size_t d_tile_bytes;       // d_tiles[0..1] + d_inv_l, allocated back to back
     size_t m_line_bytes;       // one mask row
     size_t m_buf_slot_bytes;   // one dma_cache slot = align_up(Br * m_line_bytes, 4096)
     size_t col_vec_bytes;
@@ -149,7 +149,12 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
     const size_t k_tile_size  = hex_align_up(Bc   * DK   * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
     const size_t v_tile_size  = hex_align_up(Bc   * DV   * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
     const size_t s_tile_size  = hex_align_up(g_br * Bc   * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
-    const size_t d_tile_size  = hex_align_up(g_br * g_br * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
+
+    // The rescale matrices are diagonal: the HMX kernels only ever load the g_br/32
+    // tiles that sit on the diagonal, so store just those, packed back to back with
+    // a stride of one tile.  The old [g_br, g_br] square layout allocated g_br/32
+    // times more than it used, which is also why a second D buffer was unaffordable.
+    const size_t d_tile_size = (g_br / HMX_FP16_TILE_N_ROWS) * HTP_FA_HMX_TILE_SIZE;
 
     const size_t q_dma_size   = hex_align_up(g_br * DK * (is_q_fp32 ? sizeof(float) : sizeof(__fp16)), 128);
     const size_t k_dma_size   = hex_align_up(Bc * hex_round_up(DK * sizeof(__fp16), 128), 128);
@@ -167,7 +172,8 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
     VTCM_LAYOUT_ALLOC(off, off_q_tiles,       q_tile_size);
     VTCM_LAYOUT_ALLOC(off, off_o_tiles[0],    o_tile_size);
     VTCM_LAYOUT_ALLOC(off, off_o_tiles[1],    o_tile_size);
-    VTCM_LAYOUT_ALLOC(off, off_d_tiles,       d_tile_size);
+    VTCM_LAYOUT_ALLOC(off, off_d_tiles[0],    d_tile_size);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_d_tiles[1], d_tile_size, pipeline);
     VTCM_LAYOUT_ALLOC(off, off_d_inv_l,       d_tile_size);
 
     // Group B & C share start offset (Group B tiles must be 2KB aligned)
@@ -213,7 +219,10 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
     L->o_tile_bytes        = o_tile_size;
     L->col_vec_bytes       = col_vec_size;
     L->s_tile_bytes        = s_tile_size;
-    L->d_tile_bytes        = d_tile_size;
+    // Measured from the actual offsets rather than assumed to be N * d_tile_size, so
+    // that inserting a region between them (or adding padding to VTCM_LAYOUT_ALLOC)
+    // cannot silently leave the tail of the run unzeroed.
+    L->d_tile_bytes        = (L->off_d_inv_l + d_tile_size) - L->off_d_tiles[0];
     L->m_line_bytes        = m_line_size;
     L->m_buf_slot_bytes    = m_buf_slot;
     L->row_buf_stride      = row_vec_size / 128;

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

@@ -329,6 +329,7 @@ typedef struct {
     uint64_t nb3;
     int32_t  n_past;
     int32_t  n_dims;
+    int32_t  n_offs;
     int32_t  n_ctx_orig;
     float    freq_base;
     float    freq_scale;
@@ -341,6 +342,7 @@ typedef struct {
     int32_t  sect_2;
     int32_t  sect_3;
     bool     src2;
+    bool     inplace;
 } ggml_metal_kargs_rope;
 
 typedef struct {

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

@@ -3884,6 +3884,11 @@ int ggml_metal_op_rope(ggml_metal_op_t ctx, int idx) {
     const int sect_2 = ((const int32_t *) op->op_params)[13];
     const int sect_3 = ((const int32_t *) op->op_params)[14];
 
+    const int n_offs = ((const int32_t *) op->op_params)[15];
+
+    // when dst aliases src0, the channels outside the rotated window already hold the correct data
+    const bool inplace = op->data == op->src[0]->data;
+
     ggml_metal_kargs_rope args = {
         /*.ne00        =*/ ne00,
         /*.ne01        =*/ ne01,
@@ -3903,6 +3908,7 @@ int ggml_metal_op_rope(ggml_metal_op_t ctx, int idx) {
         /*.nb3         =*/ nb3,
         /*.n_past      =*/ n_past,
         /*.n_dims      =*/ n_dims,
+        /*.n_offs      =*/ n_offs,
         /*.n_ctx_orig  =*/ n_ctx_orig,
         /*.freq_base   =*/ freq_base,
         /*.freq_scale  =*/ freq_scale,
@@ -3915,6 +3921,7 @@ int ggml_metal_op_rope(ggml_metal_op_t ctx, int idx) {
         /* sect_2      =*/ sect_2,
         /* sect_3      =*/ sect_3,
         /* src2        =*/ op->src[2] != nullptr,
+        /* inplace     =*/ inplace,
     };
 
     auto pipeline = ggml_metal_library_get_pipeline_rope(lib, op);

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

@@ -656,13 +656,13 @@ void dequantize_q5_1_t4(device const block_q5_1 * xb, short il, thread type4 & r
 
 template <typename type4x4>
 void dequantize_q8_0(device const block_q8_0 *xb, short il, thread type4x4 & reg) {
-    device const int8_t * qs = ((device const int8_t *)xb->qs);
+    device const packed_char4 * qs = (device const packed_char4 *) xb->qs;
     const float d = xb->d;
 
     float4x4 reg_f;
 
-    for (int i = 0; i < 16; i++) {
-        reg_f[i/4][i%4] = (qs[i + 16*il] * d);
+    for (int i = 0; i < 4; ++i) {
+        reg_f[i] = float4(qs[4*il + i]) * d;
     }
 
     reg = (type4x4) reg_f;
@@ -670,12 +670,10 @@ void dequantize_q8_0(device const block_q8_0 *xb, short il, thread type4x4 & reg
 
 template <typename type4>
 void dequantize_q8_0_t4(device const block_q8_0 *xb, short il, thread type4 & reg) {
-    device const int8_t * qs = ((device const int8_t *)xb->qs);
+    device const packed_char4 * qs = (device const packed_char4 *) xb->qs;
     const float d = xb->d;
 
-    for (int i = 0; i < 4; i++) {
-        reg[i] = (qs[4*(il%4) + i + 16*(il/4)] * d);
-    }
+    reg = (type4) (float4(qs[il]) * d);
 }
 
 template <typename type4x4>
@@ -4688,14 +4686,15 @@ kernel void kernel_rope_norm(
     float sin_theta;
 
     for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
-        if (i0 < args.n_dims) {
-            const int ic = i0/2;
+        if (i0 >= args.n_offs && i0 < args.n_offs + args.n_dims) {
+            const int iw = i0 - args.n_offs; // relative idx
+            const int ic = iw/2;
 
-            const float theta = theta_base * pow(args.freq_base, inv_ndims*i0);
+            const float theta = theta_base * pow(args.freq_base, inv_ndims*iw);
 
             const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
 
-            rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
+            rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, iw, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
 
             device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
             device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + i0*args.nb0);
@@ -4706,6 +4705,10 @@ kernel void kernel_rope_norm(
             dst_data[0] = x0*cos_theta - x1*sin_theta;
             dst_data[1] = x0*sin_theta + x1*cos_theta;
         } else {
+            if (args.inplace) {
+                continue;
+            }
+
             device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
             device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + i0*args.nb0);
 
@@ -4741,17 +4744,18 @@ kernel void kernel_rope_neox(
     float sin_theta;
 
     for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
-        if (i0 < args.n_dims) {
-            const int ic = i0/2;
+        if (i0 >= args.n_offs && i0 < args.n_offs + args.n_dims) {
+            const int iw = i0 - args.n_offs; // relative idx
+            const int ic = iw/2;
 
-            const float theta = theta_base * pow(args.freq_base, inv_ndims*i0);
+            const float theta = theta_base * pow(args.freq_base, inv_ndims*iw);
 
             const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
 
-            rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
+            rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, iw, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
 
-            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + ic*args.nb00);
-            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + ic*args.nb0);
+            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + (args.n_offs + ic)*args.nb00);
+            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + (args.n_offs + ic)*args.nb0);
 
             const float x0 = src[0];
             const float x1 = src[args.n_dims/2];
@@ -4759,6 +4763,10 @@ kernel void kernel_rope_neox(
             dst_data[0]             = x0*cos_theta - x1*sin_theta;
             dst_data[args.n_dims/2] = x0*sin_theta + x1*cos_theta;
         } else {
+            if (args.inplace) {
+                continue;
+            }
+
             device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
             device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + i0*args.nb0);
 
@@ -4793,8 +4801,9 @@ kernel void kernel_rope_multi(
     float sin_theta;
 
     for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
-        if (i0 < args.n_dims) {
-            const int ic = i0/2;
+        if (i0 >= args.n_offs && i0 < args.n_offs + args.n_dims) {
+            const int iw = i0 - args.n_offs; // relative idx
+            const int ic = iw/2;
 
             // mrope theta calculations
             // note: the rest is the same as kernel_rope_neox
@@ -4827,14 +4836,14 @@ kernel void kernel_rope_multi(
             }
             // end of mrope
 
-            const float theta = theta_base * pow(args.freq_base, inv_ndims*i0);
+            const float theta = theta_base * pow(args.freq_base, inv_ndims*iw);
 
             const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
 
-            rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
+            rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, iw, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
 
-            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + ic*args.nb00);
-            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + ic*args.nb0);
+            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + (args.n_offs + ic)*args.nb00);
+            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + (args.n_offs + ic)*args.nb0);
 
             const float x0 = src[0];
             const float x1 = src[args.n_dims/2];
@@ -4842,6 +4851,10 @@ kernel void kernel_rope_multi(
             dst_data[0]             = x0*cos_theta - x1*sin_theta;
             dst_data[args.n_dims/2] = x0*sin_theta + x1*cos_theta;
         } else {
+            if (args.inplace) {
+                continue;
+            }
+
             device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
             device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + i0*args.nb0);
 

ggml/src/ggml-opencl/CMakeLists.txt

@@ -202,6 +202,7 @@ set(GGML_OPENCL_KERNELS
     sqr
     sqrt
     ssm_conv
+    ssm_scan
     gated_delta_net
     sub
     sum_rows

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

@@ -866,6 +866,9 @@ struct ggml_backend_opencl_context {
     // [size_idx][kda][tgpp] where size_idx: 0=S_V=16, 1=32, 2=64, 3=128; kda: 0 or 1.
     // tgpp 0 = TG variant (COLS_PER_LANE_GROUP=1), tgpp 1 = prefill variant (COLS_PER_LANE_GROUP=4).
     cl_kernel kernel_gated_delta_net_f32[4][2][2] = {};
+    cl_kernel kernel_ssm_scan_f32_mamba2_d128 = nullptr;
+    cl_kernel kernel_ssm_scan_f32_mamba2_d256 = nullptr;
+
     cl_kernel kernel_timestep_embedding;
     cl_kernel kernel_gemv_moe_q4_0_f32_ns, kernel_gemm_moe_q4_0_f32_ns, kernel_gemm_moe_q4_0_f32_ns_bin;
     cl_kernel kernel_gemm_moe_q8_0_f32_ns;
@@ -3154,6 +3157,24 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
         GGML_LOG_CONT(".");
     }
 
+    // ssm_scan (Mamba-2 fused per-token recurrent step; d_state in {128, 256})
+    {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+        const std::string kernel_src {
+            #include "ssm_scan.cl.h"
+        };
+#else
+        const std::string kernel_src = read_file("ssm_scan.cl");
+#endif
+        cl_program prog =
+            build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
+
+        CL_CHECK((backend_ctx->kernel_ssm_scan_f32_mamba2_d128 = clCreateKernel(prog, "kernel_ssm_scan_f32_mamba2_d128", &err), err));
+        CL_CHECK((backend_ctx->kernel_ssm_scan_f32_mamba2_d256 = clCreateKernel(prog, "kernel_ssm_scan_f32_mamba2_d256", &err), err));
+        CL_CHECK(clReleaseProgram(prog));
+        GGML_LOG_CONT(".");
+    }
+
     // gated_delta_net: one kernel per (S_V, KDA, tgpp) triple.
     {
     #ifdef GGML_OPENCL_EMBED_KERNELS
@@ -7301,6 +7322,23 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
                    (op->src[0]->type == GGML_TYPE_F16 && op->src[1]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32);
         case GGML_OP_SSM_CONV:
             return (op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32);
+        case GGML_OP_SSM_SCAN: {
+            // Mamba-2 fused per-token scan. Requires src3->ne[0] == 1 (scalar
+            // A per head); d_state in {128, 256}; all sources f32. Falls back
+            // to CPU otherwise (incl. Mamba-1 element-wise A).
+            for (int i = 0; i < 6; ++i) {
+                if (op->src[i]->type != GGML_TYPE_F32) {
+                    return false;
+                }
+            }
+            if (op->type != GGML_TYPE_F32) {
+                return false;
+            }
+            const int K = ggml_get_op_params_i32(op, 0);
+            const int d_state = (int) op->src[0]->ne[0];
+            const bool is_mamba2 = (op->src[3]->ne[0] == 1);
+            return is_mamba2 && (d_state == 128 || d_state == 256) && (K == 1);
+        }
         case GGML_OP_GATED_DELTA_NET:
             {
                 // Match the Vulkan backend: only F32 -> F32, S_v in {16, 32, 64, 128}.
@@ -7376,6 +7414,9 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
         case GGML_OP_DIAG_MASK_INF:
             return op->ne[3] == 1;
         case GGML_OP_ROPE: {
+            if (((const int32_t *) op->op_params)[15] != 0) {
+                return false; // FIXME: support ggml_rope_set_offset
+            }
             const int mode = ((const int32_t *) op->op_params)[2];
             const bool is_mrope = mode & GGML_ROPE_TYPE_MROPE;
             const bool is_vision = mode == GGML_ROPE_TYPE_VISION;
@@ -12257,6 +12298,103 @@ static void ggml_cl_mean(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_ssm_scan(ggml_backend_t backend, ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0]; // s
+    const ggml_tensor * src1 = dst->src[1]; // x
+    const ggml_tensor * src2 = dst->src[2]; // dt
+    const ggml_tensor * src3 = dst->src[3]; // A
+    const ggml_tensor * src4 = dst->src[4]; // B
+    const ggml_tensor * src5 = dst->src[5]; // C
+    const ggml_tensor * src6 = dst->src[6]; // ids
+
+    GGML_ASSERT(src0 && src1 && src2 && src3 && src4 && src5 && src6 && dst);
+
+    ggml_backend_opencl_context * backend_ctx = (ggml_backend_opencl_context *) backend->context;
+
+    ggml_tensor_extra_cl * e0 = (ggml_tensor_extra_cl *) src0->extra;
+    ggml_tensor_extra_cl * e1 = (ggml_tensor_extra_cl *) src1->extra;
+    ggml_tensor_extra_cl * e2 = (ggml_tensor_extra_cl *) src2->extra;
+    ggml_tensor_extra_cl * e3 = (ggml_tensor_extra_cl *) src3->extra;
+    ggml_tensor_extra_cl * e4 = (ggml_tensor_extra_cl *) src4->extra;
+    ggml_tensor_extra_cl * e5 = (ggml_tensor_extra_cl *) src5->extra;
+    ggml_tensor_extra_cl * e6 = (ggml_tensor_extra_cl *) src6->extra;
+    ggml_tensor_extra_cl * ed = (ggml_tensor_extra_cl *) dst->extra;
+
+    cl_ulong o0 = e0->offset + src0->view_offs;
+    cl_ulong o1 = e1->offset + src1->view_offs;
+    cl_ulong o2 = e2->offset + src2->view_offs;
+    cl_ulong o3 = e3->offset + src3->view_offs;
+    cl_ulong o4 = e4->offset + src4->view_offs;
+    cl_ulong o5 = e5->offset + src5->view_offs;
+    cl_ulong o6 = e6->offset + src6->view_offs;
+    cl_ulong od = ed->offset + dst->view_offs;
+
+    const int d_state  = (int) src0->ne[0];
+    const int head_dim = (int) src0->ne[1];
+    const int n_head   = (int) src1->ne[1];
+    const int n_group  = (int) src4->ne[1];
+    const int n_tokens = (int) src1->ne[2];
+    const int n_seqs   = (int) src1->ne[3];
+
+    // Mirror CPU ref: s_off = ggml_nelements(src1) * sizeof(float)
+    const cl_ulong s_off_bytes = (cl_ulong) ggml_nelements(src1) * sizeof(float);
+
+    cl_kernel kernel = (d_state == 128)
+        ? backend_ctx->kernel_ssm_scan_f32_mamba2_d128
+        : backend_ctx->kernel_ssm_scan_f32_mamba2_d256;
+    GGML_ASSERT(kernel != nullptr);
+
+    cl_ulong s0_nb2 = src0->nb[2];
+    cl_ulong s0_nb3 = src0->nb[3];
+    cl_ulong x_nb2  = src1->nb[2];
+    cl_ulong x_nb3  = src1->nb[3];
+    cl_ulong dt_nb1 = src2->nb[1];
+    cl_ulong dt_nb2 = src2->nb[2];
+    cl_ulong A_nb1  = src3->nb[1];
+    cl_ulong B_nb2  = src4->nb[2];
+    cl_ulong B_nb3  = src4->nb[3];
+    cl_ulong C_nb2  = src5->nb[2];
+    cl_ulong C_nb3  = src5->nb[3];
+
+    CL_CHECK(clSetKernelArg(kernel,  0, sizeof(cl_mem),   &e0->data_device));
+    CL_CHECK(clSetKernelArg(kernel,  1, sizeof(cl_ulong), &o0));
+    CL_CHECK(clSetKernelArg(kernel,  2, sizeof(cl_mem),   &e1->data_device));
+    CL_CHECK(clSetKernelArg(kernel,  3, sizeof(cl_ulong), &o1));
+    CL_CHECK(clSetKernelArg(kernel,  4, sizeof(cl_mem),   &e2->data_device));
+    CL_CHECK(clSetKernelArg(kernel,  5, sizeof(cl_ulong), &o2));
+    CL_CHECK(clSetKernelArg(kernel,  6, sizeof(cl_mem),   &e3->data_device));
+    CL_CHECK(clSetKernelArg(kernel,  7, sizeof(cl_ulong), &o3));
+    CL_CHECK(clSetKernelArg(kernel,  8, sizeof(cl_mem),   &e4->data_device));
+    CL_CHECK(clSetKernelArg(kernel,  9, sizeof(cl_ulong), &o4));
+    CL_CHECK(clSetKernelArg(kernel, 10, sizeof(cl_mem),   &e5->data_device));
+    CL_CHECK(clSetKernelArg(kernel, 11, sizeof(cl_ulong), &o5));
+    CL_CHECK(clSetKernelArg(kernel, 12, sizeof(cl_mem),   &e6->data_device));
+    CL_CHECK(clSetKernelArg(kernel, 13, sizeof(cl_ulong), &o6));
+    CL_CHECK(clSetKernelArg(kernel, 14, sizeof(cl_mem),   &ed->data_device));
+    CL_CHECK(clSetKernelArg(kernel, 15, sizeof(cl_ulong), &od));
+    CL_CHECK(clSetKernelArg(kernel, 16, sizeof(cl_ulong), &s0_nb2));
+    CL_CHECK(clSetKernelArg(kernel, 17, sizeof(cl_ulong), &s0_nb3));
+    CL_CHECK(clSetKernelArg(kernel, 18, sizeof(cl_ulong), &x_nb2));
+    CL_CHECK(clSetKernelArg(kernel, 19, sizeof(cl_ulong), &x_nb3));
+    CL_CHECK(clSetKernelArg(kernel, 20, sizeof(cl_ulong), &dt_nb1));
+    CL_CHECK(clSetKernelArg(kernel, 21, sizeof(cl_ulong), &dt_nb2));
+    CL_CHECK(clSetKernelArg(kernel, 22, sizeof(cl_ulong), &A_nb1));
+    CL_CHECK(clSetKernelArg(kernel, 23, sizeof(cl_ulong), &B_nb2));
+    CL_CHECK(clSetKernelArg(kernel, 24, sizeof(cl_ulong), &B_nb3));
+    CL_CHECK(clSetKernelArg(kernel, 25, sizeof(cl_ulong), &C_nb2));
+    CL_CHECK(clSetKernelArg(kernel, 26, sizeof(cl_ulong), &C_nb3));
+    CL_CHECK(clSetKernelArg(kernel, 27, sizeof(cl_ulong), &s_off_bytes));
+    CL_CHECK(clSetKernelArg(kernel, 28, sizeof(int),      &head_dim));
+    CL_CHECK(clSetKernelArg(kernel, 29, sizeof(int),      &n_head));
+    CL_CHECK(clSetKernelArg(kernel, 30, sizeof(int),      &n_group));
+    CL_CHECK(clSetKernelArg(kernel, 31, sizeof(int),      &n_tokens));
+
+    size_t global_work_size[] = { (size_t)n_head * head_dim * 64, (size_t)n_seqs, 1 };
+    size_t local_work_size[]  = { 64, 1, 1 };
+
+    backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+}
+
 static void ggml_cl_ssm_conv(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
     GGML_ASSERT(src0);
     GGML_ASSERT(src0->extra);
@@ -24743,6 +24881,14 @@ bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor
             }
             func = ggml_cl_ssm_conv;
             break;
+        case GGML_OP_SSM_SCAN:
+            if (!any_on_device) {
+                return false;
+            }
+            // SSM_SCAN has 7 source tensors, so it cannot use the standard
+            // (src0, src1, dst) func signature. Dispatch directly and return.
+            ggml_cl_ssm_scan(backend, tensor);
+            return true;
         case GGML_OP_GATED_DELTA_NET:
             if (!any_on_device) {
                 return false;

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

@@ -0,0 +1,216 @@
+// Mamba2 fused SSM scan kernel. One workgroup per (head, dim, seq); WG size =
+// 64 threads. Each thread owns c_factor = d_state/64 state elements in
+// private registers; the state stays resident across the n_tokens t-loop
+//
+// References:
+//   ggml/src/ggml-cuda/ssm-scan.cu:117 ssm_scan_f32_group
+//   ggml/src/ggml-cpu/ops.cpp:9368 ggml_compute_forward_ssm_scan_f32
+
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_khr_subgroups
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#endif
+
+#if defined(cl_qcom_reqd_sub_group_size)
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+#define REQD_SUBGROUP_SIZE_64 __attribute__((qcom_reqd_sub_group_size("half")))
+#else
+#define REQD_SUBGROUP_SIZE_64
+#endif
+
+inline float softplus_f32(float x) {
+    return (x <= 20.0f) ? log(1.0f + exp(x)) : x;
+}
+
+// d_state = 128 (most Mamba-2 models, e.g. mamba2-2.7B, Codestral-Mamba).
+// WG = 64 threads, each holds 2 state elements (tid and tid+64).
+REQD_SUBGROUP_SIZE_64
+kernel void kernel_ssm_scan_f32_mamba2_d128(
+    global const char * src0_base, ulong src0_off,
+    global const char * src1_base, ulong src1_off,
+    global const char * src2_base, ulong src2_off,
+    global const char * src3_base, ulong src3_off,
+    global const char * src4_base, ulong src4_off,
+    global const char * src5_base, ulong src5_off,
+    global const char * src6_base, ulong src6_off,
+    global       char * dst_base,  ulong dst_off,
+    ulong s0_nb2, ulong s0_nb3,
+    ulong x_nb2,  ulong x_nb3,
+    ulong dt_nb1, ulong dt_nb2,
+    ulong A_nb1,
+    ulong B_nb2,  ulong B_nb3,
+    ulong C_nb2,  ulong C_nb3,
+    ulong s_off_bytes,
+    int   head_dim, int n_head, int n_group, int n_tokens
+) {
+    const int d_state = 128;
+
+    const int tid     = (int) get_local_id(0);
+    const int wg_x    = (int) get_group_id(0);
+    const int seq_id  = (int) get_group_id(1);
+
+    const int head_id = wg_x / head_dim;
+    const int dim_id  = wg_x - head_id * head_dim;
+    const int g       = head_id / (n_head / n_group);
+
+    src0_base += src0_off;
+    src1_base += src1_off;
+    src2_base += src2_off;
+    src3_base += src3_off;
+    src4_base += src4_off;
+    src5_base += src5_off;
+    src6_base += src6_off;
+    dst_base  += dst_off;
+
+    const int seq_slot = ((global const int *) src6_base)[seq_id];
+
+    const ulong state_base_off = (ulong)seq_slot * s0_nb3 + (ulong)head_id * s0_nb2
+                                + (ulong)dim_id * d_state * sizeof(float);
+    global const float * s0_warp = (global const float *)(src0_base + state_base_off);
+    const ulong state_out_off = (ulong)seq_id * s0_nb3 + (ulong)head_id * s0_nb2
+                              + (ulong)dim_id * d_state * sizeof(float);
+    global float * s_warp = (global float *)(dst_base + s_off_bytes + state_out_off);
+
+    global const char * x_seq  = src1_base + (ulong)seq_id * x_nb3;
+    global const char * dt_seq = src2_base + (ulong)seq_id * dt_nb2;
+    global const char * B_seq  = src4_base + (ulong)seq_id * B_nb3 + (ulong)g * d_state * sizeof(float);
+    global const char * C_seq  = src5_base + (ulong)seq_id * C_nb3 + (ulong)g * d_state * sizeof(float);
+
+    const ulong y_dim_total = (ulong)n_head * head_dim;
+    global float * y_seq = (global float *)dst_base
+                           + (ulong)seq_id * (ulong)n_tokens * y_dim_total;
+
+    const float A_val = ((global const float *)src3_base)[(ulong)head_id * A_nb1 / sizeof(float)];
+
+    // c_factor = 2: each thread owns 2 state elements (tid and tid+64).
+    float state0 = s0_warp[tid];
+    float state1 = s0_warp[tid + 64];
+
+    for (int t = 0; t < n_tokens; ++t) {
+        const float dt_h        = ((global const float *)(dt_seq + (ulong)t * dt_nb1))[head_id];
+        const float dt_softplus = softplus_f32(dt_h);
+        const float dA          = exp(dt_softplus * A_val);
+        const float x_val       = ((global const float *)(x_seq + (ulong)t * x_nb2))[(ulong)head_id * head_dim + dim_id];
+        const float x_dt        = x_val * dt_softplus;
+
+        const float B0 = ((global const float *)(B_seq + (ulong)t * B_nb2))[tid];
+        const float B1 = ((global const float *)(B_seq + (ulong)t * B_nb2))[tid + 64];
+        const float C0 = ((global const float *)(C_seq + (ulong)t * C_nb2))[tid];
+        const float C1 = ((global const float *)(C_seq + (ulong)t * C_nb2))[tid + 64];
+
+        state0 = state0 * dA + B0 * x_dt;
+        state1 = state1 * dA + B1 * x_dt;
+        const float partial = state0 * C0 + state1 * C1;
+
+        const float sum = sub_group_reduce_add(partial);
+        if (tid == 0) {
+            y_seq[(ulong)t * y_dim_total + (ulong)head_id * head_dim + dim_id] = sum;
+        }
+    }
+
+    s_warp[tid]      = state0;
+    s_warp[tid + 64] = state1;
+}
+
+// d_state = 256 (Falcon-H1). WG = 64 threads, each holds 4 state elements.
+REQD_SUBGROUP_SIZE_64
+kernel void kernel_ssm_scan_f32_mamba2_d256(
+    global const char * src0_base, ulong src0_off,
+    global const char * src1_base, ulong src1_off,
+    global const char * src2_base, ulong src2_off,
+    global const char * src3_base, ulong src3_off,
+    global const char * src4_base, ulong src4_off,
+    global const char * src5_base, ulong src5_off,
+    global const char * src6_base, ulong src6_off,
+    global       char * dst_base,  ulong dst_off,
+    ulong s0_nb2, ulong s0_nb3,
+    ulong x_nb2,  ulong x_nb3,
+    ulong dt_nb1, ulong dt_nb2,
+    ulong A_nb1,
+    ulong B_nb2,  ulong B_nb3,
+    ulong C_nb2,  ulong C_nb3,
+    ulong s_off_bytes,
+    int   head_dim, int n_head, int n_group, int n_tokens
+) {
+    const int d_state = 256;
+
+    const int tid     = (int) get_local_id(0);
+    const int wg_x    = (int) get_group_id(0);
+    const int seq_id  = (int) get_group_id(1);
+
+    const int head_id = wg_x / head_dim;
+    const int dim_id  = wg_x - head_id * head_dim;
+    const int g       = head_id / (n_head / n_group);
+
+    src0_base += src0_off;
+    src1_base += src1_off;
+    src2_base += src2_off;
+    src3_base += src3_off;
+    src4_base += src4_off;
+    src5_base += src5_off;
+    src6_base += src6_off;
+    dst_base  += dst_off;
+
+    const int seq_slot = ((global const int *) src6_base)[seq_id];
+
+    const ulong state_base_off = (ulong)seq_slot * s0_nb3 + (ulong)head_id * s0_nb2
+                                + (ulong)dim_id * d_state * sizeof(float);
+    global const float * s0_warp = (global const float *)(src0_base + state_base_off);
+    const ulong state_out_off = (ulong)seq_id * s0_nb3 + (ulong)head_id * s0_nb2
+                              + (ulong)dim_id * d_state * sizeof(float);
+    global float * s_warp = (global float *)(dst_base + s_off_bytes + state_out_off);
+
+    global const char * x_seq  = src1_base + (ulong)seq_id * x_nb3;
+    global const char * dt_seq = src2_base + (ulong)seq_id * dt_nb2;
+    global const char * B_seq  = src4_base + (ulong)seq_id * B_nb3 + (ulong)g * d_state * sizeof(float);
+    global const char * C_seq  = src5_base + (ulong)seq_id * C_nb3 + (ulong)g * d_state * sizeof(float);
+
+    const ulong y_dim_total = (ulong)n_head * head_dim;
+    global float * y_seq = (global float *)dst_base
+                           + (ulong)seq_id * (ulong)n_tokens * y_dim_total;
+
+    const float A_val = ((global const float *)src3_base)[(ulong)head_id * A_nb1 / sizeof(float)];
+
+    // c_factor = 4: each thread owns 4 state elements.
+    float state0 = s0_warp[tid];
+    float state1 = s0_warp[tid + 64];
+    float state2 = s0_warp[tid + 128];
+    float state3 = s0_warp[tid + 192];
+
+    for (int t = 0; t < n_tokens; ++t) {
+        const float dt_h        = ((global const float *)(dt_seq + (ulong)t * dt_nb1))[head_id];
+        const float dt_softplus = softplus_f32(dt_h);
+        const float dA          = exp(dt_softplus * A_val);
+        const float x_val       = ((global const float *)(x_seq + (ulong)t * x_nb2))[(ulong)head_id * head_dim + dim_id];
+        const float x_dt        = x_val * dt_softplus;
+
+        global const float * B_t = (global const float *)(B_seq + (ulong)t * B_nb2);
+        global const float * C_t = (global const float *)(C_seq + (ulong)t * C_nb2);
+
+        const float B0 = B_t[tid];
+        const float B1 = B_t[tid + 64];
+        const float B2 = B_t[tid + 128];
+        const float B3 = B_t[tid + 192];
+        const float C0 = C_t[tid];
+        const float C1 = C_t[tid + 64];
+        const float C2 = C_t[tid + 128];
+        const float C3 = C_t[tid + 192];
+
+        state0 = state0 * dA + B0 * x_dt;
+        state1 = state1 * dA + B1 * x_dt;
+        state2 = state2 * dA + B2 * x_dt;
+        state3 = state3 * dA + B3 * x_dt;
+        const float partial = state0 * C0 + state1 * C1 + state2 * C2 + state3 * C3;
+
+        const float sum = sub_group_reduce_add(partial);
+        if (tid == 0) {
+            y_seq[(ulong)t * y_dim_total + (ulong)head_id * head_dim + dim_id] = sum;
+        }
+    }
+
+    s_warp[tid]       = state0;
+    s_warp[tid + 64]  = state1;
+    s_warp[tid + 128] = state2;
+    s_warp[tid + 192] = state3;
+}

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

@@ -1227,6 +1227,10 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
         const int32_t * op_params = op->op_params;
         const int n_dims = op_params[1];
         const int mode = op_params[2];
+        if (op_params[15] != 0) {
+            // FIXME: support ggml_rope_set_offset
+            return true;
+        }
         if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX && mode != GGML_ROPE_TYPE_IMROPE) {
             // GGML_LOG_WARN("OpenVINO backend does not support ROPE with mode %d\n", mode);
             return true;

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

@@ -109,7 +109,14 @@ int g_ggml_sycl_enable_host_pinned_mem = 1;
 static ggml_sycl_device_info ggml_sycl_init() {
     ggml_sycl_device_info info = {};
 
-    info.device_count = dpct::dev_mgr::instance().device_count();
+    // Do not hard crash when there exists no SYCL devices.
+    // We want to allow the user to use non-SYCL tools when SYCL is compiled (such as llama-quantize)
+    try {
+        info.device_count = dpct::dev_mgr::instance().device_count();
+    } catch (sycl::exception const &exc) {
+        GGML_LOG_INFO("%s: no SYCL device available: %s\n", __func__, exc.what());
+        info.device_count = 0;
+    }
     if (info.device_count == 0) {
         GGML_LOG_ERROR("%s: failed to initialize: %s\n", GGML_SYCL_NAME, __func__);
         return info;
@@ -6235,6 +6242,8 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
         }
         case GGML_OP_ROPE:
         case GGML_OP_ROPE_BACK:
+            // FIXME: support ggml_rope_set_offset
+            return ((const int32_t *) op->op_params)[15] == 0;
         case GGML_OP_IM2COL:
         case GGML_OP_IM2COL_3D:
         case GGML_OP_UPSCALE:

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

@@ -913,6 +913,7 @@ struct vk_device_struct {
     vk_pipeline pipeline_quantize_q8_1_x4;
 
     vk_pipeline pipeline_dequant[GGML_TYPE_COUNT];
+    vk_pipeline pipeline_dequant_transpose[GGML_TYPE_COUNT]; // fused dequant+transpose for FA quant-KV
     vk_pipeline pipeline_dequant_mul_mat_vec_f32_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols];
     vk_pipeline pipeline_dequant_mul_mat_vec_f16_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols];
     vk_pipeline pipeline_dequant_mul_mat_vec_id_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT];
@@ -962,6 +963,7 @@ struct vk_device_struct {
     vk_pipeline pipeline_cpy_f32_quant[GGML_TYPE_COUNT];
     vk_pipeline pipeline_cpy_quant_f32[GGML_TYPE_COUNT];
     vk_pipeline pipeline_cpy_transpose_16, pipeline_cpy_transpose_32;
+    vk_pipeline pipeline_cpy_transpose_02_16, pipeline_cpy_transpose_02_32;
     // [src0 0=fp32,1=fp16][dst]
     vk_pipeline pipeline_set_rows_i32[2][GGML_TYPE_COUNT];
     vk_pipeline pipeline_set_rows_i64[2][GGML_TYPE_COUNT];
@@ -1644,6 +1646,7 @@ struct vk_op_rope_push_constants {
     uint32_t rope_mode;
     uint32_t nrows;
     uint32_t n_dims;
+    uint32_t n_offs;
     float freq_scale;
     float freq_base;
     float ext_factor;
@@ -3382,10 +3385,10 @@ static void ggml_vk_queue_command_pools_cleanup(vk_device& device) {
     // Arbitrary frequency to cleanup/reuse command buffers
     static constexpr uint32_t cleanup_frequency = 10;
 
-    if (device->compute_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) {
+    if (device->compute_queue && device->compute_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) {
         ggml_vk_command_pool_cleanup(device, device->compute_queue->cmd_pool);
     }
-    if (device->transfer_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) {
+    if (device->transfer_queue && device->transfer_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) {
         ggml_vk_command_pool_cleanup(device, device->transfer_queue->cmd_pool);
     }
 }
@@ -5389,6 +5392,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_0], "dequant_q5_0", dequant_q5_0_len, dequant_q5_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_1], "dequant_q5_1", dequant_q5_1_len, dequant_q5_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q8_0], "dequant_q8_0", dequant_q8_0_len, dequant_q8_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q8_0], "dequant_q8_0_transpose", dequant_q8_0_transpose_len, dequant_q8_0_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_K], "dequant_q2_k", dequant_q2_k_len, dequant_q2_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_TQ2_0], "dequant_tq2_0", dequant_tq2_0_len, dequant_tq2_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q3_K], "dequant_q3_k", dequant_q3_k_len, dequant_q3_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1);
@@ -5525,6 +5529,8 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
 
     ggml_vk_create_pipeline(device, device->pipeline_cpy_transpose_32, "cpy_transpose_32", cpy_transpose_32_len, cpy_transpose_32_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_cpy_transpose_16, "cpy_transpose_16", cpy_transpose_16_len, cpy_transpose_16_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_cpy_transpose_02_32, "cpy_transpose_02_32", cpy_transpose_02_32_len, cpy_transpose_02_32_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_cpy_transpose_02_16, "cpy_transpose_02_16", cpy_transpose_02_16_len, cpy_transpose_02_16_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1);
 
     ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q1_0], "cpy_f32_q1_0", cpy_f32_q1_0_len, cpy_f32_q1_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q2_0], "cpy_f32_q2_0", cpy_f32_q2_0_len, cpy_f32_q2_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
@@ -8931,6 +8937,18 @@ static vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const
         }
     }
 
+    // Same, for a 0<->2 swap: src dim2 is the innermost dimension.
+    bool transpose02 = dst && !contig && src->nb[2] == ggml_type_size(to) &&
+                       ggml_is_contiguous(dst) && ggml_are_same_shape(dst, src);
+
+    if (transpose02 && src->type == to) {
+        if (ggml_type_size(to) == 4) {
+            return ctx->device->pipeline_cpy_transpose_02_32;
+        } else if (ggml_type_size(to) == 2) {
+            return ctx->device->pipeline_cpy_transpose_02_16;
+        }
+    }
+
     if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_F32) {
         if (contig) {
             return ctx->device->pipeline_contig_cpy_f32_f32;
@@ -10807,9 +10825,32 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
 
     const bool f32acc = !ctx->device->fp16 || dst->op_params[3] == GGML_PREC_F32 || k->type == GGML_TYPE_BF16;
 
+    // dequant K/V once into an f16 scratch, reordered KV layout so FA can read without a stride
+    auto is_dense_kv_cache = [](const ggml_tensor * t) {
+        return t->nb[0] == ggml_type_size(t->type) &&
+               t->nb[2] == ggml_row_size(t->type, t->ne[0]) &&
+               t->nb[1] == t->nb[2] * t->ne[2] &&
+               t->nb[3] == t->nb[1] * t->ne[1];
+    };
+    const bool k_quant = k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_BF16 && k->type != GGML_TYPE_F32;
+    const bool v_quant = v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_BF16 && v->type != GGML_TYPE_F32;
+    const bool use_dequant_kv = k_quant && v_quant && neq1 >= 64 &&
+                                is_dense_kv_cache(k) && is_dense_kv_cache(v) &&
+                                (uint64_t)ggml_nelements(k) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange &&
+                                (uint64_t)ggml_nelements(v) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange &&
+                                ctx->device->pipeline_dequant_transpose[k->type] != nullptr &&
+                                ctx->device->pipeline_dequant_transpose[v->type] != nullptr &&
+                                // coopmat2 path does not benefit from the f16 scratch
+                                !ctx->device->coopmat2 &&
+                                // Intel Xe1 regresses, see PR 25494
+                                (ctx->device->vendor_id != VK_VENDOR_ID_INTEL ||
+                                 (ctx->device->coopmat_support && ctx->device->architecture != vk_device_architecture::INTEL_XE1));
+    const ggml_type k_type_eff = use_dequant_kv ? GGML_TYPE_F16 : k->type;
+    const ggml_type v_type_eff = use_dequant_kv ? GGML_TYPE_F16 : v->type;
+
     // For scalar/coopmat1 FA, we can use the "large" size to accommodate qga.
     // For coopmat2 FA, we always use the small size (which is still pretty large for gqa).
-    vk_fa_tuning_params tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, 512, KV, k->type, v->type, f32acc);
+    vk_fa_tuning_params tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, 512, KV, k_type_eff, v_type_eff, f32acc);
     const uint32_t max_gqa = std::min(tuning_params.block_rows, 32u);
 
     if (N <= 8 && qk_ratio > 1 && qk_ratio <= max_gqa &&
@@ -10822,7 +10863,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
         workgroups_y /= gqa_ratio;
     }
 
-    tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k->type, v->type, f32acc);
+    tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k_type_eff, v_type_eff, f32acc);
 
     const uint32_t q_stride = (uint32_t)(nbq1 / ggml_type_size(q->type));
     uint32_t k_stride = (uint32_t)(nbk1 / ggml_type_size(k->type));
@@ -10836,6 +10877,17 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
         v_stride /= 4;
     }
 
+    uint32_t nbk2_eff = (uint32_t)nbk2, nbk3_eff = (uint32_t)nbk3;
+    uint32_t nbv2_eff = (uint32_t)nbv2, nbv3_eff = (uint32_t)nbv3;
+    if (use_dequant_kv) {
+        k_stride = HSK;
+        v_stride = HSV;
+        nbk2_eff = (uint32_t)((uint64_t)HSK * KV * sizeof(ggml_fp16_t));
+        nbk3_eff = (uint32_t)((uint64_t)HSK * KV * nek2 * sizeof(ggml_fp16_t));
+        nbv2_eff = (uint32_t)((uint64_t)HSV * KV * sizeof(ggml_fp16_t));
+        nbv3_eff = (uint32_t)((uint64_t)HSV * KV * nev2 * sizeof(ggml_fp16_t));
+    }
+
     const uint32_t alignment = tuning_params.block_cols;
     bool aligned = (KV % alignment) == 0 &&
                    // the "aligned" shader variant will forcibly align strides, for performance
@@ -10862,7 +10914,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
     bool use_mask_opt = mask && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16
                         && (ctx->device->architecture != vk_device_architecture::AMD_GCN || HSK > 256 || HSV > 256);
     vk_fa_pipeline_state fa_pipeline_state = get_fa_pipeline_state(ctx->device, tuning_params, HSK, HSV, aligned, f32acc,
-                                                                   mask != nullptr, use_mask_opt, logit_softcap != 0, k->type, v->type);
+                                                                   mask != nullptr, use_mask_opt, logit_softcap != 0, k_type_eff, v_type_eff);
 
     vk_pipeline pipeline = nullptr;
 
@@ -10966,6 +11018,34 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
     vk_subbuffer sinks_buf = sinks ? ggml_vk_tensor_subbuffer(ctx, sinks) : q_buf;
     vk_subbuffer mask_opt_buf = use_mask_opt ? ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0) : q_buf;
 
+    if (use_dequant_kv) {
+        const uint64_t fp = sizeof(ggml_fp16_t);
+        const uint64_t k_f16_sz = (uint64_t)ggml_nelements(k) * fp;
+        const uint64_t v_f16_sz = (uint64_t)ggml_nelements(v) * fp;
+        if (ctx->prealloc_size_x < k_f16_sz + v_f16_sz) {
+            ctx->prealloc_size_x = k_f16_sz + v_f16_sz;
+            ggml_vk_preallocate_buffers(ctx, subctx);
+        }
+        vk_pipeline tr_k = ctx->device->pipeline_dequant_transpose[k->type];
+        vk_pipeline tr_v = ctx->device->pipeline_dequant_transpose[v->type];
+        ggml_pipeline_request_descriptor_sets(ctx, tr_k, 1);
+        ggml_pipeline_request_descriptor_sets(ctx, tr_v, 1);
+        if (ctx->prealloc_x_need_sync) {
+            ggml_vk_sync_buffers(ctx, subctx);
+        }
+        vk_subbuffer k_dst = vk_subbuffer{ ctx->prealloc_x, 0,        k_f16_sz };
+        vk_subbuffer v_dst = vk_subbuffer{ ctx->prealloc_x, k_f16_sz, v_f16_sz };
+        const uint32_t k_nel = (uint32_t)ggml_nelements(k);
+        const uint32_t v_nel = (uint32_t)ggml_nelements(v);
+        { const std::vector<uint32_t> pc = { (uint32_t)HSK, (uint32_t)nek2, (uint32_t)KV, 0, k_nel };
+          ggml_vk_dispatch_pipeline(ctx, subctx, tr_k, { k_buf, k_dst }, pc, { k_nel, 1, 1 }); }
+        { const std::vector<uint32_t> pc = { (uint32_t)HSV, (uint32_t)nev2, (uint32_t)KV, 0, v_nel };
+          ggml_vk_dispatch_pipeline(ctx, subctx, tr_v, { v_buf, v_dst }, pc, { v_nel, 1, 1 }); }
+        ggml_vk_sync_buffers(ctx, subctx);
+        k_buf = k_dst;
+        v_buf = v_dst;
+    }
+
     uint32_t mask_n_head_log2 = ((sinks != nullptr) << 24) | n_head_log2;
 
     if (use_mask_opt)
@@ -10995,8 +11075,8 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
                                               (uint32_t)nev2, (uint32_t)nev3,
                                               nem1, nem2, nem3,
                                               q_stride, (uint32_t)nbq2, (uint32_t)nbq3,
-                                              k_stride, (uint32_t)nbk2, (uint32_t)nbk3,
-                                              v_stride, (uint32_t)nbv2, (uint32_t)nbv3,
+                                              k_stride, nbk2_eff, nbk3_eff,
+                                              v_stride, nbv2_eff, nbv3_eff,
                                               scale, max_bias, logit_softcap,
                                               mask_n_head_log2, m0, m1,
                                               gqa_ratio, split_kv, split_k };
@@ -11038,6 +11118,10 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
                                     {q_buf, k_buf, v_buf, mask_buf, sinks_buf, dst_buf, mask_opt_buf},
                                     pc, { workgroups_x, workgroups_y, workgroups_z });
     }
+
+    if (use_dequant_kv) {
+        ctx->prealloc_x_need_sync = true;
+    }
 }
 
 static vk_conv_shapes ggml_vk_conv_select_shape(ggml_backend_vk_context * ctx, uint32_t K, uint32_t NPQ) {
@@ -12192,7 +12276,16 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
                 elements = { ne, 1, 1 };
             }
 
-            if (pipeline == ctx->device->pipeline_cpy_transpose_32 ||
+            if (pipeline == ctx->device->pipeline_cpy_transpose_02_32 ||
+                pipeline == ctx->device->pipeline_cpy_transpose_02_16) {
+                // 32x32 tiles over dims 0 and 2; dim1 and dim3 are the batch
+                elements[0] = (uint32_t)CEIL_DIV(dst->ne[0], 32);
+                elements[1] = (uint32_t)CEIL_DIV(dst->ne[2], 32);
+                elements[2] = (uint32_t)(dst->ne[1]*dst->ne[3]);
+                elements[0] = std::min(elements[0], ctx->device->properties.limits.maxComputeWorkGroupCount[0]);
+                elements[1] = std::min(elements[1], ctx->device->properties.limits.maxComputeWorkGroupCount[1]);
+                elements[2] = std::min(elements[2], ctx->device->properties.limits.maxComputeWorkGroupCount[2]);
+            } else if (pipeline == ctx->device->pipeline_cpy_transpose_32 ||
                 pipeline == ctx->device->pipeline_cpy_transpose_16) {
                 // 32x32 tiles
                 elements[0] = (uint32_t)CEIL_DIV(dst->ne[0], 32);
@@ -13120,6 +13213,7 @@ static uint32_t ggml_vk_rms_partials_size(ggml_backend_vk_context * ctx, const g
 static vk_op_rope_push_constants ggml_vk_make_rope_constants(const ggml_tensor *dst, const ggml_tensor *src0, const bool has_ff, bool backprop, const uint32_t set_rows_stride) {
     const int n_dims        = ((const int32_t *) dst->op_params)[1];
     const int mode          = ((const int32_t *) dst->op_params)[2];
+    const int n_offs        = ((const int32_t *) dst->op_params)[15];
     // const int n_ctx         = ((const int32_t *) dst->op_params)[3];
     const int n_ctx_orig    = ((const int32_t *) dst->op_params)[4];
     const float freq_base   = ((const float *)   dst->op_params)[5];
@@ -13149,7 +13243,7 @@ static vk_op_rope_push_constants ggml_vk_make_rope_constants(const ggml_tensor *
     uint32_t nb13 = dst->nb[3] / ggml_type_size(dst->type);
 
     vk_op_rope_push_constants rope {
-        (uint32_t)mode, (uint32_t)ggml_nrows(src0), (uint32_t)n_dims, freq_scale,
+        (uint32_t)mode, (uint32_t)ggml_nrows(src0), (uint32_t)n_dims, (uint32_t)n_offs, freq_scale,
         freq_base, ext_factor, attn_factor, {corr_dims[0], corr_dims[1]}, theta_scale, has_ff,
         { sections[0], sections[1], sections[2], sections[3] }, is_imrope, backprop, set_rows_stride,
 
@@ -19195,6 +19289,10 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
                     tensor_clone = ggml_rope_ext_back(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], n_dims, mode, n_ctx_orig_ggml, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
                 }
             }
+            const int n_offs = ((int32_t *) tensor->op_params)[15];
+            if (n_offs != 0) {
+                tensor_clone = ggml_rope_set_offset(tensor_clone, n_offs);
+            }
         } else if (tensor->op == GGML_OP_UNARY) {
             switch (ggml_get_unary_op(tensor)) {
             case GGML_UNARY_OP_EXP:

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

@@ -0,0 +1,61 @@
+#version 450
+
+#include "types.glsl"
+#include "generic_unary_head.glsl"
+
+// workgroup does 32x32 tile, but uses 32x8 threads
+#define TILE_DIM 32
+layout(local_size_x = 32, local_size_y = 8, local_size_z = 1) in;
+
+// +1 padding avoids shared-memory bank conflicts on the transposed read
+shared uint sh[TILE_DIM][TILE_DIM + 1];
+
+void iter(uvec3 wg_id) {
+    const uint tile_i0 = wg_id.x;   // tiles dst ne10 (== src ne00)
+    const uint tile_i2 = wg_id.y;   // tiles dst ne12 (== src ne02)
+
+    const uint tid_col = gl_LocalInvocationID.x;
+    const uint tid_row = gl_LocalInvocationID.y;
+
+    const uint i1 = wg_id.z % p.ne11;
+    const uint i3 = wg_id.z / p.ne11;
+    const uint i01 = i1;
+    const uint i03 = i3;
+
+    [[unroll]] for (uint y = 0; y < 4; ++y) {
+        const uint i00 = tile_i0 * TILE_DIM + tid_row + 8 * y;
+        const uint i02 = tile_i2 * TILE_DIM + tid_col;
+        if (i00 < p.ne00 && i01 < p.ne01 && i02 < p.ne02 && i03 < p.ne03) {
+            const uint src_idx = i00 * p.nb00 + i01 * p.nb01 + i02 * p.nb02 + i03 * p.nb03;
+            sh[tid_row + 8 * y][tid_col] = uint(data_a[get_aoffset() + src_idx]);
+        }
+    }
+
+    barrier();
+
+    [[unroll]] for (uint y = 0; y < 4; ++y) {
+        const uint i0 = tile_i0 * TILE_DIM + tid_col;
+        const uint i2 = tile_i2 * TILE_DIM + tid_row + 8 * y;
+        if (i0 < p.ne10 && i1 < p.ne11 && i2 < p.ne12 && i3 < p.ne13) {
+            const uint dst_idx = i0 * p.nb10 + i1 * p.nb11 + i2 * p.nb12 + i3 * p.nb13;
+            data_d[get_doffset() + dst_idx] = D_TYPE(sh[tid_col][tid_row + 8 * y]);
+        }
+    }
+}
+
+#define CEIL_DIV(a, b) (((a) + (b) - 1) / (b))
+
+void main() {
+    bool need_barrier = false;
+    for (uint z = gl_WorkGroupID.z; z < p.ne11 * p.ne13; z += gl_NumWorkGroups.z) {
+        for (uint y = gl_WorkGroupID.y; y < CEIL_DIV(p.ne12, TILE_DIM); y += gl_NumWorkGroups.y) {
+            for (uint x = gl_WorkGroupID.x; x < CEIL_DIV(p.ne10, TILE_DIM); x += gl_NumWorkGroups.x) {
+                if (need_barrier) {
+                    barrier();
+                }
+                need_barrier = true;
+                iter(uvec3(x, y, z));
+            }
+        }
+    }
+}

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

@@ -18,7 +18,18 @@ void main() {
         return;
     }
 
+#ifdef DEQUANT_TRANSPOSE
+    // read [HS, NH, KV, NS], write [HS, KV, NH, NS]
+    const uint HS = p.M, NH = p.K, KVn = p.stride_a;
+    const uint e0 = ib * 32;
+    const uint b_idx = (e0 % HS)
+                     + ((e0 / (HS * NH)) % KVn) * HS
+                     + ((e0 / HS) % NH) * (HS * KVn)
+                     + (e0 / (HS * NH * KVn)) * (HS * KVn * NH)
+                     + 16 * il;
+#else
     const uint b_idx = 1024*i + 32*ir + 16*il;
+#endif
 
     const float d = float(data_a[ib].d);
 

ggml/src/ggml-vulkan/vulkan-shaders/rope_funcs.glsl

@@ -50,19 +50,21 @@ void rope_norm(const uint i0, const uint i1, const uint i2, const uint i3, rope_
     }
     idst += p.d_offset;
 
-    if (i0 >= p.n_dims) {
+    if (i0 < p.n_offs || i0 >= p.n_offs + p.n_dims) {
         rope_data_d[idst + 0] = ROPE_D_TYPE(rope_data_a[ix + 0]);
         rope_data_d[idst + 1] = ROPE_D_TYPE(rope_data_a[ix + 1]);
 
         return;
     }
 
-    const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, i0/2.0f);
+    const uint iw = i0 - p.n_offs; // relative idx
 
-    const float freq_factor = p.has_ff != 0 ? rope_data_ff[i0/2] : 1.0f;
+    const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, iw/2.0f);
+
+    const float freq_factor = p.has_ff != 0 ? rope_data_ff[iw/2] : 1.0f;
 
     float cos_theta, sin_theta;
-    rope_yarn(theta_base / freq_factor, i0, cos_theta, sin_theta, p);
+    rope_yarn(theta_base / freq_factor, iw, cos_theta, sin_theta, p);
 
     const float x0 = float(rope_data_a[ix + 0]);
     const float x1 = float(rope_data_a[ix + 1]);
@@ -87,25 +89,28 @@ void rope_neox(const uint i0, const uint i1, const uint i2, const uint i3, rope_
     }
     idst += p.d_offset;
 
-    if (i0 >= p.n_dims) {
+    if (i0 < p.n_offs || i0 >= p.n_offs + p.n_dims) {
         rope_data_d[idst + i0/2 + 0] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 0]);
         rope_data_d[idst + i0/2 + 1] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 1]);
 
         return;
     }
 
-    const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, i0/2.0f);
+    const uint iw = i0 - p.n_offs; // relative idx
 
-    const float freq_factor = p.has_ff != 0 ? rope_data_ff[i0/2] : 1.0f;
+    const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, iw/2.0f);
+
+    const float freq_factor = p.has_ff != 0 ? rope_data_ff[iw/2] : 1.0f;
 
     float cos_theta, sin_theta;
-    rope_yarn(theta_base / freq_factor, i0, cos_theta, sin_theta, p);
+    rope_yarn(theta_base / freq_factor, iw, cos_theta, sin_theta, p);
 
-    const float x0 = float(rope_data_a[ix + 0]);
-    const float x1 = float(rope_data_a[ix + p.n_dims/2]);
+    // idst/ix point at channel i0/2; the first channel of the rotated pair is p.n_offs + iw/2 = i0/2 + p.n_offs/2
+    const float x0 = float(rope_data_a[ix + p.n_offs/2 + 0]);
+    const float x1 = float(rope_data_a[ix + p.n_offs/2 + p.n_dims/2]);
 
-    rope_data_d[idst + 0]          = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta);
-    rope_data_d[idst + p.n_dims/2] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta);
+    rope_data_d[idst + p.n_offs/2 + 0]          = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta);
+    rope_data_d[idst + p.n_offs/2 + p.n_dims/2] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta);
 }
 
 
@@ -125,53 +130,56 @@ void rope_multi(const uint i0, const uint i1, const uint i2, const uint i3, rope
     }
     idst += p.d_offset;
 
-    if (i0 >= p.n_dims) {
+    if (i0 < p.n_offs || i0 >= p.n_offs + p.n_dims) {
         rope_data_d[idst + i0/2 + 0] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 0]);
         rope_data_d[idst + i0/2 + 1] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 1]);
 
         return;
     }
 
+    const uint iw = i0 - p.n_offs; // relative idx
+
     const int sect_dims = p.sections[0] + p.sections[1] + p.sections[2] + p.sections[3];
     const int sec_w = p.sections[1] + p.sections[0];
-    const uint sector = (i0 / 2) % sect_dims;
+    const uint sector = (iw / 2) % sect_dims;
 
     float theta_base = 0.0;
     if (p.is_imrope != 0) {
         if (sector % 3 == 1 && sector < 3 * p.sections[1]) {
-            theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, i0/2.0f);
+            theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, iw/2.0f);
         } else if (sector % 3 == 2 && sector < 3 * p.sections[2]) {
-            theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, i0/2.0f);
+            theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, iw/2.0f);
         } else if (sector % 3 == 0 && sector < 3 * p.sections[0]) {
-            theta_base = rope_data_pos[i2]*pow(p.theta_scale, i0/2.0f);
+            theta_base = rope_data_pos[i2]*pow(p.theta_scale, iw/2.0f);
         } else {
-            theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, i0/2.0f);
+            theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, iw/2.0f);
         }
     } else {
         if (sector < p.sections[0]) {
-            theta_base = rope_data_pos[i2]*pow(p.theta_scale, i0/2.0f);
+            theta_base = rope_data_pos[i2]*pow(p.theta_scale, iw/2.0f);
         }
         else if (sector >= p.sections[0] && sector < sec_w) {
-            theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, i0/2.0f);
+            theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, iw/2.0f);
         }
         else if (sector >= sec_w && sector < sec_w + p.sections[2]) {
-            theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, i0/2.0f);
+            theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, iw/2.0f);
         }
         else if (sector >= sec_w + p.sections[2]) {
-            theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, i0/2.0f);
+            theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, iw/2.0f);
         }
     }
 
-    const float freq_factor = p.has_ff != 0 ? rope_data_ff[i0/2] : 1.0f;
+    const float freq_factor = p.has_ff != 0 ? rope_data_ff[iw/2] : 1.0f;
 
     float cos_theta, sin_theta;
-    rope_yarn(theta_base / freq_factor, i0, cos_theta, sin_theta, p);
+    rope_yarn(theta_base / freq_factor, iw, cos_theta, sin_theta, p);
 
-    const float x0 = float(rope_data_a[ix + 0]);
-    const float x1 = float(rope_data_a[ix + p.n_dims/2]);
+    // idst/ix point at channel i0/2; the first channel of the rotated pair is p.n_offs + iw/2 = i0/2 + p.n_offs/2
+    const float x0 = float(rope_data_a[ix + p.n_offs/2 + 0]);
+    const float x1 = float(rope_data_a[ix + p.n_offs/2 + p.n_dims/2]);
 
-    rope_data_d[idst + 0]          = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta);
-    rope_data_d[idst + p.n_dims/2] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta);
+    rope_data_d[idst + p.n_offs/2 + 0]          = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta);
+    rope_data_d[idst + p.n_offs/2 + p.n_dims/2] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta);
 }
 
 void rope_vision(const uint i0, const uint i1, const uint i2, const uint i3, rope_params p) {

ggml/src/ggml-vulkan/vulkan-shaders/rope_params.glsl

@@ -5,6 +5,7 @@ struct rope_params {
     uint rope_mode;
     uint nrows;
     uint n_dims;
+    uint n_offs;
     float freq_scale;
     float freq_base;
     float ext_factor;

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

@@ -780,6 +780,10 @@ void process_shaders() {
         if (tname != "f16" && tname != "bf16") {
             string_to_spv("dequant_" + tname, "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}}));
         }
+        // Fused dequant+transpose variant for FA quant-KV (per-head-contiguous f16 scratch).
+        if (tname == "q8_0") {
+            string_to_spv("dequant_" + tname + "_transpose", "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}, {"DEQUANT_TRANSPOSE", "1"}}));
+        }
 
         shader = (tname == "f32" || tname == "f16" || tname == "bf16") ? "get_rows.comp" : "get_rows_quant.comp";
 
@@ -826,6 +830,8 @@ void process_shaders() {
 
     string_to_spv("cpy_transpose_16", "copy_transpose.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}});
     string_to_spv("cpy_transpose_32", "copy_transpose.comp", {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}});
+    string_to_spv("cpy_transpose_02_16", "copy_transpose_02.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}});
+    string_to_spv("cpy_transpose_02_32", "copy_transpose_02.comp", {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}});
 
     for (std::string t : {"q1_0", "q2_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {
         string_to_spv("cpy_f32_" + t, "copy_to_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"S_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});

ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp

@@ -954,10 +954,11 @@ struct ggml_webgpu_mul_mat_vec_pipeline_key {
     int       vectorized;
     uint32_t  num_cols;
     bool      use_mmvq;
+    bool      src_overlap;
 
     bool operator==(const ggml_webgpu_mul_mat_vec_pipeline_key & other) const {
         return src0_type == other.src0_type && src1_type == other.src1_type && vectorized == other.vectorized &&
-               num_cols == other.num_cols && use_mmvq == other.use_mmvq;
+               num_cols == other.num_cols && use_mmvq == other.use_mmvq && src_overlap == other.src_overlap;
     }
 };
 
@@ -969,6 +970,7 @@ struct ggml_webgpu_mul_mat_vec_pipeline_key_hash {
         ggml_webgpu_hash_combine(seed, key.vectorized);
         ggml_webgpu_hash_combine(seed, key.num_cols);
         ggml_webgpu_hash_combine(seed, key.use_mmvq);
+        ggml_webgpu_hash_combine(seed, key.src_overlap);
         return seed;
     }
 };
@@ -977,6 +979,7 @@ struct ggml_webgpu_mul_mat_vec_shader_decisions {
     uint32_t wg_size;
     uint32_t outputs_per_wg;
     uint32_t vec_size;
+    bool     src_overlap = false;
 };
 
 struct ggml_webgpu_quantize_q8_pipeline_key {
@@ -998,10 +1001,11 @@ struct ggml_webgpu_mul_mat_pipeline_key {
     ggml_type src1_type;
     int       vectorized;
     int       use_subgroup_matrix;
+    bool      src_overlap;
 
     bool operator==(const ggml_webgpu_mul_mat_pipeline_key & other) const {
         return src0_type == other.src0_type && src1_type == other.src1_type && vectorized == other.vectorized &&
-               use_subgroup_matrix == other.use_subgroup_matrix;
+               use_subgroup_matrix == other.use_subgroup_matrix && src_overlap == other.src_overlap;
     }
 };
 
@@ -1012,6 +1016,7 @@ struct ggml_webgpu_mul_mat_pipeline_key_hash {
         ggml_webgpu_hash_combine(seed, key.src1_type);
         ggml_webgpu_hash_combine(seed, key.vectorized);
         ggml_webgpu_hash_combine(seed, key.use_subgroup_matrix);
+        ggml_webgpu_hash_combine(seed, key.src_overlap);
         return seed;
     }
 };
@@ -1034,6 +1039,7 @@ struct ggml_webgpu_mul_mat_shader_decisions {
     uint32_t subgroup_matrix_n;
 
     uint32_t mul_mat_wg_size;
+    bool     src_overlap = false;
 };
 
 /** MUL_MAT_ID **/
@@ -1950,7 +1956,7 @@ class ggml_webgpu_shader_lib {
         return quantize_q8_pipelines[key];
     }
 
-    webgpu_pipeline get_mul_mat_vec_pipeline(const ggml_webgpu_shader_lib_context & context) {
+    webgpu_pipeline get_mul_mat_vec_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) {
         ggml_webgpu_mul_mat_vec_pipeline_key key = {};
         key.src0_type                            = context.src0->type;
         key.src1_type                            = context.src1->type;
@@ -1961,6 +1967,7 @@ class ggml_webgpu_shader_lib {
         key.num_cols   = context.dst->ne[1];
         key.use_mmvq =
             ggml_webgpu_can_use_mmvq(context.src0, context.src1, context.supports_dot_product, context.vendor);
+        key.src_overlap = src_overlap;
 
         auto it = mul_mat_vec_pipelines.find(key);
         if (it != mul_mat_vec_pipelines.end()) {
@@ -2068,6 +2075,11 @@ class ggml_webgpu_shader_lib {
             defines.push_back("Q8_1_T");
         }
 
+        if (key.src_overlap) {
+            defines.push_back("SRC_OVERLAP");
+            variant += "_src_overlap";
+        }
+
         defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size));
         defines.push_back(std::string("OUTPUTS_PER_WG=") + std::to_string(outputs_per_wg));
         defines.push_back(context.supports_subgroups ? "USE_SUBGROUP_REDUCTION" : "USE_WORKGROUP_REDUCTION");
@@ -2089,7 +2101,7 @@ class ggml_webgpu_shader_lib {
         return mul_mat_vec_pipelines[key];
     }
 
-    webgpu_pipeline get_mul_mat_fast_pipeline(const ggml_webgpu_shader_lib_context & context) {
+    webgpu_pipeline get_mul_mat_fast_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) {
         ggml_webgpu_mul_mat_pipeline_key key = {};
         key.src0_type                        = context.src0->type;
         key.src1_type                        = context.src1->type;
@@ -2098,6 +2110,7 @@ class ggml_webgpu_shader_lib {
                                       1 :
                                       0;
         key.use_subgroup_matrix = context.supports_subgroup_matrix;
+        key.src_overlap         = src_overlap;
 
         auto it = mul_mat_fast_pipelines.find(key);
         if (it != mul_mat_fast_pipelines.end()) {
@@ -2216,6 +2229,11 @@ class ggml_webgpu_shader_lib {
             variant += "_vectorized";
         }
 
+        if (key.src_overlap) {
+            defines.push_back("SRC_OVERLAP");
+            variant += "_src_overlap";
+        }
+
         if (!key.use_subgroup_matrix) {
             defines.push_back("WORKGROUP_SIZE_M=" + std::to_string(WEBGPU_MUL_MAT_WG_SIZE_M) + "u");
             defines.push_back("WORKGROUP_SIZE_N=" + std::to_string(WEBGPU_MUL_MAT_WG_SIZE_N) + "u");

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

@@ -1628,48 +1628,65 @@ static webgpu_encoded_op ggml_webgpu_mul_mat(webgpu_context & ctx,
     // Get or create pipeline
     webgpu_pipeline                   pipeline;
     std::vector<webgpu_dispatch_desc> dispatches;
+    const bool src_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src0, src1) && !use_mmvq;
 
     if (use_mat_vec) {
         if (use_mmvq) {
             ggml_webgpu_quantize_q8_dispatch(ctx, src0, src1, dst, dispatches);
         }
-        pipeline = ctx->shader_lib->get_mul_mat_vec_pipeline(shader_lib_ctx);
+        pipeline = ctx->shader_lib->get_mul_mat_vec_pipeline(shader_lib_ctx, src_overlap);
     } else {
-        pipeline = ctx->shader_lib->get_mul_mat_fast_pipeline(shader_lib_ctx);
+        pipeline = ctx->shader_lib->get_mul_mat_fast_pipeline(shader_lib_ctx, src_overlap);
+    }
+
+    uint32_t offset_src0   = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type));
+    uint32_t offset_src1   = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type));
+    size_t   merged_offset = 0;
+    size_t   merged_size   = 0;
+    if (src_overlap) {
+        const ggml_webgpu_merged_binding_range merged_range =
+            ggml_webgpu_tensor_merged_binding_range(ctx, { src0, src1 });
+        merged_offset = merged_range.offset;
+        merged_size   = merged_range.size;
+        offset_src0   = ggml_webgpu_tensor_merged_element_offset(src0, merged_range);
+        offset_src1   = ggml_webgpu_tensor_merged_element_offset(src1, merged_range);
     }
 
     // Build params
-    std::vector<uint32_t> params = {
-        (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)),
-        (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)),
-        (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)),
-        (uint32_t) dst->ne[0],
-        (uint32_t) dst->ne[1],
-        (uint32_t) src0->ne[0],
-        (uint32_t) (src0->nb[1] / ggml_type_size(src0->type)),
-        (uint32_t) (src1->nb[1] / ggml_type_size(src1->type)),
-        (uint32_t) (src0->nb[2] / ggml_type_size(src0->type)),
-        (uint32_t) (src1->nb[2] / ggml_type_size(src1->type)),
-        (uint32_t) (src0->nb[3] / ggml_type_size(src0->type)),
-        (uint32_t) (src1->nb[3] / ggml_type_size(src1->type)),
-        (uint32_t) src0->ne[2],
-        (uint32_t) src0->ne[3],
-        (uint32_t) (src1->ne[2] / src0->ne[2]),
-        (uint32_t) (src1->ne[3] / src0->ne[3])
-    };
+    std::vector<uint32_t> params = { offset_src0,
+                                     offset_src1,
+                                     (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)),
+                                     (uint32_t) dst->ne[0],
+                                     (uint32_t) dst->ne[1],
+                                     (uint32_t) src0->ne[0],
+                                     (uint32_t) (src0->nb[1] / ggml_type_size(src0->type)),
+                                     (uint32_t) (src1->nb[1] / ggml_type_size(src1->type)),
+                                     (uint32_t) (src0->nb[2] / ggml_type_size(src0->type)),
+                                     (uint32_t) (src1->nb[2] / ggml_type_size(src1->type)),
+                                     (uint32_t) (src0->nb[3] / ggml_type_size(src0->type)),
+                                     (uint32_t) (src1->nb[3] / ggml_type_size(src1->type)),
+                                     (uint32_t) src0->ne[2],
+                                     (uint32_t) src0->ne[3],
+                                     (uint32_t) (src1->ne[2] / src0->ne[2]),
+                                     (uint32_t) (src1->ne[3] / src0->ne[3]) };
 
     // Build bind group entries
     std::vector<wgpu::BindGroupEntry> entries = {};
-
-    entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0));
     if (use_mmvq) {
+        entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0));
         auto & mmvq_qq8_entry = dispatches[0].bind_group_entries[1];
         entries.push_back(ggml_webgpu_make_bind_group_entry(1, ggml_webgpu_tensor_buf(dst), mmvq_qq8_entry.offset,
                                                             mmvq_qq8_entry.size));
+        entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, dst));
+    } else if (src_overlap) {
+        entries.push_back(
+            ggml_webgpu_make_bind_group_entry(0, ggml_webgpu_tensor_buf(src0), merged_offset, merged_size));
+        entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, dst));
     } else {
+        entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0));
         entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, src1));
+        entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, dst));
     }
-    entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, dst));
 
     // Calculate workgroup dimensions
     uint32_t       wg_x           = 1;
@@ -4455,7 +4472,9 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
             supports_op = (op->type == GGML_TYPE_F32 && src0->type == GGML_TYPE_F32) && ggml_is_contiguous_rows(src0);
             break;
         case GGML_OP_ROPE:
-            supports_op = op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16;
+            // FIXME: support ggml_rope_set_offset
+            supports_op =
+                (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && ((const int32_t *) op->op_params)[15] == 0;
             break;
         case GGML_OP_GLU:
             switch (ggml_get_glu_op(op)) {

ggml/src/ggml-webgpu/wgsl-shaders/common_decls.tmpl

@@ -1,3 +1,7 @@
+#ifndef SRC0
+#define SRC0 src0
+#endif
+
 #ifdef BYTE_HELPERS
 fn get_byte(value: u32, index: u32) -> u32 {
     return (value >> (index * 8)) & 0xFF;
@@ -46,7 +50,7 @@ fn load_f16_as_f32_at_src(byte_offset: u32) -> f32 {
 
 #ifdef DECLARE_BYTE_LOADERS_SRC0
 fn load_u16_at_src0(byte_offset: u32) -> u32 {
-    let word = src0[byte_offset / 4u];
+    let word = SRC0[byte_offset / 4u];
     let shift = (byte_offset & 0x2u) * 8u;
     return (word >> shift) & 0xFFFFu;
 }
@@ -55,14 +59,14 @@ fn load_u16_at_src0(byte_offset: u32) -> u32 {
 // Caller extracts the 16-bit half it needs via & 0xFFFFu or >> 16u.
 // this is used in k-quants for better performance
 fn load_u32_at_src0_aligned(byte_offset: u32) -> u32 {
-    return src0[(byte_offset & ~3u) / 4u];
+    return SRC0[(byte_offset & ~3u) / 4u];
 }
 
 fn load_u32_at_src0(byte_offset: u32) -> u32 {
     let word_idx = byte_offset / 4u;
     let shift = (byte_offset & 0x3u) * 8u;
-    let lo = src0[word_idx];
-    let hi = src0[word_idx + 1u];
+    let lo = SRC0[word_idx];
+    let hi = SRC0[word_idx + 1u];
     let shifted = (lo >> shift) | (hi << (32u - shift));
     return select(shifted, lo, shift == 0u);
 }
@@ -73,7 +77,7 @@ fn load_f16_at_src0(byte_offset: u32) -> f16 {
 }
 
 fn load_f16_as_f32_at_src0(byte_offset: u32) -> f32 {
-    let word = src0[byte_offset / 4u];
+    let word = SRC0[byte_offset / 4u];
     let shift = (byte_offset & 0x2u) * 8u;
     let d_bits = (word >> shift) & 0xFFFFu;
     return unpack2x16float(d_bits)[0];

ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl

@@ -1,3 +1,10 @@
+#ifndef SRC0
+#define SRC0 src0
+#endif
+#ifndef SRC1
+#define SRC1 src1
+#endif
+
 #ifdef VEC
 #define VEC_SIZE 4
 #define SHMEM_TYPE vec4<f16>
@@ -39,7 +46,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3
         let src0_idx = batch_offset + global_m * params.stride_01 + global_k;
         let src0_val = select( // taking a slight performance hit to avoid oob
             SRC0_TYPE(0.0),
-            src0[src0_idx/VEC_SIZE],
+            SRC0[src0_idx/VEC_SIZE],
             global_m < params.m && global_k < params.k);
         store_shmem(SHMEM_TYPE(src0_val), elem_idx);
     }
@@ -57,7 +64,7 @@ fn init_shmem_src1(thread_id: u32, batch_offset: u32, offset_n: u32, k_outer: u3
         let src1_idx = batch_offset + global_n * params.stride_11 + global_k;
         let src1_val = select(
             SRC1_TYPE(0.0),
-            src1[src1_idx/VEC_SIZE],
+            SRC1[src1_idx/VEC_SIZE],
             global_n < params.n && global_k < params.k);
         store_shmem(SHMEM_TYPE(src1_val), TILE_SRC0_SHMEM + elem_idx);
     }

ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_reg_tile.wgsl

@@ -1,8 +1,12 @@
 enable f16;
 
 #define DECLARE_BYTE_LOADERS_SRC0
-#include "common_decls.tmpl"
 
+#ifdef SRC_OVERLAP
+#define SRC0 merged_src
+#define SRC1 merged_src
+#endif
+#include "common_decls.tmpl"
 #include "mul_mat_decls.tmpl"
 
 #ifdef VEC
@@ -36,11 +40,17 @@ struct MulMatParams {
     broadcast3: u32
 };
 
+#ifdef SRC_OVERLAP
+@group(0) @binding(0) var<storage, read_write> merged_src: array<SRC0_TYPE>;
+#define DST_BINDING 1
+#else
 @group(0) @binding(0) var<storage, read_write> src0: array<SRC0_TYPE>; // M rows, K columns
 @group(0) @binding(1) var<storage, read_write> src1: array<SRC1_TYPE>; // K rows, N columns (transposed)
-@group(0) @binding(2) var<storage, read_write> dst: array<DST_TYPE>; // M rows, N columns (transposed)
+#define DST_BINDING 2
+#endif
 
-@group(0) @binding(3) var<uniform> params: MulMatParams;
+@group(0) @binding(DST_BINDING) var<storage, read_write> dst: array<DST_TYPE>; // M rows, N columns (transposed)
+@group(0) @binding(DST_BINDING + 1) var<uniform> params: MulMatParams;
 
 fn get_local_n(thread_id: u32) -> u32 {
     return thread_id / WORKGROUP_SIZE_M;

ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_subgroup_matrix.wgsl

@@ -4,6 +4,10 @@ enable subgroups;
 enable chromium_experimental_subgroup_matrix;
 
 #define DECLARE_BYTE_LOADERS_SRC0
+#ifdef SRC_OVERLAP
+#define SRC0 merged_src
+#define SRC1 merged_src
+#endif
 #include "common_decls.tmpl"
 
 #include "mul_mat_decls.tmpl"
@@ -48,11 +52,17 @@ struct MulMatParams {
 };
 
 // SRC0_TYPE and SRC1_TYPE are defined in mul_mat_decls, which is included
+#ifdef SRC_OVERLAP
+@group(0) @binding(0) var<storage, read_write> merged_src: array<SRC0_TYPE>;
+#define DST_BINDING 1
+#else
 @group(0) @binding(0) var<storage, read_write> src0: array<SRC0_TYPE>; // M rows, K columns
 @group(0) @binding(1) var<storage, read_write> src1: array<SRC1_TYPE>; // K rows, N columns (transposed)
-@group(0) @binding(2) var<storage, read_write> dst: array<DST_TYPE>; // M rows, N columns (transposed)
+#define DST_BINDING 2
+#endif
 
-@group(0) @binding(3) var<uniform> params: MulMatParams;
+@group(0) @binding(DST_BINDING) var<storage, read_write> dst: array<DST_TYPE>; // M rows, N columns (transposed)
+@group(0) @binding(DST_BINDING + 1) var<uniform> params: MulMatParams;
 
 const WG_M_SG_TILE_SIZE = SUBGROUP_M * SUBGROUP_MATRIX_M * SUBGROUP_MATRIX_M_SIZE;
 const WG_N_SG_TILE_SIZE = SUBGROUP_N * SUBGROUP_MATRIX_N * SUBGROUP_MATRIX_N_SIZE;

ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec.wgsl

@@ -7,6 +7,11 @@ enable f16;
 requires packed_4x8_integer_dot_product;
 #endif
 
+#ifdef SRC_OVERLAP
+#define SRC0 merged_src
+#define SRC1 merged_src
+#endif
+
 #define DECLARE_BYTE_LOADERS_SRC0
 #include "common_decls.tmpl"
 
@@ -35,17 +40,22 @@ struct MulMatParams {
     broadcast3: u32
 };
 
+#if defined(MMVQ)
 @group(0) @binding(0) var<storage, read_write> src0: array<SRC0_TYPE>;
-
-#ifdef MMVQ
 @group(0) @binding(1) var<storage, read_write> src1q: array<q8_1>;
+#define DST_BINDING 2
+#elif defined(SRC_OVERLAP)
+@group(0) @binding(0) var<storage, read_write> merged_src: array<SRC0_TYPE>;
+#define DST_BINDING 1
 #else
+@group(0) @binding(0) var<storage, read_write> src0: array<SRC0_TYPE>;
 @group(0) @binding(1) var<storage, read_write> src1: array<SRC1_TYPE>;
+#define DST_BINDING 2
 #endif
 
-@group(0) @binding(2) var<storage, read_write> dst: array<f32>;
+@group(0) @binding(DST_BINDING) var<storage, read_write> dst: array<f32>;
 // "mul_mat_vec_acc.tmpl" requires params.k, params.m, params.stride_01
-@group(0) @binding(3) var<uniform> params: MulMatParams;
+@group(0) @binding(DST_BINDING + 1) var<uniform> params: MulMatParams;
 
 // Flattened as [row][thread] to keep each row's reduction contiguous in memory.
 var<workgroup> partial_sums: array<f32, OUTPUTS_PER_WG * WG_SIZE>;

ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec_acc.tmpl

@@ -1,3 +1,10 @@
+#ifndef SRC0
+#define SRC0 src0
+#endif
+#ifndef SRC1
+#define SRC1 src1
+#endif
+
 #ifdef U32_DEQUANT_HELPERS
 #define SRC0_TYPE u32
 
@@ -43,13 +50,13 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
     for (var k = thread_id; k < k_vec; k += WG_SIZE) {
         var x_vals: array<SRC1_TYPE, NUM_COLS>;
         for (var col = 0u;col < NUM_COLS;col += 1) {
-            x_vals[col] = src1[src1_idx_base_vec + col * (params.stride_11 / VEC_SIZE) + k];
+            x_vals[col] = SRC1[src1_idx_base_vec + col * (params.stride_11 / VEC_SIZE) + k];
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
             let output_row = row_base + row;
             if (output_row < params.m) {
                 let src0_idx = (src0_batch_offset + output_row * params.stride_01) / VEC_SIZE + k;
-                let w = src0[src0_idx];
+                let w = SRC0[src0_idx];
                 for (var col = 0u;col < NUM_COLS;col += 1) {
                     acc[col][row] += inner_dot(w, x_vals[col]);
                 }
@@ -76,7 +83,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < ELEMS_PER_THREAD; i++) {
-                x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
+                x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -116,8 +123,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
-                x_block[col][i]     = f32(src1[x_base + col * params.stride_11 + i]);
-                x_block[col][i + 4] = f32(src1[x_base + col * params.stride_11 + i + 16]);
+                x_block[col][i]     = f32(SRC1[x_base + col * params.stride_11 + i]);
+                x_block[col][i + 4] = f32(SRC1[x_base + col * params.stride_11 + i + 16]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -160,8 +167,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
-                x_block[col][i]     = f32(src1[x_base + col * params.stride_11 + i]);
-                x_block[col][i + 4] = f32(src1[x_base + col * params.stride_11 + i + 16]);
+                x_block[col][i]     = f32(SRC1[x_base + col * params.stride_11 + i]);
+                x_block[col][i + 4] = f32(SRC1[x_base + col * params.stride_11 + i + 16]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -205,8 +212,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
-                x_block[col][i]     = f32(src1[x_base + col * params.stride_11 + i]);
-                x_block[col][i + 4] = f32(src1[x_base + col * params.stride_11 + i + 16]);
+                x_block[col][i]     = f32(SRC1[x_base + col * params.stride_11 + i]);
+                x_block[col][i + 4] = f32(SRC1[x_base + col * params.stride_11 + i + 16]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -253,8 +260,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
-                x_block[col][i]     = f32(src1[x_base + col * params.stride_11 + i]);
-                x_block[col][i + 4] = f32(src1[x_base + col * params.stride_11 + i + 16]);
+                x_block[col][i]     = f32(SRC1[x_base + col * params.stride_11 + i]);
+                x_block[col][i + 4] = f32(SRC1[x_base + col * params.stride_11 + i + 16]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -302,7 +309,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < ELEMS_PER_THREAD; i++) {
-                x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
+                x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -347,7 +354,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < ELEMS_PER_THREAD; i++) {
-                x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
+                x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -409,10 +416,10 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, 16>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < 4u; i++) {
-                x_block[col][i]       = f32(src1[x_base + col * params.stride_11 + i]);
-                x_block[col][i + 4u]  = f32(src1[x_base + col * params.stride_11 + 32u + i]);
-                x_block[col][i + 8u]  = f32(src1[x_base + col * params.stride_11 + 64u + i]);
-                x_block[col][i + 12u] = f32(src1[x_base + col * params.stride_11 + 96u + i]);
+                x_block[col][i]       = f32(SRC1[x_base + col * params.stride_11 + i]);
+                x_block[col][i + 4u]  = f32(SRC1[x_base + col * params.stride_11 + 32u + i]);
+                x_block[col][i + 8u]  = f32(SRC1[x_base + col * params.stride_11 + 64u + i]);
+                x_block[col][i + 12u] = f32(SRC1[x_base + col * params.stride_11 + 96u + i]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -518,8 +525,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, 16>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < 8u; i++) {
-                x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
-                x_block[col][i + 8u] = f32(src1[x_base + col * params.stride_11 + 32u + i]);
+                x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
+                x_block[col][i + 8u] = f32(SRC1[x_base + col * params.stride_11 + 32u + i]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -610,10 +617,10 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         for (var col = 0u; col < NUM_COLS;col += 1) {
             let col_base = x_base + col * params.stride_11;
             for (var i = 0u; i < 4u; i++) {
-                x_block[col][i]       = f32(src1[col_base + i]);
-                x_block[col][i + 4u]  = f32(src1[col_base + 32u + i]);
-                x_block[col][i + 8u]  = f32(src1[col_base + 128u + i]);
-                x_block[col][i + 12u] = f32(src1[col_base + 160u + i]);
+                x_block[col][i]       = f32(SRC1[col_base + i]);
+                x_block[col][i + 4u]  = f32(SRC1[col_base + 32u + i]);
+                x_block[col][i + 8u]  = f32(SRC1[col_base + 128u + i]);
+                x_block[col][i + 12u] = f32(SRC1[col_base + 160u + i]);
             }
         }
 
@@ -713,10 +720,10 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         for (var col = 0u; col < NUM_COLS;col += 1) {
             let col_base = x_base + col * params.stride_11;
             for (var i = 0u; i < 4u; i++) {
-                x_block[col][i]       = f32(src1[col_base + i]);
-                x_block[col][i + 4u]  = f32(src1[col_base + 32u + i]);
-                x_block[col][i + 8u]  = f32(src1[col_base + 128u + i]);
-                x_block[col][i + 12u] = f32(src1[col_base + 160u + i]);
+                x_block[col][i]       = f32(SRC1[col_base + i]);
+                x_block[col][i + 4u]  = f32(SRC1[col_base + 32u + i]);
+                x_block[col][i + 8u]  = f32(SRC1[col_base + 128u + i]);
+                x_block[col][i + 12u] = f32(SRC1[col_base + 160u + i]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -823,10 +830,10 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         for (var col = 0u; col < NUM_COLS;col += 1) {
             let col_base = x_base + col * params.stride_11;
             for (var l = 0u; l < 4u; l++) {
-                x_block[col][l]       = f32(src1[col_base + l]);
-                x_block[col][l + 4u]  = f32(src1[col_base + 32u + l]);
-                x_block[col][l + 8u]  = f32(src1[col_base + 64u + l]);
-                x_block[col][l + 12u] = f32(src1[col_base + 96u + l]);
+                x_block[col][l]       = f32(SRC1[col_base + l]);
+                x_block[col][l + 4u]  = f32(SRC1[col_base + 32u + l]);
+                x_block[col][l + 8u]  = f32(SRC1[col_base + 64u + l]);
+                x_block[col][l + 12u] = f32(SRC1[col_base + 96u + l]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -899,7 +906,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, 16>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < 16u; i++) {
-                x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
+                x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -960,7 +967,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, 16>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < 16u; i++) {
-                x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
+                x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1039,7 +1046,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, 16>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < 16u; i++) {
-                x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
+                x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1101,7 +1108,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, 16>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < 16u; i++) {
-                x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
+                x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1168,7 +1175,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, 16>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < 16u; i++) {
-                x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
+                x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1234,7 +1241,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, 16>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < 16u; i++) {
-                x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
+                x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1302,7 +1309,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, 16>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < 16u; i++) {
-                x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
+                x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1367,8 +1374,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < ELEMS_PER_THREAD / 2u; i++) {
-                x_block[col][i]     = f32(src1[x_base + col * params.stride_11 + i]);
-                x_block[col][i + 4u] = f32(src1[x_base + col * params.stride_11 + i + 16u]);
+                x_block[col][i]     = f32(SRC1[x_base + col * params.stride_11 + i]);
+                x_block[col][i + 4u] = f32(SRC1[x_base + col * params.stride_11 + i + 16u]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1418,7 +1425,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, 16>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < 16u; i++) {
-                x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
+                x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1476,8 +1483,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
-                x_block[col][i]     = f32(src1[x_base + col * params.stride_11 + i]);
-                x_block[col][i + 4] = f32(src1[x_base + col * params.stride_11 + i + 16]);
+                x_block[col][i]     = f32(SRC1[x_base + col * params.stride_11 + i]);
+                x_block[col][i + 4] = f32(SRC1[x_base + col * params.stride_11 + i + 16]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1521,8 +1528,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
         var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
         for (var col = 0u; col < NUM_COLS;col += 1) {
             for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
-                x_block[col][i]     = f32(src1[x_base + col * params.stride_11 + i]);
-                x_block[col][i + 8] = f32(src1[x_base + col * params.stride_11 + i + 8]);
+                x_block[col][i]     = f32(SRC1[x_base + col * params.stride_11 + i]);
+                x_block[col][i + 8] = f32(SRC1[x_base + col * params.stride_11 + i + 8]);
             }
         }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {

ggml/src/ggml.c

@@ -4200,7 +4200,7 @@ static struct ggml_tensor * ggml_rope_impl(
 
     struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a);
 
-    int32_t params[15] = { /*n_past*/ 0, n_dims, mode, /*n_ctx*/ 0, n_ctx_orig };
+    int32_t params[16] = { /*n_past*/ 0, n_dims, mode, /*n_ctx*/ 0, n_ctx_orig };
     memcpy(params +  5, &freq_base,    sizeof(float));
     memcpy(params +  6, &freq_scale,   sizeof(float));
     memcpy(params +  7, &ext_factor,   sizeof(float));
@@ -4212,6 +4212,8 @@ static struct ggml_tensor * ggml_rope_impl(
     } else {
         memset(params + 11, 0,         sizeof(int32_t) * GGML_MROPE_SECTIONS);
     }
+    params[15] = 0; // n_offs, set via ggml_rope_set_offset()
+
     ggml_set_op_params(result, params, sizeof(params));
 
     result->op     = GGML_OP_ROPE;
@@ -4422,6 +4424,20 @@ struct ggml_tensor * ggml_rope_multi_back(
     result->op = GGML_OP_ROPE_BACK;
     return result;
 }
+
+struct ggml_tensor * ggml_rope_set_offset(
+        struct ggml_tensor  * a,
+        int                   n_offs) {
+    GGML_ASSERT(a->op == GGML_OP_ROPE || a->op == GGML_OP_ROPE_BACK);
+    GGML_ASSERT(n_offs >= 0);
+
+    const int32_t mode = ggml_get_op_params_i32(a, 2);
+    GGML_ASSERT(mode != GGML_ROPE_TYPE_VISION);
+
+    ggml_set_op_params_i32(a, 15, n_offs);
+    return a;
+}
+
 // ggml_clamp
 
 struct ggml_tensor * ggml_clamp(

gguf-py/gguf/constants.py

@@ -208,6 +208,7 @@ class Keys:
         SHARED_KV_LAYERS             = "{arch}.attention.shared_kv_layers"
         SLIDING_WINDOW_PATTERN       = "{arch}.attention.sliding_window_pattern"
         TEMPERATURE_SCALE            = "{arch}.attention.temperature_scale"
+        ROPE_PATTERN                 = "{arch}.attention.rope_pattern"
 
         class Indexer:
             HEAD_COUNT = "{arch}.attention.indexer.head_count"
@@ -549,6 +550,7 @@ class MODEL_ARCH(IntEnum):
     GRANITE_MOE      = auto()
     GRANITE_HYBRID   = auto()
     GRANITE_SWITCH   = auto()
+    GRANITE_SWA      = auto()
     CHAMELEON        = auto()
     WAVTOKENIZER_DEC = auto()
     PLM              = auto()
@@ -1265,6 +1267,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
     MODEL_ARCH.GRANITE_MOE:      "granitemoe",
     MODEL_ARCH.GRANITE_HYBRID:   "granitehybrid",
     MODEL_ARCH.GRANITE_SWITCH:   "graniteswitch",
+    MODEL_ARCH.GRANITE_SWA:      "granite_swa",
     MODEL_ARCH.CHAMELEON:        "chameleon",
     MODEL_ARCH.WAVTOKENIZER_DEC: "wavtokenizer-dec",
     MODEL_ARCH.PLM:              "plm",
@@ -4152,6 +4155,31 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.FFN_DOWN,
         MODEL_TENSOR.FFN_UP,
     ],
+    MODEL_ARCH.GRANITE_SWA: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_SINKS,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        # MoE (GraniteMoeSWA)
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_GATE_UP_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        # Shared expert - gate+up kept fused in FFN_UP_SHEXP (LLM_FFN_SWIGLU)
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+    ],
     MODEL_ARCH.CHAMELEON: [
         MODEL_TENSOR.TOKEN_EMBD,
         MODEL_TENSOR.OUTPUT_NORM,

gguf-py/gguf/gguf_writer.py

@@ -824,6 +824,9 @@ class GGUFWriter:
         else:
             self.add_array(key, value)
 
+    def add_rope_pattern(self, value: Sequence[bool]) -> None:
+        self.add_array(Keys.Attention.ROPE_PATTERN.format(arch=self.arch), value)
+
     def add_dense_features_dims(self, dense:str, in_f:int, out_f:int) -> None:
         self.add_uint32(Keys.LLM.DENSE_FEAT_IN_SIZE.format(arch=self.arch, dense=dense), in_f)
         self.add_uint32(Keys.LLM.DENSE_FEAT_OUT_SIZE.format(arch=self.arch, dense=dense), out_f)

gguf-py/gguf/tensor_mapping.py

@@ -458,6 +458,7 @@ class TensorNameMap:
             "transformer.decoder_layer.{bid}.router",           # Grok
             "transformer.blocks.{bid}.ffn.router.layer",        # dbrx
             "model.layers.{bid}.block_sparse_moe.router.layer", # granitemoe
+            "model.layers.{bid}.block_sparse_moe.router",       # granite_swa
             "model.layers.{bid}.feed_forward.router",           # llama4 jamba
             "encoder.layers.{bid}.mlp.router.layer",            # nomic-bert-moe
             "model.layers.{bid}.mlp.router",                    # openai-moe

src/llama-arch.cpp

@@ -102,6 +102,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
     { LLM_ARCH_GRANITE_MOE,      "granitemoe"       },
     { LLM_ARCH_GRANITE_HYBRID,   "granitehybrid"    },
     { LLM_ARCH_GRANITE_SWITCH,   "graniteswitch"    },
+    { LLM_ARCH_GRANITE_SWA,      "granite_swa"      },
     { LLM_ARCH_CHAMELEON,        "chameleon"        },
     { LLM_ARCH_WAVTOKENIZER_DEC, "wavtokenizer-dec" },
     { LLM_ARCH_PLM,              "plm"              },
@@ -261,6 +262,8 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT,       "%s.attention.relative_buckets_count"       },
     { LLM_KV_ATTENTION_SLIDING_WINDOW,               "%s.attention.sliding_window"               },
     { LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN,       "%s.attention.sliding_window_pattern"       },
+    { LLM_KV_ATTENTION_ROPE_PATTERN,                 "%s.attention.rope_pattern"                 },
+
     { LLM_KV_ATTENTION_SCALE,                        "%s.attention.scale"                        },
     { LLM_KV_ATTENTION_OUTPUT_SCALE,                 "%s.attention.output_scale"                 },
     { LLM_KV_ATTENTION_VALUE_SCALE,                  "%s.attention.value_scale"                  },

src/llama-arch.h

@@ -107,6 +107,7 @@ enum llm_arch {
     LLM_ARCH_GRANITE_MOE,
     LLM_ARCH_GRANITE_HYBRID,
     LLM_ARCH_GRANITE_SWITCH,
+    LLM_ARCH_GRANITE_SWA,
     LLM_ARCH_CHAMELEON,
     LLM_ARCH_WAVTOKENIZER_DEC,
     LLM_ARCH_PLM,
@@ -267,6 +268,8 @@ enum llm_kv {
     LLM_KV_ATTENTION_SLIDING_WINDOW,
     LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN,
     LLM_KV_ATTENTION_SCALE,
+    LLM_KV_ATTENTION_ROPE_PATTERN,
+
     LLM_KV_ATTENTION_OUTPUT_SCALE,
     LLM_KV_ATTENTION_VALUE_SCALE,
     LLM_KV_ATTENTION_TEMPERATURE_LENGTH,

src/llama-hparams.cpp

@@ -291,7 +291,11 @@ bool llama_hparams::has_rope(uint32_t il) const {
         return false;
     }
 
-    return true;
+    if (il < n_layer_all) {
+        return rope_pattern[il] != 0;
+    }
+
+    GGML_ABORT("%s: il (%u) out of bounds (n_layer_all: %u)\n", __func__, il, n_layer_all);
 }
 
 uint32_t llama_hparams::n_layer() const {

src/llama-hparams.h

@@ -144,6 +144,10 @@ struct llama_hparams {
 
     std::array<int, 4> rope_sections;
 
+    // Per-layer RoPE enable flags (1 = use RoPE, 0 = NoPE)
+    // by default, all layers use RoPE (controlled by rope_finetuned)
+    std::array<uint32_t, LLAMA_MAX_LAYERS> rope_pattern;
+
     // Sliding Window Attention (SWA)
     llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;
     // the size of the sliding window (0 - no SWA)

src/llama-model-saver.cpp

@@ -30,6 +30,7 @@ bool llama_model_saver_supports_arch(llm_arch arch) {
         case LLM_ARCH_MUSE_GLIMMER:
         case LLM_ARCH_MELLUM:
         case LLM_ARCH_LAGUNA:
+        case LLM_ARCH_GRANITE_SWA:
             return false;
         default:
             return true;
@@ -272,6 +273,7 @@ void llama_model_saver::add_kv_from_model() {
     add_kv(LLM_KV_ATTENTION_VALUE_RESIDUAL_MIX_LORA_RANK, hparams.n_lora_value_res_mix);
     add_kv(LLM_KV_ATTENTION_GATE_LORA_RANK,          hparams.n_lora_gate);
     add_kv(LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT,  hparams.n_rel_attn_bkts);
+    add_kv(LLM_KV_ATTENTION_ROPE_PATTERN,            hparams.rope_pattern, true);
     add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW,          hparams.n_swa);
     // add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN,  ???);
     add_kv(LLM_KV_ATTENTION_SCALE,                   hparams.f_attention_scale);

src/llama-model.cpp

@@ -246,6 +246,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
             return new llama_model_minicpm(params);
         case LLM_ARCH_GRANITE_HYBRID:
             return new llama_model_granite_hybrid(params);
+        case LLM_ARCH_GRANITE_SWA:
+            return new llama_model_granite_swa(params);
         case LLM_ARCH_CHAMELEON:
             return new llama_model_chameleon(params);
         case LLM_ARCH_WAVTOKENIZER_DEC:
@@ -1157,6 +1159,7 @@ void llama_model_base::load_hparams(llama_model_loader & ml) {
     std::fill(hparams.n_ff_arr.begin(),      hparams.n_ff_arr.end(),      0);
 
     std::fill(hparams.rope_sections.begin(), hparams.rope_sections.end(), 0);
+    std::fill(hparams.rope_pattern.begin(),  hparams.rope_pattern.end(), 1);
     std::fill(hparams.is_swa_impl.begin(),   hparams.is_swa_impl.end(), 0);
     std::fill(hparams.is_recr_impl.begin(),  hparams.is_recr_impl.end(),  llm_arch_is_recurrent(ml.get_arch()) ? 1 : 0);
     std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 0);
@@ -2639,6 +2642,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
         case LLM_ARCH_GRANITE_MOE:
         case LLM_ARCH_GRANITE_HYBRID:
         case LLM_ARCH_GRANITE_SWITCH:
+        case LLM_ARCH_GRANITE_SWA:
         case LLM_ARCH_CHAMELEON:
         case LLM_ARCH_BAILINGMOE:
         case LLM_ARCH_BAILINGMOE3:

src/models/deepseek32.cpp

@@ -10,8 +10,6 @@ void llama_model_deepseek32::load_arch_hparams(llama_model_loader & ml) {
     ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false);
 
     // MoE parameters
-    ml.get_key(LLM_KV_EXPERT_COUNT,                hparams.n_expert);
-    ml.get_key(LLM_KV_EXPERT_USED_COUNT,           hparams.n_expert_used);
     ml.get_key(LLM_KV_EXPERT_SHARED_COUNT,         hparams.n_expert_shared);
     ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT,   hparams.n_layer_dense_lead, false);
     ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE,        hparams.expert_weights_scale, false);

src/models/glm-dsa.cpp

@@ -32,8 +32,6 @@ void llama_model_glm_dsa::load_arch_hparams(llama_model_loader & ml) {
     ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false);
 
     // MoE parameters
-    ml.get_key(LLM_KV_EXPERT_COUNT,                hparams.n_expert);
-    ml.get_key(LLM_KV_EXPERT_USED_COUNT,           hparams.n_expert_used);
     ml.get_key(LLM_KV_EXPERT_SHARED_COUNT,         hparams.n_expert_shared);
     ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT,   hparams.n_layer_dense_lead, false);
     ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE,        hparams.expert_weights_scale, false);

src/models/glm4-moe.cpp

@@ -6,8 +6,6 @@ void llama_model_glm4_moe::load_arch_hparams(llama_model_loader & ml) {
     ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false);
 
     // MoE parameters
-    ml.get_key(LLM_KV_EXPERT_COUNT,                hparams.n_expert);
-    ml.get_key(LLM_KV_EXPERT_USED_COUNT,           hparams.n_expert_used);
     ml.get_key(LLM_KV_EXPERT_SHARED_COUNT,         hparams.n_expert_shared);
     ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT,   hparams.n_layer_dense_lead, false);
     ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE,        hparams.expert_weights_scale, false);

src/models/granite-hybrid.cpp

@@ -16,7 +16,8 @@ void llama_model_granite_hybrid::load_arch_hparams(llama_model_loader & ml) {
     // Granite uses rope_finetuned as a switch for rope, so default to true
     bool rope_finetuned = true;
     ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false);
-    hparams.rope_finetuned = rope_finetuned;
+    hparams.rope_finetuned = rope_finetuned; // needed for round trip save
+    std::fill(hparams.rope_pattern.begin(), hparams.rope_pattern.end(), rope_finetuned);
 
     // A layer is recurrent IFF the n_head_kv value is set to 0
     for (uint32_t i = 0; i < hparams.n_layer(); ++i) {
@@ -147,7 +148,7 @@ llama_model_granite_hybrid::graph::graph(const llama_model & model, const llm_gr
 
     // Positional embeddings populated if rope enabled
     ggml_tensor * inp_pos = nullptr;
-    if (hparams.rope_finetuned) {
+    if (hparams.has_rope(0)) {
         inp_pos = build_inp_pos();
     }
 
@@ -206,8 +207,7 @@ ggml_tensor * llama_model_granite_hybrid::graph::build_attention_layer(ggml_tens
                                                               const int                 il) {
     auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il);
 
-    const bool use_rope = hparams.rope_finetuned;
-    if (use_rope) {
+    if (hparams.has_rope(il)) {
         ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
         Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
                              ext_factor, attn_factor, beta_fast, beta_slow);

src/models/granite-moe.cpp

@@ -7,11 +7,6 @@ void llama_model_granite_moe::load_arch_hparams(llama_model_loader & ml) {
     ml.get_key(LLM_KV_EMBEDDING_SCALE,             hparams.f_embedding_scale, false);
     ml.get_key(LLM_KV_ATTENTION_SCALE,             hparams.f_attention_scale, false);
 
-    // Granite uses rope_finetuned as a switch for rope, so default to true
-    bool rope_finetuned = true;
-    ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false);
-    hparams.rope_finetuned = rope_finetuned;
-
     switch (hparams.n_layer()) {
         case 32: type = LLM_TYPE_3B; break;
         case 40: type = LLM_TYPE_3B; break;

src/models/granite-swa.cpp

@@ -0,0 +1,319 @@
+#include "models.h"
+
+#include <sstream>
+
+void llama_model_granite_swa::load_arch_hparams(llama_model_loader & ml) {
+    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+    ml.get_key(LLM_KV_LOGIT_SCALE,                 hparams.f_logit_scale);
+    ml.get_key(LLM_KV_RESIDUAL_SCALE,              hparams.f_residual_scale, false);
+    ml.get_key(LLM_KV_EMBEDDING_SCALE,             hparams.f_embedding_scale, false);
+    ml.get_key(LLM_KV_ATTENTION_SCALE,             hparams.f_attention_scale, false);
+
+    // MoE expert configuration
+    ml.get_key(LLM_KV_EXPERT_COUNT,                hparams.n_expert, false);
+    ml.get_key(LLM_KV_EXPERT_USED_COUNT,           hparams.n_expert_used, false);
+
+     // iSWA configuration
+    ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl);
+    ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
+    hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
+
+    // Granite4 Vision uses array deepstack_mapping
+    ml.get_arr(LLM_KV_DEEPSTACK_MAPPING, hparams.deepstack_mapping_arr, false);
+
+    // Count the unique deepstack input indices
+    std::unordered_set<uint32_t> unique_deepstack_idxs;
+    for (const auto val : hparams.deepstack_mapping_arr) {
+        if (val >= 0) {
+            unique_deepstack_idxs.insert(val);
+        }
+    }
+    hparams.n_deepstack_layers = unique_deepstack_idxs.size();
+
+    // Ensure all values are valid (avoid overflow attacks)
+    for (const auto val : unique_deepstack_idxs) {
+        if (val > hparams.n_deepstack_layers) {
+            std::stringstream ss;
+            ss << "Invalid deepstack index: " << val << " > " << hparams.n_deepstack_layers;
+            throw std::runtime_error(ss.str());
+        }
+    }
+
+    // Per-layer RoPE pattern (optional)
+    ml.get_arr(LLM_KV_ATTENTION_ROPE_PATTERN, hparams.rope_pattern, false);
+
+    switch (hparams.n_layer()) {
+        case 32: type = LLM_TYPE_3B; break;
+        case 40: type = LLM_TYPE_3B; break;
+        // Add additional layer/vocab/etc checks here for other model sizes
+        default: type = LLM_TYPE_UNKNOWN;
+    }
+
+    // For Granite MoE Shared
+    ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, /* required */ false);
+}
+
+void llama_model_granite_swa::load_arch_tensors(llama_model_loader &) {
+    LLAMA_LOAD_LOCALS;
+
+    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+    // output
+    output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
+    output      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
+
+    // if output is NULL, init from the input tok embed
+    if (output == NULL) {
+        output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
+    }
+
+    for (int i = 0; i < n_layer; ++i) {
+        auto & layer = layers[i];
+
+        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+
+        create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);
+        layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);
+
+        // optional bias tensors
+        layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED);
+
+        // Per-layer attention sinks for iSWA
+        layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0);
+
+        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+
+        if (hparams.rope_scaling_type_train == LLAMA_ROPE_SCALING_TYPE_LONGROPE) {
+            layer.rope_long  = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_LONG,  "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
+            layer.rope_short = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_SHORT, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
+        }
+        else {
+            layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
+        }
+
+        if (n_expert == 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);
+            layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd,   n_ff}, 0);
+
+            // optional MLP bias
+            layer.ffn_gate_b = create_tensor(tn(LLM_TENSOR_FFN_GATE, "bias", i), {n_ff}, TENSOR_NOT_REQUIRED);
+            layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED);
+            layer.ffn_up_b   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "bias", i), {n_ff}, TENSOR_NOT_REQUIRED);
+        } else {
+            layer.ffn_gate_inp  = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,  "weight", i), {n_embd, n_expert}, 0);
+            layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {  n_ff, n_embd, n_expert}, 0);
+            create_tensor_gate_up_exps(layer, i, n_embd, n_ff, n_expert, 0);
+
+            // For Granite MoE Shared - gate+up kept fused in ffn_up_shexp (see LLM_FFN_SWIGLU below)
+            if (hparams.n_ff_shexp > 0) {
+                layer.ffn_up_shexp   = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), {n_embd, 2*hparams.n_ff_shexp}, 0);
+                layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {hparams.n_ff_shexp, n_embd}, 0);
+            }
+        }
+    }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_granite_swa::build_arch_graph(const llm_graph_params & params) const {
+    return std::make_unique<graph>(*this, params);
+}
+
+llama_model_granite_swa::graph::graph(
+    const llama_model & model,
+    const llm_graph_params & params)
+    : llm_graph_context(params) {
+
+    const int64_t n_embd_head = hparams.n_embd_head_v();
+
+    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+    GGML_ASSERT(n_embd_head == n_rot);
+
+    ggml_tensor * cur;
+    ggml_tensor * inpL;
+
+    inpL = build_inp_embd(model.tok_embd);
+
+    // inp_pos - built only if rope enabled
+    ggml_tensor * inp_pos = build_inp_pos();
+    auto * inp_attn = build_attn_inp_kv_iswa();
+
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+    for (int il = 0; il < n_layer; ++il) {
+
+        // Granite Vision 4.1 deepstack: inject the projector stream that
+        // targets decoder layer `il` before the decoder runs.
+        // NOTE: skip the first deepstack layer since that's inpL
+        const auto & deepstack_emb_idx = hparams.deepstack_mapping_arr[il];
+        if (il > 0 && deepstack_emb_idx >= 0) {
+            ggml_tensor * ds = ggml_view_2d(ctx0,
+                res->t_inp_embd, n_embd, n_tokens,
+                res->t_inp_embd->nb[1],
+                deepstack_emb_idx * n_embd * sizeof(float));
+            inpL = ggml_add(ctx0, inpL, ds);
+            cb(inpL, "deepstack_in", il);
+        }
+
+        ggml_tensor * inpSA = inpL;
+
+        // norm
+        cur = build_norm(inpL,
+                model.layers[il].attn_norm, NULL,
+                LLM_NORM_RMS, il);
+        cb(cur, "attn_norm", il);
+
+        // self-attention
+        cur = build_attention_layer(
+            cur, inp_pos, inp_attn,
+            model, n_embd_head, il);
+
+        if (il == n_layer - 1 && inp_out_ids) {
+            cur   = ggml_get_rows(ctx0,   cur, inp_out_ids);
+            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+        }
+        // ffn
+        cur = build_layer_ffn(cur, inpSA, model, il);
+
+        // input for next layer
+        inpL = cur;
+    }
+    cur = inpL;
+
+    cur = build_norm(cur,
+            model.output_norm, NULL,
+            LLM_NORM_RMS, -1);
+
+    cb(cur, "result_norm", -1);
+    res->t_embd = cur;
+
+    // lm_head
+    cur = build_lora_mm(model.output, cur, model.output_s);
+
+    // For Granite architectures - scale logits
+    cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale);
+    cb(cur, "result_output", -1);
+    res->t_logits = cur;
+
+    ggml_build_forward_expand(gf, cur);
+}
+
+ggml_tensor * llama_model_granite_swa::graph::build_attention_layer(
+          ggml_tensor                  * cur,
+          ggml_tensor                  * inp_pos,
+          llm_graph_input_attn_kv_iswa * inp_attn,
+    const llama_model                  & model,
+    const int64_t                        n_embd_head,
+    const int                            il) {
+
+    auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
+            n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il);
+
+    const bool use_rope = hparams.has_rope(il);
+    if (use_rope) {
+        ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
+        Qcur = ggml_rope_ext(
+                ctx0, Qcur, inp_pos, rope_factors,
+                n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+                ext_factor, attn_factor, beta_fast, beta_slow
+                );
+
+        Kcur = ggml_rope_ext(
+                ctx0, Kcur, inp_pos, rope_factors,
+                n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+                ext_factor, attn_factor, beta_fast, beta_slow
+                );
+    }
+
+    cb(Qcur, "Qcur", il);
+    cb(Kcur, "Kcur", il);
+    cb(Vcur, "Vcur", il);
+
+    const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
+
+    // Pass layer.attn_sinks to build_attn for sink-based attention modulation
+    cur = build_attn(inp_attn,
+            model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s,
+            Qcur, Kcur, Vcur, nullptr, model.layers[il].attn_sinks, nullptr, kq_scale, il);
+    cb(cur, "attn_out", il);
+    return cur;
+}
+
+ggml_tensor * llama_model_granite_swa::graph::build_layer_ffn(
+          ggml_tensor * cur,
+          ggml_tensor * inpSA,
+    const llama_model & model,
+    const int           il) {
+
+    // For Granite architectures - scale residual
+    if (hparams.f_residual_scale) {
+        cur = ggml_scale(ctx0, cur, hparams.f_residual_scale);
+    }
+    ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+    cb(ffn_inp, "ffn_inp", il);
+
+    // feed-forward network (non-MoE)
+    if (model.layers[il].ffn_gate_inp == nullptr) {
+
+        cur = build_norm(ffn_inp,
+                model.layers[il].ffn_norm, NULL,
+                LLM_NORM_RMS, il);
+        cb(cur, "ffn_norm", il);
+
+        cur = build_ffn(cur,
+                model.layers[il].ffn_up,   model.layers[il].ffn_up_b,   NULL,
+                model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, NULL,
+                model.layers[il].ffn_down, model.layers[il].ffn_down_b, NULL,
+                NULL,
+                LLM_FFN_SILU, LLM_FFN_PAR, il);
+        cb(cur, "ffn_out", il);
+
+    } else {
+        // MoE branch
+        cur = build_norm(ffn_inp,
+                model.layers[il].ffn_norm, NULL,
+                LLM_NORM_RMS, il);
+        cb(cur, "ffn_norm", il);
+
+        ggml_tensor * moe_out = build_moe_ffn(cur,
+                model.layers[il].ffn_gate_inp,
+                model.layers[il].ffn_up_exps,
+                model.layers[il].ffn_gate_exps,
+                model.layers[il].ffn_down_exps,
+                nullptr,
+                n_expert, n_expert_used,
+                LLM_FFN_SILU, true,
+                hparams.expert_weights_scale,
+                LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX,
+                il,
+                nullptr, model.layers[il].ffn_gate_up_exps);
+        cb(moe_out, "ffn_moe_out", il);
+
+        // For Granite MoE Shared - gate+up kept fused in ffn_up_shexp
+        if (hparams.n_ff_shexp > 0) {
+            ggml_tensor * ffn_shexp = build_ffn(cur,
+                model.layers[il].ffn_up_shexp,   NULL, NULL,
+                NULL,                            NULL, NULL,
+                model.layers[il].ffn_down_shexp, NULL, NULL,
+                NULL,
+                LLM_FFN_SWIGLU, LLM_FFN_SEQ, il);
+            cb(ffn_shexp, "ffn_shexp", il);
+
+            cur = ggml_add(ctx0, moe_out, ffn_shexp);
+            cb(cur, "ffn_out", il);
+        } else {
+            cur = moe_out;
+        }
+    }
+
+    // For Granite architectures - scale residual
+    if (hparams.f_residual_scale) {
+        cur = ggml_scale(ctx0, cur, hparams.f_residual_scale);
+    }
+    cur = ggml_add(ctx0, cur, ffn_inp);
+    cb(cur, "ffn_out", il);
+
+    cur = build_cvec(cur, il);
+    cb(cur, "l_out", il);
+
+    return cur;
+}

src/models/granite-switch.cpp

@@ -11,7 +11,8 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) {
 
     bool rope_finetuned = true;
     ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false);
-    hparams.rope_finetuned = rope_finetuned;
+    hparams.rope_finetuned = rope_finetuned; // needed for round trip save
+    std::fill(hparams.rope_pattern.begin(), hparams.rope_pattern.end(), rope_finetuned);
 
     switch (hparams.n_layer()) {
         case 40: type = hparams.n_embd == 4096 ? LLM_TYPE_8B : LLM_TYPE_3B; break;
@@ -254,7 +255,7 @@ llama_model_granite_switch::graph::graph(
     cb(inpL, "inp_embd", -1);
 
     ggml_tensor * inp_pos = nullptr;
-    if (hparams.rope_finetuned) {
+    if (hparams.has_rope(0)) {
         inp_pos = build_inp_pos();
     }
     auto * inp_attn = build_attn_inp_kv();
@@ -361,7 +362,7 @@ ggml_tensor * llama_model_granite_switch::graph::build_attention_layer(
     Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
     Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
 
-    if (hparams.rope_finetuned) {
+    if (hparams.has_rope(il)) {
         ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
         Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors,
                 n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,

src/models/granite.cpp

@@ -33,7 +33,8 @@ void llama_model_granite::load_arch_hparams(llama_model_loader & ml) {
     // Granite uses rope_finetuned as a switch for rope, so default to true
     bool rope_finetuned = true;
     ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false);
-    hparams.rope_finetuned = rope_finetuned;
+    hparams.rope_finetuned = rope_finetuned; // needed for round trip save
+    std::fill(hparams.rope_pattern.begin(), hparams.rope_pattern.end(), rope_finetuned);
 
     switch (hparams.n_layer()) {
         case 32: type = LLM_TYPE_3B; break;
@@ -127,7 +128,7 @@ llama_model_granite::graph::graph(
 
     // inp_pos - built only if rope enabled
     ggml_tensor * inp_pos = nullptr;
-    if (hparams.rope_finetuned) {
+    if (hparams.has_rope(0)) {
         inp_pos = build_inp_pos();
     }
     auto * inp_attn = build_attn_inp_kv();
@@ -203,8 +204,7 @@ ggml_tensor * llama_model_granite::graph::build_attention_layer(
     auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
             n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il);
 
-    const bool use_rope = hparams.rope_finetuned;
-    if (use_rope) {
+    if (hparams.has_rope(il)) {
         ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
         Qcur = ggml_rope_ext(
                 ctx0, Qcur, inp_pos, rope_factors,

src/models/models.h

@@ -1719,6 +1719,34 @@ struct llama_model_granite_hybrid : public llama_model_base {
 };
 
 
+struct llama_model_granite_swa : public llama_model_base {
+    llama_model_granite_swa(const struct llama_model_params & params) : llama_model_base(params) {}
+    void load_arch_hparams(llama_model_loader & ml) override;
+    void load_arch_tensors(llama_model_loader & ml) override;
+
+    struct graph : public llm_graph_context {
+        graph(const llama_model & model, const llm_graph_params & params);
+
+    private:
+        ggml_tensor * build_attention_layer(
+                  ggml_tensor                  * cur,
+                  ggml_tensor                  * inp_pos,
+                  llm_graph_input_attn_kv_iswa * inp_attn,
+            const llama_model                  & model,
+            const int64_t                        n_embd_head,
+            const int                            il);
+
+        ggml_tensor * build_layer_ffn(
+                  ggml_tensor * cur,
+                  ggml_tensor * inpSA,
+            const llama_model & model,
+            const int           il);
+    };
+
+    std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
+};
+
+
 struct llama_model_chameleon : public llama_model_base {
     llama_model_chameleon(const struct llama_model_params & params) : llama_model_base(params) {}
     void load_arch_hparams(llama_model_loader & ml) override;

tests/test-backend-ops.cpp

@@ -3061,28 +3061,36 @@ struct test_cpy : public test_case {
 };
 
 // GGML_OP_CONT
+// permute = {0, 0, 0, 0} means no permutation: the source is transposed (or
+// view-sliced). A non-identity permute applies ggml_permute before ggml_cont.
 struct test_cont : public test_case {
     const ggml_type type;
     const std::array<int64_t, 4> ne;
     bool use_view_slice;
+    const std::array<int64_t, 4> permute;
 
     std::string vars() override {
-        return VARS_TO_STR3(type, ne, use_view_slice);
+        return VARS_TO_STR4(type, ne, use_view_slice, permute);
     }
 
     test_cont(ggml_type type = GGML_TYPE_F32,
             std::array<int64_t, 4> ne = {10, 10, 10, 1},
-            bool use_view_slice = false)
-        : type(type), ne(ne), use_view_slice(use_view_slice) {}
+            bool use_view_slice = false,
+            std::array<int64_t, 4> permute = {0, 0, 0, 0})
+        : type(type), ne(ne), use_view_slice(use_view_slice), permute(permute) {}
 
     ggml_tensor * build_graph(ggml_context * ctx) override {
         ggml_tensor * src = ggml_new_tensor(ctx, type, 4, ne.data());
         ggml_set_param(src);
         ggml_set_name(src, "src");
 
+        const bool permuted = permute[0] != 0 || permute[1] != 0 || permute[2] != 0 || permute[3] != 0;
 
         ggml_tensor * dst;
-        if (use_view_slice) {
+        if (permuted) {
+            dst = ggml_permute(ctx, src, permute[0], permute[1], permute[2], permute[3]);
+            ggml_set_name(dst, "src_permuted");
+        } else if (use_view_slice) {
             dst = ggml_view_4d(ctx, src, src->ne[0], 1, src->ne[2], src->ne[3],
                 src->nb[1], src->nb[2], src->nb[3], src->nb[0] * (src->ne[1] - 1));
             ggml_set_name(dst, "src_view_slice");
@@ -4470,9 +4478,10 @@ struct test_mul_mat : public test_case {
     const std::array<int64_t, 4> per; // permutation of dimensions
     const int64_t k_v; // size of k in memory, resulting in a non-contiguous view for k_v > k, no view for k_v == 0
     const uint32_t o; // number of outputs
+    const bool src_overlap; // a and b are overlapping views of the same tensor
 
     std::string vars() override {
-        return VARS_TO_STR10(type_a, type_b, m, n, k, bs, nr, per, k_v, o);
+        return VARS_TO_STR11(type_a, type_b, m, n, k, bs, nr, per, k_v, o, src_overlap);
     }
 
     double max_nmse_err() override {
@@ -4501,8 +4510,8 @@ struct test_mul_mat : public test_case {
             std::array<int64_t, 2> bs = {10, 10},
             std::array<int64_t, 2> nr = {2, 2},
             std::array<int64_t, 4> per = {0, 1, 2, 3},
-            int64_t k_v = 0, uint32_t o = 1)
-        : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), per(per), k_v(k_v), o(o) {}
+            int64_t k_v = 0, uint32_t o = 1, bool src_overlap = false)
+        : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), per(per), k_v(k_v), o(o), src_overlap(src_overlap) {}
 
     ggml_tensor * build_graph(ggml_context * ctx) override {
         // C^T = A * B^T: (k, m) * (k, n) => (m, n)
@@ -4535,6 +4544,18 @@ struct test_mul_mat : public test_case {
             b = ggml_permute(ctx, b, per[0], per[1], per[2], per[3]);
             ggml_set_name(a, "a_permuted");
             ggml_set_name(b, "b_permuted");
+        } else if (src_overlap) {
+            GGML_ASSERT(type_a == type_b);
+            GGML_ASSERT(k_v == 0);
+
+            // a and b are interleaved views of the same tensor: (e.g. fused QKV in MiniMax-01)
+            ggml_tensor * base = ggml_new_tensor_4d(ctx, type_a, 2*k, std::max(m, n), bs[0]*nr[0], bs[1]*nr[1]);
+            ggml_set_name(base, "base");
+
+            a = ggml_view_4d(ctx, base, k, m, bs[0],       bs[1],       base->nb[1], base->nb[2], base->nb[3], 0);
+            b = ggml_view_4d(ctx, base, k, n, bs[0]*nr[0], bs[1]*nr[1], base->nb[1], base->nb[2], base->nb[3], k*ggml_type_size(type_a));
+            ggml_set_name(a, "a");
+            ggml_set_name(b, "b");
         } else {
             const int64_t k_physical = k_v == 0 ? k : k_v;
             a = ggml_new_tensor_4d(ctx, type_a, k_physical, m, bs[0],       bs[1]);
@@ -5331,24 +5352,27 @@ struct test_rope : public test_case {
     int v; // view (1 : non-contiguous a)
     bool forward;
     bool inplace;
+    int n_offs; // offset of the rotated dims window, set via ggml_rope_set_offset()
 
     std::string vars() override {
         // forward can be inferred from the op, does not need to be printed
-        return VARS_TO_STR11(type, ne_a, n_dims, mode, n_ctx, fs, ef, af, ff, v, inplace);
+        return VARS_TO_STR12(type, ne_a, n_dims, mode, n_ctx, fs, ef, af, ff, v, inplace, n_offs);
     }
 
     test_rope(ggml_type type = GGML_TYPE_F32,
             std::array<int64_t, 4> ne_a = {10, 5, 3, 1},
             int n_dims = 10, int mode = GGML_ROPE_TYPE_NORMAL, int n_ctx = 512, float fs = 1.0f,
-            float ef = 0.0f, float af = 0.0f, bool ff = false, int v = 0, bool forward = true, bool inplace = false)
-        : type(type), ne_a(ne_a), n_dims(n_dims), mode(mode), n_ctx(n_ctx), fs(fs), ef(ef), af(af), ff(ff), v(v), forward(forward), inplace(inplace) {}
+            float ef = 0.0f, float af = 0.0f, bool ff = false, int v = 0, bool forward = true, bool inplace = false,
+            int n_offs = 0)
+        : type(type), ne_a(ne_a), n_dims(n_dims), mode(mode), n_ctx(n_ctx), fs(fs), ef(ef), af(af), ff(ff), v(v), forward(forward), inplace(inplace), n_offs(n_offs) {}
 
     ggml_tensor * build_graph(ggml_context * ctx) override {
         ggml_tensor * a;
         if (v & 1) {
             auto ne = ne_a; ne[0] *= 2; ne[1] *= 4; ne[2] *= 3;
             a = ggml_new_tensor(ctx, type, 4, ne.data());
-            if (forward) {
+            if (forward && n_offs == 0) {
+                // FIXME: support gradients with n_offs > 0
                 ggml_set_param(a);
             }
             ggml_set_name(a, "a");
@@ -5361,7 +5385,8 @@ struct test_rope : public test_case {
             // non-aligned buffer offset, which exercises backends' alignment paths.
             auto ne = ne_a; ne[0] *= 2;
             a = ggml_new_tensor(ctx, type, 4, ne.data());
-            if (forward) {
+            if (forward && n_offs == 0) {
+                // FIXME: support gradients with n_offs > 0
                 ggml_set_param(a);
             }
             ggml_set_name(a, "a");
@@ -5372,7 +5397,8 @@ struct test_rope : public test_case {
             ggml_set_name(a, "view_of_a");
         } else {
             a = ggml_new_tensor(ctx, type, 4, ne_a.data());
-            if (forward) {
+            if (forward && n_offs == 0) {
+                // FIXME: support gradients with n_offs > 0
                 ggml_set_param(a);
             }
             ggml_set_name(a, "a");
@@ -5433,6 +5459,9 @@ struct test_rope : public test_case {
                 out = ggml_rope_ext_back(ctx, a, pos, freq, n_dims, mode, 0, 10000.0f, fs, ef, af, 1.0f, 1.0f);
             }
         }
+        if (n_offs != 0) {
+            out = ggml_rope_set_offset(out, n_offs);
+        }
         ggml_set_name(out, "out");
 
         return out;
@@ -7055,9 +7084,10 @@ struct test_flash_attn_ext : public test_case {
     const ggml_type type_K;
     const ggml_type type_V;
     std::array<int32_t, 4> permute;
+    const bool kv_view; // create K/V as views of a larger buffer (like a KV cache)
 
     std::string vars() override {
-        return VARS_TO_STR14(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute);
+        return VARS_TO_STR15(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view);
     }
 
     double max_nmse_err() override {
@@ -7073,9 +7103,10 @@ struct test_flash_attn_ext : public test_case {
 
     test_flash_attn_ext(int64_t hsk = 128, int64_t hsv = 128, int64_t nh = 32, std::array<int64_t, 2> nr23 = {1, 1}, int64_t kv = 96, int64_t nb = 8,
                         bool mask = true, bool sinks = false, float max_bias = 0.0f, float logit_softcap = 0.0f, ggml_prec prec = GGML_PREC_F32,
-                        ggml_type type_K = GGML_TYPE_F16, ggml_type type_V = GGML_TYPE_F16, std::array<int32_t, 4> permute = {0, 1, 2, 3})
+                        ggml_type type_K = GGML_TYPE_F16, ggml_type type_V = GGML_TYPE_F16, std::array<int32_t, 4> permute = {0, 1, 2, 3},
+                        bool kv_view = true)
         : hsk(hsk), hsv(hsv), nh(nh), nr23(nr23), kv(kv), nb(nb), mask(mask), sinks(sinks), max_bias(max_bias), logit_softcap(logit_softcap), prec(prec),
-          type_K(type_K), type_V(type_V), permute(permute) {}
+          type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view) {}
 
     ggml_tensor * build_graph(ggml_context * ctx) override {
         const int64_t hsk_padded = GGML_PAD(hsk, ggml_blck_size(type_K));
@@ -7103,7 +7134,7 @@ struct test_flash_attn_ext : public test_case {
         ggml_tensor * q = create_permuted(GGML_TYPE_F32, hsk_padded, nb, nh*nr23[0], nr23[1], false);
         ggml_set_name(q, "q");
 
-        ggml_tensor * k = create_permuted(type_K,        hsk_padded, kv, nh,         nr23[1], true); // the K tensor is usually a view of the K cache
+        ggml_tensor * k = create_permuted(type_K,        hsk_padded, kv, nh,         nr23[1], kv_view); // the K tensor is usually a view of the K cache
         ggml_set_name(k, "k");
 
         ggml_tensor * v = nullptr;
@@ -7117,7 +7148,7 @@ struct test_flash_attn_ext : public test_case {
             //   - https://github.com/ggml-org/llama.cpp/pull/18986
             v = ggml_view_4d(ctx, k, hsv_padded, kv, nh, nr23[1], k->nb[1], k->nb[2], k->nb[3], 0);
         } else {
-            v = create_permuted(type_V,        hsv_padded, kv, nh,         nr23[1], true); // the V tensor is usually a view of the V cache
+            v = create_permuted(type_V,        hsv_padded, kv, nh,         nr23[1], kv_view); // the V tensor is usually a view of the V cache
         }
         ggml_set_name(v, "v");
 
@@ -8892,6 +8923,20 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         }
     }
 
+    for (ggml_type type_dst : { GGML_TYPE_F32, GGML_TYPE_F16 }) {
+        for (std::array<int64_t, 4> ne : std::initializer_list<std::array<int64_t, 4>>{
+                {10, 10, 10, 1}, {33, 5, 7, 1}, {64, 3, 65, 1}, {2, 3, 5, 7},
+                // large, tile-aligned and tile-unaligned, matching the perf cases
+                {1024, 64, 64, 1}, {2304, 64, 64, 1}, {1000, 33, 65, 1} }) {
+            for (std::array<int64_t, 4> perm : std::initializer_list<std::array<int64_t, 4>>{
+                    {2, 1, 0, 3},   // 0<->2 swap
+                    {1, 2, 0, 3},   // 3-cycle
+                    {0, 2, 1, 3} }) {
+                test_cases.emplace_back(new test_cont(type_dst, ne, false, perm));
+            }
+        }
+    }
+
     auto add_test_bin_bcast = [&](ggml_type type, std::array<int64_t, 4> ne, std::array<int, 4> nr, bool perm1 = false, bool src_overlap = false) {
         for (auto op : {ggml_add, ggml_sub, ggml_mul, ggml_div}) {
             test_cases.emplace_back(new test_bin_bcast(op, type, ne, nr, 1, perm1, src_overlap));
@@ -9243,6 +9288,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 1056, 1, 67,  {1,  1}, {4, 1}, {0, 2, 1, 3}));
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 16, 32, 32, { 1,  1}, {1, 1}, {0, 1, 2, 3}, 64, 3));
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 64, 77, 77, {12,1}, {1,1}));
+    test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 32, 4, 96, {3, 2}, {1, 1}, {0, 1, 2, 3}, 0, 1, true));
 
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 576, 512, 576, {1,1}, {1,1}));
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 1, 2048, 8192, {1,  1}, {1, 1}));
@@ -9585,6 +9631,20 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         }
     }
 
+    // rotated dims window at an offset (ggml_rope_set_offset), not supported for vision mode
+    for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_F16}) {
+        for (bool fw : {true, false}) { // fw == forward
+            for (bool ff : {false, true}) {
+                test_cases.emplace_back(new test_rope(type, {128, 32, 2, 1}, 32, GGML_ROPE_TYPE_NORMAL, 512, 1.4245f, 0.7465f, 1.4245f, ff, 0, fw, false, 32));
+                test_cases.emplace_back(new test_rope(type, {128, 32, 2, 1}, 32, GGML_ROPE_TYPE_NEOX,   512, 1.4245f, 0.7465f, 1.4245f, ff, 0, fw, false, 32));
+                test_cases.emplace_back(new test_rope(type, {128, 12, 2, 1}, 24, GGML_ROPE_TYPE_MROPE,  512, 1.4245f, 0.7465f, 1.4245f, ff, 0, fw, false, 32));
+                test_cases.emplace_back(new test_rope(type, {128, 12, 2, 1}, 24, GGML_ROPE_TYPE_IMROPE, 512, 1.4245f, 0.7465f, 1.4245f, ff, 0, fw, false, 32));
+            }
+        }
+        // inplace with an offset
+        test_cases.emplace_back(new test_rope(type, {128, 32, 2, 1}, 32, GGML_ROPE_TYPE_NEOX, 512, 1.4245f, 0.7465f, 1.4245f, false, 0, true, true, 32));
+    }
+
     for (int v : { 0, 1, 2, 3 }) {
         for (int dim : { 0, 1, 2, 3, }) {
             test_cases.emplace_back(new test_concat(GGML_TYPE_F32, {11, 12, 13, 14}, 7, dim, v));
@@ -9883,6 +9943,12 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
                                                         GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
     }
 
+    // dense-allocated (non-view) quant K/V at batch >= 64, in cache and native layouts
+    test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false));
+    test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {4, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false));
+    test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false));
+    test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, false));
+
     test_cases.emplace_back(new test_cross_entropy_loss     (GGML_TYPE_F32, {   10, 5, 4, 3}));
     test_cases.emplace_back(new test_cross_entropy_loss     (GGML_TYPE_F32, {30000, 1, 1, 1}));
     test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, {   10, 5, 4, 3}));
@@ -10042,6 +10108,17 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
         }
     }
 
+    // CONT of a 0<->2 permute at DeepSeek-V4 lightning-indexer shapes:
+    // indexer_kq is [n_kv, n_tokens, n_head=64] and gets ggml_cont(ggml_permute(.., 2,1,0,3)).
+    for (int64_t n_kv : { 1024, 1280, 2048, 2304 }) {
+        test_cases.emplace_back(new test_cont(
+            GGML_TYPE_F32, {n_kv, 64, 64, 1}, false, {2, 1, 0, 3}));
+    }
+    for (int64_t n_kv : { 2048, 2304 }) {
+        test_cases.emplace_back(new test_cont(
+            GGML_TYPE_F32, {n_kv, 512, 64, 1}, false, {2, 1, 0, 3}));
+    }
+
     // Conv2d: K=CRS=NPQ=4096 matmul performance
     uint32_t                        iwh_idx  = 0;
     uint32_t                        kwh_idx  = 1;

tests/test-llama-archs.cpp

@@ -197,7 +197,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
         ms.add_kv(LLM_KV_ROPE_FREQ_BASE_SWA,              10000.0f);
         // SWA pattern: every 5th layer is full attention (matches E2B layer_types)
         ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(5));
-    } else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_MUSE_GLIMMER) {
+    } else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_GRANITE_SWA) {
         std::vector<uint32_t> pattern;
         pattern.reserve(n_layer);
         for (uint32_t il = 0; il < n_layer; il++) {
@@ -456,7 +456,7 @@ static bool arch_supported(const llm_arch arch) {
 
     // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI.
 #ifdef GGML_USE_WEBGPU
-    if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MINIMAX_01) {
+    if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA) {
         return false;
     }
 #endif // GGML_USE_WEBGPU

tests/test-mtmd-impl.cpp

@@ -1,9 +1,12 @@
 #include "testing.h"
 
 #include "mtmd-image.h"
+#include "mtmd-internal.h"
 
 #include <iostream>
+#include <stdexcept>
 #include <string>
+#include <tuple>
 #include <utility>
 #include <vector>
 
@@ -67,6 +70,73 @@ MAKE_TEST(test_image_preprocessor_lfm2) {
     }
 }
 
+//
+// mtmd temporal merge
+//
+
+MAKE_TEST(test_temporal_merge_grouping) {
+    std::vector<mtmd::bitmap_ptr> pool; // keeps the bitmaps alive until the end of the test
+
+    // spec chars:
+    //   v = video frame, w = video frame of another size, a = audio, i = plain image, t = text
+    auto make_parts = [&pool](const std::string & spec) {
+        std::vector<mtmd_input_part> parts;
+        for (char c : spec) {
+            if (c == 't') {
+                parts.push_back({ "hello", nullptr });
+                continue;
+            }
+            mtmd_bitmap * bm = nullptr;
+            switch (c) {
+                case 'v': bm = mtmd_bitmap_init(100, 100, nullptr);   break;
+                case 'w': bm = mtmd_bitmap_init(200, 200, nullptr);   break;
+                case 'a': bm = mtmd_bitmap_init_from_audio(100, nullptr); break;
+                case 'i': bm = mtmd_bitmap_init(100, 100, nullptr);   break;
+                default: throw std::runtime_error(std::string("unknown spec char: ") + c);
+            }
+            mtmd_bitmap_set_mergeable(bm, c != 'i');
+            pool.emplace_back(bm);
+            parts.push_back({ "", bm });
+        }
+        return parts;
+    };
+
+    // { parts, n_merge, expected size of each group }
+    const std::vector<std::tuple<std::string, int, std::string>> cases = {
+        { "vv",   2, "2"    },
+        { "vvv",  2, "21"   },
+        { "vvvv", 2, "22"   },
+        { "vvi",  2, "21"   },
+        { "tvvt", 2, "2"    },
+        { "vtv",  2, "11"   }, // text in between breaks the merge
+        { "vw",   2, "11"   }, // different sizes cannot be merged
+        { "aa",   2, "11"   }, // audio is never merged
+        { "ii",   2, "11"   }, // two unrelated images must stay separated
+        { "iv",   2, "11"   },
+        { "vi",   2, "11"   },
+        { "vv",   1, "11"   }, // model without temporal merge
+    };
+
+    for (const auto & [spec, n_merge, expected] : cases) {
+        auto parts  = make_parts(spec);
+        auto groups = mtmd_group_mergeable_bitmaps(parts, n_merge);
+
+        std::string actual;
+        for (const auto & group : groups) {
+            actual += std::to_string(group.size());
+        }
+
+        const std::string name = "\"" + spec + "\" with n_merge=" + std::to_string(n_merge);
+        t.assert_equal("groups for " + name, expected, actual);
+
+        size_t n_bitmap_parts = 0;
+        for (const auto & p : parts) {
+            n_bitmap_parts += p.bitmap != nullptr ? 1 : 0;
+        }
+        t.assert_equal("remaining bitmap parts for " + name, groups.size(), n_bitmap_parts);
+    }
+}
+
 //
 // main
 //

tools/mtmd/CMakeLists.txt

@@ -17,6 +17,7 @@ add_library(mtmd
             mtmd-audio.cpp
             mtmd-image.cpp
             mtmd.h
+            mtmd-internal.h
             mtmd-helper.cpp
             mtmd-helper-gen.cpp
             mtmd-helper-common.h

tools/mtmd/README-dev.md

@@ -21,7 +21,7 @@ A typical pipeline of the core libmtmd is as follows:
 - A bitmap (RGB image or PCM audio) is created
 - Bitmap and the text prompt is provided to `mtmd_tokenize()` that breaks the input into chunks
     - The tokenizer function first expands a "lazy" bitmap if it finds one. Typically, this is used by video, so that one media token corresponds to one input bitmap
-    - For models that support "fused" temporal frames like Qwen-VL, the tokenizer tries to merge pair of consecutive frames into one batch
+    - For models that support "fused" temporal frames like Qwen-VL, the tokenizer tries to merge pair of consecutive frames into one batch. Only bitmaps marked by `mtmd_bitmap_set_mergeable()` are merged
     - The preprocessor will then be called, which produces a list of chunks
     - Depending on the model itself, special tokens will be injected to separate image chunks (i.e. llava-uhd-style models)
 - Multiple bitmaps may be batched together to form a larger `mtmd_batch()`

tools/mtmd/mtmd-helper.cpp

@@ -727,7 +727,9 @@ struct mtmd_helper_video {
 
         LOG_DBG("%s: frame %d read OK\n", __func__, current_frame);
         current_frame++;
-        return mtmd_bitmap_init(info.width, info.height, frame_buf.data());
+        mtmd_bitmap * frame = mtmd_bitmap_init(info.width, info.height, frame_buf.data());
+        mtmd_bitmap_set_mergeable(frame, true);
+        return frame;
     }
 
     int32_t read_next(mtmd_bitmap ** out_bitmap, char ** out_text) {

tools/mtmd/mtmd-internal.h

@@ -0,0 +1,19 @@
+#pragma once
+
+#include "mtmd.h"
+
+#include <string>
+#include <vector>
+
+// !!! Internal header, to be used by mtmd and its unit tests only !!!
+
+#define MTMD_INTERNAL_HEADER
+
+// bitmap is null for text parts
+struct mtmd_input_part {
+    std::string text;
+    const mtmd_bitmap * bitmap;
+};
+
+// [QWEN_VIDEO] merged parts are erased from `parts`, so one group always maps to one part
+std::vector<std::vector<const mtmd_bitmap *>> mtmd_group_mergeable_bitmaps(std::vector<mtmd_input_part> & parts, int n_merge);

tools/mtmd/mtmd.cpp

@@ -1,6 +1,7 @@
 #include "clip.h"
 #include "clip-impl.h"
 #include "mtmd.h"
+#include "mtmd-internal.h"
 #include "mtmd-audio.h"
 #include "mtmd-image.h"
 #include "debug/mtmd-debug.h"
@@ -149,6 +150,7 @@ struct mtmd_bitmap {
     uint32_t ny = 0;
     std::string id; // optional user-defined id, for ex: can be set to image hash, useful for KV cache tracking
     bool is_audio = false; // true if the bitmap is audio
+    bool mergeable = false; // [QWEN_VIDEO] set only on frames of the same video
 
     // lazy-loaded bitmap
     mtmd_bitmap_lazy_callback lazy_callback = nullptr;
@@ -186,7 +188,9 @@ struct mtmd_bitmap {
 
     bool can_merge_with(const mtmd_bitmap & other) const {
         // [QWEN_VIDEO] can (temporal) merge if both are images with same size
-        return !is_audio && !other.is_audio && nx == other.nx && ny == other.ny;
+        return mergeable && other.mergeable
+            && !is_audio && !other.is_audio
+            && nx == other.nx && ny == other.ny;
     }
 
   private:
@@ -1076,6 +1080,25 @@ void mtmd_free(mtmd_context * ctx) {
     delete ctx;
 }
 
+std::vector<std::vector<const mtmd_bitmap *>> mtmd_group_mergeable_bitmaps(std::vector<mtmd_input_part> & parts, int n_merge) {
+    std::vector<std::vector<const mtmd_bitmap *>> output;
+    for (size_t i = 0; i < parts.size(); i++) {
+        if (parts[i].bitmap == nullptr) {
+            continue; // text part
+        }
+        const bool has_next = n_merge > 1 && i + 1 < parts.size() && parts[i + 1].bitmap != nullptr;
+        if (has_next && parts[i].bitmap->can_merge_with(*parts[i + 1].bitmap)) {
+            LOG_DBG("%s: merging 2 frames at part index %zu and %zu\n", __func__, i, i + 1);
+            output.push_back({parts[i].bitmap, parts[i + 1].bitmap});
+            parts.erase(parts.begin() + i + 1);
+            continue;
+        }
+        LOG_DBG("%s: no merging for part index %zu\n", __func__, i);
+        output.push_back({parts[i].bitmap});
+    }
+    return output;
+}
+
 struct mtmd_tokenizer {
     mtmd_context * ctx;
 
@@ -1084,10 +1107,7 @@ struct mtmd_tokenizer {
     bool parse_special;
     const llama_vocab * vocab;
 
-    struct part {
-        std::string text;
-        const mtmd_bitmap * bitmap;
-    };
+    using part = mtmd_input_part;
     std::vector<part> parts;
     // these will be freed when mtmd_tokenizer finishes
     std::vector<mtmd::bitmap> bm_from_lazy; // TODO @ngxson : refactor, free bm_from_lazy progressively
@@ -1192,34 +1212,7 @@ struct mtmd_tokenizer {
             GGML_ASSERT(n_merge_frames <= 2 && "we only support merging maximum 2 images for now; open an issue if this model supports merging more");
         }
 
-        // Build merged_bitmaps: each entry is a group of 1 or 2 bitmaps.
-        // For consecutive mergeable bitmap parts, merge them and collapse the second part out of this->parts.
-        std::vector<std::vector<const mtmd_bitmap *>> merged_bitmaps;
-        if (n_merge_frames > 1) {
-            for (size_t i = 0; i < parts.size(); ++i) {
-                if (parts[i].bitmap == nullptr) {
-                    continue;
-                }
-                if (i + 1 < parts.size() && parts[i + 1].bitmap != nullptr) {
-                    const mtmd_bitmap * bm_a = parts[i].bitmap;
-                    const mtmd_bitmap * bm_b = parts[i + 1].bitmap;
-                    if (bm_a->can_merge_with(*bm_b)) {
-                        LOG_DBG("%s: merging 2 frames at part index %zu and %zu\n", __func__, i, i + 1);
-                        merged_bitmaps.push_back({bm_a, bm_b});
-                        parts.erase(parts.begin() + i + 1); // collapse the second bitmap part
-                        continue;
-                    }
-                }
-                LOG_DBG("%s: no merging for part index %zu\n", __func__, i);
-                merged_bitmaps.push_back({parts[i].bitmap});
-            }
-        } else {
-            for (const auto & p : parts) {
-                if (p.bitmap != nullptr) {
-                    merged_bitmaps.push_back({p.bitmap});
-                }
-            }
-        }
+        auto merged_bitmaps = mtmd_group_mergeable_bitmaps(parts, n_merge_frames);
 
         size_t i_bm = 0;
         for (const auto & p : parts) {
@@ -2200,6 +2193,10 @@ void mtmd_bitmap_set_id(mtmd_bitmap * bitmap, const char * id) {
     }
 }
 
+void mtmd_bitmap_set_mergeable(mtmd_bitmap * bitmap, bool mergeable) {
+    bitmap->mergeable = mergeable;
+}
+
 mtmd_bitmap * mtmd_bitmap_init_lazy(mtmd_context * ctx,
                                     const char * id,
                                     void * user_data,

tools/mtmd/mtmd.h

@@ -154,7 +154,8 @@ MTMD_API const char * mtmd_get_marker(const mtmd_context * ctx);
 //     length of data must be nx * ny * 3
 //     the data is in RGBRGBRGB... format
 //     note: some video-capable models (i.e. qwen-vl) can merge consecutive bitmaps
-//           into one chunk, mtmd_tokenize() will automatically handle this
+//           into one chunk; mtmd_tokenize() handles this, but remember to set
+//           mtmd_bitmap_set_mergeable(true) for every frame
 // if bitmap is audio:
 //     length of data must be n_samples * sizeof(float)
 //     the data is in float format (PCM F32)
@@ -175,6 +176,8 @@ MTMD_API void                  mtmd_bitmap_free       (mtmd_bitmap * bitmap);
 // these getters/setters are dedicated functions, so you can for example calculate the hash of the image based on mtmd_bitmap_get_data()
 MTMD_API const char * mtmd_bitmap_get_id(const mtmd_bitmap * bitmap);
 MTMD_API void         mtmd_bitmap_set_id(mtmd_bitmap * bitmap, const char * id);
+// if true, this bitmap can be merged (temporal merge) with an adjacent mergeable bitmap by certain video input models
+MTMD_API void         mtmd_bitmap_set_mergeable(mtmd_bitmap * bitmap, bool mergeable);
 
 // mtmd_bitmap lazy
 //

tools/server/README-dev.md

@@ -291,6 +291,36 @@ The flow for downloading a new model:
 - If a stop request comes in, the router asks the child process to stop (same mechanism as running a model in child process)
 - Otherwise, upon completion, we call `load_models()` to refresh the list of models
 
+### Sleep mode
+
+Sleep mode was initially introduced in PR [#18228](https://github.com/ggml-org/llama.cpp/pull/18228). The main idea is to have:
+- `server_queue` keeping track of the idle timeout
+- When the timeout is detected, `server_queue` signals to `server_context_impl` that it should go into sleep
+- `server_context_impl` frees all `llama_context` and `mtmd_context`
+
+Compared to simply exiting the whole process, this approach allows accessing some read-only endpoints during sleep, while also handling wakeup-on-request. Any inference request will wake the server up.
+
+Call stack on entering sleeping:
+- `server_queue::start_loop` (main thread) sees no task for `idle_sleep_ms` --> `sleeping = true`
+- `cb0(true)` --> `server_routes::update_cached_responses`
+    - snapshots `/props`, `/models` and metrics; the model is still alive here
+- `cb1(true)` --> `server_context_impl::handle_sleeping_state`
+    - `callback_state(SERVER_STATE_SLEEPING)` --> reported to router in child mode
+    - `destroy()` --> frees `llama_context` and `mtmd_context`
+- `condition_tasks.wait` until `req_stop_sleeping`
+
+Call stack on waking up:
+- `server_res_generator` constructor (HTTP thread) --> `server_queue::wait_until_no_sleep`
+    - sets `req_stop_sleeping = true`, then waits until `sleeping == false`
+- `server_queue::start_loop` (main thread) wakes up
+- `cb1(false)` --> `server_context_impl::handle_sleeping_state`
+    - `load_model()`, which then emits `callback_state(SERVER_STATE_READY)`
+- `cb0(false)` --> `server_routes::update_cached_responses`
+    - nothing to do, the cache is only read during sleep
+- `sleeping = false` --> `notify_all` unblocks the HTTP thread, the request is handled as usual
+
+Endpoints created with `create_response(true)` (`/health`, `/props`, `/models`, `/metrics`) skip `wait_until_no_sleep`, so they answer from the cached responses instead of waking the server.
+
 ### Notable Related PRs
 
 - Initial server implementation: https://github.com/ggml-org/llama.cpp/pull/1443

tools/server/README.md

@@ -1759,6 +1759,7 @@ The precedence rule for preset options is as follows:
 We also offer additional options that are exclusive to presets (these aren't treated as command-line arguments):
 - `load-on-startup` (boolean): Controls whether the model loads automatically when the server starts
 - `stop-timeout` (int, seconds): After requested unload, wait for this many seconds before forcing termination (default: 10)
+- `dedup-cache-models` (boolean): When the preset uses `hf-repo` pointing to a model that is already downloaded, hide the corresponding cached model entry from `GET /models` (the preset entry remains visible). Set it in the `[*]` section to apply to all presets.
 
 ### Routing requests
 
@@ -2070,6 +2071,7 @@ Note that the following endpoints are exempt from being considered as incoming t
 - `GET /health`
 - `GET /props`
 - `GET /models`
+- `GET /metrics`
 
 ## More examples
 

tools/server/server-context.cpp

@@ -818,6 +818,14 @@ public:
         }
     }
 
+    server_metrics get_metrics() const {
+        return metrics;
+    }
+
+    void reset_metrics_bucket() {
+        metrics.reset_bucket();
+    }
+
 private:
     // note: accessing these fields outside of this class is not thread-safe
     // use server_context methods instead
@@ -898,6 +906,10 @@ private:
     void handle_sleeping_state(bool new_state) {
         GGML_ASSERT(sleeping != new_state);
         if (new_state) {
+            if (callback_state) {
+                callback_state(SERVER_STATE_SLEEPING, {});
+                // note: for sleeping == false, event is emitted by load_model()
+            }
             SRV_INF("%s", "server is entering sleeping state\n");
             destroy();
         } else {
@@ -2290,8 +2302,8 @@ private:
 
     // returns false to decline the task, it is offered again after the decode is done
     bool process_single_task(server_task && task, bool is_yielding) {
-        // while yielding, an encode / decode is running and only accessing metrics is safe
-        if (is_yielding && task.type != SERVER_TASK_TYPE_METRICS) {
+        // while yielding, an encode / decode is running and only reading the server state is safe
+        if (is_yielding && task.type != SERVER_TASK_TYPE_METRICS && task.type != SERVER_TASK_TYPE_SLOT_GET) {
             SRV_DBG("decoding, decline task, id_task = %d\n", task.id);
             return false;
         }
@@ -2417,28 +2429,17 @@ private:
                 } break;
             case SERVER_TASK_TYPE_METRICS:
                 {
-                    json slots_data = json::array();
-
-                    int n_idle_slots       = 0;
                     int n_processing_slots = 0;
 
                     for (server_slot & slot : slots) {
-                        json slot_data = slot.to_json(slots_debug == 0);
-
                         if (slot.is_processing()) {
                             n_processing_slots++;
-                        } else {
-                            n_idle_slots++;
                         }
-
-                        slots_data.push_back(slot_data);
                     }
-                    SRV_DBG("n_idle_slots = %d, n_processing_slots = %d\n", n_idle_slots, n_processing_slots);
+                    SRV_DBG("n_processing_slots = %d\n", n_processing_slots);
 
                     auto res = std::make_unique<server_task_result_metrics>();
                     res->id                  = task.id;
-                    res->slots_data          = std::move(slots_data);
-                    res->n_idle_slots        = n_idle_slots;
                     res->n_processing_slots  = n_processing_slots;
                     res->n_tasks_deferred    = queue_tasks.queue_tasks_deferred_size();
                     res->metrics             = metrics;
@@ -2446,6 +2447,28 @@ private:
                     if (task.metrics_reset_bucket) {
                         metrics.reset_bucket();
                     }
+                    queue_results.send(std::move(res));
+                } break;
+            case SERVER_TASK_TYPE_SLOT_GET:
+                {
+                    json slots_data = json::array();
+
+                    int n_idle_slots = 0;
+
+                    for (server_slot & slot : slots) {
+                        if (!slot.is_processing()) {
+                            n_idle_slots++;
+                        }
+
+                        slots_data.push_back(slot.to_json(slots_debug == 0));
+                    }
+                    SRV_DBG("n_idle_slots = %d\n", n_idle_slots);
+
+                    auto res = std::make_unique<server_task_result_slots>();
+                    res->id           = task.id;
+                    res->slots_data   = std::move(slots_data);
+                    res->n_idle_slots = n_idle_slots;
+
                     queue_results.send(std::move(res));
                 } break;
             case SERVER_TASK_TYPE_SLOT_SAVE:
@@ -4142,12 +4165,6 @@ struct server_res_generator : server_res_spipe {
 
 void server_context::set_state_callback(server_state_callback_t callback) {
     impl->callback_state = std::move(callback);
-    impl->queue_tasks.on_sleeping_state([this](bool sleeping) {
-        if (sleeping) {
-            impl->callback_state(SERVER_STATE_SLEEPING, {});
-        }
-        // for sleeping == false, event is emitted by load_model()
-    });
 }
 
 //
@@ -4431,6 +4448,119 @@ server_routes::server_routes(const common_params & params, server_context & ctx_
           queue_tasks(ctx_server.impl->queue_tasks),
           queue_results(ctx_server.impl->queue_results) {
     init_routes();
+
+    // note: this must be registered before load_model()
+    //       so that on sleep phase, the callback is called before ctx is destroyed
+    queue_tasks.on_sleeping_state([this](bool is_sleeping) {
+        update_cached_responses(is_sleeping);
+    });
+}
+
+static json get_res_model_info(const server_context_meta & meta) {
+    // note: do NOT use ctx_server here, otherwise it's not possible to use this during sleep
+
+    return {
+        {"id",       meta.model_name},
+        {"aliases",  meta.model_aliases},
+        {"tags",     meta.model_tags},
+        {"object",   "model"},
+        {"created",  std::time(0)},
+        {"owned_by", "llamacpp"},
+        {"meta",     {
+            {"vocab_type",  meta.model_vocab_type},
+            {"n_vocab",     meta.model_vocab_n_tokens},
+            {"n_ctx",       meta.slot_n_ctx},
+            {"n_ctx_train", meta.model_n_ctx_train},
+            {"n_embd",      meta.model_n_embd_inp},
+            {"n_params",    meta.model_n_params},
+            {"size",        meta.model_size},
+            {"ftype",       meta.model_ftype},
+        }},
+    };
+}
+
+static json get_res_models(const server_context_meta & meta) {
+    // note: do NOT use ctx_server here, otherwise it's not possible to use this during sleep
+
+    return {
+        {"models", {
+            {
+                {"name",  meta.model_name},
+                {"model", meta.model_name},
+                {"modified_at", ""},
+                {"size", ""},
+                {"digest", ""}, // dummy value, llama.cpp does not support managing model file's hash
+                {"type", "model"},
+                {"description", ""},
+                {"tags", {""}},
+                {"capabilities", meta.has_mtmd ? json({"completion","multimodal"}) : json({"completion"})},
+                {"parameters", ""},
+                {"details", {
+                    {"parent_model", ""},
+                    {"format", "gguf"},
+                    {"family", ""},
+                    {"families", {""}},
+                    {"parameter_size", ""},
+                    {"quantization_level", ""}
+                }}
+            }
+        }},
+        {"object", "list"},
+        {"data", {
+            get_res_model_info(meta),
+        }}
+    };
+}
+
+static json get_res_props(const server_context_meta & meta, const common_params & params, bool is_sleeping) {
+    // note: do NOT use ctx_server here, otherwise it's not possible to use this during sleep
+
+    task_params tparams;
+    tparams.sampling = params.sampling;
+    json default_generation_settings_for_props = json {
+        { "params", tparams.to_json(true) },
+        { "n_ctx",  meta.slot_n_ctx },
+    };
+
+    std::string tmpl_default = common_chat_templates_source(meta.chat_params.tmpls.get(), "");
+    std::string tmpl_tools   = common_chat_templates_source(meta.chat_params.tmpls.get(), "tool_use");
+
+    json props = {
+        { "default_generation_settings", default_generation_settings_for_props },
+        { "total_slots",                 params.n_parallel },
+        { "model_alias",                 meta.model_name },
+        { "model_ftype",                 meta.model_ftype },
+        { "model_path",                  meta.model_path },
+        { "modalities",                  json {
+            {"vision", meta.has_inp_image},
+            {"video",  meta.has_inp_video},
+            {"audio",  meta.has_inp_audio},
+        } },
+        { "media_marker",                get_media_marker() },
+        { "endpoint_slots",              params.endpoint_slots },
+        { "endpoint_props",              params.endpoint_props },
+        { "endpoint_metrics",            params.endpoint_metrics },
+        { "ui",                          params.ui },
+        { "ui_settings",                 meta.json_ui_settings },
+        { "chat_template",               tmpl_default },
+        { "chat_template_caps",          meta.chat_template_caps },
+        { "bos_token",                   meta.bos_token_str },
+        { "eos_token",                   meta.eos_token_str },
+        { "build_info",                  meta.build_info },
+        { "is_sleeping",                 is_sleeping },
+        { "cors_proxy_enabled",          params.ui_mcp_proxy },
+    };
+    if (params.use_jinja) {
+        if (!tmpl_tools.empty()) {
+            props["chat_template_tool_use"] = tmpl_tools;
+        }
+    }
+
+    return props;
+}
+
+json server_routes::get_model_info() const {
+    return get_res_model_info(*meta);
 }
 
 void server_routes::init_routes() {
@@ -4451,41 +4581,64 @@ void server_routes::init_routes() {
     };
 
     this->get_metrics = [this](const server_http_req & req) {
-        auto res = create_response();
+        auto res = create_response(true);
         if (!params.endpoint_metrics) {
             res->error(format_error_response("This server does not support metrics endpoint. Start it with `--metrics`", ERROR_TYPE_NOT_SUPPORTED));
             return res;
         }
 
-        // request slots data using task queue
-        {
-            server_task task(SERVER_TASK_TYPE_METRICS);
-            task.id = res->rd.get_new_id();
+        // render response using cached_metrics
+        auto use_cached_metrics = [&]() {
+            std::unique_lock<std::mutex> lock(mutex_cache);
+            res->headers["Process-Start-Time-Unix"] = std::to_string(cached_metrics.t_start);
+            server_task_result_metrics tmp;
+            tmp.metrics = cached_metrics;
+            res->content_type = "text/plain; version=0.0.4";
+            res->status = 200;
+            res->data = tmp.to_metrics();
             // the gauges are averaged over the window between two scrapes
-            task.metrics_reset_bucket = true;
-            res->rd.post_task(std::move(task), true); // high-priority task
+            cached_metrics.reset_bucket();
+            should_reset_buckets = true;
+        };
+
+        if (queue_tasks.is_sleeping()) {
+            use_cached_metrics();
+
+        } else {
+            // request slots data using task queue
+            {
+                server_task task(SERVER_TASK_TYPE_METRICS);
+                task.id = res->rd.get_new_id();
+                // the gauges are averaged over the window between two scrapes
+                task.metrics_reset_bucket = true;
+                res->rd.post_task(std::move(task), true); // high-priority task
+            }
+
+            // a task posted right before sleeping is never processed, do not wait for it
+            auto result = res->rd.next([&]{
+                return req.should_stop() || queue_tasks.is_sleeping();
+            });
+            if (!result) {
+                if (!req.should_stop()) {
+                    use_cached_metrics();
+                }
+                return res;
+            }
+
+            if (result->is_error()) {
+                res->error(result->to_json());
+                return res;
+            }
+
+            auto res_task = dynamic_cast<server_task_result_metrics*>(result.get());
+            GGML_ASSERT(res_task != nullptr);
+
+            res->headers["Process-Start-Time-Unix"] = std::to_string(res_task->metrics.t_start);
+            res->content_type = "text/plain; version=0.0.4";
+            res->status = 200;
+            res->data = res_task->to_metrics();
         }
 
-        // get the result
-        auto result = res->rd.next(req.should_stop);
-        if (!result) {
-            // connection was closed
-            GGML_ASSERT(req.should_stop());
-            return res;
-        }
-
-        if (result->is_error()) {
-            res->error(result->to_json());
-            return res;
-        }
-
-        auto res_task = dynamic_cast<server_task_result_metrics*>(result.get());
-        GGML_ASSERT(res_task != nullptr);
-
-        res->headers["Process-Start-Time-Unix"] = std::to_string(res_task->metrics.t_start);
-        res->content_type = "text/plain; version=0.0.4";
-        res->status = 200;
-        res->data = res_task->to_metrics();
         return res;
     };
 
@@ -4498,7 +4651,7 @@ void server_routes::init_routes() {
 
         // request slots data using task queue
         {
-            server_task task(SERVER_TASK_TYPE_METRICS);
+            server_task task(SERVER_TASK_TYPE_SLOT_GET);
             task.id = res->rd.get_new_id();
             res->rd.post_task(std::move(task), true); // high-priority task
         }
@@ -4516,7 +4669,7 @@ void server_routes::init_routes() {
             return res;
         }
 
-        auto * res_task = dynamic_cast<server_task_result_metrics*>(result.get());
+        auto * res_task = dynamic_cast<server_task_result_slots*>(result.get());
         GGML_ASSERT(res_task != nullptr);
 
         // optionally return "fail_on_no_slot" error
@@ -4566,53 +4719,13 @@ void server_routes::init_routes() {
 
     this->get_props = [this](const server_http_req &) {
         auto res = create_response(true);
-
-        // this endpoint can be accessed during sleeping
-        // the next LOC is to avoid someone accidentally use ctx_server
-        bool ctx_server; // do NOT delete this line
-        GGML_UNUSED(ctx_server);
-
-        task_params tparams;
-        tparams.sampling = params.sampling;
-        json default_generation_settings_for_props = json {
-            { "params", tparams.to_json(true) },
-            { "n_ctx",  meta->slot_n_ctx },
-        };
-
-        std::string tmpl_default = common_chat_templates_source(meta->chat_params.tmpls.get(), "");
-        std::string tmpl_tools   = common_chat_templates_source(meta->chat_params.tmpls.get(), "tool_use");
-
-        json props = {
-            { "default_generation_settings", default_generation_settings_for_props },
-            { "total_slots",                 params.n_parallel },
-            { "model_alias",                 meta->model_name },
-            { "model_ftype",                 meta->model_ftype },
-            { "model_path",                  meta->model_path },
-            { "modalities",                  json {
-                {"vision", meta->has_inp_image},
-                {"video",  meta->has_inp_video},
-                {"audio",  meta->has_inp_audio},
-            } },
-            { "media_marker",                get_media_marker() },
-            { "endpoint_slots",              params.endpoint_slots },
-            { "endpoint_props",              params.endpoint_props },
-            { "endpoint_metrics",            params.endpoint_metrics },
-            { "ui",                          params.ui },
-            { "ui_settings",                 meta->json_ui_settings },
-            { "chat_template",               tmpl_default },
-            { "chat_template_caps",          meta->chat_template_caps },
-            { "bos_token",                   meta->bos_token_str },
-            { "eos_token",                   meta->eos_token_str },
-            { "build_info",                  meta->build_info },
-            { "is_sleeping",                 queue_tasks.is_sleeping() },
-            { "cors_proxy_enabled",          params.ui_mcp_proxy },
-        };
-        if (params.use_jinja) {
-            if (!tmpl_tools.empty()) {
-                props["chat_template_tool_use"] = tmpl_tools;
-            }
+        // note: do NOT use ctx_server here, this endpoint must be accessible during sleep
+        if (queue_tasks.is_sleeping()) {
+            std::unique_lock<std::mutex> lock(mutex_cache);
+            res->ok(cached_props);
+        } else {
+            res->ok(get_res_props(*meta, params, false));
         }
-        res->ok(props);
         return res;
     };
 
@@ -4874,42 +4987,13 @@ void server_routes::init_routes() {
 
     this->get_models = [this](const server_http_req &) {
         auto res = create_response(true);
-
-        // this endpoint can be accessed during sleeping
-        // the next LOC is to avoid someone accidentally use ctx_server
-        bool ctx_server; // do NOT delete this line
-        GGML_UNUSED(ctx_server);
-
-        json models = {
-            {"models", {
-                {
-                    {"name",  meta->model_name},
-                    {"model", meta->model_name},
-                    {"modified_at", ""},
-                    {"size", ""},
-                    {"digest", ""}, // dummy value, llama.cpp does not support managing model file's hash
-                    {"type", "model"},
-                    {"description", ""},
-                    {"tags", {""}},
-                    {"capabilities", meta->has_mtmd ? json({"completion","multimodal"}) : json({"completion"})},
-                    {"parameters", ""},
-                    {"details", {
-                        {"parent_model", ""},
-                        {"format", "gguf"},
-                        {"family", ""},
-                        {"families", {""}},
-                        {"parameter_size", ""},
-                        {"quantization_level", ""}
-                    }}
-                }
-            }},
-            {"object", "list"},
-            {"data", {
-                get_model_info(),
-            }}
-        };
-
-        res->ok(models);
+        // note: do NOT use ctx_server here, this endpoint must be accessible during sleep
+        if (queue_tasks.is_sleeping()) {
+            std::unique_lock<std::mutex> lock(mutex_cache);
+            res->ok(cached_models);
+        } else {
+            res->ok(get_res_models(*meta));
+        }
         return res;
     };
 
@@ -5119,27 +5203,6 @@ void server_routes::init_routes() {
     };
 }
 
-json server_routes::get_model_info() const {
-    return json {
-        {"id",       meta->model_name},
-        {"aliases",  meta->model_aliases},
-        {"tags",     meta->model_tags},
-        {"object",   "model"},
-        {"created",  std::time(0)},
-        {"owned_by", "llamacpp"},
-        {"meta",     {
-            {"vocab_type",  meta->model_vocab_type},
-            {"n_vocab",     meta->model_vocab_n_tokens},
-            {"n_ctx",       meta->slot_n_ctx},
-            {"n_ctx_train", meta->model_n_ctx_train},
-            {"n_embd",      meta->model_n_embd_inp},
-            {"n_params",    meta->model_n_params},
-            {"size",        meta->model_size},
-            {"ftype",       meta->model_ftype},
-        }},
-    };
-}
-
 std::unique_ptr<server_res_generator> server_routes::handle_slots_save(const server_http_req & req, int id_slot) {
     auto res = create_response();
     const json request_data = json::parse(req.body);
@@ -5388,3 +5451,24 @@ std::unique_ptr<server_res_generator> server_routes::handle_count_tokens(const l
     res->ok(response);
     return res;
 }
+
+void server_routes::update_cached_responses(bool is_sleeping) {
+    // caller is task_queue, so ctx_server can be accessed without holding locks
+    std::unique_lock<std::mutex> lock(mutex_cache);
+
+    if (is_sleeping) {
+        cached_models  = get_res_models(*meta);
+        cached_props   = get_res_props(*meta, params, true);
+        cached_metrics = ctx_server.get_metrics();
+
+        should_reset_buckets = false;
+
+        SRV_DBG("%s\n", "cached responses updated");
+
+    } else if (should_reset_buckets) {
+        // a scrape during sleep already reported these buckets
+        ctx_server.reset_metrics_bucket();
+
+        should_reset_buckets = false;
+    }
+}

tools/server/server-context.h

@@ -8,6 +8,7 @@
 
 #include <cstddef>
 #include <memory>
+#include <mutex>
 #include <set>
 
 struct server_context_impl; // private implementation
@@ -174,9 +175,19 @@ private:
     std::unique_ptr<const server_context_meta> meta;
 
     const common_params & params;
-    const server_context_impl & ctx_server;
+    server_context_impl & ctx_server;
 
     server_queue & queue_tasks;
     server_response & queue_results;
     std::unique_ptr<server_res_generator> create_response(bool bypass_sleep = false);
+
+    // cached responses, to be used during sleep
+    std::mutex     mutex_cache;
+    json           cached_models  = nullptr;
+    json           cached_props   = nullptr;
+    server_metrics cached_metrics;
+    // set when a scrape during sleep already reported the throughput buckets
+    bool           should_reset_buckets = false;
+    // call right before sleep to update the cached responses
+    void update_cached_responses(bool is_sleeping);
 };

tools/server/server-http.cpp

@@ -198,8 +198,6 @@ bool server_http_context::init(const common_params & params) {
         std::unordered_set<std::string> endpoints {
             "/health",
             "/v1/health",
-            "/models",
-            "/v1/models",
         };
         endpoints.insert(frontend_paths.begin(), frontend_paths.end());
         return endpoints;

tools/server/server-models.cpp

@@ -555,6 +555,40 @@ void server_models::load_models() {
         return source_map.count(name) ? source_map.at(name) : SERVER_MODEL_SOURCE_PRESET;
     };
 
+    // hide cache models whose resolved file is already used by a preset with dedup-cache-models enabled
+    std::set<std::string> hidden_models;
+    {
+        std::set<std::string> preset_paths;
+        for (const auto & [name, preset] : custom_presets) {
+            std::string val;
+            if (!preset.get_option(COMMON_ARG_PRESET_DEDUP_CACHE_MODELS, val) || !common_arg_utils::is_truthy(val)) {
+                continue;
+            }
+            std::string hf_repo;
+            if (!preset.get_option("LLAMA_ARG_HF_REPO", hf_repo) || hf_repo.empty()) {
+                continue;
+            }
+            std::string hf_file;
+            preset.get_option("LLAMA_ARG_HF_FILE", hf_file);
+            std::string path = common_download_resolve_path(hf_repo, hf_file);
+            if (!path.empty()) {
+                preset_paths.insert(path);
+            }
+        }
+        if (!preset_paths.empty()) {
+            for (const auto & [name, preset] : cached_models) {
+                if (get_source(name) != SERVER_MODEL_SOURCE_CACHE) {
+                    continue; // merged with another source, not a pure cache entry
+                }
+                std::string path = common_download_resolve_path(name);
+                if (!path.empty() && preset_paths.count(path)) {
+                    SRV_INF("hiding cache model name=%s (deduplicated by a preset)\n", name.c_str());
+                    hidden_models.insert(name);
+                }
+            }
+        }
+    }
+
     // Helpers that read `mapping` - must be called while holding the lock.
     std::unordered_set<std::string> custom_names;
     for (const auto & [name, preset] : custom_presets) custom_names.insert(name);
@@ -590,6 +624,11 @@ void server_models::load_models() {
             }
         }
     };
+    auto apply_hidden = [&]() {
+        for (auto & [name, inst] : mapping) {
+            inst.meta.hidden = hidden_models.count(name) > 0;
+        }
+    };
     // update_args() injects HOST/PORT/ALIAS, so strip them before comparing presets
     auto preset_options_for_compare = [](common_preset p) {
         p.unset_option("LLAMA_ARG_HOST");
@@ -630,6 +669,7 @@ void server_models::load_models() {
             add_model(std::move(meta));
         }
         apply_stop_timeout();
+        apply_hidden();
         log_available_models();
 
         std::vector<std::string> models_to_load;
@@ -806,6 +846,7 @@ void server_models::load_models() {
         }
 
         apply_stop_timeout();
+        apply_hidden();
 
         // clear reload flag before unlocking for autoload - load() blocks on !is_reloading,
         // so clearing it here (while still locked) prevents a deadlock in the autoload calls below
@@ -1929,6 +1970,9 @@ void server_models_routes::init_routes() {
         auto all_models = models.get_all_meta();
         std::time_t t = std::time(0);
         for (const auto & meta : all_models) {
+            if (meta.hidden) {
+                continue; // cache model deduplicated by a preset
+            }
             json status {
                 {"value",  server_model_status_to_string(meta.status)},
                 {"args",   meta.args},

tools/server/server-models.h

@@ -84,6 +84,7 @@ struct server_model_meta {
     int exit_code = 0; // exit code of the model instance process (only valid if status == FAILED)
     int stop_timeout = 0; // seconds to wait before force-killing the model instance during shutdown
     mtmd_caps multimodal; // multimodal capabilities
+    bool hidden = false; // hidden from GET /models, but still accept if requested
 
     bool is_ready() const {
         return status == SERVER_MODEL_STATUS_LOADED;

tools/server/server-queue.cpp

@@ -3,6 +3,7 @@
 
 #include "log.h"
 
+#include <algorithm>
 #include <chrono>
 #include <thread>
 
@@ -20,6 +21,10 @@
 // server_queue
 //
 
+static bool task_resets_idle_timer(server_task_type type) {
+    return type != SERVER_TASK_TYPE_METRICS;
+}
+
 int server_queue::post(server_task && task, bool front) {
     std::unique_lock<std::mutex> lock(mutex_tasks);
     GGML_ASSERT(task.id != -1);
@@ -27,20 +32,24 @@ int server_queue::post(server_task && task, bool front) {
     if (task.type == SERVER_TASK_TYPE_CANCEL) {
         cleanup_pending_task(task.id_target);
     }
-    const int task_id = task.id;
+    const int  task_id     = task.id;
+    const bool reset_timer = task_resets_idle_timer(task.type);
     QUE_DBG("new task, id = %d, front = %d\n", task_id, front);
     if (front) {
         queue_tasks.push_front(std::move(task));
     } else {
         queue_tasks.push_back(std::move(task));
     }
-    time_last_task = ggml_time_ms();
+    if (reset_timer) {
+        time_last_task = ggml_time_ms();
+    }
     condition_tasks.notify_one();
     return task_id;
 }
 
 int server_queue::post(std::vector<server_task> && tasks, bool front) {
     std::unique_lock<std::mutex> lock(mutex_tasks);
+    bool reset_timer = false;
     for (auto & task : tasks) {
         if (task.id == -1) {
             task.id = id++;
@@ -49,6 +58,7 @@ int server_queue::post(std::vector<server_task> && tasks, bool front) {
         if (task.type == SERVER_TASK_TYPE_CANCEL) {
             cleanup_pending_task(task.id_target);
         }
+        reset_timer |= task_resets_idle_timer(task.type);
         QUE_DBG("new task, id = %d/%d, front = %d\n", task.id, (int) tasks.size(), front);
         if (front) {
             queue_tasks.push_front(std::move(task));
@@ -56,7 +66,9 @@ int server_queue::post(std::vector<server_task> && tasks, bool front) {
             queue_tasks.push_back(std::move(task));
         }
     }
-    time_last_task = ggml_time_ms();
+    if (reset_timer) {
+        time_last_task = ggml_time_ms();
+    }
     condition_tasks.notify_one();
     return 0;
 }
@@ -294,11 +306,14 @@ void server_queue::start_loop(int64_t idle_sleep_ms) {
         QUE_DBG("%s", "update slots\n");
 
         // this will run the main inference process for all slots
+        const int64_t t_update_slots = ggml_time_ms();
         callback_update_slots();
         {
             // update_slots() may take a while to finish, we need to make sure it's not counted as idle
+            // shift instead of reset, so that non-task_resets_idle_timer tasks do not delay the sleep
             std::unique_lock<std::mutex> lock(mutex_tasks);
-            time_last_task = ggml_time_ms();
+            const int64_t now = ggml_time_ms();
+            time_last_task = std::min(now, time_last_task + (now - t_update_slots));
         }
 
         QUE_DBG("%s", "waiting for new tasks\n");
@@ -312,7 +327,10 @@ void server_queue::start_loop(int64_t idle_sleep_ms) {
             if (should_sleep()) {
                 QUE_INF("%s", "entering sleeping state\n");
                 sleeping = true;
-                callback_sleeping_state(true);
+                // Call order cb0 -> cb1 -> cb{N}
+                for (auto & cb : callback_sleeping_state) {
+                    cb(true);
+                }
                 req_stop_sleeping = false;
                 // wait until we are requested to exit sleeping state
                 condition_tasks.wait(lock, [&]{
@@ -323,7 +341,10 @@ void server_queue::start_loop(int64_t idle_sleep_ms) {
                 }
                 QUE_INF("%s", "exiting sleeping state\n");
                 req_stop_sleeping = false;
-                callback_sleeping_state(false);
+                // Call order cb{N} -> cb1 -> cb0
+                for (size_t i = callback_sleeping_state.size(); i > 0; i--) {
+                    callback_sleeping_state[i - 1](false);
+                }
                 sleeping = false;
                 time_last_task = ggml_time_ms();
                 condition_tasks.notify_all(); // notify wait_until_no_sleep()

tools/server/server-queue.h

@@ -44,7 +44,7 @@ private:
     // callback functions
     std::function<bool(server_task &&, bool)> callback_new_task;
     std::function<void(void)>                 callback_update_slots;
-    std::function<void(bool)>                 callback_sleeping_state;
+    std::vector<std::function<void(bool)>>    callback_sleeping_state;
 
 public:
     ~server_queue() { worker_stop(); }
@@ -86,6 +86,7 @@ public:
      *
      * Sleeping procedure (disabled if idle_sleep_ms < 0):
      * - If there is no task after idle_sleep_ms, enter sleeping state
+     *   note: metrics tasks are processed as usual, but do not reset the idle timer
      * - Call callback_sleeping_state(true)
      * - Wait until req_stop_sleeping is set to true
      * - Call callback_sleeping_state(false)
@@ -127,18 +128,12 @@ public:
     }
 
     // Register callback for sleeping state change; multiple callbacks are allowed
-    // note: when entering sleeping state, the callback is called AFTER sleeping is set to true
-    //       when leaving sleeping state, the callback is called BEFORE sleeping is set to false
+    // for example: register order cb0, cb1, cb2
+    // entering sleep: queue.sleeping = true --> cb0(true) --> cb1(true) --> cb2(true)
+    // leaving sleep: cb2(false) --> cb1(false) --> cb0(false) --> queue.sleeping = false
+    // note: caller will hold mutex_tasks while calling the callbacks
     void on_sleeping_state(std::function<void(bool)> callback) {
-        if (callback_sleeping_state) {
-            auto prev_callback = std::move(callback_sleeping_state);
-            callback_sleeping_state = [prev_callback, callback](bool sleeping) {
-                prev_callback(sleeping);
-                callback(sleeping);
-            };
-        } else {
-            callback_sleeping_state = std::move(callback);
-        }
+        callback_sleeping_state.push_back(std::move(callback));
     }
 
 private:

tools/server/server-task.cpp

@@ -1512,10 +1512,15 @@ json server_task_result_error::to_json() {
 //
 // server_task_result_metrics
 //
-json server_task_result_metrics::to_json() {
+json server_task_result_slots::to_json() {
     return slots_data;
 }
 
+json server_task_result_metrics::to_json() {
+    // not used, /metrics renders prometheus text via to_metrics()
+    return json{};
+}
+
 // metrics definition: https://prometheus.io/docs/practices/naming/#metric-names
 std::string server_task_result_metrics::to_metrics() {
     const std::vector<metric_item> counters = {

tools/server/server-task.h

@@ -22,6 +22,7 @@ enum server_task_type {
     SERVER_TASK_TYPE_CONTROL,
     SERVER_TASK_TYPE_NEXT_RESPONSE,
     SERVER_TASK_TYPE_METRICS,
+    SERVER_TASK_TYPE_SLOT_GET,
     SERVER_TASK_TYPE_SLOT_SAVE,
     SERVER_TASK_TYPE_SLOT_RESTORE,
     SERVER_TASK_TYPE_SLOT_ERASE,
@@ -489,22 +490,16 @@ struct server_task_result_error : server_task_result {
     virtual json to_json() override;
 };
 
+// used by /metrics API
 struct server_task_result_metrics : server_task_result {
     // these are immediate stats, not accumulated (server_metrics is cumulative)
-    int n_idle_slots;
-    int n_processing_slots;
-    int n_tasks_deferred;
+    int n_processing_slots = 0;
+    int n_tasks_deferred = 0;
 
     server_metrics metrics;
 
-    // while we can also use std::vector<server_slot> this requires copying the slot object which can be quite messy
-    // therefore, we use json to temporarily store the slot.to_json() result
-    json slots_data = json::array();
-
-    // used by /slots API
     virtual json to_json() override;
 
-    // used by /metrics API
     struct metric_item {
         std::string name;
         std::string description;
@@ -513,6 +508,17 @@ struct server_task_result_metrics : server_task_result {
     std::string to_metrics();
 };
 
+// used by /slots API
+struct server_task_result_slots : server_task_result {
+    int n_idle_slots = 0;
+
+    // while we can also use std::vector<server_slot> this requires copying the slot object which can be quite messy
+    // therefore, we use json to temporarily store the slot.to_json() result
+    json slots_data = json::array();
+
+    virtual json to_json() override;
+};
+
 struct server_task_result_slot_save_load : server_task_result {
     std::string filename;
     bool is_save; // true = save, false = load

tools/server/server.cpp

@@ -235,8 +235,8 @@ int llama_server(common_params & params, int argc, char ** argv) {
     ctx_http.get ("/metrics",                  ex_wrapper(routes.get_metrics));
     ctx_http.get ("/props",                    ex_wrapper(routes.get_props));
     ctx_http.post("/props",                    ex_wrapper(routes.post_props));
-    ctx_http.get ("/models",                   ex_wrapper(routes.get_models)); // public endpoint (no API key check)
-    ctx_http.get ("/v1/models",                ex_wrapper(routes.get_models)); // public endpoint (no API key check)
+    ctx_http.get ("/models",                   ex_wrapper(routes.get_models));
+    ctx_http.get ("/v1/models",                ex_wrapper(routes.get_models));
     ctx_http.post("/completion",               ex_wrapper(routes.post_completions)); // legacy
     ctx_http.post("/completions",              ex_wrapper(routes.post_completions));
     ctx_http.post("/v1/completions",           ex_wrapper(routes.post_completions_oai));

tools/server/tests/unit/test_router.py

@@ -63,14 +63,16 @@ def test_router_chat_completion_stream(model: str, success: bool):
         assert content == ""
 
 
-def _get_model_ids(is_reload: bool) -> set[str]:
-    res = server.make_request("GET", "/models" + ("?reload=1" if is_reload else ""))
+def _get_model_ids(is_reload: bool, headers: dict | None = None) -> set[str]:
+    res = server.make_request(
+        "GET", "/models" + ("?reload=1" if is_reload else ""), headers=headers
+    )
     assert res.status_code == 200
     return {item["id"] for item in res.body.get("data", [])}
 
 
-def _get_model_status(model_id: str) -> str:
-    res = server.make_request("GET", "/models")
+def _get_model_status(model_id: str, headers: dict | None = None) -> str:
+    res = server.make_request("GET", "/models", headers=headers)
     assert res.status_code == 200
     for item in res.body.get("data", []):
         if item.get("id") == model_id or item.get("model") == model_id:
@@ -78,11 +80,11 @@ def _get_model_status(model_id: str) -> str:
     raise AssertionError(f"Model {model_id} not found in /models response")
 
 
-def _wait_for_model_status(model_id: str, desired: set[str], timeout: int = 60) -> str:
+def _wait_for_model_status(model_id: str, desired: set[str], timeout: int = 60, headers: dict | None = None) -> str:
     deadline = time.time() + timeout
     last_status = None
     while time.time() < deadline:
-        last_status = _get_model_status(model_id)
+        last_status = _get_model_status(model_id, headers=headers)
         if last_status in desired:
             return last_status
         time.sleep(0.01)
@@ -100,7 +102,7 @@ def _load_model_and_wait(
     assert load_res.status_code == 200
     assert isinstance(load_res.body, dict)
     assert load_res.body.get("success") is True
-    _wait_for_model_status(model_id, {"loaded"}, timeout=timeout)
+    _wait_for_model_status(model_id, {"loaded"}, timeout=timeout, headers=headers)
 
 
 def test_router_unload_model():
@@ -406,6 +408,59 @@ def test_router_reload_models():
         os.remove(preset_path)
 
 
+def test_router_dedup_cache_models():
+    """dedup-cache-models hides the cache entry backing a preset from GET /models"""
+    global server
+
+    preset_path = os.path.join(TMP_DIR, "test_dedup.ini")
+    cache_id = "ggml-org/test-model-stories260K:F32"
+
+    with open(preset_path, "w") as f:
+        f.write(
+            "[model-dedup]\n"
+            "hf-repo = ggml-org/test-model-stories260K\n"
+            "dedup-cache-models = 1\n"
+        )
+
+    server.models_preset = preset_path
+    server.start()
+
+    try:
+        ids = _get_model_ids(is_reload=False)
+        assert "model-dedup" in ids
+        assert cache_id not in ids, "cache model should be hidden by dedup"
+        # other cache models are unaffected
+        assert "ggml-org/tinygemma3-GGUF:Q8_0" in ids
+
+        # the hidden model is only hidden from the listing, it can still be used
+        res = server.make_request("POST", "/tokenize", data={"model": cache_id, "content": "hello"})
+        assert res.status_code == 200
+
+        # disabling the flag brings the cache entry back on reload
+        with open(preset_path, "w") as f:
+            f.write(
+                "[model-dedup]\n"
+                "hf-repo = ggml-org/test-model-stories260K\n"
+            )
+        ids = _get_model_ids(is_reload=True)
+        assert cache_id in ids
+
+        # the flag also works from the global section
+        with open(preset_path, "w") as f:
+            f.write(
+                "[*]\n"
+                "dedup-cache-models = 1\n"
+                "\n"
+                "[model-dedup]\n"
+                "hf-repo = ggml-org/test-model-stories260K\n"
+            )
+        ids = _get_model_ids(is_reload=True)
+        assert "model-dedup" in ids
+        assert cache_id not in ids, "cache model should be hidden by global dedup"
+    finally:
+        os.remove(preset_path)
+
+
 def test_router_remote_preset():
     global server
     server.model_hf_repo = "ggml-org/test-preset-ci"

tools/server/tests/unit/test_security.py

@@ -15,7 +15,7 @@ def create_server():
     server.api_key = TEST_API_KEY
 
 
-@pytest.mark.parametrize("endpoint", ["/health", "/models"])
+@pytest.mark.parametrize("endpoint", ["/health"])
 def test_access_public_endpoint(endpoint: str):
     global server
     server.start()

tools/server/tests/unit/test_sleep.py

@@ -11,6 +11,35 @@ def create_server():
     server = ServerPreset.tinyllama2()
 
 
+def is_sleeping(server: ServerProcess) -> bool:
+    res = server.make_request("GET", "/props")
+    assert res.status_code == 200
+    return res.body["is_sleeping"]
+
+
+def wait_for_sleep(server: ServerProcess, timeout: float = 10.0):
+    start = time.time()
+    while time.time() - start < timeout:
+        if is_sleeping(server):
+            return
+        time.sleep(0.1)
+    raise TimeoutError("server did not go to sleep")
+
+
+def fetch_metrics(server: ServerProcess) -> str:
+    res = server.make_request("GET", "/metrics")
+    assert res.status_code == 200
+    assert isinstance(res.body, str)
+    return res.body
+
+
+def get_metric(text: str, name: str) -> float:
+    prefix = f"llamacpp:{name} "
+    values = [ln for ln in text.splitlines() if ln.startswith(prefix)]
+    assert len(values) == 1, f"{name} not found in metrics"
+    return float(values[0][len(prefix):])
+
+
 def test_server_sleep():
     global server
     server.sleep_idle_seconds = 1
@@ -25,6 +54,10 @@ def test_server_sleep():
     res = server.make_request("GET", "/props")
     assert res.status_code == 200
     assert res.body["is_sleeping"] == True
+    res = server.make_request("GET", "/models")
+    assert res.status_code == 200
+    assert len(res.body["data"]) == 1
+    assert res.body["data"][0]["id"] == server.model_alias
 
     # make a generation request to wake up the server
     res = server.make_request("POST", "/completion", data={
@@ -37,3 +70,58 @@ def test_server_sleep():
     res = server.make_request("GET", "/props")
     assert res.status_code == 200
     assert res.body["is_sleeping"] == False
+
+
+def test_server_sleep_read_only_endpoints():
+    global server
+    server.sleep_idle_seconds = 1
+    server.server_metrics = True
+    server.start()
+
+    res = server.make_request("POST", "/completion", data={
+        "n_predict": 4,
+        "prompt": "Hello",
+    })
+    assert res.status_code == 200
+
+    # the first scrape resets the throughput buckets, so that the second one reports
+    # the same zero rates as the snapshot taken on entering sleep
+    fetch_metrics(server)
+    metrics_awake = fetch_metrics(server)
+    assert get_metric(metrics_awake, "tokens_predicted_total") > 0
+
+    wait_for_sleep(server)
+
+    # during sleep, metrics are served from the snapshot taken right before sleeping
+    assert fetch_metrics(server) == metrics_awake
+
+    # scraping /metrics must not wake the server up
+    assert is_sleeping(server)
+
+
+def test_server_sleep_metrics_buckets():
+    global server
+    server.sleep_idle_seconds = 1
+    server.server_metrics = True
+    server.start()
+
+    res = server.make_request("POST", "/completion", data={
+        "n_predict": 8,
+        "prompt": "Hello",
+    })
+    assert res.status_code == 200
+
+    wait_for_sleep(server)
+
+    # the first scrape reports the throughput of the last generation
+    assert get_metric(fetch_metrics(server), "predicted_tokens_seconds") > 0
+
+    # nothing runs while sleeping, so the next scrapes report an empty window
+    assert get_metric(fetch_metrics(server), "predicted_tokens_seconds") == 0
+    assert is_sleeping(server)
+
+    # waking up must not report the buckets again
+    res = server.make_request("POST", "/tokenize", data={"content": "Hello"})
+    assert res.status_code == 200
+    assert is_sleeping(server) == False
+    assert get_metric(fetch_metrics(server), "predicted_tokens_seconds") == 0

tools/ui/src/lib/stores/init.ts

@@ -1,10 +1,15 @@
 /**
- * Explicit store initialization, called once from the root layout.
+ * Explicit store initialization, run once and shared by every caller.
  *
  * Order matters: migrations run first because they rename and rewrite
  * localStorage keys, so every store that reads localStorage initializes
  * only after they complete. Constructors and module-level side effects
  * stay empty so import order can no longer change startup behavior.
+ *
+ * The returned promise resolves once the persisted state is in memory, which
+ * route loads await before reading settings: they run ahead of the root layout
+ * script. The conversation list loads in the background, awaited by the chat
+ * page that renders it.
  */
 
 // direct imports, not via the barrel, to avoid circular deps
@@ -16,19 +21,20 @@ import { versionStore } from './version.svelte';
 import { browser } from '$app/environment';
 import { MigrationService } from '$lib/services/migration.service';
 
-let started = false;
+let startup: Promise<void> | null = null;
 
-export async function initStores(): Promise<void> {
-	if (!browser || started) return;
+export function initStores(): Promise<void> {
+	if (!browser) return Promise.resolve();
 
-	started = true;
+	startup ??= (async () => {
+		await MigrationService.runAllMigrations();
 
-	await MigrationService.runAllMigrations();
+		settingsStore.initialize();
+		permissionsStore.initialize();
+		toolsStore.initialize();
+		void versionStore.initialize();
+		void conversationsStore.init();
+	})();
 
-	settingsStore.initialize();
-	permissionsStore.initialize();
-	toolsStore.initialize();
-	void versionStore.initialize();
-
-	await conversationsStore.init();
+	return startup;
 }

tools/ui/src/lib/stores/settings.svelte.ts

@@ -358,17 +358,24 @@ class SettingsStore {
 		// UI settings are the admin's defaults for new users: applied once on
 		// the first visit, never on later loads, so the user's config can
 		// diverge. "Reset to Default" is the explicit way back to the baseline.
+		// A first visit config carries factory values only, so a key that
+		// already diverges here was set by the user before the baseline could
+		// be reached, through the API key splash, and stays theirs.
 		if (uiSettings && this.isFirstVisit) {
 			this.isFirstVisit = false;
 
 			for (const [key, value] of Object.entries(uiSettings)) {
-				if (!this.userOverrides.has(key) && value !== undefined) {
-					setConfigValue(this.config, key, value);
+				if (value === undefined || this.userOverrides.has(key)) continue;
 
-					// theme lives in mode-watcher, not just in config -> propagate
-					if (key === SETTINGS_KEYS.THEME) {
-						setMode(value as ColorMode);
-					}
+				if (getConfigValue(this.config, key) !== getConfigValue(SETTING_CONFIG_DEFAULT, key)) {
+					continue;
+				}
+
+				setConfigValue(this.config, key, value);
+
+				// theme lives in mode-watcher, not just in config -> propagate
+				if (key === SETTINGS_KEYS.THEME) {
+					setMode(value as ColorMode);
 				}
 			}
 		}

tools/ui/src/routes/(chat)/+page.ts

@@ -1,6 +1,10 @@
 import type { PageLoad } from './$types';
+import { initStores } from '$lib/stores/init';
 import { validateApiKey } from '$lib/utils';
 
 export const load: PageLoad = async ({ fetch }) => {
+	// loads run before the root layout script, so the stored API key reaches
+	// the probe only once the settings store has read localStorage
+	await initStores();
 	await validateApiKey(fetch);
 };

tools/ui/src/routes/(chat)/chat/[id]/+page.ts

@@ -1,6 +1,10 @@
 import type { PageLoad } from './$types';
+import { initStores } from '$lib/stores/init';
 import { validateApiKey } from '$lib/utils';
 
 export const load: PageLoad = async ({ fetch }) => {
+	// loads run before the root layout script, so the stored API key reaches
+	// the probe only once the settings store has read localStorage
+	await initStores();
 	await validateApiKey(fetch);
 };

tools/ui/tests/client/ui-settings-sync.svelte.test.ts

@@ -47,6 +47,19 @@ describe('server ui_settings application semantics', () => {
 		expect(stored.apiKey).toBe('sk-user-key');
 	});
 
+	it('keeps a value the user sets before the baseline is reachable', () => {
+		settingsStore.initialize();
+		// the splash is the only way in when the server runs with --api-key,
+		// so the first user write lands before the first successful /props
+		settingsStore.updateConfig('apiKey', 'sk-user-key');
+		mockProps({ apiKey: 'admin-placeholder', theme: 'dark' });
+
+		settingsStore.syncWithServerDefaults();
+
+		expect(settingsStore.config.apiKey).toBe('sk-user-key');
+		expect(settingsStore.config.theme).toBe('dark');
+	});
+
 	it('Reset to Default reapplies the full baseline, api key included', () => {
 		settingsStore.initialize();
 		settingsStore.updateConfig('theme', 'light');