Skip to content

Deepseek 4: -sm tensor #26490

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

Open
am17an wants to merge 2 commits into
base: master
Choose a base branch
from
+249 −24
Open

Deepseek 4: -sm tensor #26490

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
878c275
DSV4: sm tensor
am17an Jul 31, 2026
3bbafa2
set coarser granularity for head splits
am17an Aug 3, 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
161 changes: 151 additions & 10 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,7 +590,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};
};

Expand Down Expand Up @@ -745,14 +756,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) {
Expand Down Expand Up @@ -817,7 +847,12 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
ggml_backend_meta_split_state split_state;
switch (tensor->op) {
case GGML_OP_NONE: {
split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
if (tensor->view_src != nullptr) {
// full-tensor view created with ggml_view_tensor, transparent for the split state
split_state = ggml_backend_meta_get_split_state(stc, tensor->view_src, assume_sync);
} else {
split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
}
} break;
case GGML_OP_DUP: {
split_state = handle_generic(src_ss, /*scalar_only =*/ true);
Expand Down Expand Up @@ -920,7 +955,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 +1019,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 +1106,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 +1292,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 +1627,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 @@ -1011,7 +1011,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
60 changes: 57 additions & 3 deletions src/llama-model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -352,9 +352,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 @@ -373,8 +377,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 @@ -431,6 +438,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");
Expand Down Expand Up @@ -618,9 +651,26 @@ 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 {hparams.n_head(il) / hparams.dsv4_o_group_count};
}
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);
// 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) ||
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 @@ -651,7 +701,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
Loading