fix: chunk MiniMax-H3 QKV projections within backend buffer limits (#2286)

This commit is contained in:
leejet authored and GitHub committed 2026-11-08 00:04:37 +08:00
1 parent fdfc913c51
commit b1f5a51a1e
4 files changed
+105 -34

No files matched your search

+24 -1
View File
@@ -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);
}
}
+4 -4
View File
@@ -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) &&