Commit e321bdad7 for llama.cpp
commit e321bdad7e8c6ebf5675b641369544dc3fa3a54b
Author: Jeff Bolz <jbolz@nvidia.com>
Date: Sun Oct 25 04:13:48 2026 -0500
vulkan: Add a larger tile size for conv3d (coopmat2 only) (#29920)
* vulkan: Add a larger tile size for conv3d (coopmat2 only)
This adds a larger tile size, in practice it is only for coopmat2 where
we don't need shared memory to store the output (would use more than 64KB).
Add some additional representative backend tests. At least, this is what codex
came up with from looking at sd.cpp Wan VAE.
* only apply large tile size for newer driver versions
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-types.h b/ggml/src/ggml-vulkan/ggml-vulkan-types.h
index 00d20e843..5cf46c890 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h
+++ b/ggml/src/ggml-vulkan/ggml-vulkan-types.h
@@ -382,6 +382,7 @@ enum vk_conv_shapes {
CONV_SHAPE_64x32,
CONV_SHAPE_32x256,
CONV_SHAPE_64x128,
+ CONV_SHAPE_256x128,
CONV_SHAPE_COUNT,
};
@@ -397,6 +398,7 @@ inline vk_conv_block_size vk_conv_block_sizes[CONV_SHAPE_COUNT] = {
{ 64, 32, 32 }, // CONV_SHAPE_64x32
{ 32, 256, 16 }, // CONV_SHAPE_32x256
{ 64, 128, 16 }, // CONV_SHAPE_64x128
+ { 256, 128, 16 }, // CONV_SHAPE_256x128
};
enum dmmv_wg_sizes {
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index 2376c1948..33936acae 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -3810,7 +3810,7 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
auto shmem_req = [&](uint32_t pad, bool csh_store, bool fp16_shmem) {
const uint32_t elem_size = fp16_shmem ? (uint32_t)sizeof(uint16_t) : (uint32_t)sizeof(float);
- const uint32_t csh_elems = csh_store ? conv2d_BS.K * conv2d_BS.NPQ : 0u;
+ const uint32_t csh_elems = csh_store ? conv2d_BS.K * conv2d_BS.NPQ : (fp16_shmem ? 8u : 0u);
return (conv2d_BS.K * (conv2d_BS.CRS + pad) + conv2d_BS.CRS * (conv2d_BS.NPQ + pad) + csh_elems) * elem_size;
};
@@ -3838,9 +3838,9 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
}
// stage cm2 accumulator through shmem for coalesced global stores;
- // skipped on 128x128 where the extra Csh footprint hurts occupancy.
+ // skipped on large tiles where the extra Csh footprint hurts occupancy.
// cm1 always uses the staged path.
- uint32_t conv2d_csh_store = (device->coopmat2 && s != CONV_SHAPE_128x128) ? 1u : 0u;
+ uint32_t conv2d_csh_store = (device->coopmat2 && s != CONV_SHAPE_128x128 && s != CONV_SHAPE_256x128) ? 1u : 0u;
if (conv2d_use_cm1) {
conv2d_csh_store = 1;
}
@@ -8720,6 +8720,26 @@ static vk_conv_shapes ggml_vk_conv_select_shape(ggml_backend_vk_context * ctx, u
}
}
+static vk_conv_shapes ggml_vk_conv3d_select_shape(ggml_backend_vk_context * ctx, uint32_t K, uint32_t NPQ) {
+ const auto & block = vk_conv_block_sizes[CONV_SHAPE_256x128];
+ const uint32_t shader_core_count = ctx->device->shader_core_count;
+ // Slowdowns were observed on DGX Spark with drivers older than R615.
+ const bool driver_supported = ctx->device->driver_id != vk::DriverId::eNvidiaProprietary ||
+ (ctx->device->properties.driverVersion >> 22) >= 615;
+ if (ctx->device->coopmat2 && shader_core_count > 0 && K % block.K == 0 && driver_supported) {
+ const uint64_t n_tiles = (uint64_t)(K / block.K) * CEIL_DIV((uint64_t)NPQ, block.NPQ);
+ // Use the largest cm2 padding and include the dummy Csh[8].
+ const uint32_t pad = 8, csh_elems = 8;
+ const uint32_t shmem = (block.K * (block.CRS + pad) + block.CRS * (block.NPQ + pad) + csh_elems) * sizeof(uint16_t);
+ // Reuse each loaded input across 256 output channels instead of 128.
+ // Require full channel tiles and at least as many workgroups as GPU cores.
+ if (n_tiles >= shader_core_count && shmem <= ctx->device->properties.limits.maxComputeSharedMemorySize) {
+ return CONV_SHAPE_256x128;
+ }
+ }
+ return ggml_vk_conv_select_shape(ctx, K, NPQ);
+}
+
static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * dst, ggml_op op) {
switch (op) {
case GGML_OP_GET_ROWS:
@@ -9389,7 +9409,7 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
const uint32_t IC = (uint32_t)ggml_get_op_params_i32(dst, 9);
const uint32_t N = (uint32_t)ggml_get_op_params_i32(dst, 10);
const uint32_t NPQ = N * dst->ne[2] * dst->ne[1] * dst->ne[0];
- const vk_conv_shapes shape = ggml_vk_conv_select_shape(ctx, OC, NPQ);
+ const vk_conv_shapes shape = ggml_vk_conv3d_select_shape(ctx, OC, NPQ);
const uint32_t KW = (uint32_t)src0->ne[0];
const uint32_t KH = (uint32_t)src0->ne[1];
@@ -9702,7 +9722,7 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
case GGML_OP_CONV_3D:
if constexpr (std::is_same_v<PC, vk_op_conv3d_push_constants>) {
const uint32_t NPQ = pc.N * pc.OD * pc.OH * pc.OW;
- const vk_conv_shapes shape = ggml_vk_conv_select_shape(ctx, pc.OC, NPQ);
+ const vk_conv_shapes shape = ggml_vk_conv3d_select_shape(ctx, pc.OC, NPQ);
const uint32_t NPQ_blocks = CEIL_DIV(NPQ, vk_conv_block_sizes[shape].NPQ);
elements = { pc.OC, NPQ_blocks, 1 };
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 1d649d6ca..1eee812e6 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -10493,6 +10493,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
};
for (ggml_type kernel_type : {GGML_TYPE_F32, GGML_TYPE_F16}) {
+ test_cases.emplace_back(new test_conv_3d(1, 4, 6, 78, 54, 512, 3, 3, 3, 1, 1, 1, 0, 0, 0, 1, 1, 1, kernel_type));
+ test_cases.emplace_back(new test_conv_3d(2, 16, 4, 64, 64, 256, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, kernel_type));
for (int N : {1, 2}) {
for (int IC : {1, 3}) {
for (int OC : {1, 4}) {
@@ -12593,6 +12595,15 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
{1, 320, 8, 76, 52, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},
{1, 1280, 8, 76, 52, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},
{1, 320, 8, 152, 104, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+ // Wan VAE pads spatial borders and temporal history before conv3d; operator padding is zero.
+ {1, 384, 3, 40, 28, 384, 3, 3, 3, 1, 1, 1, 0, 0, 0, 1, 1, 1},
+ {1, 192, 4, 78, 54, 192, 3, 3, 3, 1, 1, 1, 0, 0, 0, 1, 1, 1},
+ {1, 96, 6, 154, 106, 96, 3, 3, 3, 1, 1, 1, 0, 0, 0, 1, 1, 1},
+ {1, 96, 6, 306, 210, 3, 3, 3, 3, 1, 1, 1, 0, 0, 0, 1, 1, 1},
+ {1, 1024, 3, 40, 28, 1024, 3, 3, 3, 1, 1, 1, 0, 0, 0, 1, 1, 1},
+ {1, 256, 6, 154, 106, 256, 3, 3, 3, 1, 1, 1, 0, 0, 0, 1, 1, 1},
+ {1, 384, 3, 38, 26, 768, 3, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1},
+ {1, 192, 2, 76, 52, 96, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1},
#if 0
// too slow on some devices
{1, 1280, 8, 152, 104, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},