vulkan : run small-M mul_mat as a swapped mat-vec

[Nathanw1014]

Flash-Next's hyper-connection inject projection is a matmul with only 4 output rows (f32 m=4 n=2048 k=10240, 95 times per prompt batch). On Vulkan it went through the regular matmul kernel, padded out to a full tile, and cost about 145 ms per 2048-token batch.

This extends the existing m=1 trick (ggml-org#28457 upstream) to small m: run it as a mat-vec with the operands swapped, then copy the small transposed result into place. Only applies to plain 2-D cases with at most 8 output rows and more than 8 columns, so decode and everything else keep their current paths.

On a Strix Halo (gfx1151, RADV) that op drops to 54 ms, and Flash-Next prefill goes up about 4% at depth 0 and 3.4% at 12k (measured with #115 and #119 applied). Decode is unchanged, and 27B and 35B never hit this path. Divergence from a CPU reference is the same as before. New test cases cover the exact shape and a few neighbours; skipping the copy makes them fail, so they do exercise the route.

Written with Claude, tested by me.

Assisted-by: Claude Opus 5.5

[gaetan-puleo]

Independent verification on a second Strix Halo (strix-halo-002, gfx1151, RADV)

Built test-backend-ops at head 4fa5cc3dc in the kyuz0 RADV container. Baseline is the unmodified mul_mat code from the same base (I added the two new perf cases to a pre-#129 build so the harness and shape match exactly).

Performance (isolated op, GGML_NLOOP=3, -b Vulkan0)

shape (type_a=f32,type_b=f32)	old path	this PR	speedup
m=4, n=2048, k=10240	2006.9 µs (83.6 GFLOPS)	413.9 µs (405 GFLOPS)	4.85x
m=4, n=512, k=10240	468.6 µs (89.5 GFLOPS)	25.1 µs (1.67 TFLOPS)	18.6x

Both runs report OK; the baseline numbers are the same harness on the same machine, so the route is genuinely taken. (I did not reproduce the end-to-end prefill %; the per-op gain is large and in the direction you describe.)

Correctness

Targeted runs of the new shapes on Vulkan0, all pass:

• m=4,n=2048,k=10240 → 1/1
• m in {2,3,4,5,8} x n in {9,16,509,2048}, k=2048 → 20/20
• m=4, n in {1,4,8,9,512,2048}, k=10240 → 6/6 (n <= 8 stays on the plain mat-vec, n = 9 enters the new path)
• m=4,n=509,k=2051 (F16/F32, F16/F16, batched bs=[2,1]) → 3/3
• m=9, n in {9,16,509,2048}, k=2048 → 4/4 (boundary, falls through)

Full suite: test-backend-ops test -b Vulkan0 -o MUL_MAT → 7443/7446. The three failures are all [CONCAT] ERR = ... on m=128 shapes (type_a=bf16,n=9/16,k=5120 and type_a=q8_0,n=4,k=5120). m=128 > mul_mat_vec_max_cols (8), so they cannot come from this change, and I reproduced them identically on the unmodified baseline. So this PR introduces no regression across the MUL_MAT suite.

Code review

The transpose is wired up correctly. The mat-vec writes data_d[j*batch_stride_d + row] with batch_stride_d = ne20 = n (the vector index is slow), so the scratch is [m, n] row-major; dst_t (nb[0]=n*4, nb[1]=4) describes exactly that source, and ggml_vk_cpy_to_contiguous writes pc.nb10=1, nb11=m — a contiguous [m, n] with nb0=4, i.e. the real destination. The scratch is sized to ggml_nbytes(dst) = m*n*4, which the kernel fills exactly. The prealloc_split_k growth, the pre-write need_sync check and the trailing prealloc_split_k_need_sync = true match the existing scratch pattern. The m=1 path is unchanged (transpose_d is false), the non-swap column-chunking path never sets it, and decode/speculative (n <= 8) stays on the old branches.

Nits, none blocking:

• The dispatch guard does not test dst->type, but the new code GGML_ASSERT(dst->type == GGML_TYPE_F32 ...). That is safe today since ggml_mul_mat always allocates an F32 result, but it would abort (not fall through) if that ever changed. Adding dst->type == GGML_TYPE_F32 to the branch condition would make it a clean fallback.
• The branch allows src1 to be BF16, which makes the BF16 tensor the mat-vec matrix after the swap. That allowance already existed for the m=1 route, but there is no type_b=BF16 case anywhere in the suite, so this specific combination has no dedicated test. Same for a quantized src1. Worth a case or two if you want the new small-m route exercised for those types.
• The copy passes dst = nullptr to ggml_vk_get_cpy_pipeline, so it takes the generic strided shader rather than the fast transpose pipeline. Irrelevant at a few KB, just noting it.

CI

All three red jobs are pre-existing and unrelated: ubuntu-llvmpipe fails the same test-recurrent-state-rollback (x2) and test-save-load-state as master b79812d59 (confirmed identical), windows is the test-qsa-fast.cpp setenv error fixed by #138, and metal hit the 24h timeout.

LGTM. Correct on gfx1151, no MUL_MAT regressions, and a 4.8–18x improvement on the targeted small-M shapes.

[Nathanw1014]

Thanks for the careful check, especially building the baseline harness for the perf cases. Pushed 109119a2f:

• The dispatch condition now requires an F32 dst when M > 1, so anything else falls through to the regular paths instead of asserting.
• On BF16 / quantized src1: I tried adding cases, and the CPU reference reports them as not supported, so test-backend-ops cannot check that combination. Rather than leave an untestable route, the M > 1 path is now limited to f32/f16 src1 (the graphs that hit it are f32 activations anyway). The M = 1 route is unchanged.

I left the copy on the generic strided shader: it moves a few KB.

Re-checked on gfx1151, stock Mesa 25.2.8: -o MUL_MAT 7443/7446, the same three [CONCAT] failures you saw on the baseline. m=4,n=2048,k=10240 417 µs and m=4,n=512,k=10240 28 µs, so the route is still taken.

[gaetan-puleo]

Code review of 109119a (by reading the code; not built or run here)

I found no correctness problems. LGTM, with a few optional follow-ups below.

Verified

• Transposed write: with the inputs swapped, NUM_COLS = m, batch_n is true and stride_batch_d = n. The shader writes data_d[j*batch_stride_d + d_offset + first_row + n] with d_offset = 0 (mul_mat_vec_base.glsl:124), so the scratch holds the [m, n] result transposed (j*n + i). dst_t (nb0 = 4n, nb1 = 4) describes that layout exactly. ggml_vk_cpy_to_contiguous applies no misalign offset (the push constants are zero-initialised and init_pushconst_tensor_offsets is never called), so reading prealloc_split_k at offset 0 is safe.
• Scratch buffer: sized exactly to ggml_nbytes(dst). Growing it mid-graph submits and waits before the old buffer is destroyed. The sync order (check need_sync before writing, barrier before the copy, set need_sync after) matches the existing split_k matmul. The other prealloc_split_k users all stay within a single node, and the graph-level tracking still sees dst as written.
• Kernel selection: Q8_1 only accepts quantized A, so with A = F32/F16 quantize_y is always false and prealloc_y and its cache are untouched. The swapped call returns before the column-chunk logic. Pipelines exist for every {F32, F16} × {F32, F16} combination at NUM_COLS 2..8.
• Routing: m = 1 is unchanged. n ≤ 8, batched, fused, misaligned and non-F32-dst cases all fall through. MUL_MAT+ADD fusion requires ggml_nrows(mul) == 1, so this shape can't be fused.
• Tests: exactly 25 of the new cases take the route. The m = 9, n ≤ 8 and bs={2,1} cases correctly don't.
• CI: the llvmpipe failures (test-recurrent-state-rollback x2, test-save-load-state) abort in ggml_backend_sched_alloc_splits and also fail on an unrelated PR, so they're not caused by this one. Note that the llvmpipe job builds test-backend-ops but does not run it (ctest -L main), so the new cases have only been run on gfx1151.

Optional follow-ups

1. Test gap: the route accepts src0 F32 × src1 F16, which uses the f16_f32 mat-vec kernel after the swap, but none of the new cases cover it. Adding test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F16, 4, 509, 2051, {1, 1}, {1, 1}) would close that.
2. Unmeasured shapes: the route has no lower bound on k (or upper bound on n), and perf was only measured at m = 4 with k = 2048/10240. The 3/5/6-column slowdown noted near ggml_vk_mmv_good_cols was measured on quantized weights only, so it doesn't prove anything about F32/F16 A, but m = 3/5/6/8 are unmeasured. It would be worth adding perf cases such as m=3/5/8, n=2048, k=10240 and a small-k one like m=4, n=2048, k=64, or a minimum k in the routing condition.
3. Perf numbers: m=4, n=512, k=10240 at ~25 µs is about 21 MB / 25 µs ≈ 840 GB/s, well above Strix Halo DRAM bandwidth. The activations likely stay in the 32 MB MALL across loop iterations, so the 18.6x speedup is probably cache-inflated. The n = 2048 number (~4.8x, ~200 GB/s) is the realistic one and matches the end-to-end 145 → 54 ms.
4. Nits:
   • ggml_vk_cpy_to_contiguous already ends with a barrier, so prealloc_split_k_need_sync = true after it is redundant and adds an extra barrier at the next scratch user. Harmless.
   • With dst = nullptr, the generic strided copy is used rather than pipeline_cpy_transpose_32. Irrelevant at a few KB.
   • On non-AMD devices the workgroup-size heuristic sees n as M for the swapped call. Untested on NVIDIA and Intel, but not a correctness issue.
   • The perf logger labels this route MUL_MAT rather than MUL_MAT_VEC (it checks n <= mul_mat_vec_max_cols). This was already the case for the m = 1 swapped route.