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
src/llama-arch.cpp
-1
View File
@@ -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:
src/llama-kv-cache-dsv4.cpp
+1
View File
@@ -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);
}
src/llama-model-saver.cpp
+13 -4
View File
@@ -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);
src/llama-model.cpp
+67 -3
View File
@@ -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};
src/models/dflash.cpp
+4 -3
View File
@@ -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];