vulkan: Fix OOB A reads in MUL_MAT_VEC for odd sizes

There was a TODO to fix the OOB reads from the A matrix which we do
here.

It is within performance noise (+<0.1%) in tg128 for
Qwen3.5-9B:BF16 on Intel BMG.

Author: TheBlueMatt

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

@@ -5,6 +5,9 @@
 #include "types.glsl"
 
 #if defined(DATA_A_F32)
+FLOAT_TYPE dequantize1(uint ib, uint iqs, uint a_offset) {
+    return data_a[a_offset + ib];
+}
 vec2 dequantize(uint ib, uint iqs, uint a_offset) {
     return vec2(data_a[a_offset + ib], data_a[a_offset + ib + 1]);
 }
@@ -20,6 +23,9 @@ vec4 dequantize4_2aligned(uint ib, uint iqs, uint a_offset) {
 #endif
 
 #if defined(DATA_A_F16)
+FLOAT_TYPE dequantize1(uint ib, uint iqs, uint a_offset) {
+    return data_a[a_offset + ib];
+}
 vec2 dequantize(uint ib, uint iqs, uint a_offset) {
     return vec2(data_a[a_offset + ib], data_a[a_offset + ib + 1]);
 }
@@ -35,6 +41,9 @@ vec4 dequantize4_2aligned(uint ib, uint iqs, uint a_offset) {
 #endif
 
 #if defined(DATA_A_BF16)
+FLOAT_TYPE dequantize1(uint ib, uint iqs, uint a_offset) {
+    return bf16_to_fp32(data_a[a_offset + ib]);
+}
 vec2 dequantize(uint ib, uint iqs, uint a_offset) {
     return vec2(bf16_to_fp32(data_a[a_offset + ib]), bf16_to_fp32(data_a[a_offset + ib + 1]));
 }

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

@@ -16,15 +16,11 @@ layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
 
 uint a_offset, b_offset, d_offset, y_offset;
 
-vec4 load_b(const uint j, const uint iybs, const uint iqs, const bool lastiter) {
-    // Check if the second of the pair of elements is OOB, and don't fetch B or
-    // accumulate it. We still fetch a pair of elements for A, which is fine for
-    // quantized formats since they'll be within the same block. We should
-    // probably skip fetching the second element for F16/F32, but as of now we
-    // still do.
-    const bool OOB_y = lastiter && (iybs + iqs + y_offset >= p.ncols);
-    const bool OOB_z = lastiter && (iybs + iqs + y_offset*2 >= p.ncols);
-    const bool OOB_w = lastiter && (iybs + iqs + y_offset*3 >= p.ncols);
+vec4 load_b(const uint j, const uint iybs, const uint iqs, const bool lastiter, out bool OOB_y, out bool OOB_z, out bool OOB_w) {
+    // Check if the latter elements are OOB, and don't fetch B or accumulate it.
+    OOB_y = lastiter && (iybs + iqs + y_offset >= p.ncols);
+    OOB_z = lastiter && (iybs + iqs + y_offset*2 >= p.ncols);
+    OOB_w = lastiter && (iybs + iqs + y_offset*3 >= p.ncols);
 
     if (!OOB_w) {
         return vec4(FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs]),
@@ -46,14 +42,6 @@ vec4 load_b(const uint j, const uint iybs, const uint iqs, const bool lastiter)
     }
 }
 
-vec4 load_b_aligned(const uint j, const uint iybs, const uint iqs, const bool lastiter) {
-    const bool OOB_w = lastiter && (iybs + iqs + y_offset*3 >= p.ncols);
-    if (!OOB_w)
-        return data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4];
-    else
-        return load_b(j, iybs, iqs, lastiter);
-}
-
 void iter(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const uint num_rows, const uint tid, const uint i, bool lastiter)
 {
     [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) {
@@ -63,6 +51,8 @@ void iter(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const
 
 #if K_PER_ITER == 8
 #if QUANT_R == 2
+        // Note that we end up fetching bogus elements here, but its fine as they'll be
+        // within an accessible block.
         const vec4 bv02 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4]);
         const vec4 bv13 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs + y_offset) / 4]);
         const vec4 bv0 = vec4(bv02.x, bv13.x, bv02.y, bv13.y);
@@ -72,7 +62,11 @@ void iter(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const
         const vec4 bv1 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4 + 1]);
 #endif
 #else
-        const vec4 b = load_b(j, iybs, iqs, lastiter);
+        bool OOB_y;
+        bool OOB_z;
+        bool OOB_w;
+
+        const vec4 b = load_b(j, iybs, iqs, lastiter, OOB_y, OOB_z, OOB_w);
 #endif
         uint ibi = first_row*p.ncols;
         [[unroll]] for (uint n = 0; n < num_rows; ++n) {
@@ -98,24 +92,38 @@ void iter(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const
 
             temp[j][n] += rowtmp;
 #else
-            const vec4 v = dequantize4(ib, iqs, a_offset);
-
-            // matrix multiplication
-            temp[j][n] += dot(v, b);
+            if (!OOB_w) {
+                const vec4 v = dequantize4(ib, iqs, a_offset);
+                temp[j][n] += dot(v, b);
+            } else if (!OOB_z) {
+                const vec2 v0 = dequantize(ib, iqs, a_offset);
+                const FLOAT_TYPE v1 = dequantize1(ib + 2/QUANT_R, iqs, a_offset);
+                const vec3 v = vec3(v0.x, v0.y, v1);
+                const vec3 b0 = vec3(b.x, b.y, b.z);
+                temp[j][n] += dot(v, b0);
+            } else if (!OOB_y) {
+                const vec2 v0 = dequantize(ib, iqs, a_offset);
+                const vec2 b0 = vec2(b.x, b.y);
+                temp[j][n] += dot(v0, b0);
+            } else {
+                const FLOAT_TYPE v = dequantize1(ib, iqs, a_offset);
+                temp[j][n] = fma(v, b.x, temp[j][n]);
+            }
 #endif
         }
     }
 }
 
 #if defined(DATA_A_F32) || defined(DATA_A_F16) || defined(DATA_A_BF16)
-void iter_aligned_nonquant(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const uint num_rows, const uint tid, const uint i, bool lastiter)
+void iter_aligned_nonquant(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const uint num_rows, const uint tid, const uint i)
 {
     [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) {
         const uint col = i*BLOCK_SIZE + K_PER_ITER*tid;
         const uint iqs = 0; // quant index
         const uint iybs = col; // y block start index
 
-        const vec4 b = load_b_aligned(j, iybs, iqs, lastiter);
+        const vec4 b = data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4];
+
         uint ibi = first_row*p.ncols;
         [[unroll]] for (uint n = 0; n < num_rows; ++n) {
             const uint ib = (ibi + col)/QUANT_K; // block index
@@ -136,7 +144,7 @@ void compute_outputs(const uint32_t first_row, const uint32_t num_rows) {
     get_offsets(a_offset, b_offset, d_offset);
     const bool is_aligned_nonquant =
         p.batch_stride_b % 4 == 0 && b_offset % 4 == 0 &&
-        p.ncols % 2 == 0 && BLOCK_SIZE % 4 == 0 &&
+        p.ncols % 4 == 0 && BLOCK_SIZE % 4 == 0 &&
         K_PER_ITER == 4;
 
     y_offset = QUANT_R == 1 ? 1 : QUANT_K/2;
@@ -170,7 +178,7 @@ void compute_outputs(const uint32_t first_row, const uint32_t num_rows) {
         while (i < unrolled_iters) {
             // Manually partially unroll the loop
             [[unroll]] for (uint k = 0; k < unroll_count; ++k) {
-                iter_aligned_nonquant(temp, first_row, num_rows, tid, i*K_PER_ITER, false);
+                iter_aligned_nonquant(temp, first_row, num_rows, tid, i*K_PER_ITER);
                 i++;
             }
         }
@@ -201,7 +209,7 @@ void compute_outputs(const uint32_t first_row, const uint32_t num_rows) {
         while (i < unrolled_iters && is_aligned_nonquant) {
             // Manually partially unroll the loop
             [[unroll]] for (uint k = 0; k < unroll_count; ++k) {
-                iter_aligned_nonquant(temp, first_row, num_rows, tid, i*K_PER_ITER, false);
+                iter_aligned_nonquant(temp, first_row, num_rows, tid, i*K_PER_ITER);
                 i++;
             }
         }
@@ -221,7 +229,7 @@ void compute_outputs(const uint32_t first_row, const uint32_t num_rows) {
 #if K_PER_ITER == 4
     if (is_aligned_nonquant) {
         while (i < num_iters) {
-            iter_aligned_nonquant(temp, first_row, num_rows, tid, i*K_PER_ITER, true);
+            iter_aligned_nonquant(temp, first_row, num_rows, tid, i*K_PER_ITER);
             i++;
         }
     } else {