Public/koboldcpp
1
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-09-06 13:01:21 +02:00
Code Issues 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 03:20:21 -05:00
committed by GitHub
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
tests/test-backend-ops.cpp
+37 -4
View File
@@ -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");
@@ -8905,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));
@@ -10056,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;