vulkan: tiled transpose for 0<->2 permuted CONT (#26585)

* vulkan: tiled transpose for 0<->2 permuted CONT

-ggml_vk_get_cpy_pipeline only routed to the tiled shared-memory transpose
shader when dim1 was the innermost dimension, i.e. ggml_transpose (a 0<->1
swap). A 0<->2 swap -- ggml_cont(ggml_permute(x, 2, 1, 0, 3)) -- fell back to
the generic per-element strided copy, whose source reads stride by ne0*ne1
elements: one cache line per lane.

-DeepSeek-V4's lightning indexer performs exactly that permute on a
[n_kv, n_tokens, n_head] tensor. On Vulkan/RADV gfx1151 it ran at ~1-9 GB/s of
a ~200 GB/s part and accounted for 43% of total prefill time.

-Add copy_transpose_02.comp, mirroring copy_transpose.comp but tiling over dst
dims (0, 2) with dims 1 and 3 as the batch, so reads walk src dim2 and writes
walk dst dim0 -- both contiguous. The selection condition additionally requires
a non-contiguous source and a contiguous destination so it cannot take cases
the contiguous-copy shader already handles.

-test-backend-ops only exercised ggml_transpose for CONT, so the strided path
was untested. Add test_cont_permute covering (2,1,0,3), (1,2,0,3) and (0,2,1,3)
over f32/f16 at tile-aligned, tile-unaligned and large shapes. The large shapes
are in the eval set rather than only in perf because perf mode does not verify
results.

-Measured on gfx1151, ne=[n_kv,64,64,1], perm=(2,1,0,3), f32:

  n_kv=1024:   9.08 ->  579.85 GB/s
  n_kv=1280:  20.03 ->  153.71 GB/s
  n_kv=2048:   7.11 ->   91.68 GB/s
  n_kv=2304:  16.24 ->   86.49 GB/s

-The ~2.2x penalty previously seen at power-of-two n_kv (destination-stride
aliasing) is gone. End to end, DeepSeek-V4-Flash IQ3_XXS prefill on a 9k-token
prompt goes from 56.33 t/s to 103.74 t/s (+84%).

-Note: at n_tokens=512 a single slow-path dispatch takes ~273 ms and looping it
in perf mode can trip the GPU watchdog, so the perf cases use n_tokens=64.

* tests: fold test_cont_permute into test_cont, add L2-exceeding perf shapes

Review feedback: test_cont gains a permute parameter ({0,0,0,0} = none),
matching test_mul_mat's pattern, and the separate struct is gone. Perf
adds [n_kv, 512, 64, 1] variants (~0.5 GB per run) that exceed GPU L2,
since the 64-token shapes fit in cache on large parts and read above
memory bandwidth.

* tests: trim perf-case comment to the two-line summary

* vulkan: trim comments on the 0<->2 transpose path

Drop the shader file header, the read/write block comments and the
rationale prose in the CONT test cases. Keep the tile-shape and
bank-conflict notes and the permute parameter documentation.

---------

Co-authored-by: Kevin Hopper <no-reply@maestro.press>

.github/workflows/docker.yml

@@ -394,6 +394,11 @@ jobs:
     name: Create shared tags from digests
     needs: [prepare_matrices, push_to_registry, create_tag]
     runs-on: ubuntu-24.04
+    permissions:
+      contents: read
+      packages: write
+      id-token: write
+      attestations: write
     strategy:
       fail-fast: false
       matrix:
@@ -428,6 +433,7 @@ jobs:
           password: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Create tags from digests
+        id: create_tags
         shell: bash
         run: |
           set -euo pipefail
@@ -439,6 +445,7 @@ jobs:
           SRC_TAG="${{ needs.create_tag.outputs.source_tag }}"
           BUILD_DATE="${{ steps.build_date.outputs.date }}"
           COMMIT_SHA="${{ steps.checkout.outputs.commit }}"
+          echo "image_repo=${IMAGE_REPO}" >> "$GITHUB_OUTPUT"
           TAGS="${{ matrix.config.tag }}"
           ARCHES="${{ matrix.config.arches }}"
           DIGEST_GLOB="/tmp/digests/*.tsv"
@@ -505,6 +512,16 @@ jobs:
 
               echo "Creating ${merged_versioned_tag} from ${refs[*]}"
               docker buildx imagetools create "${annotations[@]}" --tag "${merged_versioned_tag}" "${refs[@]}"
+
+              if [[ "$tag_name" == "${TAGS%% *}" ]]; then
+                  local digest
+                  digest="$(docker buildx imagetools inspect "${merged_versioned_tag}" --format '{{.Manifest.Digest}}')"
+                  if [[ ! "$digest" =~ ^sha256:[0-9a-f]{64}$ ]]; then
+                      echo "Invalid digest for ${merged_versioned_tag}: ${digest}" >&2
+                      exit 1
+                  fi
+                  echo "${image_type}_digest=${digest}" >> "$GITHUB_OUTPUT"
+              fi
           }
 
           for tag in $TAGS; do
@@ -528,3 +545,24 @@ jobs:
           done
         env:
           GITHUB_REPOSITORY_OWNER: '${{ github.repository_owner }}'
+
+      - name: Attest full image
+        if: ${{ matrix.config.full }}
+        uses: actions/attest@v4
+        with:
+          subject-name: ${{ steps.create_tags.outputs.image_repo }}
+          subject-digest: ${{ steps.create_tags.outputs.full_digest }}
+
+      - name: Attest light image
+        if: ${{ matrix.config.light }}
+        uses: actions/attest@v4
+        with:
+          subject-name: ${{ steps.create_tags.outputs.image_repo }}
+          subject-digest: ${{ steps.create_tags.outputs.light_digest }}
+
+      - name: Attest server image
+        if: ${{ matrix.config.server }}
+        uses: actions/attest@v4
+        with:
+          subject-name: ${{ steps.create_tags.outputs.image_repo }}
+          subject-digest: ${{ steps.create_tags.outputs.server_digest }}

.github/workflows/release.yml

@@ -1569,6 +1569,8 @@ jobs:
     # https://docs.github.com/en/actions/security-for-github-actions/security-guides/automatic-token-authentication#modifying-the-permissions-for-the-github_token
     permissions:
         contents: write # for creating release
+        id-token: write
+        attestations: write
 
     runs-on: ubuntu-slim
 
@@ -1662,6 +1664,11 @@ jobs:
         run: |
           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
+        uses: actions/attest@v4
+        with:
+          subject-path: 'release/*'
+
       - name: Create and push git tag
         run: |
           TAG="${{ steps.tag.outputs.name }}"

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

@@ -962,6 +962,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];
@@ -5525,6 +5526,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 +8934,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;
@@ -12192,7 +12207,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);

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/vulkan-shaders-gen.cpp

@@ -826,6 +826,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;

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++) {

gguf-py/gguf/gguf_reader.py

@@ -32,6 +32,10 @@ from gguf.constants import (
     GGUFEndian,
 )
 
+# limits mirroring ggml/src/gguf.cpp (not part of gguf.h)
+GGUF_MAX_STRING_LENGTH  = 1024 * 1024 * 1024
+GGUF_MAX_ARRAY_ELEMENTS = 1024 * 1024 * 1024
+
 logger = logging.getLogger(__name__)
 
 READER_SUPPORTED_VERSIONS = [2, GGUF_VERSION]
@@ -167,6 +171,10 @@ class GGUFReader:
         offs += self._push_field(ReaderField(offs, 'GGUF.tensor_count', [temp_counts[:1]], [0], [GGUFValueType.UINT64]))
         offs += self._push_field(ReaderField(offs, 'GGUF.kv_count', [temp_counts[1:]], [0], [GGUFValueType.UINT64]))
         tensor_count, kv_count = temp_counts
+        if tensor_count > GGUF_MAX_ARRAY_ELEMENTS:
+            raise ValueError(f'Tensor count {tensor_count} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}')
+        if kv_count > GGUF_MAX_ARRAY_ELEMENTS:
+            raise ValueError(f'KV count {kv_count} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}')
         offs = self._build_fields(offs, kv_count)
 
         # Build Tensor Info Fields
@@ -217,6 +225,10 @@ class GGUFReader:
 
     def _get_str(self, offset: int) -> tuple[npt.NDArray[np.uint64], npt.NDArray[np.uint8]]:
         slen = self._get(offset, np.uint64)
+        if int(slen[0]) > GGUF_MAX_STRING_LENGTH:
+            raise ValueError(f'String length {int(slen[0])} exceeds maximum {GGUF_MAX_STRING_LENGTH}')
+        if offset + 8 + int(slen[0]) > self.data.nbytes:
+            raise ValueError(f'String length {int(slen[0])} exceeds remaining file size {self.data.nbytes - offset - 8}')
         return slen, self._get(offset + 8, np.uint8, slen[0])
 
     def _get_field_parts(
@@ -241,6 +253,8 @@ class GGUFReader:
             raw_itype = self._get(offs, np.uint32)
             offs += int(raw_itype.nbytes)
             alen = self._get(offs, np.uint64)
+            if int(alen[0]) > GGUF_MAX_ARRAY_ELEMENTS:
+                raise ValueError(f'Array length {int(alen[0])} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}')
             offs += int(alen.nbytes)
             aparts: list[npt.NDArray[Any]] = [raw_itype, alen]
             data_idxs: list[int] = []

tests/CMakeLists.txt

@@ -116,8 +116,6 @@ function(llama_build_and_test source)
     set_property(TEST ${TEST_TARGET} PROPERTY LABELS ${LLAMA_TEST_LABEL})
 endfunction()
 
-llama_build_and_test(test-unicode.cpp)
-
 # build test-tokenizer-0 target once and add many tests
 llama_build(test-tokenizer-0.cpp)
 
@@ -154,6 +152,7 @@ llama_build(test-recurrent-state-rollback.cpp)
 
 if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
     # these tests are disabled on Windows because they use internal functions not exported with LLAMA_API (when building with shared libraries)
+    llama_build_and_test(test-unicode.cpp)
     llama_build_and_test(test-sampling.cpp)
     llama_build_and_test(test-reasoning-budget.cpp)
     llama_build_and_test(test-grammar-parser.cpp)

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]);
@@ -8892,6 +8913,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 +9278,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}));
@@ -10042,6 +10078,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

@@ -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