mirror of
https://github.com/LostRuins/koboldcpp.git
synced 2026-08-27 17:29:52 +00:00
Deepseek 4: -sm tensor (#26490)
* DSV4: sm tensor * set coarser granularity for head splits * fix dspark * add model saving for dsv4 + allow dflash to return on specific device * add comment about dsv4 seq_rm * simplify * add shared expert delayed allreduce * remove special test for dsv4
This commit is contained in:
Aman Gupta
GitHub
parent
c060ca974c
commit
bf0a29cc16
7 changed files with 314 additions and 48 deletions
|
|
@ -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<ggml_backend_meta_split_state> & 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<ggml_backend_meta_split_state> & 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<ggml_backend_meta_split_state> & 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<size_t> 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];
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<uint32_t> 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);
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -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];
|
||||
|
||||
|
|
|
|||
|
|
@ -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<uint32_t>({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<uint32_t>({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<uint32_t>({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) {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue