Kevin HopperandKevin Hopper
98d1e92c21
* 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>
62 lines
2.1 KiB
Plaintext
62 lines
2.1 KiB
Plaintext
#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));
|
|
}
|
|
}
|
|
}
|
|
}
|