From a21950ee5a2f18c5a89bbd05fb5a305b8ee36757 Mon Sep 17 00:00:00 2001 From: s3kundes Date: Fri, 17 Jul 2026 20:39:20 +0300 Subject: [PATCH 1/2] split mode tensor working --- ggml/src/ggml-backend-meta.cpp | 430 ++++++++++++++++++++++++++++----- src/llama-arch.cpp | 1 - src/llama-model.cpp | 50 ++++ 3 files changed, 423 insertions(+), 58 deletions(-) diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index a5a3a58ad054..09a2c6b927f3 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -47,6 +48,99 @@ const char * ggml_backend_meta_split_axis_name(enum ggml_backend_meta_split_axis } } +static const char * ggml_backend_meta_tensor_name(const ggml_tensor * tensor) { + return tensor && tensor->name[0] ? tensor->name : "(unnamed)"; +} + +static std::string ggml_backend_meta_shape_str(const ggml_tensor * tensor) { + if (tensor == nullptr) { + return "(null)"; + } + std::string ret = "ne={"; + for (int i = 0; i < GGML_MAX_DIMS; i++) { + if (i > 0) { + ret += ","; + } + ret += std::to_string(tensor->ne[i]); + } + ret += "}, nb={"; + for (int i = 0; i < GGML_MAX_DIMS; i++) { + if (i > 0) { + ret += ","; + } + ret += std::to_string(tensor->nb[i]); + } + ret += "}"; + return ret; +} + +static std::string ggml_backend_meta_tensor_str(const ggml_tensor * tensor) { + if (tensor == nullptr) { + return "(null)"; + } + std::string ret = std::string(ggml_backend_meta_tensor_name(tensor)) + "[" + ggml_op_name(tensor->op) + + ", " + ggml_type_name(tensor->type) + ", " + ggml_backend_meta_shape_str(tensor); + if (tensor->view_src != nullptr) { + ret += ", view_src="; + ret += ggml_backend_meta_tensor_name(tensor->view_src); + ret += ", view_offs="; + ret += std::to_string(tensor->view_offs); + } + ret += "]"; + return ret; +} + +static std::string ggml_backend_meta_split_state_str(const ggml_backend_meta_split_state & ss, size_t n_bufs) { + std::string ret = std::string(ggml_backend_meta_split_axis_name(ss.axis)) + ", n_segments=" + std::to_string(ss.n_segments); + ret += ", sums={"; + for (size_t j = 0; j < n_bufs; j++) { + if (j > 0) { + ret += ","; + } + int64_t sum = 0; + for (size_t s = 0; s < ss.n_segments; s++) { + sum += ss.ne[s*n_bufs + j] * ss.nr[s]; + } + ret += std::to_string(sum); + } + ret += "}"; + if (ss.axis >= 0 && ss.axis < GGML_MAX_DIMS) { + ret += ", segments={"; + for (size_t s = 0; s < ss.n_segments; s++) { + if (s > 0) { + ret += ";"; + } + ret += "nr=" + std::to_string(ss.nr[s]) + ":"; + for (size_t j = 0; j < n_bufs; j++) { + if (j > 0) { + ret += ","; + } + ret += std::to_string(ss.ne[s*n_bufs + j]); + } + } + ret += "}"; + } + return ret; +} + +static void ggml_backend_meta_log_split_context( + const char * label, + const ggml_tensor * tensor, + const std::vector & src_ss, + size_t n_bufs, + const std::string & detail) { + GGML_LOG_ERROR("META_SPLIT_%s: dst=%s%s%s\n", label, ggml_backend_meta_tensor_str(tensor).c_str(), + detail.empty() ? "" : ", ", detail.c_str()); + for (size_t i = 0; i < GGML_MAX_SRC; i++) { + if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { + continue; + } + GGML_LOG_ERROR("META_SPLIT_%s: src[%zu]=%s, state={%s}\n", label, i, + ggml_backend_meta_tensor_str(tensor->src[i]).c_str(), + ggml_backend_meta_split_state_str(src_ss[i], n_bufs).c_str()); + } +} + // // meta backend device // @@ -513,6 +607,42 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return true; }; + auto log_split_failure = [&](const char * label, const std::vector & src_ss, const std::string & detail = std::string()) { + ggml_backend_meta_log_split_context(label, tensor, src_ss, n_bufs, detail); + }; + + auto check_split = [&](bool ok, const char * label, const std::vector & src_ss, const std::string & detail = std::string()) { + if (!ok) { + log_split_failure(label, src_ss, detail); + } + GGML_ASSERT(ok); + }; + + auto split_trace_interesting = [&]() -> bool { + switch (tensor->op) { + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: + case GGML_OP_CONCAT: + case GGML_OP_VIEW: + case GGML_OP_RESHAPE: + case GGML_OP_CONT: + case GGML_OP_PERMUTE: + case GGML_OP_GET_ROWS: + case GGML_OP_SET_ROWS: + case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: + case GGML_OP_TOP_K: + case GGML_OP_FLASH_ATTN_EXT: + case GGML_OP_LIGHTNING_INDEXER: + case GGML_OP_DSV4_HC_COMB: + case GGML_OP_DSV4_HC_PRE: + case GGML_OP_DSV4_HC_POST: + return true; + default: + return false; + } + }; + auto handle_generic = [&](const std::vector & src_ss, bool scalar_only) -> ggml_backend_meta_split_state { ggml_backend_meta_split_state ret = {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, {1}, 1}; for (size_t i = 0; i < GGML_MAX_SRC; i++) { @@ -532,13 +662,14 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (scalar_only && ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) { ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } - GGML_ASSERT(ret.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); + check_split(ret.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN, "GENERIC_ERROR", src_ss, + "scalar_only=" + std::to_string(scalar_only ? 1 : 0)); return ret; }; // Some ops process data on a per-row bases: auto handle_per_row = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0); + check_split(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0, "PER_ROW_ERROR", src_ss); return src_ss[0]; }; @@ -552,18 +683,21 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL)))) { return src_ss[0]; // GGML_OP_ADD_ID } - GGML_ASSERT(tensor->src[2] == nullptr || src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + check_split(tensor->src[2] == nullptr || src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, + "BIN_BCAST_SRC2_ERROR", src_ss); return handle_generic(src_ss, /*scalar_only =*/ false); }; auto handle_concat = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { const ggml_backend_meta_split_axis concat_axis = ggml_backend_meta_split_axis(ggml_get_op_params_i32(tensor, 0)); if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis >= 0 && src_ss[1].axis < GGML_MAX_DIMS) { - GGML_ASSERT(concat_axis != src_ss[1].axis); + check_split(concat_axis != src_ss[1].axis, "CONCAT_AXIS_ERROR", src_ss, + "concat_axis=" + std::to_string(concat_axis)); return src_ss[1]; } if (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) { - GGML_ASSERT(concat_axis != src_ss[0].axis); + check_split(concat_axis != src_ss[0].axis, "CONCAT_AXIS_ERROR", src_ss, + "concat_axis=" + std::to_string(concat_axis)); return src_ss[0]; } if (src_ss[0].axis == src_ss[1].axis && src_ss[0].axis != concat_axis) { @@ -587,9 +721,14 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return src_ss[1]; } if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_0) { - GGML_ASSERT(split_states_equal(src_ss[0], src_ss[1])); + check_split(split_states_equal(src_ss[0], src_ss[1]), "MUL_MAT_RATIO_ERROR", src_ss); return {assume_sync ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_PARTIAL, {0}, {1}, 1}; } + if (src_ss[0].axis == src_ss[1].axis && src_ss[0].axis >= GGML_BACKEND_SPLIT_AXIS_2 && src_ss[0].axis < GGML_MAX_DIMS) { + check_split(split_states_equal(src_ss[0], src_ss[1]), "MUL_MAT_BATCH_RATIO_ERROR", src_ss); + return src_ss[0]; + } + log_split_failure("MUL_MAT_UNSUPPORTED", src_ss, "assume_sync=" + std::to_string(assume_sync ? 1 : 0)); GGML_ABORT("fatal error"); //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; }; @@ -600,7 +739,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( case GGML_BACKEND_SPLIT_AXIS_1: case GGML_BACKEND_SPLIT_AXIS_2: case GGML_BACKEND_SPLIT_AXIS_3: { - GGML_ASSERT(src_ss[0].n_segments == 1); + check_split(src_ss[0].n_segments == 1, "RESHAPE_SEGMENT_ERROR", src_ss); if (src_ss[0].axis == ggml_n_dims(tensor->src[0]) - 1 && src_ss[0].nr[0] == 1) { return {ggml_backend_meta_split_axis(ggml_n_dims(tensor) - 1), {0}, {1}, 1}; } @@ -616,12 +755,13 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return {ggml_backend_meta_split_axis(dim), {0}, {uint32_t(base_ne_out_next/base_ne_in)}, 1}; } if (base_ne_out_next > base_ne_in) { - GGML_ASSERT(src_ss[0].n_segments == 1); - GGML_ASSERT(src_ss[0].nr[0] == 1); + check_split(src_ss[0].n_segments == 1, "RESHAPE_SEGMENT_ERROR", src_ss); + check_split(src_ss[0].nr[0] == 1, "RESHAPE_REPEAT_ERROR", src_ss); return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1}; } base_ne_out = base_ne_out_next; } + log_split_failure("RESHAPE_SHAPE_MISMATCH", src_ss, "base_ne_in=" + std::to_string(base_ne_in)); GGML_ABORT("shape mismatch for %s", ggml_op_name(tensor->op)); } case GGML_BACKEND_SPLIT_AXIS_MIRRORED: @@ -629,6 +769,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return src_ss[0]; } default: { + log_split_failure("RESHAPE_UNSUPPORTED", src_ss); GGML_ABORT("fatal error"); //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } @@ -668,11 +809,13 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1}; } } + log_split_failure("VIEW_STRIDE_ERROR", src_ss); GGML_ABORT("fatal error"); } if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED || src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) { return src_ss[0]; } + log_split_failure("VIEW_PERMUTED_UNSUPPORTED", src_ss); GGML_ABORT("view of permuted tensor not implemented"); //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; }; @@ -683,7 +826,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( case GGML_BACKEND_SPLIT_AXIS_1: case GGML_BACKEND_SPLIT_AXIS_2: case GGML_BACKEND_SPLIT_AXIS_3: { - GGML_ASSERT(src_ss[0].n_segments == 1 || src_ss[0].nr[0] == 1); + check_split(src_ss[0].n_segments == 1 || src_ss[0].nr[0] == 1, "PERMUTE_SEGMENT_ERROR", src_ss); return {ggml_backend_meta_split_axis(tensor->op_params[src_ss[0].axis]), {0}, {src_ss[0].nr[0]}, 1}; } case GGML_BACKEND_SPLIT_AXIS_MIRRORED: @@ -691,6 +834,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return src_ss[0]; } default: { + log_split_failure("PERMUTE_UNSUPPORTED", src_ss); GGML_ABORT("fatal error"); //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } @@ -701,7 +845,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( switch (src_ss[0].axis) { case GGML_BACKEND_SPLIT_AXIS_0: case GGML_BACKEND_SPLIT_AXIS_1: { - GGML_ASSERT(src_ss[0].n_segments == 1 || src_ss[0].nr[0] == 1); + check_split(src_ss[0].n_segments == 1 || src_ss[0].nr[0] == 1, "TRANSPOSE_SEGMENT_ERROR", src_ss); return {ggml_backend_meta_split_axis(int(src_ss[0].axis) ^ 1), {0}, {src_ss[0].nr[0]}, 1}; } case GGML_BACKEND_SPLIT_AXIS_2: @@ -711,6 +855,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return src_ss[0]; } default: { + log_split_failure("TRANSPOSE_UNSUPPORTED", src_ss); GGML_ABORT("fatal error"); //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } @@ -725,31 +870,51 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( }; auto handle_set_rows = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_1); - GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); - GGML_ASSERT(split_states_equal(src_ss[0], src_ss[2])); + check_split(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_1, "SET_ROWS_DST_AXIS_ERROR", src_ss); + check_split(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, "SET_ROWS_SRC_AXIS_ERROR", src_ss); + check_split(split_states_equal(src_ss[0], src_ss[2]), "SET_ROWS_INDEX_RATIO_ERROR", src_ss); return src_ss[0]; }; auto handle_rope = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + check_split(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, "ROPE_POS_ERROR", src_ss); return src_ss[0]; }; auto handle_pad = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { if (src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) { - GGML_ASSERT(tensor->op_params[2*src_ss[0].axis + 0] == 0); - GGML_ASSERT(tensor->op_params[2*src_ss[0].axis + 1] == 0); + check_split(tensor->op_params[2*src_ss[0].axis + 0] == 0, "PAD_BEGIN_ERROR", src_ss); + check_split(tensor->op_params[2*src_ss[0].axis + 1] == 0, "PAD_END_ERROR", src_ss); } return src_ss[0]; }; auto handle_flash_attn_ext = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - GGML_ASSERT( src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2); - GGML_ASSERT( src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2); - GGML_ASSERT( src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2); - GGML_ASSERT(tensor->src[4] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); - GGML_ASSERT(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_0); + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && + src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && + src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + check_split(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, + "FLASH_ATTN_MASK_ERROR", src_ss); + check_split(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, + "FLASH_ATTN_SINK_ERROR", src_ss); + return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; + } + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2 && + src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && + src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + check_split(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, + "FLASH_ATTN_MASK_ERROR", src_ss); + check_split(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_0, + "FLASH_ATTN_SINK_ERROR", src_ss); + return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; + } + check_split(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2, "FLASH_ATTN_Q_AXIS_ERROR", src_ss); + check_split(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2, "FLASH_ATTN_K_AXIS_ERROR", src_ss); + check_split(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2, "FLASH_ATTN_V_AXIS_ERROR", src_ss); + check_split(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, + "FLASH_ATTN_MASK_ERROR", src_ss); + check_split(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_0, + "FLASH_ATTN_SINK_ERROR", src_ss); return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; }; @@ -771,17 +936,26 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { return src_ss[0]; } - GGML_ASSERT(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1); - GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_1); - GGML_ASSERT(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_1); - GGML_ASSERT(src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_1); - GGML_ASSERT(src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_1); + check_split(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1, "GDN_SRC0_AXIS_ERROR", src_ss); + check_split(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_1, "GDN_SRC1_AXIS_ERROR", src_ss); + check_split(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_1, "GDN_SRC2_AXIS_ERROR", src_ss); + check_split(src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_1, "GDN_SRC3_AXIS_ERROR", src_ss); + check_split(src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_1, "GDN_SRC4_AXIS_ERROR", src_ss); // state shape is [S_v, S_v, H_v, n_seqs] (s0 only); the heads dim is its own axis 2, // so a head-aligned split on the input cache lands on axis 2 here. - GGML_ASSERT(src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_2 || src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_1 || src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_0); + check_split(src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_2 || src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_1 || src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_0, + "GDN_STATE_AXIS_ERROR", src_ss); return {GGML_BACKEND_SPLIT_AXIS_0, {0}, {1}, 1}; }; + auto handle_lightning_indexer = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + check_split(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, "LIGHTNING_INDEXER_Q_AXIS_ERROR", src_ss); + check_split(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, "LIGHTNING_INDEXER_K_AXIS_ERROR", src_ss); + check_split(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, "LIGHTNING_INDEXER_WEIGHT_AXIS_ERROR", src_ss); + check_split(src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, "LIGHTNING_INDEXER_MASK_AXIS_ERROR", src_ss); + return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; + }; + auto calculate_split_state = [&]() -> ggml_backend_meta_split_state { if (ggml_nelements(tensor) == 0) { return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; @@ -793,12 +967,21 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (ret.axis >= 0 && ret.axis <= GGML_MAX_DIMS) { const int64_t granularity = ret.axis == GGML_BACKEND_SPLIT_AXIS_0 ? ggml_blck_size(tensor->type) : 1; int64_t ne_sum = 0; + std::vector empty_src_ss(GGML_MAX_SRC, {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, {1}, 1}); for (size_t s = 0; s < ret.n_segments; s++) { for (size_t j = 0; j < n_bufs; j++) { + if (ret.ne[s*n_bufs + j] % granularity != 0) { + ggml_backend_meta_log_split_context("STATIC_GRANULARITY_ERROR", tensor, empty_src_ss, n_bufs, + "state={" + ggml_backend_meta_split_state_str(ret, n_bufs) + "}, granularity=" + std::to_string(granularity)); + } GGML_ASSERT(ret.ne[s*n_bufs + j] % granularity == 0); ne_sum += ret.ne[s*n_bufs + j] * ret.nr[s]; } } + if (ne_sum != tensor->ne[ret.axis]) { + ggml_backend_meta_log_split_context("STATIC_SIZE_ERROR", tensor, empty_src_ss, n_bufs, + "state={" + ggml_backend_meta_split_state_str(ret, n_bufs) + "}, ne_sum=" + std::to_string(ne_sum)); + } GGML_ASSERT(ne_sum == tensor->ne[ret.axis]); } return ret; @@ -811,6 +994,9 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( continue; } src_ss[i] = ggml_backend_meta_get_split_state(stc, tensor->src[i], /*assume_sync =*/ true); + if (src_ss[i].axis == GGML_BACKEND_SPLIT_AXIS_UNKNOWN) { + log_split_failure("SRC_UNKNOWN", src_ss, "src=" + std::to_string(i)); + } GGML_ASSERT(src_ss[i].axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); } @@ -920,7 +1106,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( split_state = handle_rope(src_ss); } break; case GGML_OP_ROPE_BACK: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); + split_state = handle_rope(src_ss); } break; case GGML_OP_CLAMP: { split_state = handle_generic(src_ss, /*scalar_only =*/ false); @@ -984,6 +1170,9 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( case GGML_OP_GATED_DELTA_NET: { split_state = handle_gated_delta_net(src_ss); } break; + case GGML_OP_LIGHTNING_INDEXER: { + split_state = handle_lightning_indexer(src_ss); + } break; case GGML_OP_DSV4_HC_COMB: case GGML_OP_DSV4_HC_PRE: case GGML_OP_DSV4_HC_POST: { @@ -1008,6 +1197,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( split_state = handle_generic(src_ss, /*scalar_only =*/ false); } break; default: { + log_split_failure("OP_UNIMPLEMENTED", src_ss); GGML_ABORT("ggml op not implemented: %s", ggml_op_name(tensor->op)); split_state = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } break; @@ -1032,24 +1222,43 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( split_state.ne[j] *= tensor->ne[split_state.axis]; if (split_state.ne[j] != 0 || tensor->src[i]->ne[src_ss[i].axis] != 0) { const int64_t div = tensor->src[i]->ne[src_ss[i].axis] * split_state.nr[0]; + if (div == 0 || split_state.ne[j] % div != 0) { + log_split_failure("RATIO_DIV_ERROR", src_ss, + "src=" + std::to_string(i) + ", buf=" + std::to_string(j) + + ", div=" + std::to_string(div) + ", value=" + std::to_string(split_state.ne[j]) + + ", dst_state={" + ggml_backend_meta_split_state_str(split_state, n_bufs) + "}"); + } + GGML_ASSERT(div != 0); GGML_ASSERT(split_state.ne[j] % div == 0); split_state.ne[j] /= div; } } } else { - GGML_ASSERT(split_state.n_segments == 1); + check_split(split_state.n_segments == 1, "RATIO_SEGMENT_ERROR", src_ss, + "dst_state={" + ggml_backend_meta_split_state_str(split_state, n_bufs) + "}"); for (size_t j = 0; j < n_bufs; j++) { // Assert that ratio is consistent: int64_t sum = 0; for (size_t s = 0; s < src_ss[i].n_segments; s++) { sum += src_ss[i].ne[s*n_bufs + j] * src_ss[i].nr[s]; } - GGML_ASSERT(split_state.ne[j]*split_state.nr[0] * tensor->src[i]->ne[src_ss[i].axis] - == sum * tensor->ne[split_state.axis]); + const int64_t lhs = split_state.ne[j]*split_state.nr[0] * tensor->src[i]->ne[src_ss[i].axis]; + const int64_t rhs = sum * tensor->ne[split_state.axis]; + if (lhs != rhs) { + log_split_failure("RATIO_ERROR", src_ss, + "src=" + std::to_string(i) + ", buf=" + std::to_string(j) + + ", lhs=" + std::to_string(lhs) + ", rhs=" + std::to_string(rhs) + + ", dst_state={" + ggml_backend_meta_split_state_str(split_state, n_bufs) + "}"); + } + GGML_ASSERT(lhs == rhs); } } first_src_split_by_axis = false; } + if (first_src_split_by_axis) { + log_split_failure("NO_SPLIT_SOURCE", src_ss, + "dst_state={" + ggml_backend_meta_split_state_str(split_state, n_bufs) + "}"); + } GGML_ASSERT(!first_src_split_by_axis); } return split_state; @@ -1065,37 +1274,23 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (it == buf_ctx->split_state_cache.end()) { buf_ctx->split_state_cache[key].first = calculate_split_state(); memcpy(buf_ctx->split_state_cache[key].second, tensor, sizeof(buf_ctx->split_state_cache[key].second)); - if (buf_ctx->debug > 0) { + if (buf_ctx->debug > 1 || (buf_ctx->debug > 0 && split_trace_interesting())) { std::string srcs_info; for (size_t i = 0; i < GGML_MAX_SRC; i++) { - if (tensor->src[i] == nullptr) { + if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { continue; } if (!srcs_info.empty()) { srcs_info += ", "; } - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor->src[0], true); - GGML_ASSERT(split_state.n_segments == 1); - const char * axis_name = ggml_backend_meta_split_axis_name(split_state.axis); - std::string ne_info; - for (size_t j = 0; j < n_bufs; j++) { - if (!ne_info.empty()) { - ne_info += ", "; - } - ne_info += std::to_string(split_state.ne[j]) + "x" + std::to_string(split_state.nr[0]); - } - srcs_info += std::string(tensor->src[i]->name) + "[" + ggml_op_name(tensor->src[i]->op) + ", " + axis_name + ", {" + ne_info + "}]"; - } - std::string ne_info; - for (size_t j = 0; j < n_bufs; j++) { - if (!ne_info.empty()) { - ne_info += ", "; - } - const ggml_backend_meta_split_state & ss = buf_ctx->split_state_cache[key].first; - ne_info += std::to_string(ss.ne[j]) + "x" + std::to_string(ss.nr[0]); + const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor->src[i], true); + srcs_info += "src[" + std::to_string(i) + "]=" + ggml_backend_meta_tensor_str(tensor->src[i]) + + ", state={" + ggml_backend_meta_split_state_str(split_state, n_bufs) + "}"; } - GGML_LOG_DEBUG("SPLIT_STATE: {%s} -> %s[%s, %s, {%s}]\n", srcs_info.c_str(), tensor->name, ggml_op_name(tensor->op), - ggml_backend_meta_split_axis_name(buf_ctx->split_state_cache[key].first.axis), ne_info.c_str()); + const ggml_backend_meta_split_state & ss = buf_ctx->split_state_cache[key].first; + GGML_LOG_DEBUG("META_SPLIT_TRACE: dst=%s, assume_sync=%d, state={%s}, srcs={%s}\n", + ggml_backend_meta_tensor_str(tensor).c_str(), assume_sync ? 1 : 0, + ggml_backend_meta_split_state_str(ss, n_bufs).c_str(), srcs_info.c_str()); } } @@ -1132,7 +1327,16 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m const size_t n_simple_bufs = ggml_backend_meta_buffer_n_bufs(tensor->buffer); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(stc, tensor, /*assume_sync =*/ true); + std::vector empty_src_ss(GGML_MAX_SRC, {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, {1}, 1}); + if (ggml_nelements(tensor) != 0 && split_state.axis == GGML_BACKEND_SPLIT_AXIS_UNKNOWN) { + ggml_backend_meta_log_split_context("INIT_UNKNOWN", tensor, empty_src_ss, n_simple_bufs, + "state={" + ggml_backend_meta_split_state_str(split_state, n_simple_bufs) + "}"); + } GGML_ASSERT(ggml_nelements(tensor) == 0 || split_state.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); + if (split_state.n_segments > 16) { + ggml_backend_meta_log_split_context("INIT_SEGMENTS_ERROR", tensor, empty_src_ss, n_simple_bufs, + "state={" + ggml_backend_meta_split_state_str(split_state, n_simple_bufs) + "}"); + } GGML_ASSERT(split_state.n_segments <= 16); int split_dim = split_state.axis; @@ -1182,8 +1386,18 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m if (t_ij->view_src != nullptr && ggml_backend_buffer_is_meta(t_ij->view_src->buffer)) { t_ij->view_src = ggml_backend_meta_buffer_simple_tensor(tensor->view_src, j); if (t_ij->view_offs > 0 && split_dim >= 0 && split_dim < GGML_MAX_DIMS) { + if (tensor->ne[split_dim] == 0) { + ggml_backend_meta_log_split_context("INIT_VIEW_ZERO_DIM", tensor, empty_src_ss, n_simple_bufs, + "state={" + ggml_backend_meta_split_state_str(split_state, n_simple_bufs) + "}"); + } GGML_ASSERT(tensor->ne[split_dim] != 0); - const int split_dim_view_src = ggml_backend_meta_get_split_state(tensor->view_src, /*assume_sync =*/ true).axis; + const ggml_backend_meta_split_state view_src_ss = ggml_backend_meta_get_split_state(tensor->view_src, /*assume_sync =*/ true); + const int split_dim_view_src = view_src_ss.axis; + if (split_dim_view_src < 0 || split_dim_view_src >= GGML_MAX_DIMS) { + ggml_backend_meta_log_split_context("INIT_VIEW_SRC_AXIS_ERROR", tensor, empty_src_ss, n_simple_bufs, + "state={" + ggml_backend_meta_split_state_str(split_state, n_simple_bufs) + "}, view_src_state={" + + ggml_backend_meta_split_state_str(view_src_ss, n_simple_bufs) + "}"); + } GGML_ASSERT(split_dim_view_src >= 0 && split_dim_view_src < GGML_MAX_DIMS); // The offset can be internal to the data split, in those cases the view offset should not be scaled. @@ -1253,6 +1467,102 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor(ggml_backend_buffer return ggml_backend_meta_buffer_init_tensor_impl(buf_ctx->get_simple_tensor_container(tensor), tensor); } +static void ggml_backend_meta_buffer_memset_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) { + const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); + const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); + GGML_ASSERT(ggml_is_contiguous(tensor) || split_state.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + + if (split_state.n_segments != 1 || split_state.nr[0] != 1) { + GGML_ASSERT(split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS); + GGML_ASSERT(split_state.nr[0] != 0); + GGML_ASSERT(tensor->ne[3] == 1); + + std::vector simple_offsets(n_bufs, 0); + if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_0) { + GGML_ASSERT(tensor->ne[2] == 1); + + const size_t row_stride = tensor->nb[1]; + GGML_ASSERT(offset % row_stride == 0); + GGML_ASSERT(size % row_stride == 0); + const int64_t row_start = offset / row_stride; + const int64_t row_count = size / row_stride; + GGML_ASSERT(row_start + row_count <= tensor->ne[1]); + + const int64_t blck_size = ggml_blck_size(tensor->type); + for (size_t s = 0; s < split_state.n_segments; s++) { + for (size_t r = 0; r < split_state.nr[s]; r++) { + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + GGML_ASSERT(split_state.ne[s*n_bufs + j] % blck_size == 0); + const size_t nbytes = split_state.ne[s*n_bufs + j]/blck_size * tensor->nb[0]; + for (int64_t row = 0; row < row_count; row++) { + ggml_backend_tensor_memset(simple_tensor, value, + simple_offsets[j] + (row_start + row)*simple_tensor->nb[1], nbytes); + } + simple_offsets[j] += nbytes; + } + } + } + return; + } + GGML_ASSERT(split_state.axis == GGML_BACKEND_SPLIT_AXIS_1); + + const size_t row_stride = tensor->nb[2]; + GGML_ASSERT(offset % row_stride == 0); + GGML_ASSERT(size % row_stride == 0); + const int64_t row_start = offset / row_stride; + const int64_t row_count = size / row_stride; + GGML_ASSERT(row_start + row_count <= tensor->ne[2]); + + for (size_t s = 0; s < split_state.n_segments; s++) { + for (size_t r = 0; r < split_state.nr[s]; r++) { + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t nbytes = split_state.ne[s*n_bufs + j] * tensor->nb[1]; + for (int64_t row = 0; row < row_count; row++) { + ggml_backend_tensor_memset(simple_tensor, value, + simple_offsets[j] + (row_start + row)*simple_tensor->nb[2], nbytes); + } + simple_offsets[j] += nbytes; + } + } + } + return; + } + + switch (split_state.axis) { + case GGML_BACKEND_SPLIT_AXIS_0: + case GGML_BACKEND_SPLIT_AXIS_1: + case GGML_BACKEND_SPLIT_AXIS_2: { + const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; + GGML_ASSERT(offset % chunk_size_full == 0); + GGML_ASSERT(size % chunk_size_full == 0); + const int64_t i_start = offset /chunk_size_full; + const int64_t i_stop = (offset + size)/chunk_size_full; + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; + if (chunk_size_j == 0) { + continue; + } + for (int64_t i = i_start; i < i_stop; i++) { + ggml_backend_tensor_memset(simple_tensor, value, i*chunk_size_j, chunk_size_j); + } + } + } break; + case GGML_BACKEND_SPLIT_AXIS_MIRRORED: + case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + ggml_backend_tensor_memset(simple_tensor, value, offset, size); + } + } break; + default: { + GGML_ABORT("fatal error"); + } + } +} + static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); @@ -1486,7 +1796,7 @@ static const ggml_backend_buffer_i ggml_backend_meta_buffer_iface = { /* .free_buffer = */ ggml_backend_meta_buffer_free_buffer, /* .get_base = */ ggml_backend_meta_buffer_get_base, /* .init_tensor = */ ggml_backend_meta_buffer_init_tensor, - /* .memset_tensor = */ nullptr, // TODO implement + /* .memset_tensor = */ ggml_backend_meta_buffer_memset_tensor, /* .set_tensor = */ ggml_backend_meta_buffer_set_tensor, /* .get_tensor = */ ggml_backend_meta_buffer_get_tensor, /* .set_tensor_2d = */ nullptr, @@ -1517,6 +1827,9 @@ static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_bac bufs.reserve(n_simple_bufts); for (size_t i = 0; i < n_simple_bufts; i++) { bufs.push_back(ggml_backend_buft_alloc_buffer(ggml_backend_meta_buft_simple_buft(buft, i), size)); + if (bufs.back() == nullptr) { + GGML_LOG_ERROR("META_ALLOC_ERROR: simple_buft=%zu, requested_size=%zu\n", i, size); + } GGML_ASSERT(bufs.back() != nullptr); max_size = std::max(max_size, ggml_backend_buffer_get_size(bufs.back())); } @@ -1573,6 +1886,9 @@ struct ggml_backend_buffer * ggml_backend_meta_alloc_ctx_tensors_from_buft(struc t->buffer = meta_buf_ctx->bufs[i].get(); } } + if (!meta_buf_ctx->bufs[i]) { + GGML_LOG_ERROR("META_ALLOC_CTX_ERROR: simple_buft=%zu\n", i); + } GGML_ASSERT(meta_buf_ctx->bufs[i]); meta_buf->size = std::max(meta_buf->size, ggml_backend_buffer_get_size(meta_buf_ctx->bufs[i].get())); } diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 72968607db80..60fb33ce6612 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -989,7 +989,6 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_OLMOE: case LLM_ARCH_DEEPSEEK2: case LLM_ARCH_DEEPSEEK32: - case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_GLM_DSA: case LLM_ARCH_BITNET: case LLM_ARCH_T5: diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 4c10e4126f62..2f5d3f12e0f2 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -345,9 +345,17 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_qkv_bias ("blk\\.\\d*\\.attn_qkv.bias"); static const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight"); static const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*"); + static const std::regex pattern_dsv4_state ("dsv4_.*_state_(kv|score)_l\\d*"); static const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight"); static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight"); static const std::regex pattern_attn_out_bias ("blk\\.\\d*\\.attn_output.bias"); + static const std::regex pattern_attn_out_a_weight("blk\\.\\d*\\.attn_output_a\\.weight"); + static const std::regex pattern_attn_out_b_weight("blk\\.\\d*\\.attn_output_b\\.weight"); + static const std::regex pattern_attn_out_b_bias ("blk\\.\\d*\\.attn_output_b\\.bias"); + static const std::regex pattern_attn_q_a_weight ("blk\\.\\d*\\.attn_q_a\\.weight"); + static const std::regex pattern_attn_q_b_weight ("blk\\.\\d*\\.attn_q_b\\.weight"); + static const std::regex pattern_attn_q_b_bias ("blk\\.\\d*\\.attn_q_b\\.bias"); + static const std::regex pattern_attn_kv_weight ("blk\\.\\d*\\.attn_kv\\.weight"); static const std::regex pattern_attn_gate_weight("blk\\.\\d*\\.attn_gate.weight"); static const std::regex pattern_ssm_dt ("blk\\.\\d*\\.ssm_dt.bias"); @@ -440,6 +448,14 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str if (std::regex_match(tensor_name, pattern_qk_norm)) { return get_tensor_config_impl(tensor->ne[1] == 1 ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight"); } + if (ud->model->arch == LLM_ARCH_DEEPSEEK4 && + (std::regex_match(tensor_name, pattern_kv_cache) || + std::regex_match(tensor_name, pattern_dsv4_state))) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } + if (ud->model->arch == LLM_ARCH_DEEPSEEK4 && std::regex_match(tensor_name, pattern_attn_sinks)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output_a.weight"); + } if (std::regex_match(tensor_name, pattern_kv_cache) || std::regex_match(tensor_name, pattern_attn_sinks)) { return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output.weight"); } @@ -449,6 +465,26 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str if (std::regex_match(tensor_name, pattern_attn_out_bias)) { return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); } + if (ud->model->arch == LLM_ARCH_DEEPSEEK4) { + if (std::regex_match(tensor_name, pattern_attn_q_b_weight)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output_a.weight", "attn_output.weight"); + } + if (std::regex_match(tensor_name, pattern_attn_q_b_bias)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output_a.weight", "attn_output.weight"); + } + if (std::regex_match(tensor_name, pattern_attn_out_a_weight)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output_b.weight"); + } + if (std::regex_match(tensor_name, pattern_attn_out_b_weight)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0); + } + if (std::regex_match(tensor_name, pattern_attn_out_b_bias)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } + if (std::regex_match(tensor_name, pattern_attn_q_a_weight) || std::regex_match(tensor_name, pattern_attn_kv_weight)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } + } if (std::regex_match(tensor_name, pattern_attn_gate_weight)) { return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight", "ssm_out.weight"); @@ -611,9 +647,23 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str if (std::regex_match(tensor_name, pattern_attn_sinks)) { GGML_ASSERT(segments.size() == 1); + if (ud->model->arch == LLM_ARCH_DEEPSEEK4) { + return {1}; + } return {std::lcm(n_embd_q, blck_size_perf)/n_embd_q * n_gqa}; } + if (ud->model->arch == LLM_ARCH_DEEPSEEK4) { + if (std::regex_match(tensor_name, pattern_attn_q_b_weight) || std::regex_match(tensor_name, pattern_attn_q_b_bias)) { + GGML_ASSERT(segments.size() == 1); + return {hparams.n_embd_head_k(il)}; + } + if (std::regex_match(tensor_name, pattern_attn_out_a_weight) || std::regex_match(tensor_name, pattern_attn_out_b_weight)) { + GGML_ASSERT(segments.size() == 1); + return {std::lcm(hparams.dsv4_o_lora_rank, blck_size)}; + } + } + const int64_t granularity_q = std::lcm(n_embd_q, blck_size_perf); if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_q_bias)) { GGML_ASSERT(segments.size() == 1); From 3515fafd813144ea833a04c8c2dd0b2c4b110bcf Mon Sep 17 00:00:00 2001 From: s3kundes Date: Sat, 18 Jul 2026 17:12:38 +0300 Subject: [PATCH 2/2] sol refactor and free some vram --- ggml/src/ggml-backend-meta.cpp | 364 ++++++++------------------------- src/llama-model.cpp | 80 ++++---- tests/test-llama-archs.cpp | 50 ++++- 3 files changed, 170 insertions(+), 324 deletions(-) diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 09a2c6b927f3..b4cde81136d1 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -10,7 +10,6 @@ #include #include #include -#include #include #include #include @@ -48,99 +47,6 @@ const char * ggml_backend_meta_split_axis_name(enum ggml_backend_meta_split_axis } } -static const char * ggml_backend_meta_tensor_name(const ggml_tensor * tensor) { - return tensor && tensor->name[0] ? tensor->name : "(unnamed)"; -} - -static std::string ggml_backend_meta_shape_str(const ggml_tensor * tensor) { - if (tensor == nullptr) { - return "(null)"; - } - std::string ret = "ne={"; - for (int i = 0; i < GGML_MAX_DIMS; i++) { - if (i > 0) { - ret += ","; - } - ret += std::to_string(tensor->ne[i]); - } - ret += "}, nb={"; - for (int i = 0; i < GGML_MAX_DIMS; i++) { - if (i > 0) { - ret += ","; - } - ret += std::to_string(tensor->nb[i]); - } - ret += "}"; - return ret; -} - -static std::string ggml_backend_meta_tensor_str(const ggml_tensor * tensor) { - if (tensor == nullptr) { - return "(null)"; - } - std::string ret = std::string(ggml_backend_meta_tensor_name(tensor)) + "[" + ggml_op_name(tensor->op) + - ", " + ggml_type_name(tensor->type) + ", " + ggml_backend_meta_shape_str(tensor); - if (tensor->view_src != nullptr) { - ret += ", view_src="; - ret += ggml_backend_meta_tensor_name(tensor->view_src); - ret += ", view_offs="; - ret += std::to_string(tensor->view_offs); - } - ret += "]"; - return ret; -} - -static std::string ggml_backend_meta_split_state_str(const ggml_backend_meta_split_state & ss, size_t n_bufs) { - std::string ret = std::string(ggml_backend_meta_split_axis_name(ss.axis)) + ", n_segments=" + std::to_string(ss.n_segments); - ret += ", sums={"; - for (size_t j = 0; j < n_bufs; j++) { - if (j > 0) { - ret += ","; - } - int64_t sum = 0; - for (size_t s = 0; s < ss.n_segments; s++) { - sum += ss.ne[s*n_bufs + j] * ss.nr[s]; - } - ret += std::to_string(sum); - } - ret += "}"; - if (ss.axis >= 0 && ss.axis < GGML_MAX_DIMS) { - ret += ", segments={"; - for (size_t s = 0; s < ss.n_segments; s++) { - if (s > 0) { - ret += ";"; - } - ret += "nr=" + std::to_string(ss.nr[s]) + ":"; - for (size_t j = 0; j < n_bufs; j++) { - if (j > 0) { - ret += ","; - } - ret += std::to_string(ss.ne[s*n_bufs + j]); - } - } - ret += "}"; - } - return ret; -} - -static void ggml_backend_meta_log_split_context( - const char * label, - const ggml_tensor * tensor, - const std::vector & src_ss, - size_t n_bufs, - const std::string & detail) { - GGML_LOG_ERROR("META_SPLIT_%s: dst=%s%s%s\n", label, ggml_backend_meta_tensor_str(tensor).c_str(), - detail.empty() ? "" : ", ", detail.c_str()); - for (size_t i = 0; i < GGML_MAX_SRC; i++) { - if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { - continue; - } - GGML_LOG_ERROR("META_SPLIT_%s: src[%zu]=%s, state={%s}\n", label, i, - ggml_backend_meta_tensor_str(tensor->src[i]).c_str(), - ggml_backend_meta_split_state_str(src_ss[i], n_bufs).c_str()); - } -} - // // meta backend device // @@ -607,42 +513,6 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return true; }; - auto log_split_failure = [&](const char * label, const std::vector & src_ss, const std::string & detail = std::string()) { - ggml_backend_meta_log_split_context(label, tensor, src_ss, n_bufs, detail); - }; - - auto check_split = [&](bool ok, const char * label, const std::vector & src_ss, const std::string & detail = std::string()) { - if (!ok) { - log_split_failure(label, src_ss, detail); - } - GGML_ASSERT(ok); - }; - - auto split_trace_interesting = [&]() -> bool { - switch (tensor->op) { - case GGML_OP_MUL_MAT: - case GGML_OP_MUL_MAT_ID: - case GGML_OP_CONCAT: - case GGML_OP_VIEW: - case GGML_OP_RESHAPE: - case GGML_OP_CONT: - case GGML_OP_PERMUTE: - case GGML_OP_GET_ROWS: - case GGML_OP_SET_ROWS: - case GGML_OP_ROPE: - case GGML_OP_ROPE_BACK: - case GGML_OP_TOP_K: - case GGML_OP_FLASH_ATTN_EXT: - case GGML_OP_LIGHTNING_INDEXER: - case GGML_OP_DSV4_HC_COMB: - case GGML_OP_DSV4_HC_PRE: - case GGML_OP_DSV4_HC_POST: - return true; - default: - return false; - } - }; - auto handle_generic = [&](const std::vector & src_ss, bool scalar_only) -> ggml_backend_meta_split_state { ggml_backend_meta_split_state ret = {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, {1}, 1}; for (size_t i = 0; i < GGML_MAX_SRC; i++) { @@ -662,14 +532,13 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (scalar_only && ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) { ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } - check_split(ret.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN, "GENERIC_ERROR", src_ss, - "scalar_only=" + std::to_string(scalar_only ? 1 : 0)); + GGML_ASSERT(ret.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); return ret; }; // Some ops process data on a per-row bases: auto handle_per_row = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - check_split(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0, "PER_ROW_ERROR", src_ss); + GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0); return src_ss[0]; }; @@ -683,21 +552,18 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL)))) { return src_ss[0]; // GGML_OP_ADD_ID } - check_split(tensor->src[2] == nullptr || src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, - "BIN_BCAST_SRC2_ERROR", src_ss); + GGML_ASSERT(tensor->src[2] == nullptr || src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); return handle_generic(src_ss, /*scalar_only =*/ false); }; auto handle_concat = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { const ggml_backend_meta_split_axis concat_axis = ggml_backend_meta_split_axis(ggml_get_op_params_i32(tensor, 0)); if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis >= 0 && src_ss[1].axis < GGML_MAX_DIMS) { - check_split(concat_axis != src_ss[1].axis, "CONCAT_AXIS_ERROR", src_ss, - "concat_axis=" + std::to_string(concat_axis)); + GGML_ASSERT(concat_axis != src_ss[1].axis); return src_ss[1]; } if (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) { - check_split(concat_axis != src_ss[0].axis, "CONCAT_AXIS_ERROR", src_ss, - "concat_axis=" + std::to_string(concat_axis)); + GGML_ASSERT(concat_axis != src_ss[0].axis); return src_ss[0]; } if (src_ss[0].axis == src_ss[1].axis && src_ss[0].axis != concat_axis) { @@ -721,14 +587,14 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return src_ss[1]; } if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_0) { - check_split(split_states_equal(src_ss[0], src_ss[1]), "MUL_MAT_RATIO_ERROR", src_ss); + GGML_ASSERT(split_states_equal(src_ss[0], src_ss[1])); return {assume_sync ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_PARTIAL, {0}, {1}, 1}; } - if (src_ss[0].axis == src_ss[1].axis && src_ss[0].axis >= GGML_BACKEND_SPLIT_AXIS_2 && src_ss[0].axis < GGML_MAX_DIMS) { - check_split(split_states_equal(src_ss[0], src_ss[1]), "MUL_MAT_BATCH_RATIO_ERROR", src_ss); + if (src_ss[0].axis == src_ss[1].axis && src_ss[0].axis >= GGML_BACKEND_SPLIT_AXIS_2 && + src_ss[0].axis < GGML_MAX_DIMS) { + GGML_ASSERT(split_states_equal(src_ss[0], src_ss[1])); return src_ss[0]; } - log_split_failure("MUL_MAT_UNSUPPORTED", src_ss, "assume_sync=" + std::to_string(assume_sync ? 1 : 0)); GGML_ABORT("fatal error"); //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; }; @@ -739,7 +605,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( case GGML_BACKEND_SPLIT_AXIS_1: case GGML_BACKEND_SPLIT_AXIS_2: case GGML_BACKEND_SPLIT_AXIS_3: { - check_split(src_ss[0].n_segments == 1, "RESHAPE_SEGMENT_ERROR", src_ss); + GGML_ASSERT(src_ss[0].n_segments == 1); if (src_ss[0].axis == ggml_n_dims(tensor->src[0]) - 1 && src_ss[0].nr[0] == 1) { return {ggml_backend_meta_split_axis(ggml_n_dims(tensor) - 1), {0}, {1}, 1}; } @@ -755,13 +621,12 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return {ggml_backend_meta_split_axis(dim), {0}, {uint32_t(base_ne_out_next/base_ne_in)}, 1}; } if (base_ne_out_next > base_ne_in) { - check_split(src_ss[0].n_segments == 1, "RESHAPE_SEGMENT_ERROR", src_ss); - check_split(src_ss[0].nr[0] == 1, "RESHAPE_REPEAT_ERROR", src_ss); + GGML_ASSERT(src_ss[0].n_segments == 1); + GGML_ASSERT(src_ss[0].nr[0] == 1); return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1}; } base_ne_out = base_ne_out_next; } - log_split_failure("RESHAPE_SHAPE_MISMATCH", src_ss, "base_ne_in=" + std::to_string(base_ne_in)); GGML_ABORT("shape mismatch for %s", ggml_op_name(tensor->op)); } case GGML_BACKEND_SPLIT_AXIS_MIRRORED: @@ -769,7 +634,6 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return src_ss[0]; } default: { - log_split_failure("RESHAPE_UNSUPPORTED", src_ss); GGML_ABORT("fatal error"); //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } @@ -809,13 +673,11 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1}; } } - log_split_failure("VIEW_STRIDE_ERROR", src_ss); GGML_ABORT("fatal error"); } if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED || src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) { return src_ss[0]; } - log_split_failure("VIEW_PERMUTED_UNSUPPORTED", src_ss); GGML_ABORT("view of permuted tensor not implemented"); //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; }; @@ -826,7 +688,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( case GGML_BACKEND_SPLIT_AXIS_1: case GGML_BACKEND_SPLIT_AXIS_2: case GGML_BACKEND_SPLIT_AXIS_3: { - check_split(src_ss[0].n_segments == 1 || src_ss[0].nr[0] == 1, "PERMUTE_SEGMENT_ERROR", src_ss); + GGML_ASSERT(src_ss[0].n_segments == 1 || src_ss[0].nr[0] == 1); return {ggml_backend_meta_split_axis(tensor->op_params[src_ss[0].axis]), {0}, {src_ss[0].nr[0]}, 1}; } case GGML_BACKEND_SPLIT_AXIS_MIRRORED: @@ -834,7 +696,6 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return src_ss[0]; } default: { - log_split_failure("PERMUTE_UNSUPPORTED", src_ss); GGML_ABORT("fatal error"); //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } @@ -845,7 +706,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( switch (src_ss[0].axis) { case GGML_BACKEND_SPLIT_AXIS_0: case GGML_BACKEND_SPLIT_AXIS_1: { - check_split(src_ss[0].n_segments == 1 || src_ss[0].nr[0] == 1, "TRANSPOSE_SEGMENT_ERROR", src_ss); + GGML_ASSERT(src_ss[0].n_segments == 1 || src_ss[0].nr[0] == 1); return {ggml_backend_meta_split_axis(int(src_ss[0].axis) ^ 1), {0}, {src_ss[0].nr[0]}, 1}; } case GGML_BACKEND_SPLIT_AXIS_2: @@ -855,7 +716,6 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return src_ss[0]; } default: { - log_split_failure("TRANSPOSE_UNSUPPORTED", src_ss); GGML_ABORT("fatal error"); //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } @@ -870,54 +730,53 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( }; auto handle_set_rows = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - check_split(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_1, "SET_ROWS_DST_AXIS_ERROR", src_ss); - check_split(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, "SET_ROWS_SRC_AXIS_ERROR", src_ss); - check_split(split_states_equal(src_ss[0], src_ss[2]), "SET_ROWS_INDEX_RATIO_ERROR", src_ss); + GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_1); + GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(split_states_equal(src_ss[0], src_ss[2])); return src_ss[0]; }; auto handle_rope = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - check_split(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, "ROPE_POS_ERROR", src_ss); + GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); return src_ss[0]; }; auto handle_pad = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { if (src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) { - check_split(tensor->op_params[2*src_ss[0].axis + 0] == 0, "PAD_BEGIN_ERROR", src_ss); - check_split(tensor->op_params[2*src_ss[0].axis + 1] == 0, "PAD_END_ERROR", src_ss); + GGML_ASSERT(tensor->op_params[2*src_ss[0].axis + 0] == 0); + GGML_ASSERT(tensor->op_params[2*src_ss[0].axis + 1] == 0); } return src_ss[0]; }; auto handle_flash_attn_ext = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && - src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && - src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - check_split(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, - "FLASH_ATTN_MASK_ERROR", src_ss); - check_split(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, - "FLASH_ATTN_SINK_ERROR", src_ss); + GGML_ASSERT(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; } - if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2 && - src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && - src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - check_split(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, - "FLASH_ATTN_MASK_ERROR", src_ss); - check_split(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_0, - "FLASH_ATTN_SINK_ERROR", src_ss); - return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; - } - check_split(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2, "FLASH_ATTN_Q_AXIS_ERROR", src_ss); - check_split(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2, "FLASH_ATTN_K_AXIS_ERROR", src_ss); - check_split(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2, "FLASH_ATTN_V_AXIS_ERROR", src_ss); - check_split(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, - "FLASH_ATTN_MASK_ERROR", src_ss); - check_split(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_0, - "FLASH_ATTN_SINK_ERROR", src_ss); + + GGML_ASSERT(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2); + const bool kv_split = src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2 && + src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2; + const bool kv_mirrored = src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && + src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED; + GGML_ASSERT(kv_split || kv_mirrored); + GGML_ASSERT(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_0); return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; }; + auto handle_lightning_indexer = [&]( + const std::vector & src_ss) -> ggml_backend_meta_split_state { + for (size_t i = 0; i < 4; i++) { + GGML_ASSERT(src_ss[i].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } + return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; + }; + auto handle_ssm_conv = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { if (src_ss[0].axis == src_ss[1].axis) { if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { @@ -936,26 +795,17 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { return src_ss[0]; } - check_split(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1, "GDN_SRC0_AXIS_ERROR", src_ss); - check_split(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_1, "GDN_SRC1_AXIS_ERROR", src_ss); - check_split(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_1, "GDN_SRC2_AXIS_ERROR", src_ss); - check_split(src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_1, "GDN_SRC3_AXIS_ERROR", src_ss); - check_split(src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_1, "GDN_SRC4_AXIS_ERROR", src_ss); + GGML_ASSERT(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1); + GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_1); + GGML_ASSERT(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_1); + GGML_ASSERT(src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_1); + GGML_ASSERT(src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_1); // state shape is [S_v, S_v, H_v, n_seqs] (s0 only); the heads dim is its own axis 2, // so a head-aligned split on the input cache lands on axis 2 here. - check_split(src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_2 || src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_1 || src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_0, - "GDN_STATE_AXIS_ERROR", src_ss); + GGML_ASSERT(src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_2 || src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_1 || src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_0); return {GGML_BACKEND_SPLIT_AXIS_0, {0}, {1}, 1}; }; - auto handle_lightning_indexer = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - check_split(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, "LIGHTNING_INDEXER_Q_AXIS_ERROR", src_ss); - check_split(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, "LIGHTNING_INDEXER_K_AXIS_ERROR", src_ss); - check_split(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, "LIGHTNING_INDEXER_WEIGHT_AXIS_ERROR", src_ss); - check_split(src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED, "LIGHTNING_INDEXER_MASK_AXIS_ERROR", src_ss); - return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; - }; - auto calculate_split_state = [&]() -> ggml_backend_meta_split_state { if (ggml_nelements(tensor) == 0) { return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; @@ -967,21 +817,12 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (ret.axis >= 0 && ret.axis <= GGML_MAX_DIMS) { const int64_t granularity = ret.axis == GGML_BACKEND_SPLIT_AXIS_0 ? ggml_blck_size(tensor->type) : 1; int64_t ne_sum = 0; - std::vector empty_src_ss(GGML_MAX_SRC, {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, {1}, 1}); for (size_t s = 0; s < ret.n_segments; s++) { for (size_t j = 0; j < n_bufs; j++) { - if (ret.ne[s*n_bufs + j] % granularity != 0) { - ggml_backend_meta_log_split_context("STATIC_GRANULARITY_ERROR", tensor, empty_src_ss, n_bufs, - "state={" + ggml_backend_meta_split_state_str(ret, n_bufs) + "}, granularity=" + std::to_string(granularity)); - } GGML_ASSERT(ret.ne[s*n_bufs + j] % granularity == 0); ne_sum += ret.ne[s*n_bufs + j] * ret.nr[s]; } } - if (ne_sum != tensor->ne[ret.axis]) { - ggml_backend_meta_log_split_context("STATIC_SIZE_ERROR", tensor, empty_src_ss, n_bufs, - "state={" + ggml_backend_meta_split_state_str(ret, n_bufs) + "}, ne_sum=" + std::to_string(ne_sum)); - } GGML_ASSERT(ne_sum == tensor->ne[ret.axis]); } return ret; @@ -994,9 +835,6 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( continue; } src_ss[i] = ggml_backend_meta_get_split_state(stc, tensor->src[i], /*assume_sync =*/ true); - if (src_ss[i].axis == GGML_BACKEND_SPLIT_AXIS_UNKNOWN) { - log_split_failure("SRC_UNKNOWN", src_ss, "src=" + std::to_string(i)); - } GGML_ASSERT(src_ss[i].axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); } @@ -1197,7 +1035,6 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( split_state = handle_generic(src_ss, /*scalar_only =*/ false); } break; default: { - log_split_failure("OP_UNIMPLEMENTED", src_ss); GGML_ABORT("ggml op not implemented: %s", ggml_op_name(tensor->op)); split_state = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } break; @@ -1222,43 +1059,24 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( split_state.ne[j] *= tensor->ne[split_state.axis]; if (split_state.ne[j] != 0 || tensor->src[i]->ne[src_ss[i].axis] != 0) { const int64_t div = tensor->src[i]->ne[src_ss[i].axis] * split_state.nr[0]; - if (div == 0 || split_state.ne[j] % div != 0) { - log_split_failure("RATIO_DIV_ERROR", src_ss, - "src=" + std::to_string(i) + ", buf=" + std::to_string(j) + - ", div=" + std::to_string(div) + ", value=" + std::to_string(split_state.ne[j]) + - ", dst_state={" + ggml_backend_meta_split_state_str(split_state, n_bufs) + "}"); - } - GGML_ASSERT(div != 0); GGML_ASSERT(split_state.ne[j] % div == 0); split_state.ne[j] /= div; } } } else { - check_split(split_state.n_segments == 1, "RATIO_SEGMENT_ERROR", src_ss, - "dst_state={" + ggml_backend_meta_split_state_str(split_state, n_bufs) + "}"); + GGML_ASSERT(split_state.n_segments == 1); for (size_t j = 0; j < n_bufs; j++) { // Assert that ratio is consistent: int64_t sum = 0; for (size_t s = 0; s < src_ss[i].n_segments; s++) { sum += src_ss[i].ne[s*n_bufs + j] * src_ss[i].nr[s]; } - const int64_t lhs = split_state.ne[j]*split_state.nr[0] * tensor->src[i]->ne[src_ss[i].axis]; - const int64_t rhs = sum * tensor->ne[split_state.axis]; - if (lhs != rhs) { - log_split_failure("RATIO_ERROR", src_ss, - "src=" + std::to_string(i) + ", buf=" + std::to_string(j) + - ", lhs=" + std::to_string(lhs) + ", rhs=" + std::to_string(rhs) + - ", dst_state={" + ggml_backend_meta_split_state_str(split_state, n_bufs) + "}"); - } - GGML_ASSERT(lhs == rhs); + GGML_ASSERT(split_state.ne[j]*split_state.nr[0] * tensor->src[i]->ne[src_ss[i].axis] + == sum * tensor->ne[split_state.axis]); } } first_src_split_by_axis = false; } - if (first_src_split_by_axis) { - log_split_failure("NO_SPLIT_SOURCE", src_ss, - "dst_state={" + ggml_backend_meta_split_state_str(split_state, n_bufs) + "}"); - } GGML_ASSERT(!first_src_split_by_axis); } return split_state; @@ -1274,7 +1092,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (it == buf_ctx->split_state_cache.end()) { buf_ctx->split_state_cache[key].first = calculate_split_state(); memcpy(buf_ctx->split_state_cache[key].second, tensor, sizeof(buf_ctx->split_state_cache[key].second)); - if (buf_ctx->debug > 1 || (buf_ctx->debug > 0 && split_trace_interesting())) { + if (buf_ctx->debug > 0) { std::string srcs_info; for (size_t i = 0; i < GGML_MAX_SRC; i++) { if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { @@ -1283,14 +1101,29 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (!srcs_info.empty()) { srcs_info += ", "; } - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor->src[i], true); - srcs_info += "src[" + std::to_string(i) + "]=" + ggml_backend_meta_tensor_str(tensor->src[i]) + - ", state={" + ggml_backend_meta_split_state_str(split_state, n_bufs) + "}"; + const ggml_backend_meta_split_state split_state = + ggml_backend_meta_get_split_state(tensor->src[i], true); + GGML_ASSERT(split_state.n_segments == 1); + const char * axis_name = ggml_backend_meta_split_axis_name(split_state.axis); + std::string ne_info; + for (size_t j = 0; j < n_bufs; j++) { + if (!ne_info.empty()) { + ne_info += ", "; + } + ne_info += std::to_string(split_state.ne[j]) + "x" + std::to_string(split_state.nr[0]); + } + srcs_info += std::string(tensor->src[i]->name) + "[" + ggml_op_name(tensor->src[i]->op) + ", " + axis_name + ", {" + ne_info + "}]"; + } + std::string ne_info; + for (size_t j = 0; j < n_bufs; j++) { + if (!ne_info.empty()) { + ne_info += ", "; + } + const ggml_backend_meta_split_state & ss = buf_ctx->split_state_cache[key].first; + ne_info += std::to_string(ss.ne[j]) + "x" + std::to_string(ss.nr[0]); } - const ggml_backend_meta_split_state & ss = buf_ctx->split_state_cache[key].first; - GGML_LOG_DEBUG("META_SPLIT_TRACE: dst=%s, assume_sync=%d, state={%s}, srcs={%s}\n", - ggml_backend_meta_tensor_str(tensor).c_str(), assume_sync ? 1 : 0, - ggml_backend_meta_split_state_str(ss, n_bufs).c_str(), srcs_info.c_str()); + GGML_LOG_DEBUG("SPLIT_STATE: {%s} -> %s[%s, %s, {%s}]\n", srcs_info.c_str(), tensor->name, ggml_op_name(tensor->op), + ggml_backend_meta_split_axis_name(buf_ctx->split_state_cache[key].first.axis), ne_info.c_str()); } } @@ -1327,16 +1160,7 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m const size_t n_simple_bufs = ggml_backend_meta_buffer_n_bufs(tensor->buffer); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(stc, tensor, /*assume_sync =*/ true); - std::vector empty_src_ss(GGML_MAX_SRC, {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, {1}, 1}); - if (ggml_nelements(tensor) != 0 && split_state.axis == GGML_BACKEND_SPLIT_AXIS_UNKNOWN) { - ggml_backend_meta_log_split_context("INIT_UNKNOWN", tensor, empty_src_ss, n_simple_bufs, - "state={" + ggml_backend_meta_split_state_str(split_state, n_simple_bufs) + "}"); - } GGML_ASSERT(ggml_nelements(tensor) == 0 || split_state.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); - if (split_state.n_segments > 16) { - ggml_backend_meta_log_split_context("INIT_SEGMENTS_ERROR", tensor, empty_src_ss, n_simple_bufs, - "state={" + ggml_backend_meta_split_state_str(split_state, n_simple_bufs) + "}"); - } GGML_ASSERT(split_state.n_segments <= 16); int split_dim = split_state.axis; @@ -1386,18 +1210,8 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m if (t_ij->view_src != nullptr && ggml_backend_buffer_is_meta(t_ij->view_src->buffer)) { t_ij->view_src = ggml_backend_meta_buffer_simple_tensor(tensor->view_src, j); if (t_ij->view_offs > 0 && split_dim >= 0 && split_dim < GGML_MAX_DIMS) { - if (tensor->ne[split_dim] == 0) { - ggml_backend_meta_log_split_context("INIT_VIEW_ZERO_DIM", tensor, empty_src_ss, n_simple_bufs, - "state={" + ggml_backend_meta_split_state_str(split_state, n_simple_bufs) + "}"); - } GGML_ASSERT(tensor->ne[split_dim] != 0); - const ggml_backend_meta_split_state view_src_ss = ggml_backend_meta_get_split_state(tensor->view_src, /*assume_sync =*/ true); - const int split_dim_view_src = view_src_ss.axis; - if (split_dim_view_src < 0 || split_dim_view_src >= GGML_MAX_DIMS) { - ggml_backend_meta_log_split_context("INIT_VIEW_SRC_AXIS_ERROR", tensor, empty_src_ss, n_simple_bufs, - "state={" + ggml_backend_meta_split_state_str(split_state, n_simple_bufs) + "}, view_src_state={" + - ggml_backend_meta_split_state_str(view_src_ss, n_simple_bufs) + "}"); - } + const int split_dim_view_src = ggml_backend_meta_get_split_state(tensor->view_src, /*assume_sync =*/ true).axis; GGML_ASSERT(split_dim_view_src >= 0 && split_dim_view_src < GGML_MAX_DIMS); // The offset can be internal to the data split, in those cases the view offset should not be scaled. @@ -1467,9 +1281,11 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor(ggml_backend_buffer return ggml_backend_meta_buffer_init_tensor_impl(buf_ctx->get_simple_tensor_container(tensor), tensor); } -static void ggml_backend_meta_buffer_memset_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) { +static void ggml_backend_meta_buffer_memset_tensor( + ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) { const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); + const ggml_backend_meta_split_state split_state = + ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); GGML_ASSERT(ggml_is_contiguous(tensor) || split_state.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); if (split_state.n_segments != 1 || split_state.nr[0] != 1) { @@ -1505,6 +1321,7 @@ static void ggml_backend_meta_buffer_memset_tensor(ggml_backend_buffer_t buffer, } return; } + GGML_ASSERT(split_state.axis == GGML_BACKEND_SPLIT_AXIS_1); const size_t row_stride = tensor->nb[2]; @@ -1537,21 +1354,24 @@ static void ggml_backend_meta_buffer_memset_tensor(ggml_backend_buffer_t buffer, const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; GGML_ASSERT(offset % chunk_size_full == 0); GGML_ASSERT(size % chunk_size_full == 0); - const int64_t i_start = offset /chunk_size_full; - const int64_t i_stop = (offset + size)/chunk_size_full; + const int64_t i_start = offset / chunk_size_full; + const int64_t i_stop = (offset + size) / chunk_size_full; for (size_t j = 0; j < n_bufs; j++) { ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; - if (chunk_size_j == 0) { + const size_t chunk_size = simple_tensor->nb[split_state.axis + 1]; + if (chunk_size == 0) { continue; } for (int64_t i = i_start; i < i_stop; i++) { - ggml_backend_tensor_memset(simple_tensor, value, i*chunk_size_j, chunk_size_j); + ggml_backend_tensor_memset(simple_tensor, value, i*chunk_size, chunk_size); } } } break; - case GGML_BACKEND_SPLIT_AXIS_MIRRORED: case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { + GGML_ASSERT(value == 0); + [[fallthrough]]; + } + case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { for (size_t j = 0; j < n_bufs; j++) { ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); ggml_backend_tensor_memset(simple_tensor, value, offset, size); @@ -1827,9 +1647,6 @@ static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_bac bufs.reserve(n_simple_bufts); for (size_t i = 0; i < n_simple_bufts; i++) { bufs.push_back(ggml_backend_buft_alloc_buffer(ggml_backend_meta_buft_simple_buft(buft, i), size)); - if (bufs.back() == nullptr) { - GGML_LOG_ERROR("META_ALLOC_ERROR: simple_buft=%zu, requested_size=%zu\n", i, size); - } GGML_ASSERT(bufs.back() != nullptr); max_size = std::max(max_size, ggml_backend_buffer_get_size(bufs.back())); } @@ -1886,9 +1703,6 @@ struct ggml_backend_buffer * ggml_backend_meta_alloc_ctx_tensors_from_buft(struc t->buffer = meta_buf_ctx->bufs[i].get(); } } - if (!meta_buf_ctx->bufs[i]) { - GGML_LOG_ERROR("META_ALLOC_CTX_ERROR: simple_buft=%zu\n", i); - } GGML_ASSERT(meta_buf_ctx->bufs[i]); meta_buf->size = std::max(meta_buf->size, ggml_backend_buffer_get_size(meta_buf_ctx->bufs[i].get())); } diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 2f5d3f12e0f2..83870ae003ea 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -345,17 +345,13 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_qkv_bias ("blk\\.\\d*\\.attn_qkv.bias"); static const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight"); static const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*"); - static const std::regex pattern_dsv4_state ("dsv4_.*_state_(kv|score)_l\\d*"); + static const std::regex pattern_dsv4_state ("dsv4_(csa|hca|lid)_state_(kv|score)_l\\d*"); static const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight"); static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight"); static const std::regex pattern_attn_out_bias ("blk\\.\\d*\\.attn_output.bias"); static const std::regex pattern_attn_out_a_weight("blk\\.\\d*\\.attn_output_a\\.weight"); static const std::regex pattern_attn_out_b_weight("blk\\.\\d*\\.attn_output_b\\.weight"); - static const std::regex pattern_attn_out_b_bias ("blk\\.\\d*\\.attn_output_b\\.bias"); - static const std::regex pattern_attn_q_a_weight ("blk\\.\\d*\\.attn_q_a\\.weight"); - static const std::regex pattern_attn_q_b_weight ("blk\\.\\d*\\.attn_q_b\\.weight"); - static const std::regex pattern_attn_q_b_bias ("blk\\.\\d*\\.attn_q_b\\.bias"); - static const std::regex pattern_attn_kv_weight ("blk\\.\\d*\\.attn_kv\\.weight"); + static const std::regex pattern_attn_q_b_weight ("blk\\.\\d*\\.attn_q_b\\.weight"); static const std::regex pattern_attn_gate_weight("blk\\.\\d*\\.attn_gate.weight"); static const std::regex pattern_ssm_dt ("blk\\.\\d*\\.ssm_dt.bias"); @@ -374,8 +370,11 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_ffn_gate_bias ("blk\\.\\d*\\.ffn_gate(_exps)?.bias"); static const std::regex pattern_ffn_gate_up_weight("blk\\.\\d*\\.ffn_gate_up(_exps)?.weight"); static const std::regex pattern_ffn_down_weight ("blk\\.\\d*\\.ffn_down(_exps)?.weight"); - static const std::regex pattern_ffn_down_bias ("blk\\.\\d*\\.ffn_down.bias"); - static const std::regex pattern_ffn_down_exps_bias("blk\\.\\d*\\.ffn_down_exps.bias"); + static const std::regex pattern_ffn_down_bias ("blk\\.\\d*\\.ffn_down.bias"); + static const std::regex pattern_ffn_down_exps_bias ("blk\\.\\d*\\.ffn_down_exps.bias"); + static const std::regex pattern_ffn_up_shexp_weight ("blk\\.\\d*\\.ffn_up_shexp.weight"); + static const std::regex pattern_ffn_gate_shexp_weight ("blk\\.\\d*\\.ffn_gate_shexp.weight"); + static const std::regex pattern_ffn_down_shexp_weight ("blk\\.\\d*\\.ffn_down_shexp.weight"); static const std::regex pattern_output_weight("output\\.weight"); static const std::regex pattern_output_bias ("output\\.bias"); @@ -432,6 +431,32 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str }; auto get_tensor_config = [&]() -> tensor_config { + if (ud->model->arch == LLM_ARCH_DEEPSEEK4) { + if (std::regex_match(tensor_name, pattern_kv_cache) || + std::regex_match(tensor_name, pattern_dsv4_state)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } + if (std::regex_match(tensor_name, pattern_attn_sinks)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output_a.weight"); + } + if (std::regex_match(tensor_name, pattern_attn_q_b_weight)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output_a.weight"); + } + if (std::regex_match(tensor_name, pattern_attn_out_a_weight)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output_b.weight"); + } + if (std::regex_match(tensor_name, pattern_attn_out_b_weight)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0); + } + if (std::regex_match(tensor_name, pattern_ffn_up_shexp_weight) || + std::regex_match(tensor_name, pattern_ffn_gate_shexp_weight)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "ffn_down_shexp.weight"); + } + if (std::regex_match(tensor_name, pattern_ffn_down_shexp_weight)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "ffn_down_shexp.weight"); + } + } + // standard attention if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_kv_weight)) { return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight", "ssm_out.weight"); @@ -448,14 +473,6 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str if (std::regex_match(tensor_name, pattern_qk_norm)) { return get_tensor_config_impl(tensor->ne[1] == 1 ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight"); } - if (ud->model->arch == LLM_ARCH_DEEPSEEK4 && - (std::regex_match(tensor_name, pattern_kv_cache) || - std::regex_match(tensor_name, pattern_dsv4_state))) { - return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); - } - if (ud->model->arch == LLM_ARCH_DEEPSEEK4 && std::regex_match(tensor_name, pattern_attn_sinks)) { - return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output_a.weight"); - } if (std::regex_match(tensor_name, pattern_kv_cache) || std::regex_match(tensor_name, pattern_attn_sinks)) { return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output.weight"); } @@ -465,26 +482,6 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str if (std::regex_match(tensor_name, pattern_attn_out_bias)) { return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); } - if (ud->model->arch == LLM_ARCH_DEEPSEEK4) { - if (std::regex_match(tensor_name, pattern_attn_q_b_weight)) { - return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output_a.weight", "attn_output.weight"); - } - if (std::regex_match(tensor_name, pattern_attn_q_b_bias)) { - return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output_a.weight", "attn_output.weight"); - } - if (std::regex_match(tensor_name, pattern_attn_out_a_weight)) { - return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output_b.weight"); - } - if (std::regex_match(tensor_name, pattern_attn_out_b_weight)) { - return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0); - } - if (std::regex_match(tensor_name, pattern_attn_out_b_bias)) { - return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); - } - if (std::regex_match(tensor_name, pattern_attn_q_a_weight) || std::regex_match(tensor_name, pattern_attn_kv_weight)) { - return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); - } - } if (std::regex_match(tensor_name, pattern_attn_gate_weight)) { return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight", "ssm_out.weight"); @@ -654,11 +651,12 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } if (ud->model->arch == LLM_ARCH_DEEPSEEK4) { - if (std::regex_match(tensor_name, pattern_attn_q_b_weight) || std::regex_match(tensor_name, pattern_attn_q_b_bias)) { + if (std::regex_match(tensor_name, pattern_attn_q_b_weight)) { GGML_ASSERT(segments.size() == 1); return {hparams.n_embd_head_k(il)}; } - if (std::regex_match(tensor_name, pattern_attn_out_a_weight) || std::regex_match(tensor_name, pattern_attn_out_b_weight)) { + if (std::regex_match(tensor_name, pattern_attn_out_a_weight) || + std::regex_match(tensor_name, pattern_attn_out_b_weight)) { GGML_ASSERT(segments.size() == 1); return {std::lcm(hparams.dsv4_o_lora_rank, blck_size)}; } @@ -694,7 +692,11 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str // FFN if (std::regex_match(tensor_name, pattern_ffn_up_weight) || std::regex_match(tensor_name, pattern_ffn_up_bias) || std::regex_match(tensor_name, pattern_ffn_gate_weight) || std::regex_match(tensor_name, pattern_ffn_gate_bias) || - std::regex_match(tensor_name, pattern_ffn_gate_up_weight) || std::regex_match(tensor_name, pattern_ffn_down_weight)) { + std::regex_match(tensor_name, pattern_ffn_gate_up_weight) || + std::regex_match(tensor_name, pattern_ffn_down_weight) || + std::regex_match(tensor_name, pattern_ffn_up_shexp_weight) || + std::regex_match(tensor_name, pattern_ffn_gate_shexp_weight) || + std::regex_match(tensor_name, pattern_ffn_down_shexp_weight)) { const int64_t blck_size_perf = std::lcm(blck_size, 128); GGML_ASSERT(segments.size() == 1); return {blck_size_perf}; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 57d33a627d5d..f6b4b31fbc06 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -101,6 +101,11 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { n_head = 1; n_ff = 96; n_layer = 22; // hparams.n_layer_kv_from_start = 20 is hardcoded + } else if (arch == LLM_ARCH_DEEPSEEK4) { + n_embd = 256; + n_head = 8; + n_ff = 1024; + n_layer = 4; } else if (arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA @@ -154,11 +159,15 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_per_layer); } else { ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head); - ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head); + ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(1) : n_head); } ms.add_kv(LLM_KV_ATTENTION_MAX_ALIBI_BIAS, 8.0f); - if (arch == LLM_ARCH_DEEPSEEK2 + if (arch == LLM_ARCH_DEEPSEEK4) { + ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, n_embd_head); + ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, n_embd_head); + ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, n_embd_head/2); + } else if (arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_KIMI_LINEAR @@ -174,7 +183,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, 1e-5f); ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_EPS, 1e-5f); ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_GROUPS, uint32_t(8)); - ms.add_kv(LLM_KV_ATTENTION_Q_LORA_RANK, uint32_t(512)); + ms.add_kv(LLM_KV_ATTENTION_Q_LORA_RANK, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(64) : uint32_t(512)); ms.add_kv(LLM_KV_ATTENTION_KV_LORA_RANK, uint32_t(512)); ms.add_kv(LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT, uint32_t(8)); ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW, n_ctx/8); @@ -198,10 +207,24 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(2)); } - ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, uint32_t(1)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, arch == LLM_ARCH_DEEPSEEK4 ? n_head : uint32_t(1)); ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, uint32_t(64)); ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(8)); ms.add_kv(LLM_KV_ROPE_DIMENSION_SECTIONS, std::vector({n_embd_head/4, n_embd_head/4, n_embd_head/4, n_embd_head/4})); + + if (arch == LLM_ARCH_DEEPSEEK4) { + ms.add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, uint32_t(8)); + ms.add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, uint32_t(32)); + ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector({0, 0, 4, 128})); + ms.add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, 160000.0f); + ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(2)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, 1.0e-6f); + ms.add_kv(LLM_KV_HASH_LAYER_COUNT, uint32_t(0)); + ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, 10.0f); + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); + } ms.add_kv(LLM_KV_TOKENIZER_MODEL, "no_vocab"); // ms.add_kv(LLM_KV_DENSE_2_FEAT_OUT, n_embd); // ms.add_kv(LLM_KV_DENSE_3_FEAT_IN, n_embd); @@ -212,7 +235,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_EXPERT_COUNT, uint32_t(2)); ms.add_kv(LLM_KV_EXPERT_USED_COUNT, uint32_t(1)); ms.add_kv(LLM_KV_EXPERT_SHARED_COUNT, uint32_t(1)); - ms.add_kv(LLM_KV_EXPERT_GATING_FUNC, uint32_t(2)); // sigmoid + ms.add_kv(LLM_KV_EXPERT_GATING_FUNC, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(4) : uint32_t(2)); ms.add_kv(LLM_KV_EXPERT_GROUP_SCALE, 1.0f); ms.add_kv(LLM_KV_EXPERTS_PER_GROUP, uint32_t(1)); } @@ -338,6 +361,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_DEEPSEEK: case LLM_ARCH_DEEPSEEK2: case LLM_ARCH_DEEPSEEK32: + case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_GLM4_MOE: case LLM_ARCH_GLM_DSA: case LLM_ARCH_EXAONE_MOE: @@ -415,10 +439,6 @@ static bool arch_supported(const llm_arch arch) { if (arch == LLM_ARCH_DEEPSEEK2OCR) { return false; } - if (arch == LLM_ARCH_DEEPSEEK4) { - return false; - } - // FIXME some models are segfaulting with WebGPU: #ifdef GGML_USE_WEBGPU if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_KIMI_LINEAR) { @@ -598,10 +618,18 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg if (logits_cpu.empty()) { model_and_ctx_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, encode); logits_cpu = get_logits(model_and_ctx_cpu.first.get(), model_and_ctx_cpu.second.get(), tokens, encode); + if (arch == LLM_ARCH_DEEPSEEK4) { + GGML_ASSERT(llama_memory_seq_rm( + llama_get_memory(model_and_ctx_cpu.second.get()), 0, -1, -1)); + } } if (dc.split_mode != LLAMA_SPLIT_MODE_TENSOR || llm_arch_supports_sm_tensor(arch)) { model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, dc.devs, dc.split_mode, encode); logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode); + if (arch == LLM_ARCH_DEEPSEEK4) { + GGML_ASSERT(llama_memory_seq_rm( + llama_get_memory(model_and_ctx_dev.second.get()), 0, -1, -1)); + } const double nmse_val = nmse(logits_cpu, logits_dev); snprintf(nmse_str, sizeof(nmse_str), "(%.2e)", nmse_val); status_nmse = "\033[1;32mOK\033[0m"; @@ -614,7 +642,9 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg FILE * file = tmpfile(); // Can be null on Windows without administrator privileges. // FIXME: when adding a tensor to a gguf_context a copy is made, this changes the pointer which the meta backend // in turn uses to map the tensors to their simple equivalents - this is fundamentally incompatible - if (file != nullptr && llama_model_saver_supports_arch(arch) && dc.split_mode != LLAMA_SPLIT_MODE_TENSOR) { + // FIXME: DSV4 metadata is not implemented by llama_model_saver. + const bool can_roundtrip = llama_model_saver_supports_arch(arch) && arch != LLM_ARCH_DEEPSEEK4; + if (file != nullptr && can_roundtrip && dc.split_mode != LLAMA_SPLIT_MODE_TENSOR) { GGML_ASSERT(model_and_ctx_dev.first && model_and_ctx_dev.second); llama_model_saver ms = llama_model_saver(model_and_ctx_dev.first.get()); ms.add_kv_from_model();