Merged up/gate for DS4 (#2257)

* Merged up/gate for DS4

* Just in case

* WIP

* Fix CUDA silu kernel for merged up/gate with limit
This commit is contained in:
Kawrakow
2026-08-07 07:51:58 +03:00
committed by GitHub
parent da2293ded3
commit 466d9bb15f
3 changed files with 47 additions and 43 deletions
+4 -3
View File
@@ -113,7 +113,7 @@ static __global__ void fused_mul_silu_f32(const float * x, float * dst, const in
int row = i / ne0;
int j = i % ne0;
auto x_row = x + 2*row*ne0;
//dst[i] = x_row[j] * x_row[j + ne0] / (1.0f + expf(-x_row[j + ne0]));
// Note: gate is at the beginning of each row, up is offset by ne0
dst[i] = x_row[j] * x_row[j + ne0] / (1.0f + expf(-x_row[j]));
}
@@ -126,9 +126,10 @@ static __global__ void fused_mul_silu_f32(const float * x, float * dst, const in
int row = i / ne0;
int j = i % ne0;
auto x_row = x + 2*row*ne0;
float g = x_row[j + ne0] / (1.0f + expf(-x_row[j + ne0]));
// Note: gate is at the beginning of each row, up is offset by ne0
float g = x_row[j] / (1.0f + expf(-x_row[j]));
g = min(g, limit);
dst[i] = max(-limit, min(limit, x_row[j])) * g;
dst[i] = max(-limit, min(limit, x_row[j + ne0])) * g;
}
static __global__ void fused_mul_relu_f32(const float * x, const float * y, float * dst, const int k) {
+7 -36
View File
@@ -1327,8 +1327,8 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
} else {
// DSV4 uses separate up and gate expert tensors. Do not silently
// select the fork-only merged gate path for another GGUF.
GGML_ASSERT(model.layers[il].ffn_up_gate_exps == nullptr &&
"merged DSV4 MoE gate tensors use an unsupported layout");
//GGML_ASSERT(model.layers[il].ffn_up_gate_exps == nullptr &&
// "merged DSV4 MoE gate tensors use an unsupported layout");
ggml_tensor * selected_experts = nullptr;
ggml_tensor * exp_probs_b = model.layers[il].ffn_exp_probs_b;
if ((uint32_t) il < hparams.dsv4_hash_layer_count) {
@@ -1344,15 +1344,7 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
? selected_experts->ne[0]
: n_expert_used;
const int64_t dsv4_n_stream = std::max<int64_t>(1, lctx.dsv4.csa_ctx.graph_n_stream);
// Wide packed DSV4 fused/IQK MoE diverges above 1024 total tokens.
// Evaluate each active stream independently to preserve packed parity.
constexpr int64_t dsv4_moe_max_tokens = 1024;
auto build_dsv4_moe = [&](ggml_tensor * moe_cur,
ggml_tensor * moe_exp_probs_b,
ggml_tensor * moe_selected_experts) {
return llm_build_moe_ffn(ctx0, lctx, moe_cur,
auto moe_out = llm_build_moe_ffn(ctx0, lctx, cur,
model.layers[il].ffn_gate_inp,
nullptr,
model.layers[il].ffn_up_exps,
@@ -1361,37 +1353,15 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
nullptr,
model.layers[il].ffn_down_exps,
nullptr,
moe_exp_probs_b,
exp_probs_b,
n_expert, moe_n_expert_used,
LLM_FFN_SILU, hparams.expert_weights_norm,
true, hparams.expert_weights_scale,
(enum llm_expert_gating_func_type) hparams.expert_gating_func,
cb, il, gf, false, model.layers[il].ffn_up_gate_exps, nullptr, nullptr, nullptr,
moe_selected_experts);
};
selected_experts);
ggml_tensor * moe_out = nullptr;
if (dsv4_n_stream > 1 && cur->ne[1] > dsv4_moe_max_tokens &&
cur->ne[1] % dsv4_n_stream == 0) {
const int64_t n_tokens_stream = cur->ne[1]/dsv4_n_stream;
auto stream_view = [&](ggml_tensor * tensor, int64_t stream) {
if (tensor == nullptr || tensor->ne[1] != cur->ne[1]) {
return tensor;
}
return ggml_view_2d(ctx0, tensor, tensor->ne[0], n_tokens_stream,
tensor->nb[1], stream*n_tokens_stream*tensor->nb[1]);
};
for (int64_t stream = 0; stream < dsv4_n_stream; ++stream) {
ggml_tensor * stream_result = build_dsv4_moe(
stream_view(cur, stream),
stream_view(exp_probs_b, stream),
stream_view(selected_experts, stream));
moe_out = moe_out == nullptr ? stream_result : ggml_concat(ctx0, moe_out, stream_result, 1);
}
} else {
moe_out = build_dsv4_moe(cur, exp_probs_b, selected_experts);
}
ggml_build_forward_expand(gf, moe_out);
cb(moe_out, "ffn_moe_out", il);
ggml_tensor * ffn_shexp = llm_build_ffn(ctx0, lctx, nullptr, cur,
@@ -1403,6 +1373,7 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
cb(ffn_shexp, "ffn_shexp", il);
cur = ggml_add(ctx0, moe_out, ffn_shexp);
ggml_build_forward_expand(gf, cur);
}
cb(cur, "ffn_out", il);
+36 -4
View File
@@ -35,6 +35,7 @@ struct create_tensors_helper : public create_tensors_helper_interface {
bool create_std_ffn_exps(int64_t n_embd, const LLM_TN & tn, int i, int flags = 0, int n_ff_exps_input = 0,
ggml_context * ffn_ctx = nullptr);
bool create_std_ffn_exps_from_meta(const LLM_TN & tn, int i, int flags = 0, ggml_context * ffn_ctx = nullptr);
bool create_tensors() override;
@@ -2795,7 +2796,7 @@ bool create_tensors_helper::create_deepseek2_tensors(const LLM_TN & tn) {
return use_mmap_buffer;
}
bool create_tensors_helper::create_deepseek4_tensors(const LLM_TN &) {
bool create_tensors_helper::create_deepseek4_tensors(const LLM_TN & tn) {
LOADING_PRELUDE
auto create_tensor_from_meta = [&](ggml_context * ctx, const std::string & name, int flags = 0) -> ggml_tensor * {
@@ -2956,9 +2957,9 @@ bool create_tensors_helper::create_deepseek4_tensors(const LLM_TN &) {
layer.ffn_gate_inp = create_tensor_from_meta(ctx_split, layer_weight_name(i, "ffn_gate_inp"));
layer.ffn_norm = create_tensor_from_meta(ctx_split, layer_weight_name(i, "ffn_norm"));
layer.ffn_gate_exps = create_tensor_from_meta(ctx_split, layer_weight_name(i, "ffn_gate_exps"));
layer.ffn_down_exps = create_tensor_from_meta(ctx_split, layer_weight_name(i, "ffn_down_exps"));
layer.ffn_up_exps = create_tensor_from_meta(ctx_split, layer_weight_name(i, "ffn_up_exps"));
use_mmap_buffer &= !create_std_ffn_exps_from_meta(tn, i, 0);
layer.ffn_gate_shexp = create_tensor_from_meta(ctx_split, layer_weight_name(i, "ffn_gate_shexp"));
layer.ffn_down_shexp = create_tensor_from_meta(ctx_split, layer_weight_name(i, "ffn_down_shexp"));
layer.ffn_up_shexp = create_tensor_from_meta(ctx_split, layer_weight_name(i, "ffn_up_shexp"));
@@ -4244,6 +4245,37 @@ bool create_tensors_helper::create_std_ffn_exps(int64_t n_embd, const LLM_TN & t
return merged;
}
bool create_tensors_helper::create_std_ffn_exps_from_meta(const LLM_TN & tn, int i, int flags, ggml_context * ffn_ctx) {
auto & layer = model.layers[i];
if (!ffn_ctx) {
ffn_ctx = ctx_for_layer_split(i);
}
bool merged = false;
auto ug_name = tn(LLM_TENSOR_FFN_GATE_UP_EXPS, "weight", i);
auto ug_meta = ml.get_tensor_meta(ug_name.c_str());
if (ug_meta) {
layer.ffn_up_gate_exps = create_tensor(ffn_ctx, ug_name, { ug_meta->ne[0], ug_meta->ne[1], ug_meta->ne[2]}, flags);
} else {
merged = flags == 0 && ml.merge_up_gate_exps && merge_up_gate_exps(tn, i, 0);
if (!merged) {
auto u_name = tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i);
auto g_name = tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i);
auto u_meta = ml.get_tensor_meta(u_name.c_str());
auto g_meta = ml.get_tensor_meta(g_name.c_str());
GGML_ASSERT(u_meta && g_meta);
layer.ffn_up_exps = create_tensor(ffn_ctx, u_name, { u_meta->ne[0], u_meta->ne[1], u_meta->ne[2] }, 0);
layer.ffn_gate_exps = create_tensor(ffn_ctx, g_name, { g_meta->ne[0], g_meta->ne[1], g_meta->ne[2] }, 0);
}
}
auto d_name = tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i);
auto d_meta = ml.get_tensor_meta(d_name.c_str());
GGML_ASSERT(d_meta);
layer.ffn_down_exps = create_tensor(ffn_ctx, d_name, { d_meta->ne[0], d_meta->ne[1], d_meta->ne[2] }, 0);
return merged;
}
bool create_tensors_helper::merge_qkv(const LLM_TN & tn, int i, int bias, bool ignore_attn_scale) {
auto& hparams = model.hparams;
const int64_t n_head = hparams.n_head(i);