From b1f5a51a1ed98d89fef3013d2a1151a50d4b77ed Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 8 Nov 2026 00:04:37 +0800 Subject: [PATCH] fix: chunk MiniMax-H3 QKV projections within backend buffer limits (#2286) --- src/core/ggml_runner.cpp | 25 ++++++- src/core/vulkan/int8.cpp | 8 +-- src/model/common/ggml_block.hpp | 3 +- src/model/diffusion/minimax_h3.hpp | 103 +++++++++++++++++++++-------- 4 files changed, 105 insertions(+), 34 deletions(-) diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index 8c2c5ba2..2fa14b15 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -211,6 +211,22 @@ ggml_backend_t GGMLRunner::backend_for_weight(const ggml_tensor* tensor) const { return nullptr; } +static bool backend_can_allocate_op(ggml_backend_t backend, const ggml_tensor* node) { + auto buft = ggml_backend_get_default_buffer_type(backend); + size_t alignment = ggml_backend_buft_get_alignment(buft); + size_t max_size = ggml_backend_buft_get_max_size(buft); + max_size -= max_size % alignment; + if (ggml_backend_buft_get_alloc_size(buft, node) > max_size) { + return false; + } + for (auto src : node->src) { + if (src != nullptr && ggml_backend_buft_get_alloc_size(buft, src) > max_size) { + return false; + } + } + return true; +} + void GGMLRunner::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) { if (sched == nullptr || gf == nullptr) { return; @@ -235,7 +251,14 @@ void GGMLRunner::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) { continue; } - if (ggml_backend_supports_op(current, node)) { + // Operation support can allow tensors larger than the backend's allocation limit. + // Check inputs too, since cross-backend copies must also fit a single allocation. + if (!backend_can_allocate_op(current, node)) { + auto cpu = workspace_.cpu_backend(); + if (cpu != nullptr && ggml_backend_supports_op(cpu, node)) { + ggml_backend_sched_set_tensor_backend(sched, node, cpu); + } + } else if (ggml_backend_supports_op(current, node)) { ggml_backend_sched_set_tensor_backend(sched, node, current); } } diff --git a/src/core/vulkan/int8.cpp b/src/core/vulkan/int8.cpp index aad36900..2bac51bb 100644 --- a/src/core/vulkan/int8.cpp +++ b/src/core/vulkan/int8.cpp @@ -715,10 +715,10 @@ namespace sd_vulkan_i8 { const Parameters pc = {n, rows, k, op.bias ? 1u : 0u}; auto pipeline = op.kind == SDW4Kind::W4A4 ? SD_VK_PIPELINE(mul_mat_w4a4_convrot, 6, Parameters) : SD_VK_PIPELINE(mul_mat_w4a8_convrot, 8, Parameters); uint32_t tile_cols = 64, tile_rows = 32; - const bool large_w4a4 = op.kind == SDW4Kind::W4A4 && rows >= 512 && d.compute.shared_memory >= w4a4_large_shared; - const bool large_w4a8 = op.kind == SDW4Kind::W4A8 && rows >= 512 && d.compute.shared_memory >= w4a8_large_shared; - const bool decode_w4a4 = op.kind == SDW4Kind::W4A4 && d.w4a4_cm2 && k % 256 == 0 && n % 64 == 0 && rows >= 128 && - aligned(op.weight, 16) && sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot_decode_cm2) > 4; + const bool large_w4a4 = op.kind == SDW4Kind::W4A4 && rows >= 512 && d.compute.shared_memory >= w4a4_large_shared; + const bool large_w4a8 = op.kind == SDW4Kind::W4A8 && rows >= 512 && d.compute.shared_memory >= w4a8_large_shared; + const bool decode_w4a4 = op.kind == SDW4Kind::W4A4 && d.w4a4_cm2 && k % 256 == 0 && n % 64 == 0 && rows >= 128 && + aligned(op.weight, 16) && sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot_decode_cm2) > 4; const uint32_t constants[] = {64, large_w4a4 || large_w4a8 ? 128u : 64u, 64}; const char* force_dp4a = std::getenv("GGML_VK_W4_FORCE_DP4A"); if (!(force_dp4a && std::atoi(force_dp4a)) && aligned(op.input, 16) && diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index 0653ae1f..39172bc3 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -974,7 +974,8 @@ public: w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); } x = ggml_rms_norm(ctx->ggml_ctx, x, eps); - x = ggml_mul_inplace(ctx->ggml_ctx, x, w); + // The affine multiply may run on a different backend after normalization falls back. + x = ggml_mul(ctx->ggml_ctx, x, w); return x; } }; diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index b5aec439..9a40643a 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -243,8 +244,68 @@ namespace MiniMaxH3 { blocks["out_proj"] = std::make_shared(inner, hidden_size, false); } - ggml_tensor* project_qkv(GGMLRunnerContext* ctx, ggml_tensor* x) { - return std::dynamic_pointer_cast(blocks["qkv_proj"])->forward(ctx, x); + std::vector split_qkv(GGMLRunnerContext* ctx, ggml_tensor* projected) { + std::vector result; + for (int i = 0; i < 3; ++i) { + result.push_back(ggml_view_4d(ctx->ggml_ctx, projected, head_dim, heads, + projected->ne[1], projected->ne[2] * projected->ne[3], + head_dim * projected->nb[0], projected->nb[1], projected->nb[2], + i * heads * head_dim * projected->nb[0])); + } + return result; + } + + std::vector project_qkv(GGMLRunnerContext* ctx, + const std::vector& inputs, + const std::function& prepare_input = {}) { + auto projection = std::dynamic_pointer_cast(blocks["qkv_proj"]); + int64_t sequence = 0; + for (auto input : inputs) { + sequence += input->ne[1]; + } + int64_t chunk_size = sequence; + if (ctx->backend != nullptr) { + auto buft = ggml_backend_get_default_buffer_type(ctx->backend); + size_t alignment = ggml_backend_buft_get_alignment(buft); + size_t max_size = ggml_backend_buft_get_max_size(buft); + max_size -= max_size % alignment; + size_t token_bytes = ggml_row_size(GGML_TYPE_F32, 3 * heads * head_dim) * inputs[0]->ne[2] * inputs[0]->ne[3]; + chunk_size = static_cast(std::min(chunk_size, std::max(1, max_size / token_bytes))); + } + if (chunk_size == sequence) { + ggml_tensor* projected = nullptr; + for (size_t i = 0; i < inputs.size(); ++i) { + auto input = prepare_input ? prepare_input(inputs[i], i) : inputs[i]; + auto part = projection->forward(ctx, input); + projected = projected == nullptr ? part : ggml_concat(ctx->ggml_ctx, projected, part, 1); + } + return split_qkv(ctx, projected); + } + + // Split Q/K/V before concatenating tokens so the full packed QKV never exceeds the buffer limit. + std::vector chunks[3]; + for (size_t i = 0; i < inputs.size(); ++i) { + auto input = inputs[i]; + for (int64_t start = 0; start < input->ne[1]; start += chunk_size) { + auto part = ggml_ext_slice(ctx->ggml_ctx, input, 1, start, std::min(start + chunk_size, input->ne[1]), false); + if (!ggml_is_contiguous(part)) { + part = ggml_cont(ctx->ggml_ctx, part); + } + // Keep normalization/modulation adjacent to the projection's input quantization. + if (prepare_input) { + part = prepare_input(part, i); + } + auto qkv = split_qkv(ctx, projection->forward(ctx, part)); + for (int j = 0; j < 3; ++j) { + chunks[j].push_back(ggml_cont(ctx->ggml_ctx, qkv[j])); + } + } + } + std::vector result; + for (auto& parts : chunks) { + result.push_back(ggml_ext_vec_concat(ctx->ggml_ctx, parts, 2)); + } + return result; } ggml_tensor* project_output(GGMLRunnerContext* ctx, ggml_tensor* x) { @@ -252,22 +313,15 @@ namespace MiniMaxH3 { } ggml_tensor* attend(GGMLRunnerContext* ctx, - ggml_tensor* projected, + const std::vector& qkv, ggml_tensor* pe = nullptr) { - auto q_norm = std::dynamic_pointer_cast(blocks["q_norm"]); - auto k_norm = std::dynamic_pointer_cast(blocks["k_norm"]); - int64_t sequence = projected->ne[1]; - int64_t batch = projected->ne[2] * projected->ne[3]; - auto view = [&](int index) { - return ggml_view_4d(ctx->ggml_ctx, projected, head_dim, heads, sequence, batch, - head_dim * projected->nb[0], projected->nb[1], projected->nb[2], - index * heads * head_dim * projected->nb[0]); - }; - auto q = view(0); - auto k = view(1); - auto v = view(2); - q = q_norm->forward(ctx, q); - k = k_norm->forward(ctx, k); + auto q_norm = std::dynamic_pointer_cast(blocks["q_norm"]); + auto k_norm = std::dynamic_pointer_cast(blocks["k_norm"]); + auto q = qkv[0]; + auto k = qkv[1]; + auto v = qkv[2]; + q = q_norm->forward(ctx, q); + k = k_norm->forward(ctx, k); if (pe != nullptr) { q = apply_partial_rope(ctx, q, pe); k = apply_partial_rope(ctx, k, pe); @@ -290,7 +344,7 @@ namespace MiniMaxH3 { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe = nullptr) { - return project_output(ctx, attend(ctx, project_qkv(ctx, x), pe)); + return project_output(ctx, attend(ctx, project_qkv(ctx, {x}), pe)); } }; @@ -447,16 +501,9 @@ namespace MiniMaxH3 { auto modulate = [&](ggml_tensor* value, const std::vector& row, int shift, int scale) { return ggml_ext_modulate(ctx, value, row[shift], row[scale]); }; - ggml_tensor* qkv = nullptr; - for (size_t i = 0; i < segments.size(); ++i) { - auto h = modulate(norm1->forward(ctx, inputs[i]), rows[i], 0, 1); - h = attn->project_qkv(ctx, h); - if (qkv == nullptr) { - qkv = h; - continue; - } - qkv = ggml_concat(ctx->ggml_ctx, qkv, h, 1); - } + auto qkv = attn->project_qkv(ctx, inputs, [&](ggml_tensor* input, size_t i) { + return modulate(norm1->forward(ctx, input), rows[i], 0, 1); + }); auto attended = attn->attend(ctx, qkv, pe); std::vector projections; for (const auto& segment : segments) {