Skip to content

Deepseek V4: split-mode tensor #25860

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
S3kundes wants to merge 2 commits into
base: master
Choose a base branch
from
+233 −22
Draft

Deepseek V4: split-mode tensor #25860

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a21950e
split mode tensor working
S3kundes Jul 17, 2026
3515faf
sol refactor and free some vram
S3kundes Jul 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
146 changes: 138 additions & 8 deletions ggml/src/ggml-backend-meta.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -590,6 +590,11 @@ 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};
}
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];
}
Comment on lines +593 to +597

@JohannesGaessler JohannesGaessler Jul 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is incorrect. The only valid axis that would be covered here is PARTIAL and there is no sensible way to combine two PARTIAL tensors via a matrix multiplication since (a + b) * (a + b) == a*a + 2*a*b + b*b.

GGML_ABORT("fatal error");
//return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1};
};
Expand Down Expand Up @@ -745,14 +750,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};
}
Comment on lines +755 to +760

@JohannesGaessler JohannesGaessler Jul 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is indicative of incorrect assignments in llama-model.cpp. It is in principle possible to just run everything as mirrored but for an expensive operation like FA that is going to massively gimp the performance.


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) {
Expand Down Expand Up @@ -920,7 +944,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);
Expand Down Expand Up @@ -984,6 +1008,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: {
Expand Down Expand Up @@ -1068,13 +1095,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;
Expand Down Expand Up @@ -1253,6 +1281,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);
Expand Down Expand Up @@ -1486,7 +1616,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,
Expand Down
1 change: 0 additions & 1 deletion src/llama-arch.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
58 changes: 55 additions & 3 deletions src/llama-model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -345,9 +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_(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");
Expand All @@ -366,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");
Expand Down Expand Up @@ -424,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);
}
Comment on lines +435 to +438

@JohannesGaessler JohannesGaessler Jul 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is 100% the wrong way to do it and just a workaround in lieu of the correct logic.

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");
}
Comment on lines +439 to +457

@JohannesGaessler JohannesGaessler Jul 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd have to check against the compute graph (and also spin up the model with a debugger) but these may be correct. The bottom line is that attention heads should be split by GPU and combined once at the end. Mirroring them will result in duplicated KV cache and bad performance.

}

// 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");
Expand Down Expand Up @@ -611,9 +644,24 @@ 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};
}
Comment on lines +647 to +649

@JohannesGaessler JohannesGaessler Jul 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is wrong unless the model does not use GQA at all.

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)) {
GGML_ASSERT(segments.size() == 1);
return {hparams.n_embd_head_k(il)};
}
Comment on lines +654 to +657

@JohannesGaessler JohannesGaessler Jul 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This seems wrong to me but I don't know off the top of my head what the correct implementation would be.

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<int64_t>(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);
Expand Down Expand Up @@ -644,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};
Expand Down
Loading