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
2026-08-24 11:50:25 +05:30
committed by GitHub
parent c060ca974c
commit bf0a29cc16
7 changed files with 314 additions and 48 deletions
+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};