Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
* 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>
This commit is contained in:
Kevin Hopper
2026-08-19 10:20:21 +02:00
committed by GitHub
co-authored by Kevin Hopper
parent 5112b9738b
commit 98d1e92c21
4 changed files with 125 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-vulkan/ggml-vulkan.cpp
+25 -1
View File
@@ -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);