Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Deepseek 4: -sm tensor (#26490)

Browse Source 
* 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
2026-08-24 09:20:25 +03:00
committed by GitHub
parent c060ca974c
commit bf0a29cc16
7 changed files with 314 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-backend-meta.cpp
+202 -10
View File
@@ -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];