From 1c882a8ed089b452f35b42ee099089662a6262f1 Mon Sep 17 00:00:00 2001 From: Kevin Hopper <93635715+kh0pper@users.noreply.github.com> Date: Wed, 19 Aug 2026 03:20:21 -0500 Subject: [PATCH] vulkan: tiled transpose for 0<->2 permuted CONT (llama/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 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 26 +++++++- .../vulkan-shaders/copy_transpose_02.comp | 61 +++++++++++++++++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 2 + 3 files changed, 88 insertions(+), 1 deletion(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/copy_transpose_02.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 585e10d45..356f9abd7 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/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); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/copy_transpose_02.comp b/ggml/src/ggml-vulkan/vulkan-shaders/copy_transpose_02.comp new file mode 100644 index 000000000..5a3d66dab --- /dev/null +++ b/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)); + } + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 6c9f76af1..fbc2ea3ca 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/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"}});