mirror of
https://github.com/ikawrakow/ik_llama.cpp.git
synced 2026-08-12 22:29:39 +04:00
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:
@@ -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) {
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user