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

ggml-vulkan: tiled transpose fast-path for concat with transposed source

Browse Source 
The generic concat shader reads a transposed source (nb[1]==type_size) with an
uncoalesced stride, which is catastrophically slow on discrete GPUs (~5.6ms for
a 16MB concat on the RX 580 vs ~130us of memory bandwidth), a ~9% prefill
hotspot on Qwen3.5-4B (delta-net state concat).

Add a fast path for concat along dim 0 where one source is stored transposed and
the other source + dst are contiguous along dim 0: copy the contiguous source
with copy.comp and transpose the other source into the matching dst sub-region
with the existing tiled copy_transpose shader (shared-memory 32x32 transpose,
coalesced read+write). Reuses pipeline_cpy_* / pipeline_cpy_transpose_* with
custom push constants + doffset, no new shader. Falls back to the generic path
otherwise (gated on type/shape/contiguity + 16-bit doffset bound).

Assisted-by: opencode
This commit is contained in:
Lumpiasty
2026-07-22 21:04:20 +02:00
parent 1b8abd1c23
commit 42e52045dc
1 changed files with 99 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-vulkan/ggml-vulkan.cpp
+99
View File
@@ -12520,9 +12520,108 @@ static void ggml_vk_opt_step_sgd(ggml_backend_vk_context * ctx, vk_context& subc
ggml_vk_op_f32<vk_op_push_constants>(ctx, subctx, src0, src1, src2, nullptr, dst, GGML_OP_OPT_STEP_SGD, { (uint32_t)n, 0, 0.0f, 0.0f, 0.0f, 0.0f });
}
// Fast path for concat along dim 0 where one source is stored "transposed"
// (nb[1] == type_size, dim1 innermost) and the other source + dst are
// contiguous along dim 0. The generic concat shader reads the transposed
// source with a catastrophic uncoalesced stride; here we instead copy the
// contiguous source with copy.comp and transpose the other source into the
// matching dst sub-region with the tiled copy_transpose shader (shared-memory
// transpose, coalesced read+write). Mirrors the precedent in
// ggml_vk_cpy_to_contiguous: direct dispatch with custom push constants.
static void ggml_vk_concat_transpose_fastpath(ggml_backend_vk_context * ctx, vk_context& subctx,
const ggml_tensor * ctg, const ggml_tensor * trp,
ggml_tensor * dst, uint32_t off_ctg, uint32_t off_trp) {
const uint32_t ts = ggml_type_size(dst->type);
vk_pipeline pipeline_cpy = (ts == 4) ? ctx->device->pipeline_cpy_f32_f32
: ctx->device->pipeline_cpy_f16_f16;
vk_pipeline pipeline_trp = (ts == 4) ? ctx->device->pipeline_cpy_transpose_32
: ctx->device->pipeline_cpy_transpose_16;
ggml_pipeline_request_descriptor_sets(ctx, pipeline_cpy, 1);
ggml_pipeline_request_descriptor_sets(ctx, pipeline_trp, 1);
vk_subbuffer ctg_buf = ggml_vk_tensor_subbuffer(ctx, ctg, true);
vk_subbuffer trp_buf = ggml_vk_tensor_subbuffer(ctx, trp, true);
vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true);
const uint32_t a_misalign_ctg = get_misalign_bytes(ctx, ctg) / ts;
const uint32_t a_misalign_trp = get_misalign_bytes(ctx, trp) / ts;
const uint32_t d_misalign = get_misalign_bytes(ctx, dst) / ts;
// Dispatch A: contiguous copy of `ctg` into dst[off_ctg : off_ctg + ctg->ne[0], :]
if (ctg->ne[0] > 0) {
const uint32_t ne_ctg = (uint32_t) ggml_nelements(ctg);
vk_op_unary_push_constants pc = vk_op_unary_push_constants_init(ctg, dst, ne_ctg);
pc.ne10 = (uint32_t) ctg->ne[0]; // only ctg's columns in dst
pc.misalign_offsets = (a_misalign_ctg << 16) | (d_misalign + off_ctg);
init_pushconst_fastdiv(pc);
std::array<uint32_t, 3> el = ne_ctg > 262144 ? std::array<uint32_t,3>{512, 512, CEIL_DIV(ne_ctg, 262144)}
: ne_ctg > 512 ? std::array<uint32_t,3>{512, CEIL_DIV(ne_ctg, 512), 1}
: std::array<uint32_t,3>{ne_ctg, 1, 1};
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline_cpy, { ctg_buf, dst_buf }, pc, el);
}
// Dispatch B: tiled transpose of `trp` into dst[off_trp : off_trp + trp->ne[0], :]
if (trp->ne[0] > 0) {
vk_op_unary_push_constants pc = vk_op_unary_push_constants_init(trp, dst, ggml_nelements(trp));
pc.ne10 = (uint32_t) trp->ne[0]; // dst bound = trp columns (NOT dst->ne[0])
pc.misalign_offsets = (a_misalign_trp << 16) | (d_misalign + off_trp);
init_pushconst_fastdiv(pc);
std::array<uint32_t, 3> el = {
(uint32_t) CEIL_DIV(trp->ne[0], 32),
(uint32_t) CEIL_DIV(trp->ne[1], 32),
(uint32_t) (trp->ne[2] * trp->ne[3]),
};
el[0] = std::min(el[0], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[0]);
el[1] = std::min(el[1], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[1]);
el[2] = std::min(el[2], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[2]);
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline_trp, { trp_buf, dst_buf }, pc, el);
}
ggml_vk_sync_buffers(ctx, subctx);
}
static void ggml_vk_concat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
int * op_params = (int *)dst->op_params;
// Fast path: concat along dim 0 with one source "transposed" (nb[1]==type_size,
// dim1 innermost) and the other source + dst contiguous along dim 0. The generic
// shader reads the transposed source uncoalesced (~5.6ms on RX 580 vs ~130us for
// a coalesced copy); here we instead use a tiled shared-memory transpose.
auto src_transposed_2d = [&](const ggml_tensor * s) {
const uint32_t ts = ggml_type_size(s->type);
return s->nb[1] == ts // dim1 innermost
&& s->nb[0] == s->ne[1] * ts // consistent 2D transpose
&& s->ne[2] == 1 && s->ne[3] == 1;
};
const uint32_t dst_ts = ggml_type_size(dst->type);
const bool dim0_ok = (op_params[0] == 0) && (dst->nb[0] == dst_ts);
const bool types_ok = (dst_ts == 4 || dst_ts == 2)
&& (src0->type == src1->type && src0->type == dst->type)
&& (ggml_blck_size(dst->type) == 1);
const bool shapes_ok = (src0->ne[1] == src1->ne[1] && src1->ne[1] == dst->ne[1])
&& (src0->ne[2] == src1->ne[2] && src0->ne[2] == dst->ne[2])
&& (src0->ne[3] == src1->ne[3] && src0->ne[3] == dst->ne[3])
&& (src0->ne[0] + src1->ne[0] == dst->ne[0]);
// doffset is 16-bit for unary shaders; guard against overflow
const bool off_fits = ((get_misalign_bytes(ctx, dst)/dst_ts + dst->ne[0])) < 0xFFFFu;
const bool s1_trans = dim0_ok && types_ok && shapes_ok && off_fits
&& src_transposed_2d(src1) && ggml_is_contiguous(src0);
const bool s0_trans = dim0_ok && types_ok && shapes_ok && off_fits
&& src_transposed_2d(src0) && ggml_is_contiguous(src1);
if (s1_trans || s0_trans) {
const ggml_tensor * ctg = s1_trans ? src0 : src1; // contiguous source
const ggml_tensor * trp = s1_trans ? src1 : src0; // transposed source
// dst offsets (in elements) where each source region begins
const uint32_t off_ctg = s1_trans ? 0u : (uint32_t) src1->ne[0];
const uint32_t off_trp = s1_trans ? (uint32_t) src0->ne[0] : 0u;
ggml_vk_concat_transpose_fastpath(ctx, subctx, ctg, trp, dst, off_ctg, off_trp);
return;
}
const uint32_t src0_type_size = ggml_type_size(src0->type);
const uint32_t src1_type_size = ggml_type_size(src1->type);
const uint32_t dst_type_size = ggml_type_size(dst->type);