diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index ded678e68..fe58ea3bb 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -592,7 +592,18 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( 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}; } - GGML_ABORT("fatal error"); + 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]; + } + // batched matmul with the batches split across devices and a replicated activation + if (src_ss[0].axis >= GGML_BACKEND_SPLIT_AXIS_2 && src_ss[0].axis < GGML_MAX_DIMS && + src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + return src_ss[0]; + } + GGML_ABORT("unsupported mul_mat split states: node=%s src0=%s axis=%d src1=%s axis=%d", + tensor->name, tensor->src[0]->name, (int) src_ss[0].axis, tensor->src[1]->name, (int) src_ss[1].axis); //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; }; @@ -760,14 +771,33 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( }; 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[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}; + } + + 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) { @@ -938,7 +968,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); @@ -1002,6 +1032,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: { @@ -1086,13 +1119,14 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (buf_ctx->debug > 0) { 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); + 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; @@ -1271,6 +1305,108 @@ 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 = 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, chunk_size); + } + } + } break; + 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); + } + } 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); @@ -1518,7 +1654,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, @@ -1871,7 +2007,7 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, { // For MoE models it may make sense to delay the AllReduce in order to reduce I/O: - auto get_i_delayed = [&](const int i) -> int { + auto get_i_delayed_branch = [&](const int i) -> int { int id = i; // i_delayed int idr = i; // i_delayed return, last safe return value @@ -1971,6 +2107,62 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, return idr; }; + // AllReduce(a) + AllReduce(b) == AllReduce(a + b) for independent partial branches. + auto get_i_delayed = [&](const int i) -> int { + const int i_delayed = get_i_delayed_branch(i); + ggml_tensor * node = cgraph->nodes[i_delayed]; + + if (ggml_node_get_use_count(cgraph, i_delayed) != 1) { + return i_delayed; + } + + for (int id = i_delayed + 1; id < cgraph->n_nodes; id++) { + ggml_tensor * next = cgraph->nodes[id]; + if (next->view_src == node) { + return i_delayed; + } + for (int s = 0; s < GGML_MAX_SRC; s++) { + if (next->src[s] == node) { + return i_delayed; + } + } + + if (next->view_src != nullptr && next->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(next->view_src->buffer)) { + continue; + } + if (ggml_backend_meta_get_split_state(next, false).axis != GGML_BACKEND_SPLIT_AXIS_PARTIAL) { + continue; + } + + const int i_other = id; + const int i_other_delayed = get_i_delayed_branch(i_other); + ggml_tensor * other = cgraph->nodes[i_other_delayed]; + if (ggml_node_get_use_count(cgraph, i_other_delayed) != 1 || i_other_delayed + 1 >= cgraph->n_nodes) { + return i_delayed; + } + + ggml_tensor * sum = cgraph->nodes[i_other_delayed + 1]; + if (sum->op != GGML_OP_ADD || + !ggml_are_same_shape(node, other) || node->type != other->type || sum->type != node->type || + !((sum->src[0] == node && sum->src[1] == other) || + (sum->src[0] == other && sum->src[1] == node)) || + ggml_backend_meta_get_split_state(sum, false).axis != GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + return i_delayed; + } + + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + const bool compute = bcj.nodes[i]->flags & GGML_TENSOR_FLAG_COMPUTE; + const bool compute_other = bcj.nodes[i_other]->flags & GGML_TENSOR_FLAG_COMPUTE; + if (compute != compute_other) { + return i_delayed; + } + } + return i_other_delayed + 1; + } + return i_delayed; + }; + int i_start = 0; for (int i = 0; i < cgraph->n_nodes; i++) { ggml_tensor * node = cgraph->nodes[i]; diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 025f9fb54..eecf444fc 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -1060,7 +1060,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_DOTS3NOTE: case LLM_ARCH_GLM_DSA: case LLM_ARCH_BITNET: diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index 58f78e438..948d08146 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -1737,6 +1737,7 @@ void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) { kv->seq_rm(seq_id, -1, -1); if (data) { + //TODO: do not clear the kv-cache during `seq_rm`, ref: https://github.com/ggml-org/llama.cpp/pull/26490#discussion_r3798143663 for (uint32_t il : kv->get_layer_ids()) { dsv4_clear_tensor_stream(kv->get_k_storage(il), (uint32_t) seq_id); } diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 2eb5b7aaf..9adaa93f6 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -296,11 +296,18 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, hparams.dsv4_o_group_count); add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, hparams.dsv4_o_lora_rank); add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, hparams.dsv4_compress_rope_base); - add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, true); - add_kv(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); + if (model->arch == LLM_ARCH_DEEPSEEK4 || hparams.dsv4_hc_mult > 0) { + // the loader requires one compress ratio per layer, including nextn layers + const std::vector compress_ratios( + hparams.dsv4_compress_ratios.begin(), hparams.dsv4_compress_ratios.begin() + hparams.n_layer_all); + add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, compress_ratios); + } else { + add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, true); + } + add_kv(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters); - add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); - add_kv(LLM_KV_HASH_LAYER_COUNT, hparams.dsv4_hash_layer_count); + add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); + add_kv(LLM_KV_HASH_LAYER_COUNT, hparams.dsv4_hash_layer_count); const float rope_scaling_factor = hparams.rope_freq_scale_train == 1.0f ? 0.0f : 1.0f/hparams.rope_freq_scale_train; @@ -425,6 +432,8 @@ void llama_model_saver::add_tensors_from_model() { add_tensor(model->output_s); add_tensor(model->output_in_s); add_tensor(model->output_res_score); + add_tensor(model->nextn_proj_pre); + add_tensor(model->nextn_proj_post); add_tensor(model->cls); add_tensor(model->cls_b); add_tensor(model->cls_out); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 33f5661b2..c34700ff5 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -365,6 +365,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str const llama_meta_device_get_split_state_userdata * ud = (const llama_meta_device_get_split_state_userdata *) userdata; const llama_hparams & hparams = ud->model->hparams; const std::string tensor_name = tensor->name; + const bool is_dsv4 = ud->model->arch == LLM_ARCH_DEEPSEEK4 || + (ud->model->arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0); static const std::regex pattern_q_weight ("blk\\.\\d*\\.attn_q.weight"); static const std::regex pattern_kv_weight ("blk\\.\\d*\\.attn_(k|v).weight"); @@ -374,9 +376,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_(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_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"); @@ -395,8 +401,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"); @@ -453,6 +462,32 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str }; auto get_tensor_config = [&]() -> tensor_config { + if (is_dsv4) { + 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_2); + } + 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"); @@ -525,6 +560,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str // output if (std::regex_match(tensor_name, pattern_output_weight)) { + if (is_dsv4) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1); } if (std::regex_match(tensor_name, pattern_output_bias)) { @@ -649,8 +687,30 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str const int64_t granularity_head = granularity_q / hparams.n_embd_head_k(il); // for tensors with one value per head if (std::regex_match(tensor_name, pattern_attn_sinks)) { GGML_ASSERT(segments.size() == 1); + if (is_dsv4) { + return {hparams.n_head(il) / hparams.dsv4_o_group_count}; + } return {granularity_head}; } + + if (is_dsv4) { + if (std::regex_match(tensor_name, pattern_attn_q_b_weight)) { + GGML_ASSERT(segments.size() == 1); + // the grouped output projection requires each device to hold whole groups of heads + const int64_t n_head_group = hparams.n_head(il) / hparams.dsv4_o_group_count; + return {n_head_group * hparams.n_embd_head_k(il)}; + } + if (std::regex_match(tensor_name, pattern_attn_out_a_weight)) { + GGML_ASSERT(segments.size() == 1); + return {1}; + } + if (std::regex_match(tensor_name, pattern_attn_out_b_weight)) { + GGML_ASSERT(segments.size() == 1); + // the boundaries must align with wo_a's per-group split, so quant blocks must not straddle groups + GGML_ASSERT(hparams.dsv4_o_lora_rank % blck_size == 0); + return {hparams.dsv4_o_lora_rank}; + } + } if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_q_bias)) { GGML_ASSERT(segments.size() == 1); // some models have Q gate tensors, for those cases the granularity needs to be doubled: @@ -687,7 +747,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/src/models/dflash.cpp b/src/models/dflash.cpp index d3c919b35..ff40c16b2 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -117,6 +117,10 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc) output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm + // optional: reduced-vocab drafts ship their own lm head, full-vocab drafts can share the target's via ctx_other + // a draft with its own embeddings + head references no target tensors and can run on devices the target does not use (e.g. -devd with a tensor-split target) + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab_draft }, TENSOR_NOT_REQUIRED); + if (hparams.dsv4_hc_mult > 0) { const int64_t q_lora_rank = hparams.n_lora_q; const int64_t n_ff_exp = hparams.n_ff_exp; @@ -167,9 +171,6 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { return; } - // optional: reduced-vocab drafts ship their own, full-vocab drafts share the target's via ctx_other - output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab_draft }, TENSOR_NOT_REQUIRED); - for (int i = 0; i < n_layer; ++i) { auto & layer = layers[i]; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index dff8c4668..83922c53b 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -102,10 +102,11 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { n_ff = 96; n_layer = 22; // hparams.n_layer_kv_from_start = 20 is hardcoded } else if (arch == LLM_ARCH_DEEPSEEK4) { - n_embd = 128; - n_head = 1; - n_ff = 192; - n_layer = 3; // uncompressed + csa + hca, one layer of each ratio kind + // head size 64 so that GPU flash attention kernels support the model + n_embd = 512; + n_head = 8; + n_ff = 1024; + n_layer = 4; } else if (arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_LAGUNA) { n_embd = 160; // exercise per-head tensor split granularity with head size 80 } else if (arch == LLM_ARCH_QWEN3 || arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE) { @@ -175,11 +176,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_kv); + ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(1) : n_head_kv); } 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_DOTS3NOTE @@ -208,10 +213,6 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { } ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, indexer_types); } - } else if (arch == LLM_ARCH_DEEPSEEK4) { - ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, uint32_t(128)); - ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, uint32_t(128)); - ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64)); } else if (arch == LLM_ARCH_MINIMAX_M3) { // partial rotary: n_rot must not exceed the indexer key length (64) ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64)); @@ -221,7 +222,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); @@ -248,26 +249,26 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { // MSA requires one indexer head per GQA (KV) head, unlike the DSA archs where the // indexer head count is independent of the main attention head count. - if (arch == LLM_ARCH_DEEPSEEK4) { - ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 2.5f); - ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); - ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, 7.0f); - ms.add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, uint32_t(1)); - ms.add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, uint32_t(64)); - ms.add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, 10000.0f); - ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4)); - ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(4)); - ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, 1e-6f); - ms.add_kv(LLM_KV_HASH_LAYER_COUNT, uint32_t(0)); - ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector({0, 4, 128})); - } - - ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, arch == LLM_ARCH_MINIMAX_M3 ? n_head : uint32_t(1)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, arch == LLM_ARCH_MINIMAX_M3 || 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_ATTENTION_INDEXER_BLOCK_SIZE, uint32_t(4)); ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1)); 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); @@ -504,7 +505,6 @@ static bool arch_supported(const llm_arch arch) { if (arch == LLM_ARCH_DEEPSEEK2OCR) { return false; } - // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA) {