mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-10-24 21:47:33 -05:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
74bca39500 | ||
|
|
d327bf1f07 | ||
|
|
78a8ae14be | ||
|
|
f3269ef96a | ||
|
|
3f03b17724 | ||
|
|
9a0fde9967 | ||
|
|
e168e041ec | ||
|
|
ea1d7d66f3 | ||
|
|
71255fefc0 | ||
|
|
73adf89e38 | ||
|
|
df9beac715 | ||
|
|
4d32b29aaa | ||
|
|
3032209758 | ||
|
|
701fd25b3a |
@@ -105,7 +105,7 @@ jobs:
|
||||
cd build
|
||||
# Metal Paravirtual devices are difficult to support -> disable
|
||||
# ref: https://github.com/ggml-org/llama.cpp/pull/19802#issuecomment-4013704023
|
||||
ctest -L main -E "test-llama-archs|test-llama-context" --verbose --timeout 900
|
||||
ctest -L main -E "test-llama-archs|test-llama-context|test-speculative-rejection-sampling" --verbose --timeout 900
|
||||
|
||||
macos-latest-x64:
|
||||
runs-on: macos-15-intel
|
||||
|
||||
@@ -56,7 +56,7 @@ jobs:
|
||||
- name: Test
|
||||
id: ggml-ci
|
||||
run: |
|
||||
LLAMA_ARG_THREADS=$(nproc) GG_BUILD_HIGH_PERF=1 GG_BUILD_EXTRA_TESTS_0=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
|
||||
LLAMA_ARG_THREADS=$(nproc) GG_BUILD_HIGH_PERF=1 GG_BUILD_EXTRA_TESTS_0=1 bash ./ci/run.sh ~/results/llama.cpp/$RUNNER_NAME ~/mnt/llama.cpp
|
||||
|
||||
cpu-arm64-high-perf-graviton4:
|
||||
runs-on: ah-ubuntu_24_04-c8g_8x
|
||||
|
||||
@@ -79,7 +79,7 @@ jobs:
|
||||
id: ggml-ci
|
||||
run: |
|
||||
vulkaninfo --summary
|
||||
GG_BUILD_VULKAN=1 GGML_VK_DISABLE_COOPMAT2=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
|
||||
GG_BUILD_VULKAN=1 GGML_VK_DISABLE_COOPMAT2=1 bash ./ci/run.sh ~/results/llama.cpp/$RUNNER_NAME ~/mnt/llama.cpp
|
||||
|
||||
# - name: ccache-buckets-save
|
||||
# if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
|
||||
@@ -124,7 +124,7 @@ jobs:
|
||||
id: ggml-ci
|
||||
run: |
|
||||
vulkaninfo --summary
|
||||
GG_BUILD_VULKAN=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
|
||||
GG_BUILD_VULKAN=1 bash ./ci/run.sh ~/results/llama.cpp/$RUNNER_NAME ~/mnt/llama.cpp
|
||||
|
||||
# - name: ccache-buckets-save
|
||||
# if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
|
||||
|
||||
+89
-6
@@ -428,6 +428,41 @@ std::vector<common_chat_msg> common_chat_msgs_parse_oaicompat(const json & messa
|
||||
if (message.contains("tool_call_id")) {
|
||||
msg.tool_call_id = message.at("tool_call_id");
|
||||
}
|
||||
{
|
||||
json extra = json::object();
|
||||
for (const auto & [k, v] : message.items()) {
|
||||
if (k != "role" && k != "content" && k != "tool_calls" && k != "reasoning_content" &&
|
||||
k != "name" && k != "tool_call_id") {
|
||||
extra[k] = v;
|
||||
}
|
||||
}
|
||||
if (has_tool_calls) {
|
||||
// per tool call: keys besides type/function/id, and function keys besides name/arguments
|
||||
json tc_extra = json::array();
|
||||
bool any = false;
|
||||
for (const auto & tool_call : message.at("tool_calls")) {
|
||||
json e = json::object();
|
||||
for (const auto & [k, v] : tool_call.items()) {
|
||||
if (k != "type" && k != "function" && k != "id") {
|
||||
e[k] = v;
|
||||
}
|
||||
}
|
||||
for (const auto & [k, v] : tool_call.at("function").items()) {
|
||||
if (k != "name" && k != "arguments") {
|
||||
e["function"][k] = v;
|
||||
}
|
||||
}
|
||||
any = any || !e.empty();
|
||||
tc_extra.push_back(e);
|
||||
}
|
||||
if (any) {
|
||||
extra["__tool_call_extras"] = tc_extra;
|
||||
}
|
||||
}
|
||||
if (!extra.empty()) {
|
||||
msg.extra_json = extra.dump();
|
||||
}
|
||||
}
|
||||
|
||||
msgs.push_back(msg);
|
||||
}
|
||||
@@ -516,6 +551,50 @@ static json render_message_to_json(const std::vector<common_chat_msg> & msgs, co
|
||||
return messages_inp_normalizer(c).normalize(messages);
|
||||
}
|
||||
|
||||
// DeepSeek V4.1 reads non-OpenAI message keys and joins content blocks with "\n\n" around images,
|
||||
// so keep media markers as separate text parts and restore the extra keys
|
||||
static json render_message_to_json_deepseek_v4_1(const std::vector<common_chat_msg> & msgs, const jinja::caps & c) {
|
||||
json messages = json::array();
|
||||
for (const auto & msg : msgs) {
|
||||
json jmsg;
|
||||
if (msg.contains_media() && msg.content.empty()) {
|
||||
common_chat_msg typed = msg;
|
||||
for (auto & part : typed.content_parts) {
|
||||
part.type = "text";
|
||||
}
|
||||
jmsg = typed.to_json_oaicompat(/* concat_typed_text= */ false);
|
||||
} else {
|
||||
jmsg = msg.to_json_oaicompat(/* concat_typed_text= */ false);
|
||||
}
|
||||
if (!msg.extra_json.empty()) {
|
||||
const json extra = json::parse(msg.extra_json);
|
||||
for (const auto & [k, v] : extra.items()) {
|
||||
if (k == "__tool_call_extras") {
|
||||
if (!jmsg.contains("tool_calls")) {
|
||||
continue;
|
||||
}
|
||||
auto & tcs = jmsg.at("tool_calls");
|
||||
for (size_t i = 0; i < tcs.size() && i < v.size(); i++) {
|
||||
for (const auto & [tk, tv] : v[i].items()) {
|
||||
if (tk == "function") {
|
||||
for (const auto & [fk, fv] : tv.items()) {
|
||||
tcs[i]["function"][fk] = fv;
|
||||
}
|
||||
} else {
|
||||
tcs[i][tk] = tv;
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if (!jmsg.contains(k)) {
|
||||
jmsg[k] = v;
|
||||
}
|
||||
}
|
||||
}
|
||||
messages.push_back(std::move(jmsg));
|
||||
}
|
||||
return messages_inp_normalizer(c).normalize(messages);
|
||||
}
|
||||
|
||||
// DEPRECATED: only used in tests
|
||||
json common_chat_msgs_to_json_oaicompat(const std::vector<common_chat_msg> & msgs, bool concat_typed_text) {
|
||||
jinja::caps c;
|
||||
@@ -1208,11 +1287,13 @@ static common_chat_params_init_fn common_chat_template_detect_params_init(const
|
||||
// DeepSeek V3.2/V4 format detection: template defines dsml_token and uses it for tool calls.
|
||||
// The template source contains the token as a variable assignment, not as a literal in markup.
|
||||
// V3.2 names the tool call block "function_calls", V4 names it "tool_calls".
|
||||
if (src.find("dsml_token") != std::string::npos &&
|
||||
src.find("DSML") != std::string::npos &&
|
||||
(src.find("function_calls") != std::string::npos ||
|
||||
src.find("tool_calls") != std::string::npos)) {
|
||||
LOG_DBG("Using specialized template: DeepSeek V3.2/V4\n");
|
||||
// V4.1 tag names have a leading space and its template does not name the dsml_token variable.
|
||||
if ((src.find("dsml_token") != std::string::npos &&
|
||||
src.find("DSML") != std::string::npos &&
|
||||
(src.find("function_calls") != std::string::npos ||
|
||||
src.find("tool_calls") != std::string::npos)) ||
|
||||
src.find("|DSML| calls>") != std::string::npos) {
|
||||
LOG_DBG("Using specialized template: DeepSeek V3.2/V4/V4.1\n");
|
||||
return common_chat_params_init_deepseek_v3_2;
|
||||
}
|
||||
|
||||
@@ -1313,7 +1394,9 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_
|
||||
workaround::trim_all_content(trimmed_messages);
|
||||
messages_to_render = &trimmed_messages;
|
||||
}
|
||||
params.messages = render_message_to_json(*messages_to_render, tmpl.original_caps());
|
||||
params.messages = src.find("|DSML| calls>") != std::string::npos ?
|
||||
render_message_to_json_deepseek_v4_1(*messages_to_render, tmpl.original_caps()) :
|
||||
render_message_to_json(*messages_to_render, tmpl.original_caps());
|
||||
params.tool_choice = inputs.tool_choice;
|
||||
params.reasoning_format = inputs.reasoning_format;
|
||||
params.enable_thinking = inputs.enable_thinking;
|
||||
|
||||
@@ -94,6 +94,8 @@ struct common_chat_msg {
|
||||
std::string reasoning_content;
|
||||
std::string tool_name;
|
||||
std::string tool_call_id;
|
||||
// non-OpenAI message keys as a JSON object string, only for templates that read them (DeepSeek V4.1)
|
||||
std::string extra_json;
|
||||
|
||||
common_json to_json_oaicompat(bool concat_typed_text = false) const;
|
||||
|
||||
|
||||
+42
-63
@@ -299,74 +299,53 @@ static value & get_attribute(const value & val, const value & attr, value & defa
|
||||
return default_val;
|
||||
}
|
||||
|
||||
template<bool is_reject>
|
||||
static value selectattr(const func_args & args) {
|
||||
args.ensure_count(2, 4);
|
||||
args.ensure_vals<value_array, value_string, value_string, value_string>(true, true, false, false);
|
||||
template<bool is_reject, bool lookup_attr>
|
||||
static value select_or_reject(const func_args & args) {
|
||||
// example: array | select("odd") -> test_is_odd(item)
|
||||
// example: array | selectattr("status", "equalto", "active") -> test_is_equalto(item.status, "active")
|
||||
constexpr size_t test_pos = lookup_attr ? 2 : 1;
|
||||
args.ensure_count(test_pos);
|
||||
args.ensure_val<value_array>(args.get_pos(0));
|
||||
value attr;
|
||||
if constexpr (lookup_attr) {
|
||||
attr = args.get_pos(1);
|
||||
args.ensure_val<value_string>(attr);
|
||||
}
|
||||
|
||||
auto & builtins = global_builtins();
|
||||
auto test_it = builtins.end();
|
||||
if (args.count() > test_pos) {
|
||||
args.ensure_val<value_string>(args.get_pos(test_pos));
|
||||
std::string test_name = args.get_pos(test_pos)->as_string().str();
|
||||
test_it = builtins.find("test_is_" + test_name);
|
||||
if (test_it == builtins.end()) {
|
||||
throw raised_exception(args.func_name + ": unknown test '" + test_name + "'");
|
||||
}
|
||||
}
|
||||
|
||||
auto arr = args.get_pos(0)->as_array();
|
||||
auto attribute = args.get_pos(1);
|
||||
auto out = mk_val<value_array>();
|
||||
value val_default = mk_val<value_undefined>();
|
||||
|
||||
if (args.count() == 2) {
|
||||
// example: array | selectattr("active")
|
||||
for (const auto & item : arr) {
|
||||
value attr_val = get_attribute(item, attribute, val_default);
|
||||
bool is_selected = attr_val->as_bool();
|
||||
if constexpr (is_reject) is_selected = !is_selected;
|
||||
if (is_selected) out->push_back(item);
|
||||
for (const auto & item : arr) {
|
||||
value val = item;
|
||||
if constexpr (lookup_attr) {
|
||||
val = get_attribute(item, attr, val_default);
|
||||
}
|
||||
return out;
|
||||
|
||||
} else if (args.count() == 3) {
|
||||
// example: array | selectattr("equalto", "text")
|
||||
// translated to: test_is_equalto(item, "text")
|
||||
std::string test_name = args.get_pos(1)->as_string().str();
|
||||
value test_val = args.get_pos(2);
|
||||
auto & builtins = global_builtins();
|
||||
auto it = builtins.find("test_is_" + test_name);
|
||||
if (it == builtins.end()) {
|
||||
throw raised_exception("selectattr: unknown test '" + test_name + "'");
|
||||
}
|
||||
auto test_fn = it->second;
|
||||
for (const auto & item : arr) {
|
||||
bool is_selected;
|
||||
if (test_it == builtins.end()) {
|
||||
is_selected = val->as_bool();
|
||||
} else {
|
||||
func_args test_args(args.ctx);
|
||||
test_args.push_back(item); // current object
|
||||
test_args.push_back(test_val); // extra argument
|
||||
value test_result = test_fn(test_args);
|
||||
bool is_selected = test_result->as_bool();
|
||||
if constexpr (is_reject) is_selected = !is_selected;
|
||||
if (is_selected) out->push_back(item);
|
||||
test_args.push_back(val);
|
||||
for (size_t i = test_pos + 1; i < args.count(); i++) {
|
||||
test_args.push_back(args.get_pos(i));
|
||||
}
|
||||
is_selected = test_it->second(test_args)->as_bool();
|
||||
}
|
||||
return out;
|
||||
|
||||
} else if (args.count() == 4) {
|
||||
// example: array | selectattr("status", "equalto", "active")
|
||||
// translated to: test_is_equalto(item.status, "active")
|
||||
std::string test_name = args.get_pos(2)->as_string().str();
|
||||
auto extra_arg = args.get_pos(3);
|
||||
auto & builtins = global_builtins();
|
||||
auto it = builtins.find("test_is_" + test_name);
|
||||
if (it == builtins.end()) {
|
||||
throw raised_exception("selectattr: unknown test '" + test_name + "'");
|
||||
}
|
||||
auto test_fn = it->second;
|
||||
for (const auto & item : arr) {
|
||||
value attr_val = get_attribute(item, attribute, val_default);
|
||||
func_args test_args(args.ctx);
|
||||
test_args.push_back(attr_val); // attribute value
|
||||
test_args.push_back(extra_arg); // extra argument
|
||||
value test_result = test_fn(test_args);
|
||||
bool is_selected = test_result->as_bool();
|
||||
if constexpr (is_reject) is_selected = !is_selected;
|
||||
if (is_selected) out->push_back(item);
|
||||
}
|
||||
return out;
|
||||
} else {
|
||||
throw raised_exception("selectattr: invalid number of arguments");
|
||||
if constexpr (is_reject) is_selected = !is_selected;
|
||||
if (is_selected) out->push_back(item);
|
||||
}
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
@@ -1095,10 +1074,10 @@ const func_builtins & value_array_t::get_builtins() const {
|
||||
auto arr = slice(val->as_array(), start, stop, step);
|
||||
return is_val<value_tuple>(val) ? mk_val<value_tuple>(std::move(arr)) : mk_val<value_array>(std::move(arr));
|
||||
}},
|
||||
{"selectattr", selectattr<false>},
|
||||
{"select", selectattr<false>},
|
||||
{"rejectattr", selectattr<true>},
|
||||
{"reject", selectattr<true>},
|
||||
{"selectattr", select_or_reject<false, true>},
|
||||
{"select", select_or_reject<false, false>},
|
||||
{"rejectattr", select_or_reject<true, true>},
|
||||
{"reject", select_or_reject<true, false>},
|
||||
{"join", [](const func_args & args) -> value {
|
||||
args.ensure_count(1, 3);
|
||||
if (!is_val<value_array>(args.get_pos(0))) {
|
||||
|
||||
+53
-10
@@ -1,4 +1,5 @@
|
||||
#include "parsers.h"
|
||||
#include "log.h"
|
||||
|
||||
// The DeepSeek V4 reference implementation renders consecutive tool results into a single
|
||||
// user block, ordered by the tool call order of the preceding assistant message (matched
|
||||
@@ -75,13 +76,30 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
|
||||
// non-thinking generation prompt ends with a bare </think> instead of an empty
|
||||
// <think></think> pair.
|
||||
const bool is_v4 = tmpl.source().find("function_calls") == std::string::npos;
|
||||
// V4.1: DSML tag names carry a leading space, e.g. <|DSML| calls>
|
||||
const bool is_v41 = is_v4 && tmpl.source().find("|DSML| calls>") != std::string::npos;
|
||||
|
||||
std::optional<json> adjusted_messages;
|
||||
if (is_v4) {
|
||||
adjusted_messages = deepseek_v4_sort_tool_results(inputs.messages);
|
||||
}
|
||||
|
||||
auto has_tools = inputs.tools.is_array() && !inputs.tools.empty();
|
||||
// V4.1 can declare tools on a system message instead of top-level "tools": parse calls to those too.
|
||||
json tools = inputs.tools;
|
||||
if (is_v41 && !(tools.is_array() && !tools.empty()) && inputs.messages.is_array()) {
|
||||
for (const auto & msg : inputs.messages) {
|
||||
if (msg.value("role", "") == "system" && msg.contains("tools") && msg.at("tools").is_array()) {
|
||||
if (!tools.is_array()) {
|
||||
tools = json::array();
|
||||
}
|
||||
for (const auto & tool : msg.at("tools")) {
|
||||
tools.push_back(tool);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
auto has_tools = tools.is_array() && !tools.empty();
|
||||
auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object();
|
||||
auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
|
||||
auto include_grammar = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE);
|
||||
@@ -90,17 +108,39 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
|
||||
if (is_v4 && has_response_format) {
|
||||
additional_context = json{ { "response_format", inputs.json_schema } };
|
||||
}
|
||||
// V4.1 raises on anything but low/high/max or 1-100: map minimal -> low, xhigh -> max, medium -> default
|
||||
if (is_v41 && inputs.extra_context.is_object() && inputs.extra_context.contains("reasoning_effort") &&
|
||||
inputs.extra_context.at("reasoning_effort").is_string()) {
|
||||
const std::string effort = inputs.extra_context.at("reasoning_effort").get<std::string>();
|
||||
json mapped;
|
||||
if (effort == "minimal") {
|
||||
mapped = "low";
|
||||
} else if (effort == "xhigh") {
|
||||
mapped = "max";
|
||||
} else if (effort == "medium") {
|
||||
mapped = nullptr;
|
||||
}
|
||||
if (effort == "minimal" || effort == "xhigh" || effort == "medium") {
|
||||
LOG_DBG("%s: DeepSeek V4.1 reasoning_effort '%s' -> %s\n", __func__, effort.c_str(),
|
||||
mapped.is_null() ? "unset (template default)" : mapped.get<std::string>().c_str());
|
||||
if (!additional_context.has_value()) {
|
||||
additional_context = json::object();
|
||||
}
|
||||
(*additional_context)["reasoning_effort"] = mapped;
|
||||
}
|
||||
}
|
||||
|
||||
const std::string DSML = "|DSML|";
|
||||
const std::string THINK_START = "<think>";
|
||||
const std::string THINK_END = "</think>";
|
||||
const std::string TC_BLOCK = is_v4 ? "tool_calls" : "function_calls";
|
||||
const std::string FC_START = "<" + DSML + TC_BLOCK + ">";
|
||||
const std::string FC_END = "</" + DSML + TC_BLOCK + ">";
|
||||
const std::string INVOKE_START = "<" + DSML + "invoke";
|
||||
const std::string INVOKE_END = "</" + DSML + "invoke>";
|
||||
const std::string PARAM_START = "<" + DSML + "parameter";
|
||||
const std::string PARAM_END = "</" + DSML + "parameter>";
|
||||
const std::string TAG_SP = is_v41 ? " " : "";
|
||||
const std::string TC_BLOCK = is_v41 ? "calls" : is_v4 ? "tool_calls" : "function_calls";
|
||||
const std::string FC_START = "<" + DSML + TAG_SP + TC_BLOCK + ">";
|
||||
const std::string FC_END = "</" + DSML + TAG_SP + TC_BLOCK + ">";
|
||||
const std::string INVOKE_START = "<" + DSML + TAG_SP + "invoke";
|
||||
const std::string INVOKE_END = "</" + DSML + TAG_SP + "invoke>";
|
||||
const std::string PARAM_START = "<" + DSML + TAG_SP + "parameter";
|
||||
const std::string PARAM_END = "</" + DSML + TAG_SP + "parameter>";
|
||||
const std::string GEN_PROMPT = "<|Assistant|>";
|
||||
const std::string TC_SEPARATOR = "\n\n";
|
||||
|
||||
@@ -152,7 +192,7 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
|
||||
// build tool call section first since we might need it in reasoning
|
||||
auto tool_choice = p.choice();
|
||||
if (has_tool_calls) {
|
||||
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
|
||||
foreach_function(tools, [&](size_t tool_index, const json & tool) {
|
||||
const auto & function = tool.at("function");
|
||||
std::string name = function.at("name");
|
||||
|
||||
@@ -223,7 +263,10 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
|
||||
}
|
||||
|
||||
if (extract_reasoning && inputs.enable_thinking) {
|
||||
reasoning = p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END);
|
||||
// V4.1 renders a continuation without reasoning as "<|Assistant|></think>" even in thinking mode,
|
||||
// so consume the bare THINK_END too or it leaks into the content
|
||||
reasoning = is_v41 ? p.optional((THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END) | p.literal(THINK_END))
|
||||
: p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END);
|
||||
reasoning_with_tc = THINK_START +
|
||||
p.reasoning(p.until_one_of({ TC_SEPARATOR + FC_START, FC_START, THINK_END })) +
|
||||
p.space() + obligatory_tool_calls;
|
||||
|
||||
@@ -70,6 +70,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
|
||||
"LingDSparkModel": "qwen",
|
||||
"DeepseekV4ForCausalLM": "deepseek",
|
||||
"DeepseekV4DSparkModel": "deepseek",
|
||||
"DeepseekV41ForCausalLM": "deepseek41",
|
||||
"DistilBertForMaskedLM": "bert",
|
||||
"DistilBertForSequenceClassification": "bert",
|
||||
"DistilBertModel": "bert",
|
||||
@@ -326,6 +327,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
|
||||
"VllmD3Model": "vllm_d3",
|
||||
"DeepseekOCR2ForCausalLM": "deepseek",
|
||||
"DeepseekOCRForCausalLM": "deepseek",
|
||||
"DeepseekV41ForCausalLM": "deepseek41",
|
||||
"DeepseekV4ForCausalLM": "deepseek",
|
||||
"Dots3NoteForCausalLM": "dots3",
|
||||
"Dots3NoteForConditionalGeneration": "dots3",
|
||||
|
||||
@@ -0,0 +1,467 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Iterable, TYPE_CHECKING, cast
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch import Tensor
|
||||
|
||||
from .base import LazyTorchTensor, ModelBase, TextModel, gguf
|
||||
from .deepseek import DeepseekV4FlashVisionModel, DeepseekV4Model
|
||||
|
||||
|
||||
# Engram hash constants from inference/engram.py, stored in the GGUF: the token map needs the HF tokenizers
|
||||
# normalizers and the multipliers numpy's PCG64 stream, neither of which the runtime has.
|
||||
|
||||
def _is_prime(n: int) -> bool:
|
||||
if n < 2 or n % 2 == 0:
|
||||
return n == 2
|
||||
return all(n % d for d in range(3, math.isqrt(n) + 1, 2))
|
||||
|
||||
|
||||
def engram_primes(n_layers: int, max_ngram_size: int, n_heads: int, vocab_size: int) -> list[list[list[int]]]:
|
||||
"""EngramLayout.from_args: [layer][ngram order 2..max][head] bucket moduli, never reused across layers."""
|
||||
seen: set[int] = set()
|
||||
out: list[list[list[int]]] = []
|
||||
for _ in range(n_layers):
|
||||
per_ngram = []
|
||||
for _ in range(max_ngram_size - 1):
|
||||
sizes, cur = [], vocab_size - 1
|
||||
for _ in range(n_heads):
|
||||
cand = cur + 1
|
||||
while not _is_prime(cand) or cand in seen:
|
||||
cand += 1
|
||||
cur = cand
|
||||
seen.add(cur)
|
||||
sizes.append(cur)
|
||||
per_ngram.append(sizes)
|
||||
out.append(per_ngram)
|
||||
return out
|
||||
|
||||
|
||||
def engram_multipliers(layer_ids: list[int], max_ngram_size: int, compressed_vocab_size: int) -> list[list[int]]:
|
||||
"""compute_hash_multipliers: odd, bounded so that token_id * multiplier cannot overflow int64."""
|
||||
bound = max(1, (np.iinfo(np.int64).max // compressed_vocab_size) // 2)
|
||||
rows = []
|
||||
for layer_id in layer_ids:
|
||||
rng = np.random.default_rng(10007 * layer_id)
|
||||
values = rng.integers(low=0, high=bound, size=(max_ngram_size,), dtype=np.int64)
|
||||
rows.append([int(v) * 2 + 1 for v in values])
|
||||
return rows
|
||||
|
||||
|
||||
def engram_compressed_token_map(tokenizer_dir: Path) -> tuple[list[int], int]:
|
||||
"""build_compressed_token_map: token id -> compressed id (tokens that normalize alike collapse)."""
|
||||
from tokenizers import Regex, normalizers
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
|
||||
sentinel = "\ue000"
|
||||
normalizer = normalizers.Sequence([
|
||||
normalizers.NFKC(),
|
||||
normalizers.NFD(),
|
||||
normalizers.StripAccents(),
|
||||
normalizers.Lowercase(),
|
||||
normalizers.Replace(Regex(r"[ \t\r\n]+"), " "),
|
||||
normalizers.Replace(Regex(r"^ $"), sentinel),
|
||||
normalizers.Strip(),
|
||||
normalizers.Replace(sentinel, " "),
|
||||
])
|
||||
backend = tokenizer.backend_tokenizer
|
||||
key_to_new: dict[str, int] = {}
|
||||
lookup = [0] * len(tokenizer)
|
||||
for token_id in range(len(tokenizer)):
|
||||
text = backend.decode([token_id], skip_special_tokens=False)
|
||||
if "\ufffd" in text:
|
||||
key = backend.id_to_token(token_id)
|
||||
else:
|
||||
normalized = normalizer.normalize_str(text)
|
||||
key = normalized if normalized else text
|
||||
lookup[token_id] = key_to_new.setdefault(key, len(key_to_new))
|
||||
return lookup, len(key_to_new)
|
||||
|
||||
|
||||
class _StreamedBytes:
|
||||
"""Raw tensor bytes produced chunk by chunk at write time (an Engram table is ~200 GB in BF16)."""
|
||||
def __init__(self, shape: tuple[int, ...], chunks: Callable[[], Iterable[np.ndarray]]):
|
||||
self.shape = shape
|
||||
self.dtype = np.dtype(np.uint8)
|
||||
self.nbytes = int(np.prod(shape))
|
||||
self._chunks = chunks
|
||||
|
||||
def tofile(self, fout) -> None:
|
||||
written = 0
|
||||
for chunk in self._chunks():
|
||||
chunk = np.ascontiguousarray(chunk, dtype=np.uint8)
|
||||
chunk.tofile(fout)
|
||||
written += chunk.nbytes
|
||||
if written != self.nbytes:
|
||||
raise ValueError(f"streamed {written} bytes, expected {self.nbytes}")
|
||||
|
||||
def byteswap(self, inplace: bool = False):
|
||||
raise NotImplementedError("streamed tensors are written little-endian only")
|
||||
|
||||
|
||||
def _e8m0_check(scale_u8: np.ndarray, name: str) -> None:
|
||||
# e4m3 (|x| in [2^-9, 448]) * 2^(e - 127) is exact in BF16 for e in [3, 246]
|
||||
if scale_u8.size and (int(scale_u8.min()) < 3 or int(scale_u8.max()) > 246):
|
||||
raise ValueError(f"{name}: E8M0 scale exponents outside [3, 246], the dequantized values would not be exact")
|
||||
|
||||
|
||||
@ModelBase.register("DeepseekV41ForCausalLM")
|
||||
class DeepseekV41Model(DeepseekV4Model):
|
||||
"""DeepSeek-V4.1 text model, lossless: FP8 (e4m3, 32x32 ue8m0 blocks) linears dequantized to BF16 (F32 with
|
||||
--outtype f32), routed experts copied bit-exact as MXFP4, Engram tables (e4m3 rows, one e8m0 scale per 32) as the
|
||||
exact BF16 lookup (F32 with --outtype f32), router and Engram gate F32. DSpark (mtp.*) is not converted here
|
||||
(--dspark, DeepseekV41DSparkModel); the vision tower goes to an mmproj (--mmproj, DeepseekV41VisionModel).
|
||||
The Engram tables of all layers are written as one per_layer_token_embd (rows concatenated in layer order) with
|
||||
the PLE keys of qwen4exp: absolute head offsets and bucket counts, so the runtime can read it lazily."""
|
||||
model_arch = gguf.MODEL_ARCH.DEEPSEEK41
|
||||
supports_mtp_export = False
|
||||
|
||||
_ENGRAM_EMBED_RE = re.compile(r"layers\.(\d+)\.engram\.embed\.(weight|scale)$")
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.hparams.setdefault("num_hash_layers", 0) # no hash-routed layers in V4.1
|
||||
self._f32 = self.ftype == gguf.LlamaFileType.ALL_F32
|
||||
qc = self.hparams.get("quantization_config") or {}
|
||||
self._fp8_block = tuple(qc.get("weight_block_size") or (32, 32))
|
||||
if qc.get("scale_fmt", "ue8m0") != "ue8m0":
|
||||
raise ValueError(f"unsupported FP8 scale format {qc.get('scale_fmt')!r}")
|
||||
if self.hparams.get("gate_temp", 1.0) != 1.0:
|
||||
raise ValueError("gate_temp != 1 is not supported (the router logits would need dividing by it)")
|
||||
|
||||
# V4.1 does not use the DeepSeek-V4 template; a template shipped with the model is still added by set_vocab
|
||||
self.gguf_writer.remove_key(gguf.Keys.Tokenizer.CHAT_TEMPLATE)
|
||||
|
||||
def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
|
||||
# num_hidden_layers lives in text_config, which is merged only after indexing; MTP stages are skipped anyway
|
||||
return TextModel.index_tensors(self, remote_hf_model_id=remote_hf_model_id)
|
||||
|
||||
@classmethod
|
||||
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
|
||||
if item[0].startswith("vision."):
|
||||
return None
|
||||
return super().filter_tensors(item)
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
super().set_gguf_parameters()
|
||||
hp = self.hparams
|
||||
w = self.gguf_writer
|
||||
w.add_attention_kv_source_layers(hp["kv_source_layer_ids"])
|
||||
w.add_indexer_source_layers(hp["index_source_layer_ids"])
|
||||
if hp.get("candidate_source_layer_id", -1) >= 0:
|
||||
w.add_indexer_candidate_source_layer(hp["candidate_source_layer_id"])
|
||||
w.add_indexer_candidate_top_k_blocks(hp["candidate_topk_blocks"])
|
||||
w.add_indexer_candidate_block_size(hp["candidate_block_size"])
|
||||
|
||||
# GGUF cannot hold empty arrays: the Engram keys are omitted when the model has no Engram layers
|
||||
layer_ids = list(hp.get("engram_layer_ids") or [])
|
||||
if layer_ids:
|
||||
max_ngram, n_heads = hp["engram_max_ngram_size"], hp["engram_n_heads"]
|
||||
primes = engram_primes(len(layer_ids), max_ngram, n_heads, hp["engram_vocab_size"])
|
||||
for li, layer_primes in enumerate(primes):
|
||||
if sum(p for per in layer_primes for p in per) != hp["engram_num_embeddings"][li]:
|
||||
raise ValueError(f"engram layer {layer_ids[li]}: prime sum != num_embeddings")
|
||||
token_map, cvocab = engram_compressed_token_map(self.dir_model)
|
||||
if cvocab != hp["engram_compressed_vocab_size"] or len(token_map) != hp["vocab_size"]:
|
||||
raise ValueError(f"compressed token map ({len(token_map)} -> {cvocab}) does not match the config")
|
||||
# absolute row offsets into the concatenated table
|
||||
sizes = [p for layer in primes for per in layer for p in per]
|
||||
offsets = [sum(sizes[:i]) for i in range(len(sizes))]
|
||||
w.add_key_value(gguf.Keys.PerLayerEmbedding.LAYERS.format(arch=w.arch), layer_ids,
|
||||
gguf.GGUFValueType.ARRAY, gguf.GGUFValueType.UINT32)
|
||||
w.add_ple_ngram_size(max_ngram)
|
||||
w.add_ple_heads_per_ngram(n_heads)
|
||||
w.add_embedding_length_per_layer_input(hp["engram_head_dim"])
|
||||
w.add_ple_layer_multipliers([m for row in engram_multipliers(layer_ids, max_ngram, cvocab) for m in row])
|
||||
w.add_ple_head_offsets(offsets)
|
||||
w.add_ple_head_vocab_sizes(sizes)
|
||||
w.add_engram_pad_token_id(hp["engram_pad_token_id"])
|
||||
w.add_engram_token_map(token_map)
|
||||
|
||||
def dequant_model(self):
|
||||
bo, bi = self._fp8_block
|
||||
|
||||
def dequant(weight: Tensor, scale: Tensor) -> Tensor:
|
||||
out_f, in_f = weight.shape
|
||||
s = self._e8m0_to_float(scale).repeat_interleave(bo, 0)[:out_f].repeat_interleave(bi, 1)[:, :in_f]
|
||||
return weight.float() * s
|
||||
|
||||
for name in [n for n in self.model_tensors if n.endswith(".scale") and not self._ENGRAM_EMBED_RE.match(n)]:
|
||||
weight_name = name.removesuffix(".scale") + ".weight"
|
||||
if weight_name not in self.model_tensors or self.model_tensors[weight_name]().dtype not in self._float8_dtypes():
|
||||
continue
|
||||
w, s = self.model_tensors[weight_name], self.model_tensors.pop(name)
|
||||
out_f, in_f = w().shape
|
||||
if tuple(s().shape) != (-(-out_f // bo), -(-in_f // bi)):
|
||||
raise ValueError(f"{weight_name}: scale shape {tuple(s().shape)} does not match {(bo, bi)} blocks")
|
||||
_e8m0_check(LazyTorchTensor.to_eager(s()).view(torch.uint8).numpy(), name)
|
||||
self.model_tensors[weight_name] = lambda w=w, s=s: dequant(w(), s())
|
||||
self._dsv4_fp8_dequantized.add(weight_name)
|
||||
|
||||
def _tensor_part(self, name: str) -> Path:
|
||||
index = self.dir_model / "model.safetensors.index.json"
|
||||
if not index.is_file():
|
||||
return self.dir_model / "model.safetensors"
|
||||
with open(index, "r", encoding="utf-8") as f:
|
||||
return self.dir_model / json.load(f)["weight_map"][name]
|
||||
|
||||
def _write_mxfp4_experts(self, bid: int, proj: str, key: gguf.MODEL_TENSOR, n_expert: int | None = None) -> None:
|
||||
n_expert = self.hparams["n_routed_experts"] if n_expert is None else n_expert
|
||||
names = [f"layers.{bid}.ffn.experts.{e}.{proj}" for e in range(n_expert)]
|
||||
pairs = [(self.model_tensors.pop(n + ".weight"), self.model_tensors.pop(n + ".scale")) for n in names]
|
||||
rows, packed_cols = pairs[0][0]().shape
|
||||
|
||||
def chunks() -> Iterable[np.ndarray]:
|
||||
for w, s in pairs:
|
||||
yield self.repack_mxfp4_blocks(LazyTorchTensor.to_eager(w()), LazyTorchTensor.to_eager(s()))
|
||||
data = cast(np.ndarray, _StreamedBytes((len(pairs), rows, packed_cols * 2 // 32 * 17), chunks))
|
||||
self.gguf_writer.add_tensor(self.format_tensor_name(key, bid), data, raw_dtype=gguf.GGMLQuantizationType.MXFP4)
|
||||
|
||||
def _write_engram_tables(self, layer_ids: list[int]) -> None:
|
||||
# official lookup e4m3 * 2^(e8m0 - 127), exact in BF16
|
||||
parts, dim = [], None
|
||||
for bid in layer_ids:
|
||||
wn, sn = f"layers.{bid}.engram.embed.weight", f"layers.{bid}.engram.embed.scale"
|
||||
rows, d = self.model_tensors.pop(wn)().shape
|
||||
if d % 32 != 0 or tuple(self.model_tensors.pop(sn)().shape) != (rows, d // 32) or d != (dim or d):
|
||||
raise ValueError(f"{wn}: expected one e8m0 scale per 32 elements and the same width in every layer")
|
||||
if rows != self.hparams["engram_num_embeddings"][len(parts)]:
|
||||
raise ValueError(f"{wn}: {rows} rows, config says {self.hparams['engram_num_embeddings'][len(parts)]}")
|
||||
dim = d
|
||||
parts.append((wn, sn, rows, self._tensor_part(wn), self._tensor_part(sn)))
|
||||
assert dim is not None
|
||||
name = self.format_tensor_name(gguf.MODEL_TENSOR.PER_LAYER_TOKEN_EMBD)
|
||||
n_rows = sum(p[2] for p in parts)
|
||||
|
||||
def chunks() -> Iterable[np.ndarray]:
|
||||
from safetensors import safe_open
|
||||
for wn, sn, rows, wpart, spart in parts:
|
||||
with safe_open(wpart, framework="pt") as fw, safe_open(spart, framework="pt") as fs:
|
||||
ws, ss = fw.get_slice(wn), fs.get_slice(sn)
|
||||
for r0 in range(0, rows, 1 << 18):
|
||||
q = ws[r0:r0 + (1 << 18)].view(torch.uint8).reshape(-1, 32)
|
||||
e = ss[r0:r0 + (1 << 18)].view(torch.uint8).reshape(-1, 1)
|
||||
_e8m0_check(e.numpy(), sn)
|
||||
v = (q.view(torch.float8_e4m3fn).float() * torch.exp2(e.float() - 127)).reshape(-1, dim)
|
||||
yield (v if self._f32 else v.bfloat16()).view(torch.uint8).numpy()
|
||||
qt = gguf.GGMLQuantizationType.F32 if self._f32 else gguf.GGMLQuantizationType.BF16
|
||||
data = cast(np.ndarray, _StreamedBytes((n_rows, dim * (4 if self._f32 else 2)), chunks))
|
||||
self.gguf_writer.add_tensor(name, data, raw_dtype=qt)
|
||||
|
||||
def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
|
||||
if not self._dsv4_mxfp4_generated:
|
||||
for bid in range(self.block_count):
|
||||
self._write_mxfp4_experts(bid, "w1", gguf.MODEL_TENSOR.FFN_GATE_EXP)
|
||||
self._write_mxfp4_experts(bid, "w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP)
|
||||
self._write_mxfp4_experts(bid, "w3", gguf.MODEL_TENSOR.FFN_UP_EXP)
|
||||
if self.hparams.get("engram_layer_ids"):
|
||||
self._write_engram_tables(list(self.hparams["engram_layer_ids"]))
|
||||
self._dsv4_mxfp4_generated = True
|
||||
return ()
|
||||
|
||||
_LAYER_MAP: dict[str, tuple[gguf.MODEL_TENSOR, str]] = {
|
||||
"attn.indexer.wk.weight": (gguf.MODEL_TENSOR.INDEXER_ATTN_K, ".weight"),
|
||||
"attn.indexer.k_norm.weight": (gguf.MODEL_TENSOR.INDEXER_K_NORM, ".weight"),
|
||||
"engram.wkv.weight": (gguf.MODEL_TENSOR.ENGRAM_KV, ".weight"),
|
||||
}
|
||||
|
||||
def _map_dsv4_tensor_name(self, name: str, bid: int | None) -> tuple[gguf.MODEL_TENSOR, str]:
|
||||
m = re.match(r"layers\.\d+\.(.+)$", name)
|
||||
if m is not None and m.group(1) in self._LAYER_MAP:
|
||||
return self._LAYER_MAP[m.group(1)]
|
||||
return super()._map_dsv4_tensor_name(name, bid)
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
|
||||
if name.endswith(".engram.k_weight"):
|
||||
return [] # folded into engram_gate
|
||||
if name.endswith(".engram.q_weight"):
|
||||
# Engram gate weight = q_weight * k_weight (bf16 * bf16 is exact in f32)
|
||||
k = self.model_tensors[name.removesuffix("q_weight") + "k_weight"]
|
||||
return [(self.format_tensor_name(gguf.MODEL_TENSOR.ENGRAM_GATE, bid), data_torch.float() * k().float())]
|
||||
return super().modify_tensors(data_torch, name, bid)
|
||||
|
||||
def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
|
||||
if new_name.endswith((".engram_gate.weight", ".ffn_gate_inp.weight")) or name in self._dsv4_f32_tensors:
|
||||
return gguf.GGMLQuantizationType.F32
|
||||
if n_dims >= 2 and (name in self._dsv4_fp8_dequantized or name in self._dsv4_bf16_tensors):
|
||||
return gguf.GGMLQuantizationType.F32 if self._f32 else gguf.GGMLQuantizationType.BF16
|
||||
return False
|
||||
|
||||
|
||||
class DeepseekV41DSparkModel(DeepseekV41Model):
|
||||
"""DeepSeek-V4.1 DSpark speculative draft (mtp.0..2 of the checkpoint) as a DFLASH GGUF (`--dspark`).
|
||||
|
||||
Layout = the V4 DSpark draft (conversion/deepseek.py DeepseekV4DSparkModel) with V4.1's differences flagged by
|
||||
`dflash.hyper_connection.single_pass`: stages blk.0..2 (MXFP4 experts bit-exact, dense BF16), mtp.0.main_proj /
|
||||
main_norm -> fc / enc.output_norm, mtp.2.norm -> output_norm, markov_head.embed / head -> markov_w1 / markov_w2,
|
||||
confidence_head.proj -> conf_proj (F32), no hc_head. Token embeddings and the lm head are the target's (shared at
|
||||
runtime through ctx_other). target_layers are the V4.1 ids as-is: the draft reads the attention input of those
|
||||
layers, which is what llama_set_embeddings_layer_inp extracts.
|
||||
"""
|
||||
model_arch = gguf.MODEL_ARCH.DFLASH
|
||||
mtp_only = True
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
hp = self.hparams
|
||||
self._n_main = hp["num_hidden_layers"]
|
||||
self._n_stage = hp["num_nextn_predict_layers"]
|
||||
if self._n_stage <= 0 or "dspark_block_size" not in hp:
|
||||
raise ValueError("DeepSeek-V4.1 DSpark export needs num_nextn_predict_layers > 0 and the dspark_* config")
|
||||
# stages keep main-model block ids (mtp.M -> layers.n_main+M); format_tensor_name() maps them back
|
||||
self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self._n_stage)
|
||||
|
||||
def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
|
||||
hp = {**self.hparams, **self.hparams.get("text_config", {})} # text_config is merged only after indexing
|
||||
type(self)._dsv4_main_layers = hp["num_hidden_layers"]
|
||||
type(self)._dsv4_nextn_layers = hp.get("num_nextn_predict_layers", 0)
|
||||
return TextModel.index_tensors(self, remote_hf_model_id=remote_hf_model_id)
|
||||
|
||||
@classmethod
|
||||
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
|
||||
name, gen = item
|
||||
m = re.match(r"mtp\.(\d+)\.(.+)$", name)
|
||||
if m is None:
|
||||
return None # shared with the target at runtime
|
||||
stage, rest = int(m.group(1)), m.group(2)
|
||||
assert cls._dsv4_main_layers is not None
|
||||
if stage >= cls._dsv4_nextn_layers:
|
||||
raise ValueError(f"Unexpected DSpark stage {stage}")
|
||||
if rest.startswith(("markov_head.", "confidence_head.")):
|
||||
return f"dspark.{rest}", gen
|
||||
return f"layers.{cls._dsv4_main_layers + stage}.{rest}", gen
|
||||
|
||||
def format_tensor_name(self, key: gguf.MODEL_TENSOR, bid: int | None = None, suffix: str = ".weight") -> str:
|
||||
if bid is not None and bid >= self._n_main:
|
||||
bid -= self._n_main
|
||||
return super().format_tensor_name(key, bid, suffix)
|
||||
|
||||
def set_vocab(self):
|
||||
super().set_vocab()
|
||||
self.gguf_writer.add_mask_token_id(self.hparams["dspark_noise_token_id"])
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
hp = self.hparams
|
||||
n = self._n_stage
|
||||
w = self.gguf_writer
|
||||
w.add_block_count(n)
|
||||
w.add_context_length(hp["max_position_embeddings"])
|
||||
w.add_embedding_length(hp["hidden_size"])
|
||||
w.add_feed_forward_length(hp["moe_intermediate_size"])
|
||||
w.add_head_count(hp["num_attention_heads"])
|
||||
w.add_head_count_kv(1)
|
||||
w.add_key_length(hp["head_dim"])
|
||||
w.add_value_length(hp["head_dim"])
|
||||
w.add_rope_dimension_count(hp["qk_rope_head_dim"])
|
||||
# compress-ratio-0 layers: plain rope at rope_theta, no YaRN
|
||||
w.add_rope_freq_base(float(hp["rope_theta"]))
|
||||
w.add_layer_norm_rms_eps(hp["rms_norm_eps"])
|
||||
w.add_q_lora_rank(hp["q_lora_rank"])
|
||||
w.add_sliding_window(hp["sliding_window"])
|
||||
w.add_expert_count(hp["dspark_n_routed_experts"])
|
||||
w.add_expert_used_count(hp["dspark_num_experts_per_tok"])
|
||||
w.add_expert_feed_forward_length(hp["moe_intermediate_size"])
|
||||
w.add_expert_shared_count(hp["n_shared_experts"])
|
||||
w.add_expert_weights_scale(hp["routed_scaling_factor"])
|
||||
w.add_expert_weights_norm(hp["norm_topk_prob"])
|
||||
w.add_expert_gating_func(gguf.ExpertGatingFuncType.SQRTSOFTPLUS)
|
||||
w.add_swiglu_clamp_exp([hp["swiglu_limit"]] * n)
|
||||
w.add_swiglu_clamp_shexp([hp["swiglu_limit"]] * n)
|
||||
w.add_attention_output_group_count(hp["o_groups"])
|
||||
w.add_attention_output_lora_rank(hp["o_lora_rank"])
|
||||
w.add_attention_compress_ratios([0] * n)
|
||||
w.add_hyper_connection_count(hp["hc_mult"])
|
||||
w.add_hyper_connection_sinkhorn_iterations(hp["hc_sinkhorn_iters"])
|
||||
w.add_hyper_connection_epsilon(hp["hc_eps"])
|
||||
w.add_hyper_connection_single_pass(True)
|
||||
w.add_block_size(hp["dspark_block_size"])
|
||||
w.add_target_layers(hp["dspark_target_layer_ids"])
|
||||
w.add_sample_from_anchor(True)
|
||||
w.add_has_confidence_head(True)
|
||||
w.add_file_type(self.ftype)
|
||||
|
||||
def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
|
||||
# no Engram tables in the draft
|
||||
if not self._dsv4_mxfp4_generated:
|
||||
n_expert = self.hparams["dspark_n_routed_experts"]
|
||||
for bid in range(self._n_main, self._n_main + self._n_stage):
|
||||
self._write_mxfp4_experts(bid, "w1", gguf.MODEL_TENSOR.FFN_GATE_EXP, n_expert)
|
||||
self._write_mxfp4_experts(bid, "w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP, n_expert)
|
||||
self._write_mxfp4_experts(bid, "w3", gguf.MODEL_TENSOR.FFN_UP_EXP, n_expert)
|
||||
self._dsv4_mxfp4_generated = True
|
||||
return ()
|
||||
|
||||
_DSPARK_ROOT: dict[str, gguf.MODEL_TENSOR] = {
|
||||
"main_proj.weight": gguf.MODEL_TENSOR.FC,
|
||||
"main_norm.weight": gguf.MODEL_TENSOR.ENC_OUTPUT_NORM,
|
||||
"norm.weight": gguf.MODEL_TENSOR.OUTPUT_NORM,
|
||||
"dspark.markov_head.embed.weight": gguf.MODEL_TENSOR.DSPARK_MARKOV_W1,
|
||||
"dspark.markov_head.head.weight": gguf.MODEL_TENSOR.DSPARK_MARKOV_W2,
|
||||
"dspark.confidence_head.proj.weight": gguf.MODEL_TENSOR.DSPARK_CONF_PROJ,
|
||||
}
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
|
||||
if name.endswith(".ffn.gate.bias_vl"):
|
||||
return []
|
||||
if name in self._DSPARK_ROOT:
|
||||
return [(self.format_tensor_name(self._DSPARK_ROOT[name], None), data_torch)]
|
||||
m = re.match(r"layers\.\d+\.(main_proj\.weight|main_norm\.weight|norm\.weight)$", name)
|
||||
if m is not None:
|
||||
return [(self.format_tensor_name(self._DSPARK_ROOT[m.group(1)], None), data_torch)]
|
||||
return super().modify_tensors(data_torch, name, bid)
|
||||
|
||||
def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
|
||||
if new_name == "conf_proj.weight":
|
||||
return gguf.GGMLQuantizationType.F32
|
||||
return super().tensor_force_quant(name, new_name, bid, n_dims)
|
||||
|
||||
|
||||
@ModelBase.register("DeepseekV41ForCausalLM")
|
||||
@ModelBase.example("deepseek-ai/DeepSeek-V4.1-Flash-1017")
|
||||
class DeepseekV41VisionModel(DeepseekV4FlashVisionModel):
|
||||
"""DeepSeek-V4.1 vision tower (projector deepseek41v): the V4-Flash-Vision ViT + 3x3 aligner, with the V4.1 token
|
||||
block ([START] ([IMAGE]*w [NEWLINE])*h [END], no pads) and resize plan. ref: inference/vision.py,
|
||||
inference/image_processor.py of deepseek-ai/DeepSeek-V4.1-Flash-1017. The checkpoint config is HF-style
|
||||
(vision_config); image_start / image_end / image_newline are written like V4's sentinels, there is no image_pad."""
|
||||
|
||||
def get_vision_config(self) -> dict[str, Any] | None:
|
||||
vc = self.global_config.get("vision_config")
|
||||
if not vc or vc.get("num_hidden_layers", 0) == 0:
|
||||
raise ValueError("DeepseekV41VisionModel requires vision_config.num_hidden_layers > 0")
|
||||
return {
|
||||
"num_hidden_layers": vc["num_hidden_layers"],
|
||||
"hidden_size": vc["hidden_size"],
|
||||
"num_attention_heads": vc["num_attention_heads"],
|
||||
"intermediate_size": vc["intermediate_size"],
|
||||
"patch_size": vc["patch_size"],
|
||||
# dynamic resolution; only used for compat / warmup
|
||||
"image_size": vc["patch_size"] * vc["downsample_ratio"] * 16,
|
||||
"rope_theta": vc.get("rope_theta", 10000.0),
|
||||
"downsample_ratio": vc["downsample_ratio"],
|
||||
"min_pixels": vc["min_pixels"],
|
||||
"max_image_tokens": vc["max_image_tokens"],
|
||||
"max_wh_ratio": vc.get("max_wh_ratio"),
|
||||
}
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
assert self.hparams_vision is not None
|
||||
super(DeepseekV4FlashVisionModel, self).set_gguf_parameters()
|
||||
self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.DEEPSEEK41V)
|
||||
# pytorch default eps (vision.py), not the LLM rms_norm_eps (1e-20)
|
||||
self.gguf_writer.add_vision_attention_layernorm_eps(1e-6)
|
||||
self.gguf_writer.add_vision_use_silu(True)
|
||||
self.gguf_writer.add_vision_projector_scale_factor(self.hparams_vision["downsample_ratio"])
|
||||
self.gguf_writer.add_vision_min_pixels(self.hparams_vision["min_pixels"])
|
||||
# hardcoded in clip.cpp (PROJECTOR_TYPE_DEEPSEEK41V); add GGUF keys if a later model differs
|
||||
assert self.hparams_vision["max_image_tokens"] == 1024, self.hparams_vision["max_image_tokens"]
|
||||
assert self.hparams_vision["max_wh_ratio"] is None, self.hparams_vision["max_wh_ratio"]
|
||||
assert self.hparams_vision["rope_theta"] == 10000, self.hparams_vision["rope_theta"]
|
||||
@@ -267,11 +267,15 @@ def main() -> None:
|
||||
sys.exit(1)
|
||||
|
||||
if args.dspark:
|
||||
if is_mistral_format or model_architecture != "DeepseekV4ForCausalLM":
|
||||
logger.error("--dspark is only supported for DeepseekV4ForCausalLM")
|
||||
if is_mistral_format or model_architecture not in ("DeepseekV4ForCausalLM", "DeepseekV41ForCausalLM"):
|
||||
logger.error("--dspark is only supported for DeepseekV4ForCausalLM and DeepseekV41ForCausalLM")
|
||||
sys.exit(1)
|
||||
from conversion.deepseek import DeepseekV4DSparkModel
|
||||
model_class = DeepseekV4DSparkModel
|
||||
if model_architecture == "DeepseekV41ForCausalLM":
|
||||
from conversion.deepseek41 import DeepseekV41DSparkModel
|
||||
model_class = DeepseekV41DSparkModel
|
||||
else:
|
||||
from conversion.deepseek import DeepseekV4DSparkModel
|
||||
model_class = DeepseekV4DSparkModel
|
||||
|
||||
if args.mtp or args.no_mtp:
|
||||
if not model_class.supports_mtp_export:
|
||||
|
||||
@@ -926,6 +926,7 @@ extern "C" {
|
||||
enum ggml_type type);
|
||||
|
||||
// dst[i0, i1, i2] = a[i0, i1, i2] + b[i0, ids[i1, i2]]
|
||||
// an id of -1 adds nothing: dst[i0, i1, i2] = a[i0, i1, i2]
|
||||
GGML_API struct ggml_tensor * ggml_add_id(
|
||||
struct ggml_context * ctx,
|
||||
struct ggml_tensor * a,
|
||||
@@ -1508,6 +1509,7 @@ extern "C" {
|
||||
enum ggml_op_hint hint);
|
||||
|
||||
// indirect matrix multiplication
|
||||
// an id of -1 skips the slot and zeroes the matching row of the result
|
||||
//
|
||||
// [TAG_MUL_MAT_ID_DUPLICATE_ID]
|
||||
// note: this op has undefined behaviour if a given `a` is referenced more than once for a given `b`
|
||||
|
||||
@@ -774,13 +774,6 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
|
||||
return src_ss[0];
|
||||
};
|
||||
|
||||
auto handle_pool_1d = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
|
||||
if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) {
|
||||
return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
|
||||
}
|
||||
return src_ss[0];
|
||||
};
|
||||
|
||||
auto handle_flash_attn_ext = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
|
||||
GGML_ASSERT(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
|
||||
|
||||
@@ -997,14 +990,12 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
|
||||
case GGML_OP_CONV_3D:
|
||||
case GGML_OP_CONV_2D_DW:
|
||||
case GGML_OP_CONV_TRANSPOSE_2D:
|
||||
case GGML_OP_POOL_1D:
|
||||
case GGML_OP_POOL_2D:
|
||||
case GGML_OP_POOL_2D_BACK:
|
||||
case GGML_OP_UPSCALE: {
|
||||
split_state = handle_generic(src_ss, /*scalar_only =*/ true);
|
||||
} break;
|
||||
case GGML_OP_POOL_1D: {
|
||||
split_state = handle_pool_1d(src_ss);
|
||||
} break;
|
||||
case GGML_OP_PAD: {
|
||||
split_state = handle_pad(src_ss);
|
||||
} break;
|
||||
|
||||
@@ -1652,7 +1652,12 @@ static void ggml_compute_forward_mul_mat_id(
|
||||
for (int id = 0; id < n_ids; ++id) {
|
||||
const int32_t i02 = *(const int32_t *) ((const char *) ids->data + iid1*ids->nb[1] + id*ids->nb[0]);
|
||||
|
||||
assert(i02 >= 0 && i02 < n_as);
|
||||
assert(i02 == -1 || (i02 >= 0 && i02 < n_as));
|
||||
|
||||
if (i02 == -1) {
|
||||
memset((char *) dst->data + id*nb1 + iid1*nb2, 0, ne0*sizeof(float));
|
||||
continue;
|
||||
}
|
||||
|
||||
MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = (struct mmid_row_mapping) {id, iid1};
|
||||
matrix_row_counts[i02] += 1;
|
||||
|
||||
@@ -743,12 +743,17 @@ static void ggml_compute_forward_add_id_f32(
|
||||
// src1 indices
|
||||
const int i11 = *(int32_t *) ((char *) src2->data + i1*nb20 + i2*nb21);
|
||||
|
||||
GGML_ASSERT(i11 >= 0 && i11 < ne11);
|
||||
GGML_ASSERT(i11 == -1 || (i11 >= 0 && i11 < ne11));
|
||||
|
||||
ggml_vec_add_f32(ne0,
|
||||
(float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ),
|
||||
(float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01),
|
||||
(float *) ((char *) src1->data + i11*nb11));
|
||||
float * dst_row = (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 );
|
||||
const float * src0_row = (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01);
|
||||
|
||||
if (i11 == -1) {
|
||||
memcpy(dst_row, src0_row, ne0*sizeof(float));
|
||||
continue;
|
||||
}
|
||||
|
||||
ggml_vec_add_f32(ne0, dst_row, src0_row, (float *) ((char *) src1->data + i11*nb11));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -4855,7 +4855,12 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PAR
|
||||
const int32_t i02 =
|
||||
*(const int32_t *) ((const char *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]);
|
||||
|
||||
GGML_ASSERT(i02 >= 0 && i02 < n_as);
|
||||
GGML_ASSERT(i02 == -1 || (i02 >= 0 && i02 < n_as));
|
||||
|
||||
if (i02 == -1) {
|
||||
memset((char *) dst->data + id*nb1 + iid1*nb2, 0, ne0*sizeof(float));
|
||||
continue;
|
||||
}
|
||||
|
||||
MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = { id, iid1 };
|
||||
matrix_row_counts[i02] += 1;
|
||||
|
||||
@@ -682,7 +682,12 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_
|
||||
const int32_t i02 =
|
||||
*(const int32_t *) ((const char *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]);
|
||||
|
||||
GGML_ASSERT(i02 >= 0 && i02 < n_as);
|
||||
GGML_ASSERT(i02 == -1 || (i02 >= 0 && i02 < n_as));
|
||||
|
||||
if (i02 == -1) {
|
||||
memset((char *) dst->data + id * nb1 + iid1 * nb2, 0, ne0 * sizeof(float));
|
||||
continue;
|
||||
}
|
||||
|
||||
MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = { id, iid1 };
|
||||
matrix_row_counts[i02] += 1;
|
||||
|
||||
@@ -78,7 +78,7 @@ static inline float op_expm1(float x) {
|
||||
}
|
||||
|
||||
static inline float op_softplus(float x) {
|
||||
return (x > 20.0f) ? x : logf(1.0f + expf(x));
|
||||
return (x > 20.0f) ? x : log1pf(expf(x));
|
||||
}
|
||||
|
||||
static inline float op_floor(float x) {
|
||||
|
||||
@@ -76,7 +76,7 @@ __device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], co
|
||||
for (int i = 0; i < experts_per_thread; i++) {
|
||||
const int idx = lane + i * WARP_SIZE;
|
||||
const bool active = !use_limit || (idx < limit);
|
||||
vals[i] = active ? sqrtf(vals[i] > 20.0f ? vals[i] : logf(1.0f + expf(vals[i]))) : -INFINITY;
|
||||
vals[i] = active ? sqrtf(vals[i] > 20.0f ? vals[i] : log1pf(expf(vals[i]))) : -INFINITY;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -91,7 +91,7 @@ static __device__ __forceinline__ float op_expm1(float x) {
|
||||
}
|
||||
|
||||
static __device__ __forceinline__ float op_softplus(float x) {
|
||||
return (x > 20.0f) ? x : logf(1.0f + expf(x));
|
||||
return (x > 20.0f) ? x : log1pf(expf(x));
|
||||
}
|
||||
|
||||
static __device__ __forceinline__ float op_elu(float x) {
|
||||
|
||||
@@ -133,7 +133,7 @@ size_t opt_mbuf_total = 0; // total buffer space limit (
|
||||
|
||||
int opt_mm_select = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
|
||||
int opt_fa_select = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
|
||||
int opt_fa_head_split = 1; // 1 = partition flash_attn by KV heads in multicore (default on), 0 = token-based (original)
|
||||
int opt_fa_split = 1; // 1 = head-split + team-split when n_kv_heads < N_cores (default), 0 = off
|
||||
int opt_gdn_select = 2; // 2 = HMX -> HVX, 1 = HVX, 0 = CPU (unsupported)
|
||||
int opt_ar_select = 2; // 2 = fused ALLREDUCE+ADD (default), 1 = unfused ALLREDUCE, 0 = fallback to CPY+FENCE
|
||||
int opt_ar_scatter = 1; // 1 = reduce-scatter the fused ALLREDUCE+ADD (default), 0 = full reduction
|
||||
@@ -3393,7 +3393,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
|
||||
const char * str_nhmx = getenv("GGML_HEXAGON_NHMX");
|
||||
const char * str_mm_select = getenv("GGML_HEXAGON_MM_SELECT");
|
||||
const char * str_fa_select = getenv("GGML_HEXAGON_FA_SELECT");
|
||||
const char * str_fa_head_split = getenv("GGML_HEXAGON_FA_HEAD_SPLIT");
|
||||
const char * str_fa_split = getenv("GGML_HEXAGON_FA_SPLIT");
|
||||
const char * str_gdn_select = getenv("GGML_HEXAGON_GDN_SELECT");
|
||||
const char * str_ar_select = getenv("GGML_HEXAGON_AR_SELECT");
|
||||
const char * str_ar_scatter = getenv("GGML_HEXAGON_AR_SCATTER");
|
||||
@@ -3458,7 +3458,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
|
||||
opt_nhmx = str_nhmx ? atoi(str_nhmx) : opt_nhmx;
|
||||
opt_mm_select = str_mm_select ? atoi(str_mm_select) : opt_mm_select;
|
||||
opt_fa_select = str_fa_select ? atoi(str_fa_select) : opt_fa_select;
|
||||
opt_fa_head_split = str_fa_head_split ? atoi(str_fa_head_split) : opt_fa_head_split;
|
||||
opt_fa_split = str_fa_split ? atoi(str_fa_split) : opt_fa_split;
|
||||
opt_gdn_select = str_gdn_select ? atoi(str_gdn_select) : opt_gdn_select;
|
||||
opt_ar_select = str_ar_select ? atoi(str_ar_select) : opt_ar_select;
|
||||
opt_ar_scatter = str_ar_scatter ? atoi(str_ar_scatter) : opt_ar_scatter;
|
||||
|
||||
@@ -28,7 +28,7 @@ extern size_t opt_mbuf_total;
|
||||
|
||||
extern int opt_mm_select;
|
||||
extern int opt_fa_select;
|
||||
extern int opt_fa_head_split;
|
||||
extern int opt_fa_split;
|
||||
extern int opt_gdn_select;
|
||||
extern int opt_ar_select;
|
||||
extern int opt_ar_scatter;
|
||||
|
||||
@@ -397,9 +397,20 @@ static bool htp_kparams_solve_flash_attn(
|
||||
kparams->scale = scale;
|
||||
kparams->max_bias = max_bias;
|
||||
kparams->logit_softcap = logit_softcap;
|
||||
kparams->head_split = (opt_fa_head_split != 0) ? 1 : 0;
|
||||
kparams->head_split = (opt_fa_split != 0) ? 1 : 0;
|
||||
kparams->flags = 0;
|
||||
kparams->G = G;
|
||||
|
||||
// team-split only when a team has >1 member (n_kv_heads < N_cores)
|
||||
uint32_t fa_teams = 1, fa_team_size = 1;
|
||||
const uint32_t mdev_count = htp_session_mdev_count(sess);
|
||||
if (opt_fa_split && mdev_count > 1) {
|
||||
htp_fa_mdev_teams(n_kv_heads, mdev_count, kparams->head_split != 0, &fa_teams, &fa_team_size);
|
||||
if (fa_team_size > 1) {
|
||||
kparams->flags |= HTP_FA_FLAG_TEAM_SPLIT;
|
||||
}
|
||||
}
|
||||
|
||||
const uint32_t n_head = q->ne[2];
|
||||
kparams->n_head_log2 = 1u << (uint32_t) std::floor(std::log2(n_head));
|
||||
kparams->m0 = std::pow(2.0f, -(max_bias) / kparams->n_head_log2);
|
||||
@@ -409,7 +420,7 @@ static bool htp_kparams_solve_flash_attn(
|
||||
const uint32_t DK_pad = hex_round_up(DK, 64);
|
||||
const uint32_t DV_pad = hex_round_up(DV, 64);
|
||||
size_t Br = 0, Bc = 0;
|
||||
int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK_pad, DV_pad, neq1, nek1, vtcm_size, n_threads, (q->type == GGML_TYPE_F32), sinks != nullptr, n_head);
|
||||
int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK_pad, DV_pad, neq1, nek1, vtcm_size, n_threads, (q->type == GGML_TYPE_F32), sinks != nullptr, n_head, fa_team_size);
|
||||
if (ret == 0) {
|
||||
kparams->kernel_type = HTP_FA_KERNEL_HMX;
|
||||
if (logit_softcap == 0.0f) {
|
||||
|
||||
@@ -2094,7 +2094,29 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
const uint32_t dst_e_size = (dst->type == HTP_TYPE_F32) ? sizeof(float) : sizeof(__fp16);
|
||||
const bool can_split = htp_tensor_can_row_partition(dst, dst_e_size);
|
||||
|
||||
if (kparams->head_split && can_split && n_kv_heads >= mdev_count && n_kv_heads % mdev_count == 0) {
|
||||
if ((kparams->flags & HTP_FA_FLAG_TEAM_SPLIT) && can_split) {
|
||||
// KV heads split across teams, Q-blocks split within a team
|
||||
uint32_t n_teams, team_size;
|
||||
htp_fa_mdev_teams(n_kv_heads, mdev_count, kparams->head_split != 0, &n_teams, &team_size);
|
||||
|
||||
const uint32_t team = mdev_idx / team_size;
|
||||
const uint32_t member = mdev_idx - team * team_size;
|
||||
const uint32_t kv_per_team = n_kv_heads / n_teams;
|
||||
kv_head_min = team * kv_per_team;
|
||||
kv_head_max = kv_head_min + kv_per_team;
|
||||
|
||||
const uint32_t base = n_q_blocks / team_size;
|
||||
const uint32_t rem = n_q_blocks - base * team_size;
|
||||
const uint32_t b0 = member * base + MIN(member, rem);
|
||||
const uint32_t nblk = base + (member < rem ? 1 : 0);
|
||||
|
||||
if (nblk == 0) {
|
||||
return HTP_STATUS_OK;
|
||||
}
|
||||
|
||||
q_start_min = b0 * Br;
|
||||
q_start_max = MIN((b0 + nblk) * Br, neq1);
|
||||
} else if (kparams->head_split && can_split && n_kv_heads >= mdev_count && n_kv_heads % mdev_count == 0) {
|
||||
const uint32_t kv_per_core = n_kv_heads / mdev_count;
|
||||
kv_head_min = mdev_idx * kv_per_core;
|
||||
kv_head_max = kv_head_min + kv_per_core;
|
||||
|
||||
@@ -32,9 +32,20 @@ enum htp_fa_kernel_type {
|
||||
HTP_FA_KERNEL_HMX
|
||||
};
|
||||
|
||||
// Multicore: gcd(n_kv_heads, N) teams own the KV heads, team members share the Q-blocks
|
||||
#define HTP_FA_FLAG_TEAM_SPLIT 0x01
|
||||
|
||||
static inline void htp_fa_mdev_teams(uint32_t n_kv_heads, uint32_t mdev_count, bool head_split,
|
||||
uint32_t * n_teams, uint32_t * team_size) {
|
||||
const uint32_t t = (head_split && n_kv_heads > 0 && mdev_count > 1) ? hex_gcd_u32(n_kv_heads, mdev_count) : 1;
|
||||
*n_teams = t;
|
||||
*team_size = mdev_count > 1 ? mdev_count / t : 1;
|
||||
}
|
||||
|
||||
struct htp_fa_kernel_params {
|
||||
uint8_t kernel_type; // enum htp_fa_kernel_type
|
||||
uint8_t head_split; // 1 = partition by KV heads in multicore, 0 = token partition
|
||||
uint8_t flags; // HTP_FA_FLAG_*
|
||||
uint8_t n_threads; // Number of threads to run
|
||||
|
||||
// Common parameters
|
||||
@@ -321,7 +332,8 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
|
||||
size_t n_threads,
|
||||
bool is_q_fp32,
|
||||
bool has_sinks,
|
||||
size_t n_heads) {
|
||||
size_t n_heads,
|
||||
size_t q_split) {
|
||||
const size_t T = HMX_FP16_TILE_N_ROWS; // 32
|
||||
const size_t br_unit = hmx_ceil_div(T, gqa_factor);
|
||||
const size_t bc_unit = HMX_FP16_TILE_N_COLS * 2; // 64
|
||||
@@ -361,7 +373,8 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
|
||||
const size_t vecs_per_t = n_use > 0 ? (n_row_vec_cnt + n_use - 1) / n_use : 1;
|
||||
|
||||
const size_t c_iter_actual = c_iter_base + c_softmax * vecs_per_t;
|
||||
const size_t cost = q_blocks * (c_q_fixed + kv_blocks * c_iter_actual);
|
||||
const size_t q_blocks_nsp = (q_blocks + q_split - 1) / (q_split ? q_split : 1);
|
||||
const size_t cost = q_blocks_nsp * (c_q_fixed + kv_blocks * c_iter_actual);
|
||||
const size_t mn = Br * Bc;
|
||||
|
||||
if (cost < best_cost || (cost == best_cost && mn > best_mn)) {
|
||||
|
||||
@@ -1413,10 +1413,10 @@ static void hvx_mv_id_tiled(unsigned int nth, unsigned int ith, void * data) {
|
||||
|
||||
for (uint32_t ie1 = 0; ie1 < n_aids; ++ie1) { // for each expert
|
||||
const int32_t eid = *(const int32_t *) ((const uint8_t *) src2->data + ie1 * src2->nb[0]);
|
||||
if (eid < 0) {
|
||||
if (eid == -1) {
|
||||
continue;
|
||||
}
|
||||
assert(eid < (int32_t) n_ids);
|
||||
assert(eid >= 0 && eid < (int32_t) n_ids);
|
||||
|
||||
const dma_addr_t src0_row = src0->data + eid * nb02;
|
||||
const uint8_t * restrict act_col = (const uint8_t *) act_data;
|
||||
@@ -4166,10 +4166,10 @@ static inline void scan_expert_ids_n(
|
||||
const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
|
||||
for (uint32_t id = 0; id < n_ids; ++id) {
|
||||
const int32_t i02 = row_ptr[id];
|
||||
if (i02 < 0) {
|
||||
if (i02 == -1) {
|
||||
continue;
|
||||
}
|
||||
assert(i02 < n_as);
|
||||
assert(i02 >= 0 && i02 < n_as);
|
||||
|
||||
if (matrix_rows) {
|
||||
matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
|
||||
@@ -4179,6 +4179,28 @@ static inline void scan_expert_ids_n(
|
||||
}
|
||||
}
|
||||
|
||||
// the matmuls only write the rows that an expert owns, so the rows of the skipped slots are zeroed here
|
||||
static inline void zero_skipped_dst_rows(
|
||||
const struct htp_tensor * ids,
|
||||
uint32_t n_ids,
|
||||
const struct htp_tensor * dst,
|
||||
uint32_t ne0,
|
||||
size_t nb1,
|
||||
size_t nb2
|
||||
) {
|
||||
const uint8_t * ids_data = (const uint8_t *) ids->data;
|
||||
|
||||
for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
|
||||
for (uint32_t id = 0; id < n_ids; ++id) {
|
||||
const int32_t i02 = *(const int32_t *) (ids_data + iid1 * ids->nb[1] + id * ids->nb[0]);
|
||||
if (i02 != -1) {
|
||||
continue;
|
||||
}
|
||||
memset((uint8_t *) (dst->data + id * nb1 + iid1 * nb2), 0, ne0 * sizeof(float));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static inline void scan_expert_ids(
|
||||
const struct htp_tensor * ids,
|
||||
uint32_t n_ids,
|
||||
@@ -4204,10 +4226,10 @@ static inline void scan_expert_ids(
|
||||
const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
|
||||
for (uint32_t id = 0; id < n_ids; ++id) {
|
||||
const int32_t i02 = *(const int32_t *) ((const uint8_t *) row_ptr + id * ids_nb0);
|
||||
if (i02 < 0) {
|
||||
if (i02 == -1) {
|
||||
continue;
|
||||
}
|
||||
assert(i02 < n_as);
|
||||
assert(i02 >= 0 && i02 < n_as);
|
||||
|
||||
if (matrix_rows) {
|
||||
matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
|
||||
@@ -4287,6 +4309,8 @@ int op_matmul_id(struct htp_ops_context * octx) {
|
||||
scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride);
|
||||
}
|
||||
|
||||
zero_skipped_dst_rows(ids, n_ids, dst, ne0, nb1, nb2);
|
||||
|
||||
mmctx->matrix_row_counts = matrix_row_counts;
|
||||
mmctx->matrix_rows = matrix_rows;
|
||||
mmctx->mapping_stride = mapping_stride;
|
||||
|
||||
@@ -1225,40 +1225,6 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax_part(ggml_metal_library_t lib) {
|
||||
char base[256];
|
||||
char name[256];
|
||||
|
||||
snprintf(base, 256, "kernel_mul_mm_id_amax_part_f32");
|
||||
snprintf(name, 256, "%s", base);
|
||||
|
||||
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
|
||||
if (!res.pipeline) {
|
||||
res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
|
||||
}
|
||||
|
||||
res.smem = 32*sizeof(float);
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax(ggml_metal_library_t lib) {
|
||||
char base[256];
|
||||
char name[256];
|
||||
|
||||
snprintf(base, 256, "kernel_mul_mm_id_amax_f32");
|
||||
snprintf(name, 256, "%s", base);
|
||||
|
||||
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
|
||||
if (!res.pipeline) {
|
||||
res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
|
||||
}
|
||||
|
||||
res.smem = 32*sizeof(float);
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0(ggml_metal_library_t lib, int ne02, int ne20) {
|
||||
char base[256];
|
||||
char name[256];
|
||||
@@ -1283,6 +1249,15 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0(g
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_zero(ggml_metal_library_t lib) {
|
||||
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, "kernel_mul_mm_id_zero");
|
||||
if (!res.pipeline) {
|
||||
res = ggml_metal_library_compile_pipeline(lib, "kernel_mul_mm_id_zero", "kernel_mul_mm_id_zero", nullptr);
|
||||
}
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id(ggml_metal_library_t lib, const ggml_tensor * op) {
|
||||
char base[256];
|
||||
char name[256];
|
||||
@@ -1292,18 +1267,14 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id(ggml_m
|
||||
|
||||
const bool bc_inp = op->src[0]->ne[0] % 32 != 0;
|
||||
|
||||
// src1 prec [TAG_GGML_PREC]
|
||||
const bool amax = ggml_get_op_params_i32(op, 3) == GGML_PREC_F32;
|
||||
|
||||
snprintf(base, 256, "kernel_mul_mm_id_%s_%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1));
|
||||
snprintf(name, 256, "%s_bci=%d_amax=%d", base, bc_inp, amax);
|
||||
snprintf(name, 256, "%s_bci=%d", base, bc_inp);
|
||||
|
||||
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
|
||||
if (!res.pipeline) {
|
||||
ggml_metal_cv_t cv = ggml_metal_cv_init();
|
||||
|
||||
ggml_metal_cv_set_bool(cv, bc_inp, FC_MUL_MM + 0);
|
||||
ggml_metal_cv_set_bool(cv, amax, FC_MUL_MM + 6);
|
||||
|
||||
res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
|
||||
|
||||
|
||||
@@ -140,8 +140,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_mm
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0 (ggml_metal_library_t lib, int ne02, int ne20);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax(ggml_metal_library_t lib);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax_part(ggml_metal_library_t lib);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_zero (ggml_metal_library_t lib);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
|
||||
@@ -14,8 +14,6 @@
|
||||
#define N_MM_SIMD_GROUP_X 2
|
||||
#define N_MM_SIMD_GROUP_Y 2
|
||||
|
||||
#define N_MM_NPART_AMAX 256
|
||||
|
||||
#define N_ARGMAX_NPART_MAX 256
|
||||
#define N_ARGMAX_CHUNK 4096
|
||||
|
||||
@@ -573,12 +571,11 @@ typedef struct {
|
||||
} ggml_metal_kargs_mul_mm_id_map0;
|
||||
|
||||
typedef struct {
|
||||
int32_t ne00;
|
||||
int32_t ne01;
|
||||
int32_t ne02;
|
||||
uint64_t nb01;
|
||||
uint64_t nb02;
|
||||
} ggml_metal_kargs_mul_mm_id_amax;
|
||||
int32_t ne0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb21;
|
||||
} ggml_metal_kargs_mul_mm_id_zero;
|
||||
|
||||
typedef struct {
|
||||
int32_t ne00;
|
||||
|
||||
@@ -2812,15 +2812,6 @@ size_t ggml_metal_op_mul_mat_id_extra_ids(const ggml_tensor * op) {
|
||||
return ggml_type_size(GGML_TYPE_I32)*(ne20*ne21 + ne02);
|
||||
}
|
||||
|
||||
size_t ggml_metal_op_mul_mat_id_extra_amax(const ggml_tensor * op) {
|
||||
assert(op->op == GGML_OP_MUL_MAT_ID);
|
||||
|
||||
GGML_UNUSED(op);
|
||||
|
||||
// 2 scaling factors (8 bytes) + N_MM_NPART_AMAX per-threadgroup scales for stage-1
|
||||
return 8 + N_MM_NPART_AMAX*sizeof(float);
|
||||
}
|
||||
|
||||
int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
@@ -2872,39 +2863,6 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_metal_buffer_id bid_ids = bid_tpe;
|
||||
bid_ids.offs += ggml_metal_op_mul_mat_id_extra_tpe(op);
|
||||
|
||||
ggml_metal_buffer_id bid_amax = bid_ids;
|
||||
bid_amax.offs += ggml_metal_op_mul_mat_id_extra_ids(op);
|
||||
|
||||
// src1 prec [TAG_GGML_PREC]
|
||||
const bool use_amax = ggml_get_op_params_i32(op, 3) == GGML_PREC_F32;
|
||||
|
||||
// src1 rescale factors, computed before the matmul
|
||||
// ref: https://github.com/ggml-org/llama.cpp/pull/26223
|
||||
if (use_amax) {
|
||||
ggml_metal_kargs_mul_mm_id_amax args = {
|
||||
/*.ne00 =*/ ne10,
|
||||
/*.ne01 =*/ ne11,
|
||||
/*.ne02 =*/ ne12,
|
||||
/*.nb01 =*/ nb11,
|
||||
/*.nb02 =*/ nb12,
|
||||
};
|
||||
|
||||
auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id_amax_part(lib);
|
||||
|
||||
const size_t smem = pipeline.smem;
|
||||
|
||||
GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size);
|
||||
|
||||
ggml_metal_encoder_set_pipeline(enc, pipeline);
|
||||
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
|
||||
ggml_metal_encoder_set_buffer (enc, bid_src1, 1);
|
||||
ggml_metal_encoder_set_buffer (enc, bid_amax, 2);
|
||||
|
||||
ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
|
||||
|
||||
ggml_metal_encoder_dispatch_threadgroups(enc, N_MM_NPART_AMAX, 1, 1, 256, 1, 1);
|
||||
}
|
||||
|
||||
{
|
||||
ggml_metal_kargs_mul_mm_id_map0 args = {
|
||||
ne10,
|
||||
@@ -2935,20 +2893,28 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, ne02, 1, 1);
|
||||
}
|
||||
|
||||
ggml_metal_op_concurrency_reset(ctx);
|
||||
// zero the rows that no expert owns - disjoint from what kernel_mul_mm_id writes
|
||||
{
|
||||
ggml_metal_kargs_mul_mm_id_zero args = {
|
||||
/*.ne0 =*/ ne0,
|
||||
/*.nb1 =*/ nb1,
|
||||
/*.nb2 =*/ nb2,
|
||||
/*.nb21 =*/ nb21,
|
||||
};
|
||||
|
||||
if (use_amax) {
|
||||
auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id_amax(lib);
|
||||
auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id_zero(lib);
|
||||
|
||||
ggml_metal_encoder_set_pipeline(enc, pipeline);
|
||||
ggml_metal_encoder_set_buffer (enc, bid_amax, 0);
|
||||
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
|
||||
ggml_metal_encoder_set_buffer (enc, bid_src2, 1);
|
||||
ggml_metal_encoder_set_buffer (enc, bid_dst, 2);
|
||||
|
||||
ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, 32, 1, 1);
|
||||
|
||||
// the next kernel has to wait for the amax data
|
||||
ggml_metal_op_concurrency_reset(ctx);
|
||||
ggml_metal_encoder_dispatch_threadgroups(enc, ne20, ne21, 1, 32, 1, 1);
|
||||
}
|
||||
|
||||
// this barrier is always needed because the next kernel has to wait for the id maps to be computed
|
||||
ggml_metal_op_concurrency_reset(ctx);
|
||||
|
||||
{
|
||||
auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id(lib, op);
|
||||
|
||||
@@ -2977,7 +2943,6 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_metal_encoder_set_buffer (enc, bid_tpe, 3);
|
||||
ggml_metal_encoder_set_buffer (enc, bid_ids, 4);
|
||||
ggml_metal_encoder_set_buffer (enc, bid_dst, 5);
|
||||
ggml_metal_encoder_set_buffer (enc, bid_amax, 6);
|
||||
|
||||
const size_t smem = pipeline.smem;
|
||||
|
||||
|
||||
@@ -36,7 +36,6 @@ size_t ggml_metal_op_mul_mat_id_extra_tpe(const struct ggml_tensor * op);
|
||||
|
||||
// id map [n_tokens, n_expert]
|
||||
size_t ggml_metal_op_mul_mat_id_extra_ids(const struct ggml_tensor * op);
|
||||
size_t ggml_metal_op_mul_mat_id_extra_amax(const struct ggml_tensor * op);
|
||||
|
||||
// return true if we should use the FA vector kernel for this op
|
||||
bool ggml_metal_op_flash_attn_ext_use_vec(const struct ggml_tensor * op);
|
||||
|
||||
@@ -226,7 +226,6 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_
|
||||
{
|
||||
res += ggml_metal_op_mul_mat_id_extra_tpe(tensor);
|
||||
res += ggml_metal_op_mul_mat_id_extra_ids(tensor);
|
||||
res += ggml_metal_op_mul_mat_id_extra_amax(tensor);
|
||||
} break;
|
||||
case GGML_OP_FLASH_ATTN_EXT:
|
||||
{
|
||||
|
||||
@@ -185,6 +185,14 @@ kernel void kernel_add_id(
|
||||
|
||||
device float * dst_row = (device float *)((device char *)dst + i1*nb1 + i2*nb2);
|
||||
device const float * src0_row = (device const float *)((device char *)src0 + i1*args.nb01 + i2*args.nb02);
|
||||
|
||||
if (i11 == -1) {
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
dst_row[i0] = src0_row[i0];
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
device const float * src1_row = (device const float *)((device char *)src1 + i11*args.nb11);
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
|
||||
@@ -7,7 +7,6 @@ constant short FC_mul_mm_ne12 [[function_constant(FC_MUL_MM + 2)]];
|
||||
constant short FC_mul_mm_ne13 [[function_constant(FC_MUL_MM + 3)]];
|
||||
constant short FC_mul_mm_r2 [[function_constant(FC_MUL_MM + 4)]];
|
||||
constant short FC_mul_mm_r3 [[function_constant(FC_MUL_MM + 5)]];
|
||||
constant bool FC_mul_mm_id_amax [[function_constant(FC_MUL_MM + 6)]];
|
||||
constant int32_t FC_mul_mm_id_map0_ne02 [[function_constant(FC_MUL_MM + 7)]];
|
||||
|
||||
// each block_q contains 16*nl weights
|
||||
@@ -391,8 +390,10 @@ kernel void kernel_mul_mm_id_map0(
|
||||
if (i21 + tpitg < args.ne21) {
|
||||
device const int32_t * src2_i32 = (device const int32_t *) (src2 + (i21 + tpitg)*args.nb21);
|
||||
|
||||
// 0xFFFF marks a skipped slot, it never matches an expert since ne02 fits in a threadgroup
|
||||
FOR_UNROLL (short i20 = 0; i20 < ne20; i20++) {
|
||||
sids[tpitg*ne20 + i20] = src2_i32[i20];
|
||||
const int32_t id = src2_i32[i20];
|
||||
sids[tpitg*ne20 + i20] = id == -1 ? 0xFFFF : (uint16_t) id;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -506,82 +507,27 @@ kernel void kernel_mul_mm_id_map0(
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_mul_mm_id_amax_part_f32(
|
||||
constant ggml_metal_kargs_mul_mm_id_amax & args,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
uint tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort ntg[[threads_per_threadgroup]]) {
|
||||
const int nrow = args.ne01*args.ne02;
|
||||
// kernel_mul_mm_id only writes the rows that an expert owns, so the skipped rows are zeroed here
|
||||
kernel void kernel_mul_mm_id_zero(
|
||||
constant ggml_metal_kargs_mul_mm_id_zero & args,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int i20 = tgpig.x; // slot
|
||||
const int i21 = tgpig.y; // token
|
||||
|
||||
float lmax = 0.0f;
|
||||
const int32_t id = ((device const int32_t *) (src2 + i21*args.nb21))[i20];
|
||||
|
||||
for (int ir = tgpig; ir < nrow; ir += N_MM_NPART_AMAX) {
|
||||
const int i01 = ir % args.ne01;
|
||||
const int i02 = ir / args.ne01;
|
||||
|
||||
device const float * row = (device const float *) (src1 + i02*args.nb02 + i01*args.nb01);
|
||||
|
||||
for (int i00 = tiitg; i00 < args.ne00; i00 += ntg) {
|
||||
lmax = max(lmax, fabs(row[i00]));
|
||||
}
|
||||
if (id != -1) {
|
||||
return;
|
||||
}
|
||||
|
||||
float amax = simd_max(lmax);
|
||||
device float * dst_row = (device float *) (dst + i20*args.nb1 + i21*args.nb2);
|
||||
|
||||
threadgroup float * shared_amax = (threadgroup float *) shmem;
|
||||
|
||||
if (ntg > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
shared_amax[tiisg] = 0.0f;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shared_amax[sgitg] = amax;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
amax = shared_amax[tiisg];
|
||||
amax = simd_max(amax);
|
||||
}
|
||||
|
||||
if (tiitg == 0) {
|
||||
((device float *) (dst + 8))[tgpig] = amax;
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_mul_mm_id_amax_f32(
|
||||
device char * dst,
|
||||
ushort tiitg[[thread_index_in_threadgroup]]) {
|
||||
device const float * part = (device const float *) (dst + 8);
|
||||
|
||||
float amax = 0.0f;
|
||||
|
||||
for (int i = tiitg; i < N_MM_NPART_AMAX; i += N_SIMDWIDTH) {
|
||||
amax = max(amax, part[i]);
|
||||
}
|
||||
|
||||
amax = simd_max(amax);
|
||||
|
||||
if (tiitg == 0) {
|
||||
// leave a comfortable margin below the f16 max of 65504
|
||||
float scale = 1.0f;
|
||||
|
||||
// isfinite: src1 already inf/nan is not ours to fix - keep the
|
||||
// scale at 1.0 instead of turning it into a different failure
|
||||
if (isfinite(amax) && amax > 32768.0f) {
|
||||
scale = exp2(ceil(log2(amax)) - 15.0f);
|
||||
}
|
||||
|
||||
device float * d = (device float *) dst;
|
||||
|
||||
d[0] = 1.0f/scale; // exact: scale is a power of two
|
||||
d[1] = scale;
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
dst_row[i0] = 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -607,7 +553,6 @@ kernel void kernel_mul_mm_id(
|
||||
device const char * htpe,
|
||||
device const char * hids,
|
||||
device char * dst,
|
||||
device const char * amax,
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
@@ -663,10 +608,6 @@ kernel void kernel_mul_mm_id(
|
||||
// skip the upper half of the token tile when the expert did not fill it
|
||||
constexpr short NR1H = NR1/2;
|
||||
|
||||
// power-of-two rescaling
|
||||
const float s1_inv = FC_mul_mm_id_amax ? ((device const float *) amax)[0] : 1.0f;
|
||||
const float s1_scale = FC_mul_mm_id_amax ? ((device const float *) amax)[1] : 1.0f;
|
||||
|
||||
#ifndef GGML_METAL_HAS_TENSOR
|
||||
S0_8x8 ma[4];
|
||||
S1_8x8 mb[2];
|
||||
@@ -772,7 +713,7 @@ kernel void kernel_mul_mm_id(
|
||||
|
||||
const short ib = 4*sx + sy;
|
||||
|
||||
*(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) (*((device T1 *) y + i) * (T1) s1_inv) : 0;
|
||||
*(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
|
||||
}
|
||||
} else {
|
||||
const short sx = (tiitg%NL1);
|
||||
@@ -785,7 +726,7 @@ kernel void kernel_mul_mm_id(
|
||||
|
||||
const short ib = 4*sx + sy;
|
||||
|
||||
*(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)((*((device T1_2x4 *) y)) * (T1) s1_inv);
|
||||
*(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y));
|
||||
}
|
||||
#else
|
||||
// load data and store to threadgroup memory
|
||||
@@ -833,7 +774,7 @@ kernel void kernel_mul_mm_id(
|
||||
//const short lx = (tiitg/NL1)%8;
|
||||
//const short ly = i;
|
||||
|
||||
*(sb + NK*(8*sy + ly) + 8*sx + lx) = loop_k + iy + i < args.ne00 ? (S1) (*((device T1 *) y + i) * (T1) s1_inv) : 0;
|
||||
*(sb + NK*(8*sy + ly) + 8*sx + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
|
||||
}
|
||||
} else {
|
||||
const short sx = (tiitg%NL1);
|
||||
@@ -844,7 +785,7 @@ kernel void kernel_mul_mm_id(
|
||||
//const short lx = (tiitg/NL1)%8;
|
||||
//const short ly = i;
|
||||
|
||||
*(threadgroup S1_2x4 *)(sb + NK*(8*sy + ly) + 8*sx) = (S1_2x4)((*((device T1_2x4 *) y)) * (T1) s1_inv);
|
||||
*(threadgroup S1_2x4 *)(sb + NK*(8*sy + ly) + 8*sx) = (S1_2x4)(*((device T1_2x4 *) y));
|
||||
}
|
||||
#endif
|
||||
|
||||
@@ -935,12 +876,12 @@ kernel void kernel_mul_mm_id(
|
||||
|
||||
int i = tiisg;
|
||||
for (; i < nr0/4; i += 32) {
|
||||
*(D4 + i) = *(C4 + i) * s1_scale;
|
||||
*(D4 + i) = *(C4 + i);
|
||||
}
|
||||
|
||||
i = (4*(nr0/4)) + tiisg;
|
||||
for (; i < nr0; i += 32) {
|
||||
*(D + i) = *(C + i) * s1_scale;
|
||||
*(D + i) = *(C + i);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3302,7 +3302,8 @@ kernel void kernel_mul_mv_id(
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]]) {
|
||||
const int iid1 = tgpig.z/args.nei0;
|
||||
const int idx = tgpig.z%args.nei0;
|
||||
|
||||
@@ -3316,11 +3317,20 @@ kernel void kernel_mul_mv_id(
|
||||
const int64_t i1 = idx;
|
||||
const int64_t i2 = i12;
|
||||
|
||||
device char * dst_cur = dst + (i1*args.ne0 + i2*args.ne1*args.ne0)*sizeof(float);
|
||||
|
||||
if (i02 == -1) {
|
||||
device float * dst_f32 = (device float *) dst_cur;
|
||||
const ushort nth = ntg.x*ntg.y*ntg.z;
|
||||
for (int i0 = tiitg; i0 < args.ne0; i0 += nth) {
|
||||
dst_f32[i0] = 0.0f;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
device const char * src0_cur = src0s + i02*args.nb02;
|
||||
device const char * src1_cur = src1 + i11*args.nb11 + i12*args.nb12;
|
||||
|
||||
device char * dst_cur = dst + (i1*args.ne0 + i2*args.ne1*args.ne0)*sizeof(float);
|
||||
|
||||
ggml_metal_kargs_mul_mv args0 = {
|
||||
/*.ne00 =*/ args.ne00,
|
||||
/*.ne01 =*/ args.ne01,
|
||||
|
||||
@@ -1047,6 +1047,7 @@ struct ggml_backend_opencl_context {
|
||||
cl_kernel kernel_moe_reorder_b;
|
||||
cl_kernel kernel_moe_histogram, kernel_moe_scan, kernel_moe_fill, kernel_moe_scatter;
|
||||
cl_kernel kernel_moe_scatter_stable = nullptr; // deterministic slot assignment
|
||||
cl_kernel kernel_moe_zero_dst;
|
||||
cl_kernel kernel_moe_combine_f32 = nullptr; // fused router-weight mul + cross-expert sum
|
||||
cl_kernel kernel_moe_combine_bias_f32 = nullptr; // same, with the down-projection bias add folded in
|
||||
cl_kernel kernel_mul_mv_id_q4_0_f32_8x_flat;
|
||||
@@ -5547,6 +5548,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
|
||||
build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
|
||||
|
||||
CL_CHECK((backend_ctx->kernel_moe_histogram = clCreateKernel(prog, "kernel_moe_histogram", &err), err));
|
||||
CL_CHECK((backend_ctx->kernel_moe_zero_dst = clCreateKernel(prog, "kernel_moe_zero_dst", &err), err));
|
||||
CL_CHECK((backend_ctx->kernel_moe_scan = clCreateKernel(prog, "kernel_moe_scan", &err), err));
|
||||
CL_CHECK((backend_ctx->kernel_moe_fill = clCreateKernel(prog, "kernel_moe_fill", &err), err));
|
||||
CL_CHECK((backend_ctx->kernel_moe_scatter = clCreateKernel(prog, "kernel_moe_scatter", &err), err));
|
||||
@@ -26961,6 +26963,39 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,
|
||||
|
||||
GGML_ASSERT(ne00 == ne10);
|
||||
|
||||
// the matmul kernels only write the rows an expert owns, so zero the skipped rows here
|
||||
{
|
||||
cl_int err;
|
||||
|
||||
cl_buffer_region region;
|
||||
region.origin = offset2;
|
||||
region.size = nb21 * ne21;
|
||||
cl_mem buf_ids = clCreateSubBuffer(extra2->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &err);
|
||||
CL_CHECK(err);
|
||||
|
||||
region.origin = offsetd;
|
||||
region.size = (size_t) ne0 * ne1 * ne2 * sizeof(float);
|
||||
cl_mem buf_dst = clCreateSubBuffer(extrad->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &err);
|
||||
CL_CHECK(err);
|
||||
|
||||
const cl_uint ids_stride = nb21 / sizeof(int);
|
||||
|
||||
cl_kernel kernel_zero = backend_ctx->kernel_moe_zero_dst;
|
||||
CL_CHECK(clSetKernelArg(kernel_zero, 0, sizeof(cl_mem), &buf_ids));
|
||||
CL_CHECK(clSetKernelArg(kernel_zero, 1, sizeof(cl_mem), &buf_dst));
|
||||
CL_CHECK(clSetKernelArg(kernel_zero, 2, sizeof(cl_uint), &ne21));
|
||||
CL_CHECK(clSetKernelArg(kernel_zero, 3, sizeof(cl_uint), &ne20));
|
||||
CL_CHECK(clSetKernelArg(kernel_zero, 4, sizeof(cl_uint), &ids_stride));
|
||||
CL_CHECK(clSetKernelArg(kernel_zero, 5, sizeof(cl_uint), &ne0));
|
||||
|
||||
size_t zero_local[3] = {1, 1, 64};
|
||||
size_t zero_global[3] = {(size_t) ne21, (size_t) ne20, 64};
|
||||
backend_ctx->enqueue_ndrange_kernel(kernel_zero, 3, zero_global, zero_local, dst);
|
||||
|
||||
CL_CHECK(clReleaseMemObject(buf_ids));
|
||||
CL_CHECK(clReleaseMemObject(buf_dst));
|
||||
}
|
||||
|
||||
int sgs = 32; // subgroup size
|
||||
int nsg = 1; // number of subgroups
|
||||
int nrows = 1; // number of row in src1
|
||||
|
||||
@@ -34,6 +34,14 @@ kernel void kernel_add_id(
|
||||
|
||||
global float * dst_row = (global float *)((global char *)dst + i1*nb1 + i2*nb2);
|
||||
global float * src0_row = (global float *)((global char *)src0 + i1*nb01 + i2*nb02);
|
||||
|
||||
if (i11 == -1) { // skipped slot, add nothing
|
||||
for (int i0 = get_local_id(0); i0 < ne0; i0 += get_local_size(0)) {
|
||||
dst_row[i0] = src0_row[i0];
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
global float * src1_row = (global float *)((global char *)src1 + i11*nb11);
|
||||
|
||||
for (int i0 = get_local_id(0); i0 < ne0; i0 += get_local_size(0)) {
|
||||
|
||||
@@ -89,6 +89,9 @@ __kernel void kernel_gemv_moe_mxfp4_f32_ns(
|
||||
uint i11 = i20 % ne11;
|
||||
|
||||
uint expert_id = src2[i20];
|
||||
if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
|
||||
return;
|
||||
}
|
||||
uint expert_offset = expert_id * ne00 * ne01 / 32;
|
||||
|
||||
__private float sum = 0.0f; // each thread calculate partial sum of one output
|
||||
@@ -188,6 +191,9 @@ __kernel void kernel_gemv_moe_mxfp4_f32_ns_wimg(
|
||||
uint i11 = i20 % ne11;
|
||||
|
||||
uint expert_id = src2[i20];
|
||||
if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
|
||||
return;
|
||||
}
|
||||
uint expert_offset = expert_id * ne00 * ne01 / 32;
|
||||
|
||||
__private float sum = 0.0f;
|
||||
|
||||
@@ -44,6 +44,9 @@ __kernel void kernel_gemv_moe_q4_0_f32_ns(
|
||||
uint i11 = i20 % ne11;
|
||||
|
||||
uint expert_id = src2[i20];
|
||||
if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
|
||||
return;
|
||||
}
|
||||
uint expert_offset = expert_id * ne00 * ne01 / 32;
|
||||
|
||||
__private float sum = 0.0f; // each thread calculate partial sum of one output
|
||||
|
||||
@@ -45,6 +45,9 @@ __kernel void kernel_gemv_moe_q4_1_f32_ns(
|
||||
uint i11 = i20 % ne11;
|
||||
|
||||
uint expert_id = src2[i20];
|
||||
if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
|
||||
return;
|
||||
}
|
||||
uint expert_offset = expert_id * ne00 * ne01 / 32;
|
||||
|
||||
__private float sum = 0.0f; // each thread calculate partial sum of one output
|
||||
|
||||
@@ -61,6 +61,9 @@ __kernel void kernel_gemv_moe_q4_k_f32_ns(
|
||||
uint i11 = i20 % ne11;
|
||||
|
||||
uint expert_id = src2[i20];
|
||||
if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
|
||||
return;
|
||||
}
|
||||
|
||||
int num_superblocks = ne00 / QK_K;
|
||||
int num_subblocks = ne00 / 32;
|
||||
@@ -180,6 +183,9 @@ __kernel void kernel_gemv_moe_q4_k_f32_ns_wimg(
|
||||
uint i11 = i20 % ne11;
|
||||
|
||||
uint expert_id = src2[i20];
|
||||
if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
|
||||
return;
|
||||
}
|
||||
|
||||
int num_superblocks = ne00 / QK_K;
|
||||
int num_subblocks = ne00 / 32;
|
||||
|
||||
@@ -45,6 +45,9 @@ __kernel void kernel_gemv_moe_q5_0_f32_ns(
|
||||
uint i11 = i20 % ne11;
|
||||
|
||||
uint expert_id = src2[i20];
|
||||
if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
|
||||
return;
|
||||
}
|
||||
uint expert_offset = expert_id * ne00 * ne01 / 32;
|
||||
|
||||
__private float sum = 0.0f; // each thread calculate partial sum of one output
|
||||
|
||||
@@ -46,6 +46,9 @@ __kernel void kernel_gemv_moe_q5_1_f32_ns(
|
||||
uint i11 = i20 % ne11;
|
||||
|
||||
uint expert_id = src2[i20];
|
||||
if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
|
||||
return;
|
||||
}
|
||||
uint expert_offset = expert_id * ne00 * ne01 / 32;
|
||||
|
||||
__private float sum = 0.0f; // each thread calculate partial sum of one output
|
||||
|
||||
@@ -62,6 +62,9 @@ __kernel void kernel_gemv_moe_q5_k_f32_ns(
|
||||
uint i11 = i20 % ne11;
|
||||
|
||||
uint expert_id = src2[i20];
|
||||
if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
|
||||
return;
|
||||
}
|
||||
|
||||
int num_superblocks = ne00 / QK_K;
|
||||
int num_subblocks = ne00 / 32;
|
||||
|
||||
@@ -45,6 +45,9 @@ __kernel void kernel_gemv_moe_q6_k_f32_ns(
|
||||
uint i11 = i20 % ne11;
|
||||
|
||||
uint expert_id = src2[i20];
|
||||
if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
|
||||
return;
|
||||
}
|
||||
|
||||
int num_superblocks = ne00 / QK_K;
|
||||
int num_subblocks = ne00 / 32; // 8 sub-blocks of 32 per super-block
|
||||
|
||||
@@ -18,6 +18,9 @@ __kernel void kernel_moe_histogram(
|
||||
}
|
||||
|
||||
int expert_id = input[n * n_experts + k];
|
||||
if (expert_id == -1) { // skipped slot, owned by no expert
|
||||
return;
|
||||
}
|
||||
atomic_inc(&hist[expert_id]);
|
||||
}
|
||||
|
||||
@@ -63,6 +66,9 @@ __kernel void kernel_moe_scatter(
|
||||
}
|
||||
|
||||
int val = input[n * n_experts + k];
|
||||
if (val == -1) { // skipped slot, owned by no expert
|
||||
return;
|
||||
}
|
||||
|
||||
int local_slot = atomic_inc(&slot_counter[val]);
|
||||
|
||||
@@ -74,6 +80,32 @@ __kernel void kernel_moe_scatter(
|
||||
emap[tile_idx] = val;
|
||||
}
|
||||
|
||||
// the gemm only writes the rows an expert owns, so the skipped rows (negative ids) are zeroed here
|
||||
__kernel void kernel_moe_zero_dst(
|
||||
__global const int * input,
|
||||
__global float * dst,
|
||||
uint N,
|
||||
uint topK,
|
||||
uint n_experts,
|
||||
uint ne0
|
||||
) {
|
||||
uint n = get_global_id(0);
|
||||
uint k = get_global_id(1);
|
||||
|
||||
if (n >= N || k >= topK) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (input[n * n_experts + k] != -1) {
|
||||
return;
|
||||
}
|
||||
|
||||
__global float * dst_row = dst + ((ulong)n * topK + k) * ne0;
|
||||
for (uint i = get_local_id(2); i < ne0; i += get_local_size(2)) {
|
||||
dst_row[i] = 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
// Deterministic replacement for kernel_moe_scatter.
|
||||
//
|
||||
// kernel_moe_scatter takes each token's slot from atomic_inc(slot_counter[expert]),
|
||||
|
||||
@@ -172,6 +172,9 @@ kernel void kernel_mul_mv_id_mxfp4_f32(
|
||||
const int idx = get_group_id(2)%ne20;
|
||||
|
||||
int i02 = ((global int *) (src2 + iid1*nb21))[idx];
|
||||
if (i02 == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
|
||||
return;
|
||||
}
|
||||
|
||||
int i11 = idx % ne11;
|
||||
int i12 = iid1;
|
||||
|
||||
@@ -101,6 +101,9 @@ kernel void kernel_mul_mv_id_mxfp4_f32_flat(
|
||||
const int idx = get_group_id(2) % ne20;
|
||||
|
||||
uint i02 = ((global uint *) (src2 + offset2 + iid1 * nb21))[idx];
|
||||
if ((int)i02 == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
|
||||
return;
|
||||
}
|
||||
|
||||
int i11 = idx % ne11;
|
||||
|
||||
|
||||
@@ -267,6 +267,9 @@ kernel void kernel_mul_mv_id_q4_0_f32_8x_flat(
|
||||
const int idx = get_group_id(2)%ne20;
|
||||
|
||||
const int i02 = ((global int *)(src2 + iid1*nb21))[idx];
|
||||
if (i02 == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
|
||||
return;
|
||||
}
|
||||
|
||||
const int i11 = idx%ne11;
|
||||
const int i12 = iid1;
|
||||
|
||||
@@ -73,6 +73,9 @@ kernel void kernel_mul_mv_id_q8_0_f32(
|
||||
int idx = get_group_id(2)%ne20;
|
||||
|
||||
int i02 = ((global int *) (src2 + iid1*nb21))[idx];
|
||||
if (i02 == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
|
||||
return;
|
||||
}
|
||||
|
||||
int i11_ = idx % ne11;
|
||||
int i12_ = iid1;
|
||||
|
||||
@@ -72,6 +72,9 @@ kernel void kernel_mul_mv_id_q8_0_f32_flat(
|
||||
int idx = (int)get_group_id(2)%ne20;
|
||||
|
||||
int i02 = ((global int *) (src2 + iid1*nb21))[idx];
|
||||
if (i02 == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
|
||||
return;
|
||||
}
|
||||
|
||||
int i11_ = idx % ne11;
|
||||
int i12_ = iid1;
|
||||
|
||||
@@ -9,8 +9,12 @@
|
||||
#include <openvino/op/constant.hpp>
|
||||
#include <openvino/op/convert.hpp>
|
||||
#include <openvino/op/gather.hpp>
|
||||
#include <openvino/op/greater_eq.hpp>
|
||||
#include <openvino/op/maximum.hpp>
|
||||
#include <openvino/op/multiply.hpp>
|
||||
#include <openvino/op/reshape.hpp>
|
||||
#include <openvino/op/shape_of.hpp>
|
||||
#include <openvino/op/unsqueeze.hpp>
|
||||
|
||||
namespace ov {
|
||||
namespace frontend {
|
||||
@@ -52,8 +56,16 @@ OutputVector translate_add_id(const NodeContext & context) {
|
||||
ids = std::make_shared<ov::op::v0::Convert>(ids, ov::element::i32);
|
||||
}
|
||||
|
||||
// an id of -1 marks a skipped slot: gather row 0 for it and zero the bias with a mask
|
||||
auto zero = ov::op::v0::Constant::create(ids.get_element_type(), ov::Shape{}, {0});
|
||||
ov::Output<ov::Node> keep = std::make_shared<ov::op::v1::GreaterEqual>(ids, zero);
|
||||
keep = std::make_shared<ov::op::v0::Convert>(keep, bias.get_element_type());
|
||||
keep = std::make_shared<ov::op::v0::Unsqueeze>(keep, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {-1}));
|
||||
ids = std::make_shared<ov::op::v1::Maximum>(ids, zero);
|
||||
|
||||
auto gather_axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {0});
|
||||
ov::Output<ov::Node> selected_bias = std::make_shared<ov::op::v8::Gather>(bias, ids, gather_axis);
|
||||
selected_bias = std::make_shared<ov::op::v1::Multiply>(selected_bias, keep);
|
||||
selected_bias = std::make_shared<ov::op::v1::Reshape>(
|
||||
selected_bias, std::make_shared<ov::op::v3::ShapeOf>(input, ov::element::i64), false);
|
||||
|
||||
|
||||
@@ -517,6 +517,11 @@ bool socket_t::impl::rdma_recv(void * data, size_t size) {
|
||||
|
||||
int slot = (int)wc.wr_id;
|
||||
size_t got = wc.byte_len;
|
||||
if (got > rem) {
|
||||
LOG_ERROR("RDMA recv overrun: got=%zu rem=%zu\n", got, rem);
|
||||
c->post_rx(slot);
|
||||
return false;
|
||||
}
|
||||
memcpy(dst, c->rx_slot(slot), got);
|
||||
|
||||
if (!c->post_rx(slot)) return false;
|
||||
|
||||
@@ -26,6 +26,15 @@ static void add_id_kernel(
|
||||
float* dst_row = (float*)((char*)dst + i1 * nb1 + i2 * nb2);
|
||||
const float* src0_row =
|
||||
(const float*)((const char*)src0 + i1 * nb01 + i2 * nb02);
|
||||
|
||||
if (i11 == -1) { // skipped slot, add nothing
|
||||
for (int64_t i0 = item_ct1.get_local_id(2); i0 < ne0;
|
||||
i0 += item_ct1.get_local_range(2)) {
|
||||
dst_row[i0] = src0_row[i0];
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const float* src1_row = (const float*)((const char*)src1 + i11 * nb11);
|
||||
|
||||
for (int64_t i0 = item_ct1.get_local_id(2); i0 < ne0;
|
||||
|
||||
@@ -5700,6 +5700,23 @@ __dpct_inline__ static void k_copy_dst_from_contiguous(
|
||||
}
|
||||
}
|
||||
|
||||
__dpct_inline__ static void k_zero_dst_rows(
|
||||
char *__restrict__ dst_original,
|
||||
const mmid_row_mapping *__restrict__ row_mapping, int64_t ne0, size_t nb1,
|
||||
size_t nb2, const sycl::nd_item<3> &item_ct1) {
|
||||
int32_t i = item_ct1.get_group(2);
|
||||
|
||||
const int32_t i1 = row_mapping[i].i1;
|
||||
const int32_t i2 = row_mapping[i].i2;
|
||||
|
||||
float * dst_row_original = (float *)(dst_original + i1*nb1 + i2*nb2);
|
||||
|
||||
for (int j = item_ct1.get_local_id(2); j < ne0;
|
||||
j += item_ct1.get_local_range(2)) {
|
||||
dst_row_original[j] = 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
// returns true when ggml_sycl_mul_mat_id takes the fallback path that requires stream synchronization
|
||||
static bool ggml_sycl_mul_mat_id_needs_sync(const ggml_tensor * dst) {
|
||||
const ggml_tensor * src0 = dst->src[0];
|
||||
@@ -5785,14 +5802,18 @@ static void mmid_counting_sort_rows(
|
||||
int64_t n_ids, int64_t n_as, int64_t n_routed_rows,
|
||||
std::vector<int64_t> & expert_counts,
|
||||
std::vector<int64_t> & expert_row_offsets,
|
||||
std::vector<mmid_row_mapping> & routed_row_src) {
|
||||
std::vector<mmid_row_mapping> & routed_row_src,
|
||||
std::vector<mmid_row_mapping> & skipped_row_dst) {
|
||||
|
||||
// frequencies: how many routed rows each expert "owns"
|
||||
expert_counts.assign(n_as, 0);
|
||||
for (int64_t iid1 = 0; iid1 < ids->ne[1]; iid1++) {
|
||||
for (int64_t id = 0; id < n_ids; id++) {
|
||||
const int32_t row_id_i = *(const int32_t *) (ids_host + iid1*ids->nb[1] + id*ids->nb[0]);
|
||||
GGML_ASSERT(row_id_i >= 0 && row_id_i < n_as);
|
||||
GGML_ASSERT(row_id_i == -1 || (row_id_i >= 0 && row_id_i < n_as));
|
||||
if (row_id_i == -1) {
|
||||
continue;
|
||||
}
|
||||
expert_counts[row_id_i]++;
|
||||
}
|
||||
}
|
||||
@@ -5805,10 +5826,17 @@ static void mmid_counting_sort_rows(
|
||||
|
||||
std::vector<int64_t> expert_row_next = expert_row_offsets;
|
||||
routed_row_src.resize(n_routed_rows);
|
||||
skipped_row_dst.clear();
|
||||
for (int64_t iid1 = 0; iid1 < ids->ne[1]; iid1++) {
|
||||
for (int64_t id = 0; id < n_ids; id++) {
|
||||
const int32_t row_id_i = *(const int32_t *) (ids_host + iid1*ids->nb[1] + id*ids->nb[0]);
|
||||
GGML_ASSERT(row_id_i >= 0 && row_id_i < n_as);
|
||||
GGML_ASSERT(row_id_i == -1 || (row_id_i >= 0 && row_id_i < n_as));
|
||||
|
||||
// a skipped slot has no expert, its dst row is zeroed instead
|
||||
if (row_id_i == -1) {
|
||||
skipped_row_dst.push_back({(int32_t) id, (int32_t) iid1});
|
||||
continue;
|
||||
}
|
||||
|
||||
// find and validate the next free row for a given expert (row_id_i)
|
||||
const int64_t routed_row = expert_row_next[row_id_i]++;
|
||||
@@ -5817,6 +5845,7 @@ static void mmid_counting_sort_rows(
|
||||
routed_row_src[routed_row] = {(int32_t) id, (int32_t) iid1};
|
||||
}
|
||||
}
|
||||
routed_row_src.resize(expert_row_offsets[n_as]);
|
||||
}
|
||||
|
||||
static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
|
||||
@@ -5884,7 +5913,7 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
|
||||
for (int64_t iid1 = 0; iid1 < ids->ne[1]; iid1++) {
|
||||
for (int64_t id = 0; id < n_ids; id++) {
|
||||
const int32_t i02 = *(const int32_t *) (ids_host.data() + iid1*ids->nb[1] + id*ids->nb[0]);
|
||||
GGML_ASSERT(i02 >= 0 && i02 < n_as);
|
||||
GGML_ASSERT(i02 == -1 || (i02 >= 0 && i02 < n_as));
|
||||
|
||||
const int64_t i11 = id % ne11;
|
||||
const int64_t i12 = iid1;
|
||||
@@ -5892,6 +5921,12 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
|
||||
const int64_t i1 = id;
|
||||
const int64_t i2 = i12;
|
||||
|
||||
if (i02 == -1) { // skipped slot
|
||||
SYCL_CHECK(CHECK_TRY_ERROR(
|
||||
stream->memset(dst_original + i1*nb1 + i2*nb2, 0, ne0*sizeof(float))));
|
||||
continue;
|
||||
}
|
||||
|
||||
src0_row.data = src0_original + i02*nb02;
|
||||
src1_row.data = src1_original + i11*nb11 + i12*nb12;
|
||||
dst_row.data = dst_original + i1*nb1 + i2*nb2;
|
||||
@@ -5913,20 +5948,43 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
|
||||
std::vector<int64_t> expert_row_offsets;
|
||||
// the sources (slot/token pairs) of contiguous rows to guide k_copy_src1_to_contiguous
|
||||
std::vector<mmid_row_mapping> & routed_row_src = ctx.mmid_row_mapping_host;
|
||||
// the slot/token pairs of the skipped slots, whose dst rows are zeroed
|
||||
std::vector<mmid_row_mapping> skipped_row_dst;
|
||||
|
||||
mmid_counting_sort_rows(ids, ids_host.data(), n_ids, n_as, n_routed_rows,
|
||||
expert_row_counts, expert_row_offsets, routed_row_src);
|
||||
expert_row_counts, expert_row_offsets, routed_row_src, skipped_row_dst);
|
||||
|
||||
const int64_t n_valid_rows = (int64_t) routed_row_src.size();
|
||||
|
||||
ggml_sycl_pool_alloc<mmid_row_mapping> dev_row_mapping(ctx.pool(), n_routed_rows);
|
||||
SYCL_CHECK(CHECK_TRY_ERROR(
|
||||
stream->memcpy(dev_row_mapping.get(), routed_row_src.data(), n_routed_rows*sizeof(mmid_row_mapping))));
|
||||
stream->memcpy(dev_row_mapping.get(), routed_row_src.data(), n_valid_rows*sizeof(mmid_row_mapping))));
|
||||
|
||||
const unsigned int max_work_group_size = ggml_sycl_info().max_work_group_sizes[ctx.device];
|
||||
assert(max_work_group_size % (WARP_SIZE * WARP_SIZE) == 0);
|
||||
|
||||
{
|
||||
if (!skipped_row_dst.empty()) {
|
||||
ggml_sycl_pool_alloc<mmid_row_mapping> dev_skipped(ctx.pool(), skipped_row_dst.size());
|
||||
SYCL_CHECK(CHECK_TRY_ERROR(
|
||||
stream->memcpy(dev_skipped.get(), skipped_row_dst.data(), skipped_row_dst.size()*sizeof(mmid_row_mapping))));
|
||||
|
||||
sycl::range<3> block_dims(1, 1, std::min((unsigned int)ne0, max_work_group_size));
|
||||
sycl::range<3> grid_dims(1, 1, skipped_row_dst.size());
|
||||
stream->submit([&](sycl::handler &cgh) {
|
||||
mmid_row_mapping *__restrict dev_skipped_get = dev_skipped.get();
|
||||
|
||||
cgh.parallel_for(
|
||||
sycl::nd_range<3>(grid_dims * block_dims, block_dims),
|
||||
[=](sycl::nd_item<3> item_ct1) {
|
||||
k_zero_dst_rows(dst_original, dev_skipped_get, ne0, nb1, nb2, item_ct1);
|
||||
});
|
||||
});
|
||||
SYCL_CHECK(CHECK_TRY_ERROR(stream->wait()));
|
||||
}
|
||||
|
||||
if (n_valid_rows > 0) {
|
||||
sycl::range<3> block_dims(1, 1, std::min((unsigned int)ne10, max_work_group_size));
|
||||
sycl::range<3> grid_dims(1, 1, n_routed_rows);
|
||||
sycl::range<3> grid_dims(1, 1, n_valid_rows);
|
||||
stream->submit([&](sycl::handler &cgh) {
|
||||
char *__restrict src1_contiguous_get =
|
||||
src1_contiguous.get();
|
||||
@@ -5992,9 +6050,9 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
|
||||
ggml_sycl_mul_mat(ctx, &src0_row, &src1_row, &dst_row);
|
||||
}
|
||||
|
||||
{
|
||||
if (n_valid_rows > 0) {
|
||||
sycl::range<3> block_dims(1, 1, std::min((unsigned int)ne0, max_work_group_size));
|
||||
sycl::range<3> grid_dims(1, 1, n_routed_rows);
|
||||
sycl::range<3> grid_dims(1, 1, n_valid_rows);
|
||||
stream->submit([&](sycl::handler &cgh) {
|
||||
const char *__restrict dst_contiguous_get =
|
||||
dst_contiguous.get();
|
||||
|
||||
@@ -3149,7 +3149,6 @@ static void mul_mat_vec_q_moe(
|
||||
const int expert_idx = item_ct1.get_group(1);
|
||||
const int i02 = ids_dev[expert_idx];
|
||||
|
||||
const char * vx = (const char *) vx_base + (size_t) i02 * expert_weight_stride;
|
||||
const char * vy = (const char *) vy_base + (size_t) expert_idx * src1_row_stride;
|
||||
float * dst = (float *) ((char *) dst_base + (size_t) expert_idx * dst_row_stride);
|
||||
|
||||
@@ -3159,6 +3158,15 @@ static void mul_mat_vec_q_moe(
|
||||
return;
|
||||
}
|
||||
|
||||
if (i02 == -1) { // skipped slot
|
||||
if (item_ct1.get_local_id(2) == 0) {
|
||||
dst[row] = 0.0f;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const char * vx = (const char *) vx_base + (size_t) i02 * expert_weight_stride;
|
||||
|
||||
const int blocks_per_row = ncols / qk;
|
||||
constexpr int blocks_per_warp = (vdr * WARP_SIZE + qi - 1) / qi;
|
||||
|
||||
@@ -3355,7 +3363,6 @@ static void mul_mat_vec_q_moe_reorder(
|
||||
const int expert_idx = item_ct1.get_group(1);
|
||||
const int i02 = ids_dev[expert_idx];
|
||||
|
||||
const char * vx = (const char *) vx_base + (size_t) i02 * expert_weight_stride;
|
||||
const char * vy = (const char *) vy_base + (size_t) expert_idx * src1_row_stride;
|
||||
float * dst = (float *) ((char *) dst_base + (size_t) expert_idx * dst_row_stride);
|
||||
|
||||
@@ -3366,6 +3373,15 @@ static void mul_mat_vec_q_moe_reorder(
|
||||
|
||||
const auto sg = item_ct1.get_sub_group();
|
||||
|
||||
if (i02 == -1) { // skipped slot
|
||||
if (sg.leader()) {
|
||||
dst[row] = 0.0f;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const char * vx = (const char *) vx_base + (size_t) i02 * expert_weight_stride;
|
||||
|
||||
constexpr int vdr_mmvq = reorder_vec_dot_q_sycl::gtype == GGML_TYPE_Q8_0 ?
|
||||
VDR_Q8_0_Q8_1_MMVQ : block_traits::vdr_mmvq;
|
||||
const int blocks_per_row = ncols / block_traits::qk;
|
||||
|
||||
@@ -222,6 +222,9 @@ struct vk_op_count_experts_push_constants {
|
||||
uint32_t ne00L;
|
||||
uint32_t row_tile_size;
|
||||
uint32_t row_ids_offset;
|
||||
uint32_t ne0_dst;
|
||||
uint32_t stride_dst;
|
||||
uint32_t batch_stride_dst;
|
||||
};
|
||||
|
||||
struct vk_op_glu_push_constants {
|
||||
|
||||
@@ -3585,9 +3585,9 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
|
||||
ggml_vk_create_pipeline(device, device->pipeline_count_equal_i32, "count_equal_i32", count_equal_i32_len, count_equal_i32_data, "main", 3, sizeof(vk_op_push_constants), {512, 1, 1}, { device->subgroup_size }, 1);
|
||||
|
||||
if (device->subgroup_arithmetic && device->subgroup_vote && device->subgroup_require_full_support) {
|
||||
ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_subgroup_len, count_experts_subgroup_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true, true);
|
||||
ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_subgroup_len, count_experts_subgroup_data, "main", 3, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true, true);
|
||||
} else {
|
||||
ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true);
|
||||
ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 3, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true);
|
||||
}
|
||||
|
||||
// comb holds a token's 4x4 matrix in one 16-lane slice of a subgroup, so it
|
||||
@@ -7643,16 +7643,20 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
|
||||
ggml_vk_sync_buffers(ctx, subctx);
|
||||
}
|
||||
{
|
||||
// this also zeroes the dst rows of the skipped slots (negative ids)
|
||||
vk_op_count_experts_push_constants pc = { (uint32_t)nei0,
|
||||
(uint32_t)nei1,
|
||||
(uint32_t)(nbi0 / ggml_type_size(ids->type)),
|
||||
(uint32_t)(nbi1 / ggml_type_size(ids->type)),
|
||||
(uint32_t)(get_misalign_bytes(ctx, ids) / ggml_type_size(ids->type)),
|
||||
(uint32_t)n_as,
|
||||
0, 0, bn, (uint32_t)row_ids_offset };
|
||||
0, 0, bn, (uint32_t)row_ids_offset,
|
||||
(uint32_t)ne01,
|
||||
(uint32_t)(dst->nb[1] / sizeof(float)),
|
||||
(uint32_t)(dst->nb[2] / sizeof(float)) };
|
||||
init_pushconst_fastdiv(pc);
|
||||
ggml_vk_dispatch_pipeline(ctx, subctx, count_experts,
|
||||
{ d_ids, expert_count_buf }, pc,
|
||||
{ d_ids, expert_count_buf, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, pc,
|
||||
{ 1, 1, 1 });
|
||||
}
|
||||
|
||||
|
||||
@@ -27,14 +27,22 @@ void main() {
|
||||
const uint i1 = gl_WorkGroupID.x;
|
||||
const uint i2 = gl_WorkGroupID.y;
|
||||
|
||||
const uint i11 = data_c[i1 + i2 * p.s21];
|
||||
const int i11 = data_c[i1 + i2 * p.s21];
|
||||
|
||||
const uint s1 = p.ne0;
|
||||
const uint s2 = p.ne0 * p.ne1;
|
||||
|
||||
const uint d0 = i1 * s1 + i2 * s2;
|
||||
const uint a0 = i1 * p.s01 + i2 * p.s02;
|
||||
const uint b0 = i11 * p.s11;
|
||||
|
||||
if (i11 == -1) { // skipped slot, add nothing
|
||||
for (uint i0 = gl_LocalInvocationID.x; i0 < p.ne0; i0 += BLOCK_SIZE) {
|
||||
data_d[d0 + i0] = data_a[a0 + i0];
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const uint b0 = uint(i11) * p.s11;
|
||||
|
||||
for (uint i0 = gl_LocalInvocationID.x; i0 < p.ne0; i0 += BLOCK_SIZE) {
|
||||
data_d[d0 + i0] = data_a[a0 + i0] + data_b[b0 + i0];
|
||||
|
||||
@@ -23,6 +23,9 @@ layout (push_constant) uniform parameter
|
||||
uint32_t ne00L;
|
||||
uint32_t row_tile_size;
|
||||
uint32_t row_ids_offset;
|
||||
uint32_t ne0_dst;
|
||||
uint32_t stride_dst;
|
||||
uint32_t batch_stride_dst;
|
||||
} p;
|
||||
|
||||
#define BLOCK_SIZE 256
|
||||
@@ -32,6 +35,7 @@ layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
|
||||
|
||||
layout (binding = 0) readonly buffer A {uint data_a[];};
|
||||
layout (binding = 1) writeonly buffer D {uint data_d[];};
|
||||
layout (binding = 2) writeonly buffer DST {float data_dst[];};
|
||||
|
||||
shared uint vals[EXPERTS_PER_CHUNK];
|
||||
shared uint offsets[EXPERTS_PER_CHUNK];
|
||||
@@ -47,6 +51,26 @@ void write_tiles(uint expert, uint count, uint tile_begin, uint n_tiles) {
|
||||
}
|
||||
}
|
||||
|
||||
// the matmul writes only the rows that an expert owns, so zero the rows of the skipped slots
|
||||
void zero_skipped_rows() {
|
||||
const uint num_elements = p.ne00 * p.ne01;
|
||||
const uint tid = gl_LocalInvocationID.x;
|
||||
|
||||
for (uint idx = 0; idx < num_elements; idx++) {
|
||||
const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
|
||||
const uint i00 = idx - i01 * p.ne00;
|
||||
|
||||
if (int(data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00]) != -1) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const uint d_off = i00 * p.stride_dst + i01 * p.batch_stride_dst;
|
||||
for (uint i = tid; i < p.ne0_dst; i += BLOCK_SIZE) {
|
||||
data_dst[d_off + i] = 0.0f;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void main() {
|
||||
const uint num_elements = p.ne00 * p.ne01;
|
||||
const uint tid = gl_LocalInvocationID.x;
|
||||
@@ -123,4 +147,6 @@ void main() {
|
||||
if (tid == 0) {
|
||||
data_d[0] = total_tiles;
|
||||
}
|
||||
|
||||
zero_skipped_rows();
|
||||
}
|
||||
|
||||
@@ -42,6 +42,7 @@ layout (push_constant) uniform parameter
|
||||
|
||||
#ifdef MUL_MAT_ID
|
||||
uint expert_id;
|
||||
bool expert_skip; // an id of -1 skips the slot, the result row is zeroed
|
||||
#endif
|
||||
|
||||
void get_offsets(out uint a_offset, out uint b_offset, out uint d_offset) {
|
||||
@@ -63,7 +64,9 @@ void get_offsets(out uint a_offset, out uint b_offset, out uint d_offset) {
|
||||
batch_idx_a = i03 * p.ne02 + i02;
|
||||
}
|
||||
#else
|
||||
expert_id = data_ids[expert_i0 + p.expert_i1 * p.nbi1];
|
||||
const int expert = data_ids[expert_i0 + p.expert_i1 * p.nbi1];
|
||||
expert_skip = expert == -1;
|
||||
expert_id = expert_skip ? 0 : uint(expert); // clamp, so that the offsets below stay in range
|
||||
#endif
|
||||
|
||||
a_offset =
|
||||
@@ -113,6 +116,9 @@ void reduce_result(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const in uint32_t
|
||||
const uint expert_i0 = gl_GlobalInvocationID.y;
|
||||
temp[j][n] *= FLOAT_TYPE(data_fuse1[expert_i0]);
|
||||
}
|
||||
if (expert_skip) {
|
||||
temp[j][n] = FLOAT_TYPE(0);
|
||||
}
|
||||
#else
|
||||
if ((p.fusion_flags & MAT_VEC_FUSION_FLAGS_BIAS0) != 0) {
|
||||
temp[j][n] += FLOAT_TYPE(data_fuse0[j*p.batch_stride_d + d_offset + first_row + n]);
|
||||
@@ -168,6 +174,9 @@ void reduce_result(FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const in uint32_t d_offs
|
||||
const uint expert_i0 = gl_GlobalInvocationID.y;
|
||||
temp[j][n] *= FLOAT_TYPE(data_fuse1[expert_i0]);
|
||||
}
|
||||
if (expert_skip) {
|
||||
temp[j][n] = FLOAT_TYPE(0);
|
||||
}
|
||||
#else
|
||||
if ((p.fusion_flags & MAT_VEC_FUSION_FLAGS_BIAS0) != 0) {
|
||||
temp[j][n] += FLOAT_TYPE(data_fuse0[j*p.batch_stride_d + d_offset + first_row + n]);
|
||||
@@ -213,6 +222,9 @@ void reduce_result(FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const in uint32_t d_offs
|
||||
const uint expert_i0 = gl_GlobalInvocationID.y;
|
||||
tmpsh[j][n][0] *= FLOAT_TYPE(data_fuse1[expert_i0]);
|
||||
}
|
||||
if (expert_skip) {
|
||||
tmpsh[j][n][0] = FLOAT_TYPE(0);
|
||||
}
|
||||
#else
|
||||
if ((p.fusion_flags & MAT_VEC_FUSION_FLAGS_BIAS0) != 0) {
|
||||
tmpsh[j][n][0] += FLOAT_TYPE(data_fuse0[j*p.batch_stride_d + d_offset + first_row + n]);
|
||||
|
||||
@@ -1863,10 +1863,12 @@ static webgpu_encoded_op ggml_webgpu_mul_mat_id(webgpu_context & ctx,
|
||||
// params for mul_mat_id_gather.wgsl
|
||||
std::vector<uint32_t> gather_params = {
|
||||
(uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src2) / ggml_type_size(src2->type)),
|
||||
(uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)),
|
||||
param_n_expert,
|
||||
param_n_expert_used,
|
||||
param_n_tokens,
|
||||
(uint32_t) (src2->nb[1] / ggml_type_size(src2->type)),
|
||||
(uint32_t) dst->ne[0],
|
||||
};
|
||||
|
||||
const size_t dst_offset = ggml_webgpu_tensor_offset(dst);
|
||||
@@ -1888,6 +1890,8 @@ static webgpu_encoded_op ggml_webgpu_mul_mat_id(webgpu_context & ctx,
|
||||
ggml_webgpu_tensor_binding_size(ctx, src2)),
|
||||
ggml_webgpu_make_bind_group_entry(1, ggml_webgpu_tensor_buf(dst), ids_cache_align_offset, ids_cache_bind_size),
|
||||
ggml_webgpu_make_bind_group_entry(2, ggml_webgpu_tensor_buf(dst), ids_tiles_align_offset, ids_tiles_bind_size),
|
||||
ggml_webgpu_make_bind_group_entry(3, ggml_webgpu_tensor_buf(dst), ggml_webgpu_tensor_align_offset(ctx, dst),
|
||||
ggml_webgpu_tensor_binding_size(ctx, dst)),
|
||||
};
|
||||
|
||||
dispatches.push_back({
|
||||
|
||||
@@ -46,12 +46,24 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
|
||||
let i2 = wg_linear / params.ne1;
|
||||
let i1 = wg_linear % params.ne1;
|
||||
|
||||
let i11 = u32(ids[params.offset_ids + i1 * params.nb20 + i2 * params.nb21]);
|
||||
let i11 = i32(ids[params.offset_ids + i1 * params.nb20 + i2 * params.nb21]);
|
||||
|
||||
let src0_row = params.offset_src0 + i1 * params.nb01 + i2 * params.nb02;
|
||||
let src1_row = params.offset_src1 + i11 * params.nb11;
|
||||
let dst_row = params.offset_dst + i1 * params.ne0 + i2 * (params.ne0 * params.ne1);
|
||||
|
||||
if (i11 == -1) { // skipped slot, add nothing
|
||||
#ifdef INPLACE
|
||||
// src0 already holds the result
|
||||
#else
|
||||
for (var i = thread_id;i < params.ne0; i += WG_SIZE) {
|
||||
dst[dst_row + i] = src0[src0_row + i];
|
||||
}
|
||||
#endif
|
||||
return;
|
||||
}
|
||||
|
||||
let src1_row = params.offset_src1 + u32(i11) * params.nb11;
|
||||
|
||||
for (var i = thread_id;i < params.ne0; i += WG_SIZE) {
|
||||
#ifdef INPLACE
|
||||
src0[src0_row + i] = src0[src0_row + i] + src1[src1_row + i];
|
||||
|
||||
@@ -2,19 +2,24 @@ enable f16;
|
||||
|
||||
struct MulMatIdGatherParams {
|
||||
offset_ids: u32,
|
||||
offset_dst: u32,
|
||||
|
||||
n_expert: u32,
|
||||
n_expert_used: u32,
|
||||
n_tokens: u32,
|
||||
|
||||
stride_ids_1: u32,
|
||||
|
||||
m: u32,
|
||||
};
|
||||
|
||||
@group(0) @binding(0) var<storage, read_write> ids: array<i32>; // [n_expert_used, n_tokens]
|
||||
@group(0) @binding(1) var<storage, read_write> ids_slot_pos: array<u32>; // [n_expert_used * n_tokens]
|
||||
@group(0) @binding(2) var<storage, read_write> ids_tiles: array<u32>; // [n_tiles, (expert, first slot, n slots) * n_tiles]
|
||||
|
||||
@group(0) @binding(3) var<uniform> params: MulMatIdGatherParams;
|
||||
@group(0) @binding(3) var<storage, read_write> dst: array<f32>; // [m, n_expert_used, n_tokens]
|
||||
|
||||
@group(0) @binding(4) var<uniform> params: MulMatIdGatherParams;
|
||||
|
||||
// per expert: the slot count, replaced by the first slot of the expert in ids_slot_pos by the prefix sum, then advanced as its write cursor
|
||||
var<workgroup> ids_count: array<atomic<u32>, N_EXPERT>;
|
||||
@@ -36,8 +41,10 @@ fn main(@builtin(local_invocation_id) local_id: vec3<u32>) {
|
||||
for (var i = thread_id; i < ids_slots_num; i += WG_SIZE) {
|
||||
let row = i / params.n_expert_used;
|
||||
let col = i % params.n_expert_used;
|
||||
let expert = u32(ids[params.offset_ids + row * params.stride_ids_1 + col]);
|
||||
atomicAdd(&ids_count[expert], 1u);
|
||||
let expert = ids[params.offset_ids + row * params.stride_ids_1 + col];
|
||||
if (expert != -1) { // skipped slot, owned by no expert
|
||||
atomicAdd(&ids_count[u32(expert)], 1u);
|
||||
}
|
||||
}
|
||||
|
||||
workgroupBarrier();
|
||||
@@ -71,8 +78,23 @@ fn main(@builtin(local_invocation_id) local_id: vec3<u32>) {
|
||||
for (var i = thread_id; i < ids_slots_num; i += WG_SIZE) {
|
||||
let row = i / params.n_expert_used;
|
||||
let col = i % params.n_expert_used;
|
||||
let expert = u32(ids[params.offset_ids + row * params.stride_ids_1 + col]);
|
||||
let pos = atomicAdd(&ids_count[expert], 1u);
|
||||
ids_slot_pos[pos] = i;
|
||||
let expert = ids[params.offset_ids + row * params.stride_ids_1 + col];
|
||||
if (expert != -1) {
|
||||
let pos = atomicAdd(&ids_count[u32(expert)], 1u);
|
||||
ids_slot_pos[pos] = i;
|
||||
}
|
||||
}
|
||||
|
||||
// the matmul writes only the rows that an expert owns, so zero the rows of the skipped slots
|
||||
for (var i: u32 = 0; i < ids_slots_num; i += 1) {
|
||||
let row = i / params.n_expert_used;
|
||||
let col = i % params.n_expert_used;
|
||||
if (ids[params.offset_ids + row * params.stride_ids_1 + col] != -1) {
|
||||
continue;
|
||||
}
|
||||
let dst_row = params.offset_dst + i * params.m;
|
||||
for (var r = thread_id; r < params.m; r += WG_SIZE) {
|
||||
dst[dst_row + r] = 0.0f;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -70,11 +70,18 @@ fn main(
|
||||
let own_expert = ids[params.offset_ids + slot];
|
||||
let row_base = (wg_linear % output_groups) * OUTPUTS_PER_WG;
|
||||
|
||||
let src0_batch_offset = params.offset_src0 + own_expert * params.stride_02;
|
||||
// a skipped slot (id -1) keeps the zero accumulator, so the reduction below writes zeros
|
||||
// no early return: the barriers below must stay in uniform control flow
|
||||
let skipped = i32(own_expert) == -1;
|
||||
|
||||
let src0_batch_offset = params.offset_src0 + select(own_expert, 0u, skipped) * params.stride_02;
|
||||
let src1_idx_base = params.offset_src1 + (slot % params.b_ne1) * params.stride_11;
|
||||
let dst_idx_base = params.offset_dst + slot * params.m + row_base;
|
||||
|
||||
let acc = accumulate_vec_dot(thread_id, row_base, src0_batch_offset, src1_idx_base);
|
||||
var acc: array<array<f32, OUTPUTS_PER_WG>, NUM_COLS>;
|
||||
if (!skipped) {
|
||||
acc = accumulate_vec_dot(thread_id, row_base, src0_batch_offset, src1_idx_base);
|
||||
}
|
||||
|
||||
#ifdef USE_SUBGROUP_REDUCTION
|
||||
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
|
||||
|
||||
@@ -375,7 +375,12 @@ static void ggml_zendnn_compute_forward_mul_mat_id(
|
||||
for (int id = 0; id < n_ids; ++id) {
|
||||
const int32_t i02 = *(const int32_t *)((const char *)ids->data + iid1*ids->nb[1] + id*ids->nb[0]);
|
||||
|
||||
GGML_ASSERT(i02 >= 0 && i02 < n_as);
|
||||
GGML_ASSERT(i02 == -1 || (i02 >= 0 && i02 < n_as));
|
||||
|
||||
if (i02 == -1) {
|
||||
memset((char *) dst->data + id*nb1 + iid1*nb2, 0, ne0*sizeof(float));
|
||||
continue;
|
||||
}
|
||||
|
||||
if (matrix_row_counts[i02] == 0) {
|
||||
n_active_experts++;
|
||||
|
||||
@@ -3438,6 +3438,8 @@ void ggml_mul_mat_set_hint(
|
||||
in b, n_expert_used can be broadcasted to match the n_expert_used of ids
|
||||
|
||||
c ~= as[:,:,i] @ b[:,i%r,t], i = ids[e,t] for all e,t in ids
|
||||
|
||||
an id of -1 skips the slot: no matrix is selected, c[:,e,t] is set to zero
|
||||
*/
|
||||
struct ggml_tensor * ggml_mul_mat_id(
|
||||
struct ggml_context * ctx,
|
||||
|
||||
@@ -211,6 +211,7 @@ class Keys:
|
||||
VALUE_SCALE = "{arch}.attention.value_scale"
|
||||
COMPRESS_RATIOS = "{arch}.attention.compress_ratios"
|
||||
COMPRESS_ROPE_FREQ_BASE = "{arch}.attention.compress_rope_freq_base"
|
||||
KV_SOURCE_LAYERS = "{arch}.attention.kv_source_layers"
|
||||
TEMPERATURE_LENGTH = "{arch}.attention.temperature_length"
|
||||
KEY_LENGTH_MLA = "{arch}.attention.key_length_mla"
|
||||
VALUE_LENGTH_MLA = "{arch}.attention.value_length_mla"
|
||||
@@ -236,6 +237,10 @@ class Keys:
|
||||
TYPES = "{arch}.attention.indexer.types"
|
||||
KPOOL = "{arch}.attention.indexer.kpool" # GLM5-Next
|
||||
KPOOL_SELECT_TAIL = "{arch}.attention.indexer.kpool_select_tail" # GLM5-Next
|
||||
SOURCE_LAYERS = "{arch}.attention.indexer.source_layers"
|
||||
CANDIDATE_SOURCE_LAYER = "{arch}.attention.indexer.candidate_source_layer"
|
||||
CANDIDATE_TOP_K_BLOCKS = "{arch}.attention.indexer.candidate_top_k_blocks"
|
||||
CANDIDATE_BLOCK_SIZE = "{arch}.attention.indexer.candidate_block_size"
|
||||
|
||||
class HyperConnection:
|
||||
COUNT = "{arch}.hyper_connection.count"
|
||||
@@ -245,6 +250,12 @@ class Keys:
|
||||
MAGNITUDE = "{arch}.hyper_connection.magnitude"
|
||||
# absent means the mix projection is full rank (DeepSeek-V4 behaviour)
|
||||
LOW_RANK = "{arch}.hyper_connection.low_rank"
|
||||
# single-pass mHC: each sublayer collapses with the previous pre-mix (DeepSeek-V4.1)
|
||||
SINGLE_PASS = "{arch}.hyper_connection.single_pass"
|
||||
|
||||
class Engram: # deepseek41; the table uses the PerLayerEmbedding keys
|
||||
PAD_TOKEN_ID = "{arch}.engram.pad_token_id"
|
||||
TOKEN_MAP = "{arch}.engram.token_map"
|
||||
|
||||
class PerLayerEmbedding:
|
||||
LAYERS = "{arch}.ple.layers"
|
||||
@@ -636,6 +647,7 @@ class MODEL_ARCH(IntEnum):
|
||||
DEEPSEEK2OCR = auto()
|
||||
DEEPSEEK32 = auto()
|
||||
DEEPSEEK4 = auto()
|
||||
DEEPSEEK41 = auto()
|
||||
CHATGLM = auto()
|
||||
GLM4 = auto()
|
||||
GLM4_MOE = auto()
|
||||
@@ -1317,6 +1329,8 @@ class MODEL_TENSOR(IntEnum):
|
||||
DSPARK_MARKOV_W1 = auto() # markov head: prev-token embed
|
||||
DSPARK_MARKOV_W2 = auto() # markov head: bias projection
|
||||
DSPARK_CONF_PROJ = auto() # confidence head
|
||||
ENGRAM_KV = auto()
|
||||
ENGRAM_GATE = auto()
|
||||
DFLASH_ATTN_CONV_BASE = auto()
|
||||
DFLASH_ATTN_CONV_PROJ = auto()
|
||||
DFLASH_FFN_CONV_BASE = auto()
|
||||
@@ -1460,6 +1474,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
|
||||
MODEL_ARCH.DEEPSEEK2OCR: "deepseek2-ocr",
|
||||
MODEL_ARCH.DEEPSEEK32: "deepseek32",
|
||||
MODEL_ARCH.DEEPSEEK4: "deepseek4",
|
||||
MODEL_ARCH.DEEPSEEK41: "deepseek41",
|
||||
MODEL_ARCH.CHATGLM: "chatglm",
|
||||
MODEL_ARCH.GLM4: "glm4",
|
||||
MODEL_ARCH.GLM4_MOE: "glm4moe",
|
||||
@@ -2164,6 +2179,8 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
|
||||
MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1",
|
||||
MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2",
|
||||
MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj",
|
||||
MODEL_TENSOR.ENGRAM_KV: "blk.{bid}.engram_kv",
|
||||
MODEL_TENSOR.ENGRAM_GATE: "blk.{bid}.engram_gate",
|
||||
MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: "blk.{bid}.attn_conv_base",
|
||||
MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: "blk.{bid}.attn_conv_proj",
|
||||
MODEL_TENSOR.DFLASH_FFN_CONV_BASE: "blk.{bid}.ffn_conv_base",
|
||||
@@ -4358,6 +4375,46 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
|
||||
MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
|
||||
],
|
||||
MODEL_ARCH.DEEPSEEK41: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
MODEL_TENSOR.OUTPUT_NORM,
|
||||
MODEL_TENSOR.OUTPUT,
|
||||
MODEL_TENSOR.ATTN_NORM,
|
||||
MODEL_TENSOR.ATTN_SINKS,
|
||||
MODEL_TENSOR.ATTN_Q_A,
|
||||
MODEL_TENSOR.ATTN_Q_B,
|
||||
MODEL_TENSOR.ATTN_Q_A_NORM,
|
||||
MODEL_TENSOR.ATTN_KV,
|
||||
MODEL_TENSOR.ATTN_KV_NORM,
|
||||
MODEL_TENSOR.ATTN_OUT_A,
|
||||
MODEL_TENSOR.ATTN_OUT_B,
|
||||
MODEL_TENSOR.HC_ATTN_FN,
|
||||
MODEL_TENSOR.HC_ATTN_BASE,
|
||||
MODEL_TENSOR.HC_ATTN_SCALE,
|
||||
MODEL_TENSOR.HC_FFN_FN,
|
||||
MODEL_TENSOR.HC_FFN_BASE,
|
||||
MODEL_TENSOR.HC_FFN_SCALE,
|
||||
MODEL_TENSOR.ATTN_COMPRESSOR_WKV,
|
||||
MODEL_TENSOR.ATTN_COMPRESSOR_WGATE,
|
||||
MODEL_TENSOR.ATTN_COMPRESSOR_NORM,
|
||||
MODEL_TENSOR.INDEXER_ATTN_K,
|
||||
MODEL_TENSOR.INDEXER_K_NORM,
|
||||
MODEL_TENSOR.INDEXER_ATTN_Q_B,
|
||||
MODEL_TENSOR.INDEXER_PROJ,
|
||||
MODEL_TENSOR.PER_LAYER_TOKEN_EMBD,
|
||||
MODEL_TENSOR.ENGRAM_KV,
|
||||
MODEL_TENSOR.ENGRAM_GATE,
|
||||
MODEL_TENSOR.FFN_GATE_INP,
|
||||
MODEL_TENSOR.FFN_EXP_PROBS_B,
|
||||
MODEL_TENSOR.FFN_EXP_PROBS_B_VL,
|
||||
MODEL_TENSOR.FFN_NORM,
|
||||
MODEL_TENSOR.FFN_GATE_EXP,
|
||||
MODEL_TENSOR.FFN_DOWN_EXP,
|
||||
MODEL_TENSOR.FFN_UP_EXP,
|
||||
MODEL_TENSOR.FFN_GATE_SHEXP,
|
||||
MODEL_TENSOR.FFN_DOWN_SHEXP,
|
||||
MODEL_TENSOR.FFN_UP_SHEXP,
|
||||
],
|
||||
MODEL_ARCH.ERNIE4_5_MOE: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
MODEL_TENSOR.OUTPUT_NORM,
|
||||
@@ -6449,6 +6506,7 @@ class VisionProjectorType:
|
||||
DEEPSEEKOCR = "deepseekocr"
|
||||
DEEPSEEKOCR2 = "deepseekocr2"
|
||||
DEEPSEEK4V = "deepseek4v"
|
||||
DEEPSEEK41V = "deepseek41v"
|
||||
LFM2A = "lfm2a" # audio
|
||||
MUSIC_FLAMINGO = "musicflamingo" # audio
|
||||
GLM4V = "glm4v"
|
||||
|
||||
@@ -1077,6 +1077,28 @@ class GGUFWriter:
|
||||
def add_attention_compress_rope_freq_base(self, value: float) -> None:
|
||||
self.add_float32(Keys.Attention.COMPRESS_ROPE_FREQ_BASE.format(arch=self.arch), value)
|
||||
|
||||
# deepseek41
|
||||
def add_attention_kv_source_layers(self, layers: Sequence[int]) -> None:
|
||||
self.add_key_value(Keys.Attention.KV_SOURCE_LAYERS.format(arch=self.arch), list(layers), GGUFValueType.ARRAY, GGUFValueType.UINT32)
|
||||
|
||||
def add_indexer_source_layers(self, layers: Sequence[int]) -> None:
|
||||
self.add_key_value(Keys.Attention.Indexer.SOURCE_LAYERS.format(arch=self.arch), list(layers), GGUFValueType.ARRAY, GGUFValueType.UINT32)
|
||||
|
||||
def add_indexer_candidate_source_layer(self, layer: int) -> None:
|
||||
self.add_int32(Keys.Attention.Indexer.CANDIDATE_SOURCE_LAYER.format(arch=self.arch), layer)
|
||||
|
||||
def add_indexer_candidate_top_k_blocks(self, value: int) -> None:
|
||||
self.add_uint32(Keys.Attention.Indexer.CANDIDATE_TOP_K_BLOCKS.format(arch=self.arch), value)
|
||||
|
||||
def add_indexer_candidate_block_size(self, value: int) -> None:
|
||||
self.add_uint32(Keys.Attention.Indexer.CANDIDATE_BLOCK_SIZE.format(arch=self.arch), value)
|
||||
|
||||
def add_engram_pad_token_id(self, value: int) -> None:
|
||||
self.add_uint32(Keys.Engram.PAD_TOKEN_ID.format(arch=self.arch), value)
|
||||
|
||||
def add_engram_token_map(self, values: Sequence[int]) -> None:
|
||||
self.add_key_value(Keys.Engram.TOKEN_MAP.format(arch=self.arch), list(values), GGUFValueType.ARRAY, GGUFValueType.INT32)
|
||||
|
||||
def add_hyper_connection_count(self, count: int) -> None:
|
||||
self.add_uint32(Keys.HyperConnection.COUNT.format(arch=self.arch), count)
|
||||
|
||||
@@ -1092,6 +1114,9 @@ class GGUFWriter:
|
||||
def add_hyper_connection_low_rank(self, value: int) -> None:
|
||||
self.add_uint32(Keys.HyperConnection.LOW_RANK.format(arch=self.arch), value)
|
||||
|
||||
def add_hyper_connection_single_pass(self, value: bool) -> None:
|
||||
self.add_bool(Keys.HyperConnection.SINGLE_PASS.format(arch=self.arch), value)
|
||||
|
||||
def add_ple_layers(self, values: Sequence[int]) -> None:
|
||||
self.add_array(Keys.PerLayerEmbedding.LAYERS.format(arch=self.arch), values)
|
||||
|
||||
|
||||
@@ -87,6 +87,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
|
||||
{ LLM_ARCH_DEEPSEEK2OCR, "deepseek2-ocr" },
|
||||
{ LLM_ARCH_DEEPSEEK32, "deepseek32" },
|
||||
{ LLM_ARCH_DEEPSEEK4, "deepseek4" },
|
||||
{ LLM_ARCH_DEEPSEEK41, "deepseek41" },
|
||||
{ LLM_ARCH_CHATGLM, "chatglm" },
|
||||
{ LLM_ARCH_GLM4, "glm4" },
|
||||
{ LLM_ARCH_GLM4_MOE, "glm4moe" },
|
||||
@@ -311,6 +312,11 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
|
||||
{ LLM_KV_ATTENTION_INDEXER_TYPES, "%s.attention.indexer.types" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_KPOOL, "%s.attention.indexer.kpool" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, "%s.attention.indexer.kpool_select_tail" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_SOURCE_LAYERS, "%s.attention.indexer.source_layers" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_CANDIDATE_SOURCE_LAYER, "%s.attention.indexer.candidate_source_layer" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS, "%s.attention.indexer.candidate_top_k_blocks" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE, "%s.attention.indexer.candidate_block_size" },
|
||||
{ LLM_KV_ATTENTION_KV_SOURCE_LAYERS, "%s.attention.kv_source_layers" },
|
||||
{ LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, "%s.attention.output_group_count" },
|
||||
{ LLM_KV_ATTENTION_OUTPUT_LORA_RANK, "%s.attention.output_lora_rank" },
|
||||
{ LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, "%s.attention.compress_rope_freq_base" },
|
||||
@@ -323,6 +329,10 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
|
||||
{ LLM_KV_HYPER_CONNECTION_EPSILON, "%s.hyper_connection.epsilon" },
|
||||
{ LLM_KV_HYPER_CONNECTION_MAGNITUDE, "%s.hyper_connection.magnitude" },
|
||||
{ LLM_KV_HYPER_CONNECTION_LOW_RANK, "%s.hyper_connection.low_rank" },
|
||||
{ LLM_KV_HYPER_CONNECTION_SINGLE_PASS, "%s.hyper_connection.single_pass" },
|
||||
|
||||
{ LLM_KV_ENGRAM_PAD_TOKEN_ID, "%s.engram.pad_token_id" },
|
||||
{ LLM_KV_ENGRAM_TOKEN_MAP, "%s.engram.token_map" },
|
||||
|
||||
{ LLM_KV_PLE_LAYERS, "%s.ple.layers" },
|
||||
{ LLM_KV_PLE_NGRAM_SIZE, "%s.ple.ngram_size" },
|
||||
@@ -782,6 +792,8 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
|
||||
{ LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" },
|
||||
{ LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" },
|
||||
{ LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" },
|
||||
{ LLM_TENSOR_ENGRAM_KV, "blk.%d.engram_kv" },
|
||||
{ LLM_TENSOR_ENGRAM_GATE, "blk.%d.engram_gate" },
|
||||
{ LLM_TENSOR_DFLASH_ATTN_CONV_BASE, "blk.%d.attn_conv_base" },
|
||||
{ LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, "blk.%d.attn_conv_proj" },
|
||||
{ LLM_TENSOR_DFLASH_FFN_CONV_BASE, "blk.%d.ffn_conv_base" },
|
||||
@@ -1120,6 +1132,8 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
|
||||
{LLM_TENSOR_DSPARK_MARKOV_W1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}},
|
||||
{LLM_TENSOR_DSPARK_MARKOV_W2, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_DSPARK_CONF_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_ENGRAM_KV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_ENGRAM_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
|
||||
{LLM_TENSOR_DFLASH_ATTN_CONV_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
|
||||
{LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_DFLASH_FFN_CONV_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
|
||||
@@ -1239,6 +1253,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) {
|
||||
case LLM_ARCH_QWEN35MOE:
|
||||
case LLM_ARCH_QWEN4EXP:
|
||||
case LLM_ARCH_DEEPSEEK4:
|
||||
case LLM_ARCH_DEEPSEEK41:
|
||||
case LLM_ARCH_MINIMAX_01:
|
||||
return true;
|
||||
default:
|
||||
@@ -1265,6 +1280,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
|
||||
case LLM_ARCH_QWEN35MOE:
|
||||
case LLM_ARCH_QWEN4EXP:
|
||||
case LLM_ARCH_DEEPSEEK4:
|
||||
case LLM_ARCH_DEEPSEEK41:
|
||||
case LLM_ARCH_GLM5_NEXT:
|
||||
case LLM_ARCH_NEMOTRON_H:
|
||||
case LLM_ARCH_NEMOTRON_H_MOE:
|
||||
@@ -1282,6 +1298,7 @@ bool llm_arch_supports_mixed_batch(const llm_arch & arch) {
|
||||
switch (arch) {
|
||||
case LLM_ARCH_COGVLM:
|
||||
case LLM_ARCH_DEEPSEEK4:
|
||||
case LLM_ARCH_DEEPSEEK41:
|
||||
case LLM_ARCH_GRANITE_SWITCH:
|
||||
case LLM_ARCH_EAGLE3:
|
||||
case LLM_ARCH_DFLASH:
|
||||
@@ -1319,6 +1336,8 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
|
||||
case LLM_ARCH_MINIMAX_M2:
|
||||
case LLM_ARCH_MINIMAX_M3:
|
||||
case LLM_ARCH_MISTRAL4:
|
||||
case LLM_ARCH_MELLUM: // [TAG_MOE_AMAX_SCALE]
|
||||
case LLM_ARCH_FLINT: // [TAG_MOE_AMAX_SCALE]
|
||||
case LLM_ARCH_KIMI_LINEAR:
|
||||
case LLM_ARCH_BAILINGMOE3:
|
||||
case LLM_ARCH_KIMI_K3:
|
||||
|
||||
@@ -92,6 +92,7 @@ enum llm_arch {
|
||||
LLM_ARCH_DEEPSEEK2OCR,
|
||||
LLM_ARCH_DEEPSEEK32,
|
||||
LLM_ARCH_DEEPSEEK4,
|
||||
LLM_ARCH_DEEPSEEK41,
|
||||
LLM_ARCH_CHATGLM,
|
||||
LLM_ARCH_GLM4,
|
||||
LLM_ARCH_GLM4_MOE,
|
||||
@@ -316,6 +317,11 @@ enum llm_kv {
|
||||
LLM_KV_ATTENTION_INDEXER_TYPES,
|
||||
LLM_KV_ATTENTION_INDEXER_KPOOL,
|
||||
LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL,
|
||||
LLM_KV_ATTENTION_INDEXER_SOURCE_LAYERS,
|
||||
LLM_KV_ATTENTION_INDEXER_CANDIDATE_SOURCE_LAYER,
|
||||
LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS,
|
||||
LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE,
|
||||
LLM_KV_ATTENTION_KV_SOURCE_LAYERS,
|
||||
LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,
|
||||
LLM_KV_ATTENTION_OUTPUT_LORA_RANK,
|
||||
LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE,
|
||||
@@ -328,6 +334,10 @@ enum llm_kv {
|
||||
LLM_KV_HYPER_CONNECTION_EPSILON,
|
||||
LLM_KV_HYPER_CONNECTION_MAGNITUDE,
|
||||
LLM_KV_HYPER_CONNECTION_LOW_RANK,
|
||||
LLM_KV_HYPER_CONNECTION_SINGLE_PASS,
|
||||
|
||||
LLM_KV_ENGRAM_PAD_TOKEN_ID,
|
||||
LLM_KV_ENGRAM_TOKEN_MAP,
|
||||
|
||||
LLM_KV_PLE_LAYERS,
|
||||
LLM_KV_PLE_NGRAM_SIZE,
|
||||
@@ -788,6 +798,8 @@ enum llm_tensor {
|
||||
LLM_TENSOR_DSPARK_MARKOV_W1,
|
||||
LLM_TENSOR_DSPARK_MARKOV_W2,
|
||||
LLM_TENSOR_DSPARK_CONF_PROJ,
|
||||
LLM_TENSOR_ENGRAM_KV,
|
||||
LLM_TENSOR_ENGRAM_GATE,
|
||||
LLM_TENSOR_DFLASH_ATTN_CONV_BASE,
|
||||
LLM_TENSOR_DFLASH_ATTN_CONV_PROJ,
|
||||
LLM_TENSOR_DFLASH_FFN_CONV_BASE,
|
||||
|
||||
@@ -1339,8 +1339,9 @@ void llama_context::set_abort_callback(bool (*abort_callback)(void * data), void
|
||||
void llama_context::set_embeddings(bool value) {
|
||||
LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value);
|
||||
|
||||
// a pooled context computes the logits only without embeddings, so the toggle changes the graph shape
|
||||
if (cparams.embeddings != value && cparams.pooling_type != LLAMA_POOLING_TYPE_NONE) {
|
||||
// the toggle changes the graph shape (pooled: logits only; pooling NONE: the reserve graph outputs all tokens,
|
||||
// which crashed ggml_gallocr on the first /completion), so re-reserve on every toggle
|
||||
if (cparams.embeddings != value) {
|
||||
sched_need_reserve = true;
|
||||
}
|
||||
|
||||
@@ -2536,6 +2537,7 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
|
||||
model.arch == LLM_ARCH_CLEF ||
|
||||
model.arch == LLM_ARCH_QWEN4EXP ||
|
||||
model.arch == LLM_ARCH_DEEPSEEK4 ||
|
||||
model.arch == LLM_ARCH_DEEPSEEK41 ||
|
||||
(model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) ||
|
||||
model.arch == LLM_ARCH_NANBEIGE ||
|
||||
model.arch == LLM_ARCH_MINIMAX_01 ||
|
||||
@@ -4071,7 +4073,7 @@ llama_context * llama_init_from_model(
|
||||
}
|
||||
}
|
||||
|
||||
if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4) && params.type_k != params.type_v) {
|
||||
if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4 || model->arch == LLM_ARCH_DEEPSEEK41) && params.type_k != params.type_v) {
|
||||
LLAMA_LOG_ERROR("%s: model does not support different K (%s) and V (%s) cache types\n", __func__, ggml_type_name(params.type_k), ggml_type_name(params.type_v));
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
+10
-4
@@ -2010,7 +2010,7 @@ ggml_tensor * llm_graph_context::build_ffn(
|
||||
const float limit = hparams.swiglu_clamp_shexp[il];
|
||||
constexpr float eps = 1e-6f;
|
||||
if (limit > eps) {
|
||||
if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) {
|
||||
if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) {
|
||||
cur = ggml_swiglu_clamp(ctx0, cur, tmp, limit);
|
||||
} else {
|
||||
tmp = ggml_clamp(ctx0, tmp, -limit, limit);
|
||||
@@ -2340,8 +2340,13 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
|
||||
ggml_tensor * weights_sum = ggml_sum_rows(ctx0, weights); // [1, n_tokens]
|
||||
cb(weights_sum, "ffn_moe_weights_sum", il);
|
||||
|
||||
// Avoid division by zero, clamp to smallest number representable by F16
|
||||
weights_sum = ggml_clamp(ctx0, weights_sum, 6.103515625e-5, INFINITY);
|
||||
if (arch == LLM_ARCH_DEEPSEEK41) {
|
||||
// model.py Gate.forward: weights / (sum + 1e-20); the sum can fall below the F16 clamp below
|
||||
weights_sum = ggml_scale_bias(ctx0, weights_sum, 1.0f, 1e-20f);
|
||||
} else {
|
||||
// Avoid division by zero, clamp to smallest number representable by F16
|
||||
weights_sum = ggml_clamp(ctx0, weights_sum, 6.103515625e-5, INFINITY);
|
||||
}
|
||||
cb(weights_sum, "ffn_moe_weights_sum_clamped", il);
|
||||
|
||||
weights = ggml_div(ctx0, weights, weights_sum); // [n_expert_used, n_tokens]
|
||||
@@ -2431,7 +2436,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
|
||||
const float limit = hparams.swiglu_clamp_exp[il];
|
||||
constexpr float eps = 1e-6f;
|
||||
if (limit > eps) {
|
||||
if (arch == LLM_ARCH_MAPLE || arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) {
|
||||
if (arch == LLM_ARCH_MAPLE || arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) {
|
||||
cur = ggml_swiglu_clamp(ctx0, cur, up, limit);
|
||||
} else {
|
||||
up = ggml_clamp(ctx0, up, -limit, limit);
|
||||
@@ -2510,6 +2515,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
|
||||
{
|
||||
// TODO: this graph-level rescaling should only be applied when the corresponding backend does not support
|
||||
// ggml_prec_set_src(experts, GGML_PREC_BF16, 1);
|
||||
// the amax op cannot be split across devices, so these architectures are denied -sm tensor [TAG_MOE_AMAX_SCALE]
|
||||
// ref: https://github.com/ggml-org/llama.cpp/pull/30656#issuecomment-6269597922
|
||||
llm_graph_amax_src1 amax_src1 = { nullptr, nullptr };
|
||||
|
||||
|
||||
+14
-2
@@ -10,8 +10,8 @@
|
||||
// bump if necessary
|
||||
#define LLAMA_MAX_LAYERS 512
|
||||
#define LLAMA_MAX_EXPERTS 1024 // Kimi K3
|
||||
#define LLAMA_MAX_PLE_NGRAM 8 // qwen4exp
|
||||
#define LLAMA_MAX_PLE_HEADS 64 // qwen4exp
|
||||
#define LLAMA_MAX_PLE_NGRAM 8 // qwen4exp, deepseek41 (all hash modules)
|
||||
#define LLAMA_MAX_PLE_HEADS 64 // qwen4exp, deepseek41 (all hash modules)
|
||||
|
||||
enum llama_expert_gating_func_type {
|
||||
LLAMA_EXPERT_GATING_FUNC_TYPE_NONE = 0,
|
||||
@@ -326,12 +326,20 @@ struct llama_hparams {
|
||||
uint32_t dsv4_o_group_count = 0;
|
||||
uint32_t dsv4_o_lora_rank = 0;
|
||||
uint32_t dsv4_hc_mult = 0;
|
||||
bool dsv4_hc_single_pass = false; // DeepSeek-V4.1 DSpark draft: lagged pre-mix, no hc_head
|
||||
uint32_t dsv4_hc_sinkhorn_iters = 0;
|
||||
uint32_t dsv4_hash_layer_count = 0;
|
||||
float dsv4_compress_rope_base = 0.0f;
|
||||
float dsv4_hc_eps = 0.0f;
|
||||
std::array<uint32_t, LLAMA_MAX_LAYERS> dsv4_compress_ratios;
|
||||
|
||||
// deepseek41
|
||||
std::array<int32_t, LLAMA_MAX_LAYERS> dsv41_kv_src; // layer whose compressed KV this layer reads, -1 = none
|
||||
std::array<int32_t, LLAMA_MAX_LAYERS> dsv41_topk_src; // layer whose top-k this layer uses, -1 = none
|
||||
int32_t dsv41_cand_src = -1; // -1 = no candidate block selection
|
||||
uint32_t dsv41_cand_topk_blocks = 0;
|
||||
uint32_t dsv41_cand_block_size = 0;
|
||||
|
||||
// 0 = full rank (DeepSeek-V4)
|
||||
uint32_t hc_low_rank = 0;
|
||||
|
||||
@@ -346,6 +354,10 @@ struct llama_hparams {
|
||||
uint32_t ple_eos_token_id = 0;
|
||||
// the id the PLE hash stands in at image positions; 0 makes the loader fall back to EOS
|
||||
uint32_t ple_image_token_id = 0;
|
||||
// the id the KV cells record at embedding (image) positions for the n-gram history; LLAMA_TOKEN_NULL = none
|
||||
llama_token ple_embd_token = LLAMA_TOKEN_NULL;
|
||||
// hash constant sets, stored back to back: 1 shared by all PLE layers (qwen4exp), or 1 per PLE layer in layer order (deepseek41)
|
||||
uint32_t ple_n_modules = 1;
|
||||
// the file lists PLE layer indices, so this is never a per-layer gguf array and can hold one bit per layer
|
||||
std::bitset<LLAMA_MAX_LAYERS> is_ple_impl;
|
||||
// the hash multipliers reach ~2e13 and have to stay 64-bit
|
||||
|
||||
+237
-42
@@ -333,12 +333,24 @@ static void dsv4_state_read_tensor_streams(
|
||||
}
|
||||
}
|
||||
|
||||
// serialized compressed-KV layers: all except per-device mirrors (kv_src[il] != il), refilled from their source on read
|
||||
static std::vector<uint32_t> dsv4_k_cache_state_layers(const llama_kv_cache * kv, const std::vector<int32_t> & kv_src) {
|
||||
std::vector<uint32_t> res;
|
||||
for (uint32_t il : kv->get_layer_ids()) {
|
||||
if (kv_src[il] == (int32_t) il) {
|
||||
res.push_back(il);
|
||||
}
|
||||
}
|
||||
return res;
|
||||
}
|
||||
|
||||
static void dsv4_state_write_k_cache(
|
||||
llama_io_write_i & io,
|
||||
const llama_kv_cache * kv,
|
||||
llama_seq_id seq_id,
|
||||
llama_state_seq_flags flags,
|
||||
uint32_t n_rows) {
|
||||
uint32_t n_rows,
|
||||
const std::vector<int32_t> & kv_src) {
|
||||
GGML_UNUSED(flags);
|
||||
|
||||
uint32_t s0;
|
||||
@@ -347,7 +359,8 @@ static void dsv4_state_write_k_cache(
|
||||
|
||||
const uint32_t version = DSV4_K_CACHE_STATE_VER;
|
||||
const uint32_t kv_size = kv->get_size();
|
||||
const auto layer_ids = kv->get_layer_ids();
|
||||
// mirrors are never serialized, so the state format does not depend on the split
|
||||
const auto layer_ids = dsv4_k_cache_state_layers(kv, kv_src);
|
||||
const uint32_t n_layer = layer_ids.size();
|
||||
|
||||
if (n_rows > kv_size) {
|
||||
@@ -369,7 +382,8 @@ static void dsv4_state_read_k_cache(
|
||||
llama_io_read_i & io,
|
||||
llama_kv_cache * kv,
|
||||
llama_seq_id seq_id,
|
||||
llama_state_seq_flags flags) {
|
||||
llama_state_seq_flags flags,
|
||||
const std::vector<int32_t> & kv_src) {
|
||||
GGML_UNUSED(flags);
|
||||
|
||||
uint32_t version;
|
||||
@@ -399,7 +413,7 @@ static void dsv4_state_read_k_cache(
|
||||
uint32_t s0;
|
||||
dsv4_state_dst_stream_range(kv->get_n_stream(), seq_id, ns, s0);
|
||||
|
||||
const auto layer_ids = kv->get_layer_ids();
|
||||
const auto layer_ids = dsv4_k_cache_state_layers(kv, kv_src);
|
||||
if (n_layer_ref != layer_ids.size()) {
|
||||
throw std::runtime_error("DSV4 K-cache layer count mismatch");
|
||||
}
|
||||
@@ -413,6 +427,31 @@ static void dsv4_state_read_k_cache(
|
||||
|
||||
dsv4_state_read_tensor_streams(io, kv->get_k_storage(il), kv_size, n_rows_ref, s0, ns);
|
||||
}
|
||||
|
||||
// refill the mirrors from their source rows just read
|
||||
std::vector<uint8_t> buf;
|
||||
for (uint32_t il : kv->get_layer_ids()) {
|
||||
const int32_t src = kv_src[il];
|
||||
if (src == (int32_t) il) {
|
||||
continue;
|
||||
}
|
||||
ggml_tensor * t_src = kv->get_k_storage(src);
|
||||
ggml_tensor * t_dst = kv->get_k_storage(il);
|
||||
GGML_ASSERT(t_src->type == t_dst->type && t_src->ne[0] == t_dst->ne[0]);
|
||||
|
||||
const size_t row_size = ggml_row_size(t_src->type, t_src->ne[0]);
|
||||
const size_t stream_stride = (size_t) kv_size*row_size;
|
||||
const size_t size = (size_t) n_rows_ref*row_size;
|
||||
if (size == 0) {
|
||||
continue;
|
||||
}
|
||||
buf.resize(size);
|
||||
for (uint32_t s = 0; s < ns; ++s) {
|
||||
const size_t offset = (size_t) (s0 + s)*stream_stride;
|
||||
ggml_backend_tensor_get(t_src, buf.data(), offset, size);
|
||||
ggml_backend_tensor_set(t_dst, buf.data(), offset, size);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static std::string dsv4_plan_positions(const std::vector<int32_t> & values) {
|
||||
@@ -432,6 +471,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
|
||||
const llama_ubatch & ubatch,
|
||||
uint32_t ratio,
|
||||
bool overlap,
|
||||
bool pad_blocks,
|
||||
uint32_t state_size,
|
||||
uint32_t kv_size,
|
||||
uint32_t n_stream,
|
||||
@@ -548,7 +588,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
|
||||
}
|
||||
}
|
||||
|
||||
if (ratio == DSV4_CSA_RATIO && !plan.state_pos.empty()) {
|
||||
if (pad_blocks && !plan.state_pos.empty()) {
|
||||
assert(kv_size > 0);
|
||||
|
||||
// Pad each stream to the reserve plan's block count.
|
||||
@@ -603,7 +643,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
|
||||
}
|
||||
}
|
||||
|
||||
if (ratio == DSV4_HCA_RATIO && !plan.state_pos.empty() && plan.state_write_idxs.empty()) {
|
||||
if (!pad_blocks && !plan.state_pos.empty() && plan.state_write_idxs.empty()) {
|
||||
assert(kv_size > 0);
|
||||
// the last slot must not be live, or the dummy write would corrupt it;
|
||||
// a full stream implies a completed block, which implies real writes
|
||||
@@ -744,6 +784,7 @@ static std::vector<llama_kv_cache_dsv4_context::comp_plan> dsv4_build_comp_plans
|
||||
const std::vector<llama_ubatch> & ubatches,
|
||||
uint32_t ratio,
|
||||
bool overlap,
|
||||
bool pad_blocks,
|
||||
uint32_t state_size,
|
||||
uint32_t kv_size,
|
||||
uint32_t n_stream,
|
||||
@@ -755,7 +796,7 @@ static std::vector<llama_kv_cache_dsv4_context::comp_plan> dsv4_build_comp_plans
|
||||
// the first ubatch touching a seq consumes its rollback restore
|
||||
std::vector<uint32_t> rs(rs_idx);
|
||||
for (const llama_ubatch & ubatch : ubatches) {
|
||||
plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream, n_rs_seq, rs));
|
||||
plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, pad_blocks, state_size, kv_size, n_stream, n_rs_seq, rs));
|
||||
|
||||
for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
|
||||
const llama_seq_id seq_id = ubatch.seq_id_unq[s];
|
||||
@@ -1056,6 +1097,10 @@ uint32_t llama_dsv4_comp_state::get_ratio() const {
|
||||
return ratio;
|
||||
}
|
||||
|
||||
bool llama_dsv4_comp_state::is_overlap() const {
|
||||
return state_size == 2*ratio;
|
||||
}
|
||||
|
||||
uint32_t llama_dsv4_comp_state::get_state_size() const {
|
||||
return state_size;
|
||||
}
|
||||
@@ -1261,6 +1306,23 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
|
||||
v_trans, offload, swa_full, unified_raw, kv_size, n_seq_max, n_ubatch, n_pad,
|
||||
nullptr, filter_raw, reuse, nullptr);
|
||||
|
||||
// V4.1: CSA slot = ratio 2, HCA slot = ratio 1, no overlap and no indexer compressor; the index K is stored
|
||||
// next to the compressed KV ([kv | index K]) and only the KV source layers own rows
|
||||
const bool is_v41 = model.arch == LLM_ARCH_DEEPSEEK41;
|
||||
const uint32_t ratio_csa = is_v41 ? 2u : DSV4_CSA_RATIO;
|
||||
const uint32_t ratio_hca = is_v41 ? 1u : DSV4_HCA_RATIO;
|
||||
|
||||
if (is_v41) {
|
||||
const uint32_t n_embd_comp = model.hparams.n_embd_head_k() + model.hparams.indexer_head_size;
|
||||
for (llama_hparams * hp : { &hparams_csa, &hparams_hca }) {
|
||||
std::fill(hp->n_head_kv_arr.begin(), hp->n_head_kv_arr.end(), 1);
|
||||
hp->n_embd_head_k_full = n_embd_comp;
|
||||
hp->n_embd_head_v_full = n_embd_comp;
|
||||
hp->n_embd_head_k_swa = n_embd_comp;
|
||||
hp->n_embd_head_v_swa = n_embd_comp;
|
||||
}
|
||||
}
|
||||
|
||||
dsv4_make_k_only(hparams_csa);
|
||||
dsv4_make_k_only(hparams_hca);
|
||||
|
||||
@@ -1272,12 +1334,50 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
|
||||
hparams_lid.rope_type = LLAMA_ROPE_TYPE_NEOX;
|
||||
dsv4_make_k_only(hparams_lid);
|
||||
|
||||
// V4.1 cross-layer KV sharing: a compressed layer reads the storage of its kv source or, when the source is on
|
||||
// another device (layer split), of a per-device mirror fed with the source's new rows, so the whole cache never
|
||||
// crosses devices. no mirrors without offload
|
||||
const uint32_t n_layer = model.hparams.n_layer();
|
||||
comp_read.assign(n_layer, -1);
|
||||
comp_kv_src.resize(n_layer);
|
||||
for (uint32_t il = 0; il < n_layer; ++il) {
|
||||
const int32_t src = is_v41 ? model.hparams.dsv41_kv_src[il] : (int32_t) il;
|
||||
comp_kv_src[il] = src >= 0 ? src : (int32_t) il;
|
||||
if (model.hparams.dsv4_compress_ratios[il] == 0) {
|
||||
continue;
|
||||
}
|
||||
comp_read[il] = src;
|
||||
if (src == (int32_t) il || !offload || model.dev_layer(src) == model.dev_layer(il)) {
|
||||
continue;
|
||||
}
|
||||
// the first layer of this source on this device holds the mirror
|
||||
comp_read[il] = il;
|
||||
for (int32_t j = src + 1; j < (int32_t) il; ++j) {
|
||||
if (comp_read[j] == j && comp_kv_src[j] == src && model.dev_layer(j) == model.dev_layer(il)) {
|
||||
comp_read[il] = j;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// compressed-KV storage: kv sources and mirrors; compressor state: kv sources only
|
||||
const auto is_comp_owner = [&](int32_t il, uint32_t ratio) {
|
||||
return model.hparams.dsv4_compress_ratios[il] == ratio && comp_read[il] == (int32_t) il;
|
||||
};
|
||||
|
||||
const auto is_state_owner = [&](int32_t il, uint32_t ratio) {
|
||||
if (model.hparams.dsv4_compress_ratios[il] != ratio) {
|
||||
return false;
|
||||
}
|
||||
return !is_v41 || model.hparams.dsv41_kv_src[il] == il;
|
||||
};
|
||||
|
||||
const layer_filter_cb filter_csa = [&](int32_t il) {
|
||||
if (filter && !filter(il)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
return model.hparams.dsv4_compress_ratios[il] == DSV4_CSA_RATIO;
|
||||
return is_comp_owner(il, ratio_csa);
|
||||
};
|
||||
|
||||
const layer_filter_cb filter_hca = [&](int32_t il) {
|
||||
@@ -1285,52 +1385,67 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
|
||||
return false;
|
||||
}
|
||||
|
||||
return model.hparams.dsv4_compress_ratios[il] == DSV4_HCA_RATIO;
|
||||
return is_comp_owner(il, ratio_hca);
|
||||
};
|
||||
|
||||
const layer_filter_cb filter_csa_state = [&](int32_t il) {
|
||||
return (!filter || filter(il)) && is_state_owner(il, ratio_csa);
|
||||
};
|
||||
|
||||
const layer_filter_cb filter_hca_state = [&](int32_t il) {
|
||||
return (!filter || filter(il)) && is_state_owner(il, ratio_hca);
|
||||
};
|
||||
|
||||
const layer_filter_cb filter_lid = [&](int32_t il) {
|
||||
return !is_v41 && filter_csa(il);
|
||||
};
|
||||
|
||||
const bool unified_compressed = false;
|
||||
|
||||
LLAMA_LOG_INFO("%s: creating DSV4 CSA compressed KV cache, size = %u cells\n",
|
||||
__func__, dsv4_comp_size(kv_size, DSV4_CSA_RATIO));
|
||||
__func__, dsv4_comp_size(kv_size, ratio_csa));
|
||||
|
||||
kv_csa = std::make_unique<llama_kv_cache>(
|
||||
model, hparams_csa, type_k, type_v,
|
||||
v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_CSA_RATIO), 256u), n_seq_max, n_pad,
|
||||
v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, ratio_csa), 256u), n_seq_max, n_pad,
|
||||
0, LLAMA_SWA_TYPE_NONE, nullptr, filter_csa, nullptr, nullptr);
|
||||
|
||||
LLAMA_LOG_INFO("%s: creating DSV4 HCA compressed KV cache, size = %u cells\n",
|
||||
__func__, dsv4_comp_size(kv_size, DSV4_HCA_RATIO));
|
||||
__func__, dsv4_comp_size(kv_size, ratio_hca));
|
||||
|
||||
kv_hca = std::make_unique<llama_kv_cache>(
|
||||
model, hparams_hca, type_k, type_v,
|
||||
v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_HCA_RATIO), 256u), n_seq_max, n_pad,
|
||||
v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, ratio_hca), 256u), n_seq_max, n_pad,
|
||||
0, LLAMA_SWA_TYPE_NONE, nullptr, filter_hca, nullptr, nullptr);
|
||||
|
||||
LLAMA_LOG_INFO("%s: creating DSV4 lightning-indexer KV cache, size = %u cells\n",
|
||||
__func__, dsv4_comp_size(kv_size, DSV4_CSA_RATIO));
|
||||
__func__, dsv4_comp_size(kv_size, ratio_csa));
|
||||
|
||||
kv_lid = std::make_unique<llama_kv_cache>(
|
||||
model, hparams_lid, type_k, type_v,
|
||||
v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_CSA_RATIO), 256u), n_seq_max, n_pad,
|
||||
0, LLAMA_SWA_TYPE_NONE, nullptr, filter_csa, nullptr, nullptr);
|
||||
v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, ratio_csa), 256u), n_seq_max, n_pad,
|
||||
0, LLAMA_SWA_TYPE_NONE, nullptr, filter_lid, nullptr, nullptr);
|
||||
|
||||
LLAMA_LOG_INFO("%s: creating DSV4 CSA compressor state\n", __func__);
|
||||
|
||||
const uint32_t csa_state_size = is_v41 ? ratio_csa : 2*ratio_csa;
|
||||
const uint32_t csa_state_embd = is_v41 ? model.hparams.n_embd_head_k() : 2*model.hparams.n_embd_head_k();
|
||||
|
||||
csa_state = std::make_unique<llama_dsv4_comp_state>(
|
||||
model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO,
|
||||
2*model.hparams.n_embd_head_k(), n_rs_seq, "csa", filter_csa);
|
||||
model, offload, unified_compressed, n_seq_max, ratio_csa, csa_state_size,
|
||||
csa_state_embd, n_rs_seq, "csa", filter_csa_state);
|
||||
|
||||
LLAMA_LOG_INFO("%s: creating DSV4 HCA compressor state\n", __func__);
|
||||
|
||||
hca_state = std::make_unique<llama_dsv4_comp_state>(
|
||||
model, offload, unified_compressed, n_seq_max, DSV4_HCA_RATIO, DSV4_HCA_RATIO,
|
||||
model.hparams.n_embd_head_k(), n_rs_seq, "hca", filter_hca);
|
||||
model, offload, unified_compressed, n_seq_max, ratio_hca, ratio_hca,
|
||||
model.hparams.n_embd_head_k(), n_rs_seq, "hca", filter_hca_state);
|
||||
|
||||
LLAMA_LOG_INFO("%s: creating DSV4 lightning-indexer compressor state\n", __func__);
|
||||
|
||||
lid_state = std::make_unique<llama_dsv4_comp_state>(
|
||||
model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO,
|
||||
2*model.hparams.indexer_head_size, n_rs_seq, "lid", filter_csa);
|
||||
2*model.hparams.indexer_head_size, n_rs_seq, "lid", filter_lid);
|
||||
|
||||
// DSV4 attention reads compressed-K / compressor-state rows that the current
|
||||
// graph does not necessarily overwrite; uninitialized buffer contents would
|
||||
@@ -1339,12 +1454,76 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
|
||||
clear_compressed(-1, true);
|
||||
}
|
||||
|
||||
// Engram hashes the n-gram ending at each token, so a token shared by several sequences (llama-batched prompt)
|
||||
// is well-defined only if all of them agree on its predecessors; returns false (and logs) otherwise
|
||||
static bool dsv41_shared_history_ok(const llama_batch & batch, const llama_kv_cache * kv, uint32_t n_hist) {
|
||||
if (n_hist == 0 || !batch.token || !batch.n_seq_id) {
|
||||
return true;
|
||||
}
|
||||
|
||||
bool any_shared = false;
|
||||
for (int32_t i = 0; i < batch.n_tokens && !any_shared; ++i) {
|
||||
any_shared = batch.n_seq_id[i] > 1;
|
||||
}
|
||||
if (!any_shared) {
|
||||
return true;
|
||||
}
|
||||
|
||||
std::map<std::pair<llama_seq_id, llama_pos>, llama_token> in_batch;
|
||||
for (int32_t i = 0; i < batch.n_tokens; ++i) {
|
||||
for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) {
|
||||
in_batch[{ batch.seq_id[i][s], batch.pos[i] }] = batch.token[i];
|
||||
}
|
||||
}
|
||||
|
||||
const auto tok_at = [&](llama_seq_id seq_id, llama_pos p) -> llama_token {
|
||||
// batch first: the batch continues the cache
|
||||
auto it = in_batch.upper_bound({ seq_id, p });
|
||||
if (it != in_batch.begin()) {
|
||||
--it;
|
||||
if (it->first.first == seq_id) {
|
||||
return it->second;
|
||||
}
|
||||
}
|
||||
return kv ? kv->get_cells(seq_id).seq_pos_tok_le(seq_id, p) : LLAMA_TOKEN_NULL;
|
||||
};
|
||||
|
||||
for (int32_t i = 0; i < batch.n_tokens; ++i) {
|
||||
if (batch.n_seq_id[i] <= 1) {
|
||||
continue;
|
||||
}
|
||||
for (uint32_t d = 1; d <= n_hist; ++d) {
|
||||
const llama_pos p = batch.pos[i] - (llama_pos) d;
|
||||
if (p < 0) {
|
||||
break;
|
||||
}
|
||||
const llama_token t0 = tok_at(batch.seq_id[i][0], p);
|
||||
for (int32_t s = 1; s < batch.n_seq_id[i]; ++s) {
|
||||
const llama_token ts = tok_at(batch.seq_id[i][s], p);
|
||||
if (ts != t0) {
|
||||
LLAMA_LOG_ERROR("%s: token %d at pos %d is shared by seq %d and seq %d, whose Engram n-gram histories "
|
||||
"differ at pos %d (%d vs %d); decode such tokens per sequence\n", __func__,
|
||||
i, batch.pos[i], batch.seq_id[i][0], batch.seq_id[i][s], p, t0, ts);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
llama_memory_context_ptr llama_kv_cache_dsv4::init_batch(
|
||||
llama_batch_allocr & balloc,
|
||||
uint32_t n_ubatch,
|
||||
bool embd_all) {
|
||||
GGML_UNUSED(embd_all);
|
||||
|
||||
if (hparams_raw.ple_n_heads > 0 &&
|
||||
!dsv41_shared_history_ok(balloc.get_batch(), kv_raw->get_swa(), hparams_raw.ple_ngram_size - 1)) {
|
||||
return std::make_unique<llama_kv_cache_dsv4_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
|
||||
}
|
||||
|
||||
const bool raw_per_seq = kv_raw->get_base()->get_n_stream() != 1;
|
||||
const bool comp_per_seq = csa_state->get_n_stream() > 1;
|
||||
const bool has_coupled = dsv4_batch_has_coupled(balloc.get_batch());
|
||||
@@ -1475,9 +1654,9 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
|
||||
bool res = true;
|
||||
|
||||
res = res & kv_raw->seq_rm(seq_id, p0, -1);
|
||||
res = res & kv_csa->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
|
||||
res = res & kv_hca->seq_rm(seq_id, p0/DSV4_HCA_RATIO, -1);
|
||||
res = res & kv_lid->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
|
||||
res = res & kv_csa->seq_rm(seq_id, p0/csa_state->get_ratio(), -1);
|
||||
res = res & kv_hca->seq_rm(seq_id, p0/hca_state->get_ratio(), -1);
|
||||
res = res & kv_lid->seq_rm(seq_id, p0/lid_state->get_ratio(), -1);
|
||||
|
||||
return res;
|
||||
}
|
||||
@@ -1613,15 +1792,15 @@ void llama_kv_cache_dsv4::state_write(llama_io_write_i & io, llama_seq_id seq_id
|
||||
|
||||
//FIXME : note that we conflate token positions with rows, which is not true for multi-modal case.
|
||||
const uint32_t n_rows_csa = seq_id >= 0 ?
|
||||
dsv4_state_n_used_k_rows(pos_max, DSV4_CSA_RATIO, kv_csa->get_size()) : kv_csa->get_size();
|
||||
dsv4_state_n_used_k_rows(pos_max, csa_state->get_ratio(), kv_csa->get_size()) : kv_csa->get_size();
|
||||
const uint32_t n_rows_hca = seq_id >= 0 ?
|
||||
dsv4_state_n_used_k_rows(pos_max, DSV4_HCA_RATIO, kv_hca->get_size()) : kv_hca->get_size();
|
||||
dsv4_state_n_used_k_rows(pos_max, hca_state->get_ratio(), kv_hca->get_size()) : kv_hca->get_size();
|
||||
const uint32_t n_rows_lid = seq_id >= 0 ?
|
||||
dsv4_state_n_used_k_rows(pos_max, DSV4_CSA_RATIO, kv_lid->get_size()) : kv_lid->get_size();
|
||||
dsv4_state_n_used_k_rows(pos_max, lid_state->get_ratio(), kv_lid->get_size()) : kv_lid->get_size();
|
||||
|
||||
dsv4_state_write_k_cache(io, kv_csa.get(), seq_id, flags, n_rows_csa);
|
||||
dsv4_state_write_k_cache(io, kv_hca.get(), seq_id, flags, n_rows_hca);
|
||||
dsv4_state_write_k_cache(io, kv_lid.get(), seq_id, flags, n_rows_lid);
|
||||
dsv4_state_write_k_cache(io, kv_csa.get(), seq_id, flags, n_rows_csa, comp_kv_src);
|
||||
dsv4_state_write_k_cache(io, kv_hca.get(), seq_id, flags, n_rows_hca, comp_kv_src);
|
||||
dsv4_state_write_k_cache(io, kv_lid.get(), seq_id, flags, n_rows_lid, comp_kv_src);
|
||||
}
|
||||
|
||||
csa_state->state_write(io, seq_id, flags, rs_idx);
|
||||
@@ -1659,9 +1838,9 @@ void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id,
|
||||
if (!partial_only) {
|
||||
clear_compressed(seq_id, true);
|
||||
|
||||
dsv4_state_read_k_cache(io, kv_csa.get(), seq_id, flags);
|
||||
dsv4_state_read_k_cache(io, kv_hca.get(), seq_id, flags);
|
||||
dsv4_state_read_k_cache(io, kv_lid.get(), seq_id, flags);
|
||||
dsv4_state_read_k_cache(io, kv_csa.get(), seq_id, flags, comp_kv_src);
|
||||
dsv4_state_read_k_cache(io, kv_hca.get(), seq_id, flags, comp_kv_src);
|
||||
dsv4_state_read_k_cache(io, kv_lid.get(), seq_id, flags, comp_kv_src);
|
||||
}
|
||||
|
||||
csa_state->state_read(io, seq_id, flags);
|
||||
@@ -1704,6 +1883,10 @@ llama_dsv4_comp_state * llama_kv_cache_dsv4::get_lid_state() const {
|
||||
return lid_state.get();
|
||||
}
|
||||
|
||||
int32_t llama_kv_cache_dsv4::get_comp_read_layer(int32_t il) const {
|
||||
return comp_read[il];
|
||||
}
|
||||
|
||||
uint32_t llama_kv_cache_dsv4::get_n_rs_seq() const {
|
||||
return n_rs_seq;
|
||||
}
|
||||
@@ -1932,6 +2115,11 @@ void llama_kv_cache_dsv4_raw_context::set_input_kq_mask(ggml_tensor * dst, const
|
||||
kv_swa->set_input_kq_mask(dst, ubatch, causal_attn);
|
||||
}
|
||||
|
||||
void llama_kv_cache_dsv4_raw_context::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const {
|
||||
GGML_ASSERT(kv_swa != nullptr);
|
||||
kv_swa->get_prev_tokens(ubatch, n, res);
|
||||
}
|
||||
|
||||
void llama_kv_cache_dsv4_raw_context::set_input_k_rot(ggml_tensor * dst) const {
|
||||
kv_swa->set_input_k_rot(dst);
|
||||
}
|
||||
@@ -2016,7 +2204,8 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
|
||||
reserve_plans(true),
|
||||
status(llama_memory_status_combine(
|
||||
llama_memory_status_combine(ctx_raw->get_status(), ctx_csa_mem->get_status()),
|
||||
llama_memory_status_combine(ctx_hca_mem->get_status(), ctx_lid_mem->get_status()))) {
|
||||
llama_memory_status_combine(ctx_hca_mem->get_status(), ctx_lid_mem->get_status()))),
|
||||
kv(kv) {
|
||||
}
|
||||
|
||||
llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
|
||||
@@ -2041,7 +2230,8 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
|
||||
llama_memory_status_combine(ctx_raw->get_status(), ctx_csa_mem->get_status()),
|
||||
llama_memory_status_combine(ctx_hca_mem->get_status(), ctx_lid_mem->get_status())),
|
||||
this->sc_info_csa.empty() && this->sc_info_hca.empty() && this->sc_info_lid.empty() ?
|
||||
LLAMA_MEMORY_STATUS_NO_UPDATE : LLAMA_MEMORY_STATUS_SUCCESS)) {
|
||||
LLAMA_MEMORY_STATUS_NO_UPDATE : LLAMA_MEMORY_STATUS_SUCCESS)),
|
||||
kv(kv) {
|
||||
}
|
||||
|
||||
llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
|
||||
@@ -2052,13 +2242,13 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
|
||||
std::vector<llama_ubatch> ubatches,
|
||||
std::vector<llama_ubatch> ubatches_raw) :
|
||||
ubatches(std::move(ubatches)),
|
||||
plans_csa(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true,
|
||||
plans_csa(dsv4_build_comp_plans(this->ubatches, kv->get_csa_state()->get_ratio(), kv->get_csa_state()->is_overlap(), true,
|
||||
kv->get_csa_state()->get_state_size(), kv->get_csa()->get_size(), kv->get_csa_state()->get_n_stream(),
|
||||
kv->get_n_rs_seq(), kv->get_rs_idx())),
|
||||
plans_hca(dsv4_build_comp_plans(this->ubatches, DSV4_HCA_RATIO, false,
|
||||
plans_hca(dsv4_build_comp_plans(this->ubatches, kv->get_hca_state()->get_ratio(), kv->get_hca_state()->is_overlap(), false,
|
||||
kv->get_hca_state()->get_state_size(), kv->get_hca()->get_size(), kv->get_hca_state()->get_n_stream(),
|
||||
kv->get_n_rs_seq(), kv->get_rs_idx())),
|
||||
plans_lid(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true,
|
||||
plans_lid(dsv4_build_comp_plans(this->ubatches, kv->get_lid_state()->get_ratio(), kv->get_lid_state()->is_overlap(), true,
|
||||
kv->get_lid_state()->get_state_size(), kv->get_lid()->get_size(), kv->get_lid_state()->get_n_stream(),
|
||||
kv->get_n_rs_seq(), kv->get_rs_idx())),
|
||||
ctx_raw(std::make_unique<llama_kv_cache_dsv4_raw_context>(
|
||||
@@ -2086,7 +2276,8 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
|
||||
csa_state(kv->get_csa_state()),
|
||||
hca_state(kv->get_hca_state()),
|
||||
lid_state(kv->get_lid_state()),
|
||||
status(ctx_raw->get_status()) {
|
||||
status(ctx_raw->get_status()),
|
||||
kv(kv) {
|
||||
kv->reset_rs_idx_for_ubatches(this->ubatches);
|
||||
}
|
||||
|
||||
@@ -2181,6 +2372,10 @@ const llama_dsv4_comp_state * llama_kv_cache_dsv4_context::get_lid_state() const
|
||||
return lid_state;
|
||||
}
|
||||
|
||||
int32_t llama_kv_cache_dsv4_context::get_comp_read_layer(int32_t il) const {
|
||||
return kv->get_comp_read_layer(il);
|
||||
}
|
||||
|
||||
const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_csa_plan() const {
|
||||
assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
|
||||
|
||||
@@ -2222,7 +2417,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_
|
||||
}
|
||||
|
||||
reserve_plan_csa = dsv4_build_reserve_comp_plan(
|
||||
ubatch, DSV4_CSA_RATIO, true,
|
||||
ubatch, csa_state->get_ratio(), csa_state->is_overlap(),
|
||||
csa_state->get_state_size(), get_csa()->get_n_kv(), csa_state->get_n_stream(), csa_state->get_n_rs_seq());
|
||||
|
||||
return reserve_plan_csa;
|
||||
@@ -2236,7 +2431,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_
|
||||
}
|
||||
|
||||
reserve_plan_hca = dsv4_build_reserve_comp_plan(
|
||||
ubatch, DSV4_HCA_RATIO, false,
|
||||
ubatch, hca_state->get_ratio(), hca_state->is_overlap(),
|
||||
hca_state->get_state_size(), get_hca()->get_n_kv(), hca_state->get_n_stream(), hca_state->get_n_rs_seq());
|
||||
|
||||
return reserve_plan_hca;
|
||||
@@ -2250,7 +2445,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_
|
||||
}
|
||||
|
||||
reserve_plan_lid = dsv4_build_reserve_comp_plan(
|
||||
ubatch, DSV4_CSA_RATIO, true,
|
||||
ubatch, lid_state->get_ratio(), lid_state->is_overlap(),
|
||||
lid_state->get_state_size(), get_lid()->get_n_kv(), lid_state->get_n_stream(), lid_state->get_n_rs_seq());
|
||||
|
||||
return reserve_plan_lid;
|
||||
|
||||
@@ -31,6 +31,7 @@ public:
|
||||
void apply_copies(const stream_copy_info & sc_info) const;
|
||||
|
||||
uint32_t get_ratio() const;
|
||||
bool is_overlap() const; // V4 CSA: state holds 2*ratio rows of 2x width
|
||||
uint32_t get_state_size() const;
|
||||
uint32_t get_n_stream() const;
|
||||
uint32_t get_n_rs_seq() const;
|
||||
@@ -148,6 +149,9 @@ public:
|
||||
llama_dsv4_comp_state * get_hca_state() const;
|
||||
llama_dsv4_comp_state * get_lid_state() const;
|
||||
|
||||
// layer whose compressed-KV storage `il` reads, -1 = no compressed KV
|
||||
int32_t get_comp_read_layer(int32_t il) const;
|
||||
|
||||
uint32_t get_n_rs_seq() const;
|
||||
const std::vector<uint32_t> & get_rs_idx() const;
|
||||
void reset_rs_idx_for_ubatches(const std::vector<llama_ubatch> & ubatches);
|
||||
@@ -163,6 +167,12 @@ private:
|
||||
|
||||
std::vector<uint32_t> rs_idx;
|
||||
|
||||
// [n_layer] storage layer -> layer whose rows it holds (itself, or the source of a mirror)
|
||||
std::vector<int32_t> comp_kv_src;
|
||||
|
||||
// [n_layer] layer -> layer whose compressed-KV storage it reads (itself, its V4.1 kv source or a mirror), -1 = none
|
||||
std::vector<int32_t> comp_read;
|
||||
|
||||
std::unique_ptr<llama_kv_cache_iswa> kv_raw;
|
||||
std::unique_ptr<llama_kv_cache> kv_csa;
|
||||
std::unique_ptr<llama_kv_cache> kv_hca;
|
||||
@@ -214,6 +224,9 @@ public:
|
||||
void set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const;
|
||||
void set_input_k_rot(ggml_tensor * dst) const;
|
||||
|
||||
// token history of the raw (window) cells, for host-side n-gram hashing (deepseek41 Engram)
|
||||
void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const;
|
||||
|
||||
private:
|
||||
size_t i_next = 0;
|
||||
|
||||
@@ -371,6 +384,8 @@ public:
|
||||
const comp_plan & get_hca_plan(const llama_ubatch & ubatch) const;
|
||||
const comp_plan & get_lid_plan(const llama_ubatch & ubatch) const;
|
||||
|
||||
int32_t get_comp_read_layer(int32_t il) const;
|
||||
|
||||
private:
|
||||
size_t i_next = 0;
|
||||
|
||||
@@ -403,4 +418,6 @@ private:
|
||||
mutable comp_plan reserve_plan_lid;
|
||||
|
||||
const llama_memory_status status;
|
||||
|
||||
const llama_kv_cache_dsv4 * kv = nullptr;
|
||||
};
|
||||
|
||||
@@ -1168,9 +1168,7 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
|
||||
} else if (hparams.ple_n_heads > 0) {
|
||||
// embd batch (multimodal input) has no token ids, need to pad it with the correct ID for PLE layers
|
||||
// TODO @ngxson : check if we can do the same as gemma 3n / gemma 4
|
||||
ext.tok = hparams.ple_image_token_id != 0
|
||||
? (llama_token) hparams.ple_image_token_id
|
||||
: (llama_token) hparams.ple_eos_token_id;
|
||||
ext.tok = hparams.ple_embd_token;
|
||||
}
|
||||
|
||||
cells.ext_set(idx, ext);
|
||||
@@ -1849,7 +1847,7 @@ void llama_kv_cache::set_input_v_rot(ggml_tensor * dst) const {
|
||||
}
|
||||
|
||||
bool llama_kv_cache::has_cell_ext() const {
|
||||
// M-RoPE needs the 2D position, the PLE n-gram hash needs the token id
|
||||
// M-RoPE needs the 2D position, the PLE / Engram n-gram hashes need the token id
|
||||
return hparams.n_pos_per_embd() > 1 || hparams.ple_n_heads > 0;
|
||||
}
|
||||
|
||||
|
||||
@@ -414,6 +414,7 @@ namespace GGUFMeta {
|
||||
template bool llama_model_loader::get_arr<std::vector<int32_t>>(enum llm_kv kid, std::vector<int32_t> & result, bool required);
|
||||
template bool llama_model_loader::get_arr<std::array<uint32_t, LLAMA_MAX_LAYERS>>(enum llm_kv kid, std::array<uint32_t, LLAMA_MAX_LAYERS> & result, bool required);
|
||||
template bool llama_model_loader::get_arr<std::vector<uint32_t>>(enum llm_kv kid, std::vector<uint32_t> & result, bool required);
|
||||
template bool llama_model_loader::get_arr<std::vector<uint64_t>>(enum llm_kv kid, std::vector<uint64_t> & result, bool required);
|
||||
template bool llama_model_loader::get_arr<std::array<uint64_t, LLAMA_MAX_PLE_NGRAM>>(enum llm_kv kid, std::array<uint64_t, LLAMA_MAX_PLE_NGRAM> & result, bool required);
|
||||
template bool llama_model_loader::get_arr<std::array<uint64_t, LLAMA_MAX_PLE_HEADS>>(enum llm_kv kid, std::array<uint64_t, LLAMA_MAX_PLE_HEADS> & result, bool required);
|
||||
|
||||
@@ -441,6 +442,7 @@ namespace GGUFMeta {
|
||||
template bool llama_model_loader::get_key<bool> (enum llm_kv kid, bool & result, bool required);
|
||||
template bool llama_model_loader::get_key<float> (enum llm_kv kid, float & result, bool required);
|
||||
template bool llama_model_loader::get_key<uint32_t> (enum llm_kv kid, uint32_t & result, bool required);
|
||||
template bool llama_model_loader::get_key<int32_t> (enum llm_kv kid, int32_t & result, bool required);
|
||||
template bool llama_model_loader::get_key<std::string>(enum llm_kv kid, std::string & result, bool required);
|
||||
|
||||
template<>
|
||||
|
||||
@@ -8,6 +8,7 @@
|
||||
#include "llama-hparams.h"
|
||||
#include "llama-model.h"
|
||||
#include "llama-vocab.h"
|
||||
#include "models/models.h"
|
||||
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
@@ -122,6 +123,7 @@ void llama_model_saver::add_kv(const enum llm_kv key, const Container & value, c
|
||||
template void llama_model_saver::add_kv<std::vector<uint32_t>>(const enum llm_kv, const std::vector<uint32_t> &, const bool);
|
||||
template void llama_model_saver::add_kv<std::vector<float>>(const enum llm_kv, const std::vector<float> &, const bool);
|
||||
template void llama_model_saver::add_kv<std::vector<uint64_t>>(const enum llm_kv, const std::vector<uint64_t> &, const bool);
|
||||
template void llama_model_saver::add_kv<std::vector<int32_t>>(const enum llm_kv, const std::vector<int32_t> &, const bool);
|
||||
|
||||
void llama_model_saver::add_kv(const enum llm_kv key, const std::vector<std::string> & value) {
|
||||
std::vector<const char *> tmp(value.size());
|
||||
@@ -339,6 +341,32 @@ void llama_model_saver::add_kv_from_model() {
|
||||
add_kv(LLM_KV_HASH_LAYER_COUNT, hparams.dsv4_hash_layer_count);
|
||||
add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank);
|
||||
|
||||
if (model->arch == LLM_ARCH_DEEPSEEK41) {
|
||||
const auto * m41 = dynamic_cast<const llama_model_deepseek41 *>(model);
|
||||
GGML_ASSERT(m41 != nullptr);
|
||||
std::vector<uint32_t> kv_src;
|
||||
std::vector<uint32_t> idx_src;
|
||||
for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
|
||||
if (hparams.dsv41_kv_src[il] == (int32_t) il) {
|
||||
kv_src.push_back(il);
|
||||
}
|
||||
if (hparams.dsv41_topk_src[il] == (int32_t) il) {
|
||||
idx_src.push_back(il);
|
||||
}
|
||||
}
|
||||
add_kv(LLM_KV_ATTENTION_KV_SOURCE_LAYERS, kv_src);
|
||||
add_kv(LLM_KV_ATTENTION_INDEXER_SOURCE_LAYERS, idx_src);
|
||||
if (hparams.dsv41_cand_src >= 0) {
|
||||
add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_SOURCE_LAYER, hparams.dsv41_cand_src);
|
||||
add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS, hparams.dsv41_cand_topk_blocks);
|
||||
add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE, hparams.dsv41_cand_block_size);
|
||||
}
|
||||
if (hparams.ple_n_heads > 0) {
|
||||
add_kv(LLM_KV_ENGRAM_PAD_TOKEN_ID, m41->engram_pad_id);
|
||||
add_kv(LLM_KV_ENGRAM_TOKEN_MAP, m41->engram_token_map);
|
||||
}
|
||||
}
|
||||
|
||||
// the PLE group only means anything whole: write all of it or none
|
||||
if (hparams.ple_n_heads > 0) {
|
||||
std::vector<uint32_t> ple_layers;
|
||||
@@ -355,13 +383,13 @@ void llama_model_saver::add_kv_from_model() {
|
||||
add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.ple_head_dim);
|
||||
add_kv(LLM_KV_PLE_LAYER_MULTIPLIERS, std::vector<uint64_t>(
|
||||
hparams.ple_layer_multipliers.begin(),
|
||||
hparams.ple_layer_multipliers.begin() + hparams.ple_ngram_size));
|
||||
hparams.ple_layer_multipliers.begin() + hparams.ple_n_modules*hparams.ple_ngram_size));
|
||||
add_kv(LLM_KV_PLE_HEAD_OFFSETS, std::vector<uint64_t>(
|
||||
hparams.ple_head_offsets.begin(),
|
||||
hparams.ple_head_offsets.begin() + hparams.ple_n_heads));
|
||||
hparams.ple_head_offsets.begin() + hparams.ple_n_modules*hparams.ple_n_heads));
|
||||
add_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES, std::vector<uint64_t>(
|
||||
hparams.ple_head_vocab_sizes.begin(),
|
||||
hparams.ple_head_vocab_sizes.begin() + hparams.ple_n_heads));
|
||||
hparams.ple_head_vocab_sizes.begin() + hparams.ple_n_modules*hparams.ple_n_heads));
|
||||
}
|
||||
|
||||
const float rope_scaling_factor = hparams.rope_freq_scale_train == 1.0f ? 0.0f : 1.0f/hparams.rope_freq_scale_train;
|
||||
|
||||
+6
-2
@@ -212,6 +212,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
|
||||
return new llama_model_dots3note(params);
|
||||
case LLM_ARCH_DEEPSEEK4:
|
||||
return new llama_model_deepseek4(params);
|
||||
case LLM_ARCH_DEEPSEEK41:
|
||||
return new llama_model_deepseek41(params);
|
||||
case LLM_ARCH_GLM_DSA:
|
||||
return new llama_model_glm_dsa(params);
|
||||
case LLM_ARCH_MISTRAL4:
|
||||
@@ -398,7 +400,7 @@ 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 ||
|
||||
const bool is_dsv4 = ud->model->arch == LLM_ARCH_DEEPSEEK4 || ud->model->arch == LLM_ARCH_DEEPSEEK41 ||
|
||||
(ud->model->arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0);
|
||||
|
||||
static const std::regex pattern_q_weight ("blk\\.\\d*\\.attn_q.weight");
|
||||
@@ -3032,6 +3034,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
|
||||
nullptr);
|
||||
}
|
||||
} break;
|
||||
case LLM_ARCH_DEEPSEEK41:
|
||||
case LLM_ARCH_DEEPSEEK4:
|
||||
{
|
||||
GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE);
|
||||
@@ -3441,7 +3444,7 @@ int32_t llama_model_n_head_kv(const llama_model * model) {
|
||||
int32_t llama_model_n_swa(const llama_model * model) {
|
||||
// dsv4 kv-cache has SWA but it cannot be used as a rollback because of
|
||||
// other compression ratios, so we return 0 here
|
||||
if (model->arch == LLM_ARCH_DEEPSEEK4) {
|
||||
if (model->arch == LLM_ARCH_DEEPSEEK4 || model->arch == LLM_ARCH_DEEPSEEK41) {
|
||||
return 0;
|
||||
}
|
||||
return model->hparams.n_swa;
|
||||
@@ -3536,6 +3539,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
|
||||
case LLM_ARCH_DEEPSEEK2OCR:
|
||||
case LLM_ARCH_DEEPSEEK32:
|
||||
case LLM_ARCH_DEEPSEEK4:
|
||||
case LLM_ARCH_DEEPSEEK41:
|
||||
case LLM_ARCH_MUSE_GLIMMER:
|
||||
case LLM_ARCH_MUSE_SPARK:
|
||||
case LLM_ARCH_PLM:
|
||||
|
||||
@@ -605,6 +605,10 @@ struct llama_layer {
|
||||
struct ggml_tensor * indexer_kpool_gate = nullptr;
|
||||
struct ggml_tensor * indexer_kpool_ape = nullptr;
|
||||
|
||||
// deepseek41
|
||||
struct ggml_tensor * engram_kv = nullptr;
|
||||
struct ggml_tensor * engram_gate = nullptr; // q_weight*k_weight folded
|
||||
|
||||
// MSA
|
||||
struct ggml_tensor * index_q_proj = nullptr;
|
||||
struct ggml_tensor * index_k_proj = nullptr;
|
||||
|
||||
@@ -316,6 +316,9 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param
|
||||
// do not quantize the i32 token-id -> expert-id routing table (DeepSeek-V4)
|
||||
quantize &= name.find("ffn_gate_tid2eid.weight") == std::string::npos;
|
||||
|
||||
// Engram gate is an elementwise weight (ggml_mul), not a matmul operand
|
||||
quantize &= name.find("engram_gate.weight") == std::string::npos;
|
||||
|
||||
// these are very small (e.g. 4x4)
|
||||
quantize &= name.find("altup") == std::string::npos;
|
||||
quantize &= name.find("laurel") == std::string::npos;
|
||||
|
||||
+17
-10
@@ -8,14 +8,6 @@
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
static float dsv4_rope_attn_factor(float freq_scale, float ext_factor) {
|
||||
if (ext_factor == 0.0f) {
|
||||
return 1.0f;
|
||||
}
|
||||
|
||||
return 1.0f / (1.0f + 0.1f*logf(1.0f/freq_scale));
|
||||
}
|
||||
|
||||
void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {
|
||||
if (hparams.n_layer_nextn > 0) {
|
||||
const uint32_t n_layer_main = hparams.n_layer_all - hparams.n_layer_nextn;
|
||||
@@ -358,6 +350,18 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_pre(
|
||||
ggml_tensor ** post,
|
||||
ggml_tensor ** comb,
|
||||
int il) const {
|
||||
return build_hc_pre(x, build_hc_mixes(x, hc_fn, hc_scale, hc_base, post, comb, false, il), il);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_model_deepseek4::graph::build_hc_mixes(
|
||||
ggml_tensor * x,
|
||||
ggml_tensor * hc_fn,
|
||||
ggml_tensor * hc_scale,
|
||||
ggml_tensor * hc_base,
|
||||
ggml_tensor ** post,
|
||||
ggml_tensor ** comb,
|
||||
bool f32_src,
|
||||
int il) const {
|
||||
const int64_t hc = hparams.dsv4_hc_mult;
|
||||
const int64_t hc_dim = hc*n_embd;
|
||||
const int64_t hc_mix_dim = (2 + hc)*hc;
|
||||
@@ -369,6 +373,10 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_pre(
|
||||
ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt);
|
||||
ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, norm_rms_eps);
|
||||
ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm);
|
||||
if (f32_src && !ggml_is_quantized(hc_fn->type)) {
|
||||
// F32 weights (model.py hc_pre runs in fp32): request F32 inputs, the mixes are near-tie sensitive
|
||||
ggml_prec_set_src(mixes, GGML_PREC_F32, 1);
|
||||
}
|
||||
cb(mixes, "hc_mixes", il);
|
||||
|
||||
ggml_tensor * scale_pre = dsv4_view_1d(ctx0, hc_scale, 1, 0);
|
||||
@@ -404,8 +412,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_pre(
|
||||
}
|
||||
cb(*comb, "hc_comb", il);
|
||||
|
||||
ggml_tensor * result = build_hc_pre(x, pre, il);
|
||||
return result;
|
||||
return pre;
|
||||
}
|
||||
|
||||
ggml_tensor * llama_model_deepseek4::graph::build_hc_post(
|
||||
|
||||
@@ -0,0 +1,950 @@
|
||||
#include "models.h"
|
||||
#include "llama-impl.h"
|
||||
|
||||
#include "llama-kv-cache-dsv4.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cinttypes>
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
// ref: deepseek-ai/DeepSeek-V4.1-Flash-1017 inference/model.py, engram.py
|
||||
// compressed layers attend over [window | all compressed rows] with top-k as a mask. with flash attention the default
|
||||
// is sparse: fused indexer, FA bounded to window + top-k, decode-sized ubatches gather only the selected rows (same
|
||||
// keys and mask, different key order). LLAMA_DSV41_SPARSE=0 restores the reference path
|
||||
|
||||
static constexpr int64_t DSV41_ENGRAM_DEAD = -1;
|
||||
|
||||
static bool dsv41_env_on(const char * name) {
|
||||
const char * v = getenv(name);
|
||||
return !v || atoi(v) != 0;
|
||||
}
|
||||
|
||||
void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) {
|
||||
llama_model_deepseek4::load_arch_hparams(ml);
|
||||
|
||||
hparams.n_embd_out_impl = hparams.n_embd; // no hc_head
|
||||
if (hparams.indexer_top_k == 0) {
|
||||
throw std::runtime_error("DeepSeek-V4.1 indexer top_k must be > 0");
|
||||
}
|
||||
|
||||
// each compressed layer reads the KV / top-k of the latest source layer at or before it
|
||||
std::vector<uint32_t> kv_src;
|
||||
std::vector<uint32_t> idx_src;
|
||||
ml.get_arr(LLM_KV_ATTENTION_KV_SOURCE_LAYERS, kv_src);
|
||||
ml.get_arr(LLM_KV_ATTENTION_INDEXER_SOURCE_LAYERS, idx_src);
|
||||
|
||||
const uint32_t n_layer = hparams.n_layer();
|
||||
const auto out_of_range = [&](const std::vector<uint32_t> & v) {
|
||||
return std::any_of(v.begin(), v.end(), [&](uint32_t il) { return il >= n_layer; });
|
||||
};
|
||||
if (out_of_range(kv_src) || out_of_range(idx_src)) {
|
||||
throw std::runtime_error("DeepSeek-V4.1 kv/index source layer out of range");
|
||||
}
|
||||
hparams.dsv41_kv_src.fill(-1);
|
||||
hparams.dsv41_topk_src.fill(-1);
|
||||
int32_t last_kv = -1;
|
||||
int32_t last_idx = -1;
|
||||
for (uint32_t il = 0; il < n_layer; ++il) {
|
||||
const uint32_t ratio = hparams.dsv4_compress_ratios[il];
|
||||
const bool is_kv = std::find(kv_src.begin(), kv_src.end(), il) != kv_src.end();
|
||||
const bool is_idx = std::find(idx_src.begin(), idx_src.end(), il) != idx_src.end();
|
||||
if ((is_kv || is_idx) && ratio == 0) {
|
||||
throw std::runtime_error(format("DeepSeek-V4.1 source layer %u has compress ratio 0", il));
|
||||
}
|
||||
last_kv = is_kv ? (int32_t) il : last_kv;
|
||||
last_idx = is_idx ? (int32_t) il : last_idx;
|
||||
if (ratio == 0) {
|
||||
continue;
|
||||
}
|
||||
if (ratio > 2 || last_kv < 0 || last_idx < last_kv ||
|
||||
hparams.dsv4_compress_ratios[last_kv] != ratio || hparams.dsv4_compress_ratios[last_idx] != ratio) {
|
||||
throw std::runtime_error(format("DeepSeek-V4.1 layer %u: ratio %u without a matching kv/index source", il, ratio));
|
||||
}
|
||||
hparams.dsv41_kv_src[il] = last_kv;
|
||||
hparams.dsv41_topk_src[il] = last_idx;
|
||||
}
|
||||
|
||||
ml.get_key(LLM_KV_ATTENTION_INDEXER_CANDIDATE_SOURCE_LAYER, hparams.dsv41_cand_src, false);
|
||||
if (hparams.dsv41_cand_src >= 0) {
|
||||
ml.get_key(LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS, hparams.dsv41_cand_topk_blocks);
|
||||
ml.get_key(LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE, hparams.dsv41_cand_block_size);
|
||||
const int32_t cs = hparams.dsv41_cand_src;
|
||||
if ((uint32_t) cs >= n_layer || hparams.dsv41_topk_src[cs] != cs || hparams.dsv41_cand_block_size == 0 ||
|
||||
hparams.dsv41_cand_topk_blocks == 0) {
|
||||
throw std::runtime_error("DeepSeek-V4.1 invalid candidate selection (source must be an index source, block size and top-k blocks > 0)");
|
||||
}
|
||||
for (uint32_t il = cs + 1; il < n_layer; ++il) {
|
||||
if (hparams.dsv41_topk_src[il] == (int32_t) il && hparams.dsv4_compress_ratios[il] != hparams.dsv4_compress_ratios[cs]) {
|
||||
throw std::runtime_error("DeepSeek-V4.1 candidate consumers must share the source's compress ratio");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Engram: the PLE n-gram hash of qwen4exp with one set of hash constants per Engram layer, stored back to back
|
||||
std::vector<uint32_t> e_layers;
|
||||
ml.get_arr(LLM_KV_PLE_LAYERS, e_layers, false);
|
||||
if (!e_layers.empty()) {
|
||||
std::vector<uint64_t> mult;
|
||||
std::vector<uint64_t> offsets;
|
||||
std::vector<uint64_t> primes;
|
||||
ml.get_key(LLM_KV_PLE_NGRAM_SIZE, hparams.ple_ngram_size);
|
||||
ml.get_key(LLM_KV_PLE_HEADS_PER_NGRAM, hparams.ple_heads_per_ngram);
|
||||
ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.ple_head_dim);
|
||||
ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS, mult);
|
||||
ml.get_arr(LLM_KV_PLE_HEAD_OFFSETS, offsets);
|
||||
ml.get_arr(LLM_KV_PLE_HEAD_VOCAB_SIZES, primes);
|
||||
ml.get_key(LLM_KV_ENGRAM_PAD_TOKEN_ID, engram_pad_id);
|
||||
ml.get_arr(LLM_KV_ENGRAM_TOKEN_MAP, engram_token_map);
|
||||
|
||||
const size_t n_e = e_layers.size();
|
||||
const size_t n_gram = hparams.ple_ngram_size;
|
||||
const size_t n_cols = n_gram >= 2 ? (n_gram - 1)*hparams.ple_heads_per_ngram : 0;
|
||||
if (n_cols == 0 || hparams.ple_head_dim == 0 ||
|
||||
mult.size() != n_e*n_gram || primes.size() != n_e*n_cols || offsets.size() != n_e*n_cols ||
|
||||
mult.size() > LLAMA_MAX_PLE_NGRAM || primes.size() > LLAMA_MAX_PLE_HEADS ||
|
||||
engram_pad_id >= engram_token_map.size()) {
|
||||
throw std::runtime_error("DeepSeek-V4.1 engram metadata is inconsistent");
|
||||
}
|
||||
for (size_t e = 0; e < n_e; ++e) {
|
||||
// module e belongs to the e-th Engram layer
|
||||
if (e_layers[e] >= n_layer || (e > 0 && e_layers[e] <= e_layers[e - 1])) {
|
||||
throw std::runtime_error(format("DeepSeek-V4.1 engram layer %u is out of range or out of order", e_layers[e]));
|
||||
}
|
||||
hparams.is_ple_impl.set(e_layers[e]);
|
||||
}
|
||||
for (size_t c = 0; c < n_e*n_cols; ++c) {
|
||||
// a bucket count of 0 traps in the hash; row ids are int32
|
||||
if (primes[c] < 2 || offsets[c] + primes[c] > (uint64_t) INT32_MAX) {
|
||||
throw std::runtime_error(format("DeepSeek-V4.1 engram head %zu: bucket range out of range", c));
|
||||
}
|
||||
hparams.ple_head_offsets[c] = (uint32_t) offsets[c];
|
||||
hparams.ple_head_vocab_sizes[c] = (uint32_t) primes[c];
|
||||
}
|
||||
std::copy(mult.begin(), mult.end(), hparams.ple_layer_multipliers.begin());
|
||||
hparams.ple_n_heads = (uint32_t) n_cols;
|
||||
hparams.ple_n_modules = (uint32_t) n_e;
|
||||
}
|
||||
}
|
||||
|
||||
void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) {
|
||||
LLAMA_LOAD_LOCALS;
|
||||
|
||||
const int64_t q_lora_rank = hparams.n_lora_q;
|
||||
const int64_t n_ff_exp = hparams.n_ff_exp();
|
||||
const int64_t n_expert_shared = hparams.n_expert_shared;
|
||||
const int64_t n_embd_head = hparams.n_embd_head_k();
|
||||
const int64_t o_groups = hparams.dsv4_o_group_count;
|
||||
const int64_t o_lora_rank = hparams.dsv4_o_lora_rank;
|
||||
const int64_t hc_mult = hparams.dsv4_hc_mult;
|
||||
const int64_t hc_dim = hc_mult * n_embd;
|
||||
const int64_t hc_mix_dim = (2 + hc_mult) * hc_mult;
|
||||
const int64_t n_idx_embd = hparams.indexer_head_size;
|
||||
const int64_t n_idx_head = hparams.indexer_n_head;
|
||||
|
||||
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
|
||||
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
|
||||
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);
|
||||
|
||||
// all Engram tables in one lazily read tensor, like the qwen4exp n-gram table
|
||||
if (hparams.ple_n_heads > 0) {
|
||||
int64_t n_rows = 0;
|
||||
for (size_t c = 0; c < (size_t) hparams.ple_n_modules*hparams.ple_n_heads; ++c) {
|
||||
n_rows = std::max(n_rows, (int64_t) hparams.ple_head_offsets[c] + hparams.ple_head_vocab_sizes[c]);
|
||||
}
|
||||
const std::string name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str();
|
||||
if (const auto * w = ml.get_weight(name.c_str())) {
|
||||
if (w->tensor->ne[1] < n_rows) {
|
||||
throw std::runtime_error(format("%s has %" PRId64 " rows, too few for the Engram heads (%" PRId64 ")",
|
||||
name.c_str(), w->tensor->ne[1], n_rows));
|
||||
}
|
||||
n_rows = w->tensor->ne[1];
|
||||
}
|
||||
per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {(int64_t) hparams.ple_head_dim, n_rows}, TENSOR_READ_LAZY);
|
||||
}
|
||||
|
||||
for (int i = 0; i < n_layer; ++i) {
|
||||
auto & layer = layers[i];
|
||||
|
||||
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
|
||||
layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0);
|
||||
layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0);
|
||||
layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, 0);
|
||||
layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, 0);
|
||||
layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, 0);
|
||||
layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, 0);
|
||||
layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank, o_groups}, TENSOR_ALLOW_RESHAPE);
|
||||
layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, 0);
|
||||
|
||||
layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);
|
||||
layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, 0);
|
||||
layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0);
|
||||
layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);
|
||||
layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, 0);
|
||||
layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0);
|
||||
|
||||
if (hparams.dsv41_kv_src[i] == i) {
|
||||
layer.attn_comp_wkv = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WKV, "weight", i), {n_embd, n_embd_head}, 0);
|
||||
layer.attn_comp_norm = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_NORM, "weight", i), {n_embd_head}, 0);
|
||||
if (hparams.dsv4_compress_ratios[i] > 1) {
|
||||
layer.attn_comp_wgate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "weight", i), {n_embd, n_embd_head}, 0);
|
||||
}
|
||||
layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd_head, n_idx_embd}, 0);
|
||||
layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {n_idx_embd}, 0);
|
||||
}
|
||||
if (hparams.dsv41_topk_src[i] == i) {
|
||||
layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, n_idx_head * n_idx_embd}, 0);
|
||||
layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, n_idx_head}, 0);
|
||||
}
|
||||
|
||||
if (engram_index(i) >= 0) {
|
||||
layer.engram_kv = create_tensor(tn(LLM_TENSOR_ENGRAM_KV, "weight", i), {(int64_t) hparams.ple_n_heads * hparams.ple_head_dim, n_embd * (hc_mult + 1)}, 0);
|
||||
layer.engram_gate = create_tensor(tn(LLM_TENSOR_ENGRAM_GATE, "weight", i), {n_embd, hc_mult}, 0);
|
||||
}
|
||||
|
||||
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
|
||||
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
|
||||
layer.ffn_exp_probs_b_vl = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B_VL, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED);
|
||||
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
|
||||
|
||||
layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
|
||||
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);
|
||||
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
|
||||
|
||||
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
|
||||
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd}, 0);
|
||||
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
|
||||
}
|
||||
}
|
||||
|
||||
std::unique_ptr<llm_graph_context> llama_model_deepseek41::build_arch_graph(const llm_graph_params & params) const {
|
||||
if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
|
||||
throw std::runtime_error("DeepSeek-V4.1 has no MTP graph: its DSpark draft is a separate DFLASH GGUF "
|
||||
"(convert_hf_to_gguf.py --dspark), used with --spec-type draft-dspark");
|
||||
}
|
||||
return std::make_unique<graph>(*this, params);
|
||||
}
|
||||
|
||||
// Engram n-gram hash (engram.py NgramHashState.forward)
|
||||
int llama_model_deepseek41::engram_index(int il) const {
|
||||
if (hparams.ple_n_heads == 0 || !hparams.is_ple(il)) {
|
||||
return -1;
|
||||
}
|
||||
int e = 0;
|
||||
for (int j = 0; j < il; ++j) {
|
||||
e += hparams.is_ple(j);
|
||||
}
|
||||
return e;
|
||||
}
|
||||
|
||||
void llama_model_deepseek41::engram_hash(int e, const llama_ubatch & ubatch, const std::vector<llama_token> & prev, std::vector<int32_t> & out) const {
|
||||
const int64_t n_gram = hparams.ple_ngram_size;
|
||||
const int64_t n_prev = n_gram - 1;
|
||||
const int64_t n_head = hparams.ple_heads_per_ngram;
|
||||
const int64_t n_cols = hparams.ple_n_heads;
|
||||
|
||||
const uint64_t * mult = hparams.ple_layer_multipliers.data() + (size_t) e*n_gram;
|
||||
const uint32_t * primes = hparams.ple_head_vocab_sizes.data() + (size_t) e*n_cols;
|
||||
const int64_t pad = engram_token_map[engram_pad_id];
|
||||
const uint32_t * offs = hparams.ple_head_offsets.data() + (size_t) e*n_cols;
|
||||
|
||||
const auto cid = [&](llama_token t) -> int64_t {
|
||||
return t < 0 || t >= (llama_token) engram_token_map.size() ? DSV41_ENGRAM_DEAD : engram_token_map[t];
|
||||
};
|
||||
|
||||
out.resize(ubatch.n_tokens*n_cols);
|
||||
std::vector<uint64_t> toks(n_gram);
|
||||
for (int64_t i = 0; i < ubatch.n_tokens; ++i) {
|
||||
// shift 0 is the token itself; once a shift is blocked (sequence start or dead token) every longer one is too
|
||||
const bool is_embd = !ubatch.token || (ubatch.is_mixed() && ubatch.type[i]);
|
||||
int64_t c = is_embd ? DSV41_ENGRAM_DEAD : cid(ubatch.token[i]);
|
||||
bool blocked = c == DSV41_ENGRAM_DEAD;
|
||||
toks[0] = blocked ? pad : c;
|
||||
for (int64_t s = 1; s < n_gram; ++s) {
|
||||
c = ubatch.pos[i] - s < 0 ? DSV41_ENGRAM_DEAD : cid(prev[i*n_prev + (n_prev - s)]);
|
||||
blocked = blocked || c == DSV41_ENGRAM_DEAD;
|
||||
toks[s] = blocked ? pad : c;
|
||||
}
|
||||
|
||||
uint64_t rolling = toks[0]*mult[0];
|
||||
for (int64_t o = 1; o < n_gram; ++o) {
|
||||
rolling ^= toks[o]*mult[o];
|
||||
for (int64_t h = 0; h < n_head; ++h) {
|
||||
const int64_t col = (o - 1)*n_head + h;
|
||||
out[i*n_cols + col] = (int32_t) (rolling % primes[col] + offs[col]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cand_pin: +1e30 (the reference uses +inf) on the block holding each query's newest compressed row, so a partly
|
||||
// filled block always survives; a finite pin keeps fully masked blocks at -inf instead of NaN
|
||||
class llm_graph_input_dsv41 : public llm_graph_input_i {
|
||||
public:
|
||||
llm_graph_input_dsv41(const llama_model_deepseek41 & model, const llama_kv_cache_dsv4_context * mctx) :
|
||||
model(model), mctx(mctx) {}
|
||||
|
||||
void set_input(const llama_ubatch * ubatch) override;
|
||||
bool can_reuse(const llm_graph_params & params) override;
|
||||
|
||||
std::vector<ggml_tensor *> engram_ids;
|
||||
ggml_tensor * engram_keep = nullptr; // 0 for embedding (image) tokens; only in mixed ubatches
|
||||
ggml_tensor * cand_pin = nullptr;
|
||||
bool cand_csa = false;
|
||||
|
||||
const llama_model_deepseek41 & model;
|
||||
const llama_kv_cache_dsv4_context * mctx;
|
||||
|
||||
private:
|
||||
const llama_kv_cache_dsv4_context::comp_plan & cand_plan(const llama_ubatch & ubatch) const {
|
||||
return cand_csa ? mctx->get_csa_plan(ubatch) : mctx->get_hca_plan(ubatch);
|
||||
}
|
||||
};
|
||||
|
||||
void llm_graph_input_dsv41::set_input(const llama_ubatch * ubatch) {
|
||||
if (!engram_ids.empty()) {
|
||||
// a token shared by several sequences has the same history in each, so its first sequence resolves it
|
||||
std::vector<llama_token> prev;
|
||||
mctx->get_raw()->get_prev_tokens(*ubatch, model.hparams.ple_ngram_size - 1, prev);
|
||||
// hash all Engram layers first so a lazily read table gets one batched readahead before the first gather
|
||||
std::vector<std::vector<int32_t>> ids(engram_ids.size());
|
||||
std::vector<int32_t> all;
|
||||
for (size_t e = 0; e < engram_ids.size(); ++e) {
|
||||
model.engram_hash((int) e, *ubatch, prev, ids[e]);
|
||||
all.insert(all.end(), ids[e].begin(), ids[e].end());
|
||||
}
|
||||
if (model.can_prefetch.count(model.per_layer_tok_embd)) {
|
||||
llama_prefetch_rows(model.per_layer_tok_embd, all.data(), all.size());
|
||||
}
|
||||
for (size_t e = 0; e < engram_ids.size(); ++e) {
|
||||
if (engram_ids[e]->buffer) {
|
||||
ggml_backend_tensor_set(engram_ids[e], ids[e].data(), 0, ids[e].size()*sizeof(int32_t));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// embedding (image) tokens are hashed DEAD and get no Engram contribution (gate 0)
|
||||
if (engram_keep && engram_keep->buffer) {
|
||||
std::vector<float> keep(ubatch->n_tokens);
|
||||
for (uint32_t i = 0; i < ubatch->n_tokens; ++i) {
|
||||
const bool is_embd = !ubatch->token || (ubatch->is_mixed() && ubatch->type[i]);
|
||||
keep[i] = is_embd ? 0.0f : 1.0f;
|
||||
}
|
||||
ggml_backend_tensor_set(engram_keep, keep.data(), 0, keep.size()*sizeof(float));
|
||||
}
|
||||
|
||||
if (cand_pin && cand_pin->buffer) {
|
||||
const auto & plan = cand_plan(*ubatch);
|
||||
const int64_t n_blk = cand_pin->ne[0];
|
||||
const int64_t bs = model.hparams.dsv41_cand_block_size;
|
||||
GGML_ASSERT((int64_t) plan.n_visible.size() == (int64_t) ubatch->n_tokens);
|
||||
std::vector<float> pin(n_blk*ubatch->n_tokens, 0.0f);
|
||||
for (int64_t i = 0; i < (int64_t) ubatch->n_tokens; ++i) {
|
||||
const int64_t nv = plan.n_visible[i];
|
||||
if (nv > 0) {
|
||||
GGML_ASSERT((nv - 1)/bs < n_blk);
|
||||
pin[i*n_blk + (nv - 1)/bs] = 1e30f;
|
||||
}
|
||||
}
|
||||
ggml_backend_tensor_set(cand_pin, pin.data(), 0, pin.size()*sizeof(float));
|
||||
}
|
||||
}
|
||||
|
||||
bool llm_graph_input_dsv41::can_reuse(const llm_graph_params & params) {
|
||||
mctx = static_cast<const llama_kv_cache_dsv4_context *>(params.mctx);
|
||||
|
||||
bool res = true;
|
||||
for (auto * t : engram_ids) {
|
||||
res &= t->ne[0] == (int64_t) model.hparams.ple_n_heads*params.ubatch.n_tokens;
|
||||
}
|
||||
res &= (engram_keep != nullptr) == (!engram_ids.empty() && params.ubatch.embd != nullptr);
|
||||
if (engram_keep) {
|
||||
res &= engram_keep->ne[2] == (int64_t) params.ubatch.n_tokens;
|
||||
}
|
||||
if (cand_pin) {
|
||||
const auto & plan = cand_plan(params.ubatch);
|
||||
const int64_t bs = model.hparams.dsv41_cand_block_size;
|
||||
res &= plan.n_kv > 0 && cand_pin->ne[0] == (plan.n_kv + bs - 1)/bs;
|
||||
res &= cand_pin->ne[1]*cand_pin->ne[3] == (int64_t) params.ubatch.n_tokens;
|
||||
}
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_tensor * llama_model_deepseek41::graph::build_engram(
|
||||
const llama_model_deepseek41 & model,
|
||||
ggml_tensor * x,
|
||||
ggml_tensor * rows,
|
||||
ggml_tensor * keep,
|
||||
int il) const {
|
||||
const auto & layer = model.layers[il];
|
||||
const int64_t hc = hparams.dsv4_hc_mult;
|
||||
const int64_t nt = x->ne[2];
|
||||
|
||||
GGML_ASSERT(rows->ne[0] == (int64_t) (hparams.ple_head_dim*hparams.ple_n_heads) && rows->ne[1] == nt);
|
||||
|
||||
ggml_tensor * kv = build_lora_mm(layer.engram_kv, rows);
|
||||
cb(kv, "engram_kv", il);
|
||||
|
||||
ggml_tensor * key = ggml_cont(ctx0, ggml_view_3d(ctx0, kv, n_embd, hc, nt, ggml_row_size(kv->type, n_embd), kv->nb[1], 0));
|
||||
ggml_tensor * value = ggml_cont(ctx0, ggml_view_3d(ctx0, kv, n_embd, 1, nt, ggml_row_size(kv->type, n_embd), kv->nb[1],
|
||||
ggml_row_size(kv->type, hc*n_embd)));
|
||||
|
||||
// quantized copy has no binary-op kernel
|
||||
ggml_tensor * egate = ggml_is_quantized(layer.engram_gate->type) ? ggml_cast(ctx0, layer.engram_gate, GGML_TYPE_F32) : layer.engram_gate;
|
||||
ggml_tensor * dot = ggml_mul(ctx0, ggml_mul(ctx0, ggml_rms_norm(ctx0, x, norm_rms_eps), ggml_rms_norm(ctx0, key, norm_rms_eps)), egate);
|
||||
dot = ggml_scale(ctx0, ggml_sum_rows(ctx0, dot), 1.0f/sqrtf((float) n_embd));
|
||||
cb(dot, "engram_dot", il);
|
||||
|
||||
// sigmoid(copysign(sqrt(max(|dot|, 1e-6)), dot)); sign(0) = +1 as in copysign
|
||||
ggml_tensor * mag = ggml_sqrt(ctx0, ggml_clamp(ctx0, ggml_abs(ctx0, dot), 1e-6f, INFINITY));
|
||||
ggml_tensor * sign = ggml_scale_bias(ctx0, ggml_step(ctx0, ggml_neg(ctx0, dot)), -2.0f, 1.0f);
|
||||
ggml_tensor * gate = ggml_sigmoid(ctx0, ggml_mul(ctx0, mag, sign));
|
||||
if (keep) {
|
||||
gate = ggml_mul(ctx0, gate, keep);
|
||||
}
|
||||
cb(gate, "engram_gate", il);
|
||||
|
||||
return ggml_add(ctx0, x, ggml_mul(ctx0, ggml_repeat_4d(ctx0, value, n_embd, hc, nt, 1), gate));
|
||||
}
|
||||
|
||||
ggml_tensor * llama_model_deepseek41::graph::build_attention41(
|
||||
const llama_model_deepseek41 & model,
|
||||
llm_graph_input_dsv4 * inp_dsv4,
|
||||
ggml_tensor * cur,
|
||||
ggml_tensor * inp_pos,
|
||||
ggml_tensor * cand_pin,
|
||||
int il) {
|
||||
const auto & layer = model.layers[il];
|
||||
llm_graph_input_dsv4_raw * inp_attn = inp_dsv4->get_raw();
|
||||
|
||||
const int64_t n_embd_head = hparams.n_embd_head_k();
|
||||
const int64_t n_embd_head_rope = hparams.n_rot();
|
||||
const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;
|
||||
const int64_t n_groups = hparams.dsv4_o_group_count;
|
||||
const int64_t o_lora_rank = hparams.dsv4_o_lora_rank;
|
||||
const int64_t n_idx_embd = hparams.indexer_head_size;
|
||||
const int64_t n_idx_head = hparams.indexer_n_head;
|
||||
const int64_t nt = cur->ne[1];
|
||||
const float kq_scale = 1.0f/sqrtf(float(n_embd_head));
|
||||
const int64_t ratio = hparams.dsv4_compress_ratios[il];
|
||||
|
||||
const bool comp_rope = ratio != 0;
|
||||
const float freq_base_l = comp_rope ? hparams.dsv4_compress_rope_base : freq_base;
|
||||
const float freq_scale_l = comp_rope ? freq_scale : 1.0f;
|
||||
const float ext_factor_l = comp_rope ? ext_factor : 0.0f;
|
||||
const float attn_factor_l = dsv4_rope_attn_factor(freq_scale_l, ext_factor_l);
|
||||
const float beta_fast_l = comp_rope ? beta_fast : 0.0f;
|
||||
const float beta_slow_l = comp_rope ? beta_slow : 0.0f;
|
||||
const int32_t n_ctx_orig_l = comp_rope ? n_ctx_orig : 0;
|
||||
|
||||
const auto rope = [&](ggml_tensor * t, ggml_tensor * pos, int64_t n_nope) {
|
||||
t = ggml_rope_ext(ctx0, t, pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l,
|
||||
freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l);
|
||||
return ggml_rope_set_offset(t, n_nope);
|
||||
};
|
||||
|
||||
ggml_tensor * qr = build_norm(build_lora_mm(layer.wq_a, cur), layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(qr, "qr_norm", il);
|
||||
|
||||
ggml_tensor * q = ggml_reshape_3d(ctx0, build_lora_mm(layer.wq_b, qr), n_embd_head, n_head, nt);
|
||||
q = rope(q, inp_pos, n_embd_head_nope);
|
||||
cb(q, "q", il);
|
||||
|
||||
ggml_tensor * kv = build_norm(build_lora_mm(layer.wkv, cur), layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);
|
||||
kv = rope(ggml_reshape_3d(ctx0, kv, n_embd_head, 1, nt), inp_pos, n_embd_head_nope);
|
||||
cb(kv, "kv", il);
|
||||
|
||||
const auto & ci = ratio == 2 ? inp_dsv4->get_csa() : inp_dsv4->get_hca();
|
||||
const auto * comp_ctx = ratio == 2 ? inp_dsv4->mctx->get_csa() : inp_dsv4->mctx->get_hca();
|
||||
const auto * comp_state = ratio == 2 ? inp_dsv4->mctx->get_csa_state() : inp_dsv4->mctx->get_hca_state();
|
||||
|
||||
const int32_t rd = inp_dsv4->mctx->get_comp_read_layer(il);
|
||||
|
||||
if (ratio != 0) {
|
||||
GGML_ASSERT(hparams.dsv41_kv_src[il] >= 0 && (int64_t) comp_state->get_ratio() == ratio);
|
||||
|
||||
if (hparams.dsv41_kv_src[il] == il && ci.state_pos) {
|
||||
ggml_tensor * state_kv = build_lora_mm(layer.attn_comp_wkv, cur);
|
||||
ggml_tensor * state_score = ratio > 1 ? build_lora_mm(layer.attn_comp_wgate, cur) : ggml_fill(ctx0, state_kv, 0.0f);
|
||||
|
||||
ggml_tensor * restored_kv = comp_state->get_kv_all(ctx0, il);
|
||||
ggml_tensor * restored_score = comp_state->get_score_all(ctx0, il);
|
||||
if (ci.state_restore_src_idxs && ci.state_restore_dst_idxs) {
|
||||
restored_kv = comp_state->cpy_kv (ctx0, ggml_get_rows(ctx0, restored_kv, ci.state_restore_src_idxs), ci.state_restore_dst_idxs, il);
|
||||
restored_score = comp_state->cpy_score(ctx0, ggml_get_rows(ctx0, restored_score, ci.state_restore_src_idxs), ci.state_restore_dst_idxs, il);
|
||||
}
|
||||
|
||||
if (ci.state_write_idxs) {
|
||||
const int64_t n_blocks = ci.state_write_pos->ne[0];
|
||||
const int64_t n_state = comp_state->get_n_rows();
|
||||
GGML_ASSERT(ci.state_read_idxs->ne[0] == ratio*n_blocks);
|
||||
|
||||
ggml_tensor * src_kv = ggml_concat(ctx0, ggml_view_2d(ctx0, restored_kv, restored_kv->ne[0], n_state, restored_kv->nb[1], 0), state_kv, 1);
|
||||
ggml_tensor * src_score = ggml_concat(ctx0, ggml_view_2d(ctx0, restored_score, restored_score->ne[0], n_state, restored_score->nb[1], 0), state_score, 1);
|
||||
|
||||
ggml_tensor * v = ggml_get_rows(ctx0, src_kv, ci.state_read_idxs);
|
||||
ggml_tensor * s = ggml_get_rows(ctx0, src_score, ci.state_read_idxs);
|
||||
v = ggml_cont(ctx0, ggml_permute(ctx0, ggml_reshape_3d(ctx0, v, n_embd_head, ratio, n_blocks), 1, 0, 2, 3));
|
||||
s = ggml_cont(ctx0, ggml_permute(ctx0, ggml_reshape_3d(ctx0, s, n_embd_head, ratio, n_blocks), 1, 0, 2, 3));
|
||||
|
||||
ggml_tensor * comp = ggml_sum_rows(ctx0, ggml_mul(ctx0, v, ggml_soft_max(ctx0, s)));
|
||||
comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3));
|
||||
comp = build_norm(comp, layer.attn_comp_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(comp, "comp_latent", il);
|
||||
|
||||
// index K from the RoPE-free latent; a latent stands for the first token of its group
|
||||
ggml_tensor * ik = build_norm(build_lora_mm(layer.indexer_attn_k, comp), layer.indexer_k_norm, nullptr, LLM_NORM_RMS, il);
|
||||
ik = rope(ik, ci.state_write_pos, n_idx_embd - n_embd_head_rope);
|
||||
comp = rope(comp, ci.state_write_pos, n_embd_head_nope);
|
||||
|
||||
ggml_tensor * row = ggml_concat(ctx0, comp, ik, 0);
|
||||
ggml_build_forward_expand(gf, comp_ctx->cpy_k(ctx0, row, ci.state_write_idxs, il));
|
||||
comp_row_of[il] = row;
|
||||
}
|
||||
|
||||
if (ci.state_snapshot_src_idxs && ci.state_snapshot_dst_idxs) {
|
||||
ggml_tensor * snap_kv = ggml_concat(ctx0, restored_kv, state_kv, 1);
|
||||
ggml_tensor * snap_score = ggml_concat(ctx0, restored_score, state_score, 1);
|
||||
ggml_build_forward_expand(gf, comp_state->cpy_kv (ctx0, ggml_get_rows(ctx0, snap_kv, ci.state_snapshot_src_idxs), ci.state_snapshot_dst_idxs, il));
|
||||
ggml_build_forward_expand(gf, comp_state->cpy_score(ctx0, ggml_get_rows(ctx0, snap_score, ci.state_snapshot_src_idxs), ci.state_snapshot_dst_idxs, il));
|
||||
}
|
||||
ggml_build_forward_expand(gf, comp_state->cpy_kv (ctx0, ggml_get_rows(ctx0, state_kv, ci.state_persist_src_idxs), ci.state_persist_dst_idxs, il));
|
||||
ggml_build_forward_expand(gf, comp_state->cpy_score(ctx0, ggml_get_rows(ctx0, state_score, ci.state_persist_src_idxs), ci.state_persist_dst_idxs, il));
|
||||
}
|
||||
|
||||
if (rd == il && hparams.dsv41_kv_src[il] != il && comp_row_of[hparams.dsv41_kv_src[il]]) {
|
||||
// first consumer on this device: append the source's new rows to the local mirror
|
||||
ggml_build_forward_expand(gf, comp_ctx->cpy_k(ctx0, comp_row_of[hparams.dsv41_kv_src[il]], ci.state_write_idxs, il));
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor * out = nullptr;
|
||||
if (ratio == 0 || !ci.kq_mask) {
|
||||
// no compressed row visible yet: window only
|
||||
out = build_raw_attention(inp_attn, q, kv, layer.attn_sinks, kq_scale, il);
|
||||
} else {
|
||||
const int64_t n_comp = ci.kq_mask->ne[0];
|
||||
|
||||
ggml_tensor * comp_all = comp_ctx->get_k(ctx0, rd);
|
||||
GGML_ASSERT(comp_all->ne[0] == n_embd_head + n_idx_embd && n_comp <= comp_all->ne[2]);
|
||||
const int64_t n_stream = comp_all->ne[3];
|
||||
|
||||
ggml_tensor * top_k = topk_of[hparams.dsv41_topk_src[il]];
|
||||
if (hparams.dsv41_topk_src[il] == il) {
|
||||
ggml_tensor * ik = ggml_view_4d(ctx0, comp_all, n_idx_embd, 1, n_comp, n_stream,
|
||||
comp_all->nb[1], comp_all->nb[2], comp_all->nb[3], ggml_row_size(comp_all->type, n_embd_head));
|
||||
// the fused indexer reads the strided view in place, the unfused path needs it contiguous; quantized
|
||||
// caches are dequantized straight from the strided view (CUDA has no same-type copy of quantized views)
|
||||
const bool fused_idx = sparse && cparams.fused_lid && ci.kq_mask->type == GGML_TYPE_F16;
|
||||
if (ggml_is_quantized(ik->type)) {
|
||||
ik = ggml_cast(ctx0, ik, GGML_TYPE_F32);
|
||||
} else if (!fused_idx) {
|
||||
ik = ggml_cont(ctx0, ik);
|
||||
}
|
||||
|
||||
ggml_tensor * iq = ggml_reshape_3d(ctx0, build_lora_mm(layer.indexer_attn_q_b, qr), n_idx_embd, n_idx_head, nt);
|
||||
iq = rope(iq, inp_pos, n_idx_embd - n_embd_head_rope);
|
||||
iq = ggml_view_4d(ctx0, iq, iq->ne[0], iq->ne[1], nt/n_stream, n_stream, iq->nb[1], iq->nb[2], iq->nb[2]*(nt/n_stream), 0);
|
||||
|
||||
ggml_tensor * iw = ggml_scale(ctx0, build_lora_mm(layer.indexer_proj, cur), 1.0f/sqrtf(float(n_idx_embd*n_idx_head)));
|
||||
iw = ggml_view_4d(ctx0, iw, iw->ne[0], nt/n_stream, 1, n_stream, iw->nb[1], iw->nb[1]*(nt/n_stream), iw->nb[1]*(nt/n_stream), 0);
|
||||
|
||||
// chunk queries so the [n_ih, tokens, n_comp] scores stay near 2^26 elements (at most 64 chunks); the fused
|
||||
// indexer has no such intermediate and runs as one chunk
|
||||
const int64_t n_tok = nt/n_stream;
|
||||
GGML_ASSERT(ci.kq_mask->ne[1] == n_tok);
|
||||
const int64_t chunk = fused_idx ? n_tok :
|
||||
std::max({ (n_tok + 63)/64, (INT64_C(1) << 26)/(n_idx_head*n_comp*n_stream), INT64_C(1) });
|
||||
|
||||
const int32_t cand_src = hparams.dsv41_cand_src;
|
||||
const bool pick_blk = cand_src == il && cand_pin; // level one: pick the candidate blocks
|
||||
const bool use_blk = cand_src >= 0 && cand_src < il && cand_sel; // level two: score inside them
|
||||
const int64_t bs = hparams.dsv41_cand_block_size;
|
||||
const int64_t n_blk = pick_blk || use_blk ? (n_comp + bs - 1)/bs : 0;
|
||||
GGML_ASSERT(!pick_blk || cand_pin->ne[0] == n_blk);
|
||||
GGML_ASSERT(!use_blk || cand_n_comp == n_comp);
|
||||
|
||||
ggml_tensor * ik_p = fused_idx ? nullptr : ggml_permute(ctx0, ik, 0, 2, 1, 3);
|
||||
std::vector<ggml_tensor *> parts_top_k;
|
||||
std::vector<ggml_tensor *> parts_sel;
|
||||
for (int64_t t0 = 0; t0 < n_tok; t0 += chunk) {
|
||||
const int64_t nc = std::min(chunk, n_tok - t0);
|
||||
const auto rows = [&](ggml_tensor * t) {
|
||||
return ggml_view_4d(ctx0, t, t->ne[0], nc, 1, n_stream, t->nb[1], t->nb[2], t->nb[3], t0*t->nb[1]);
|
||||
};
|
||||
ggml_tensor * score = nullptr;
|
||||
if (fused_idx) {
|
||||
score = ggml_lightning_indexer(ctx0, iq, ik, iw, ci.kq_mask);
|
||||
res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, score, il});
|
||||
} else {
|
||||
ggml_tensor * iqc = ggml_view_4d(ctx0, iq, iq->ne[0], iq->ne[1], nc, n_stream, iq->nb[1], iq->nb[2], iq->nb[3], t0*iq->nb[2]);
|
||||
|
||||
ggml_tensor * kq = ggml_mul_mat(ctx0, ik_p, ggml_permute(ctx0, iqc, 0, 2, 1, 3));
|
||||
if (!ggml_is_quantized(ik->type)) {
|
||||
ggml_prec_set_src(kq, GGML_PREC_F32, 1);
|
||||
}
|
||||
kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 2, 1, 0, 3));
|
||||
score = ggml_sum_rows(ctx0, ggml_mul(ctx0, ggml_relu(ctx0, kq), rows(iw)));
|
||||
ggml_tensor * mask = rows(ci.kq_mask);
|
||||
score = ggml_add(ctx0, ggml_cont(ctx0, ggml_permute(ctx0, score, 2, 1, 0, 3)),
|
||||
mask->type == GGML_TYPE_F32 ? mask : ggml_cast(ctx0, mask, GGML_TYPE_F32));
|
||||
}
|
||||
cb(score, "idx_score", il);
|
||||
|
||||
if (pick_blk) {
|
||||
// best position per block, pin the newest block, keep the top blocks
|
||||
ggml_tensor * s2 = ggml_reshape_2d(ctx0, score, n_comp, nc*n_stream);
|
||||
if (n_blk*bs != n_comp) {
|
||||
s2 = ggml_concat(ctx0, s2, ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_blk*bs - n_comp, nc*n_stream), -INFINITY), 0);
|
||||
}
|
||||
ggml_tensor * bmax = ggml_pool_1d(ctx0, s2, GGML_OP_POOL_MAX, bs, bs, 0);
|
||||
bmax = ggml_add(ctx0, ggml_reshape_4d(ctx0, bmax, n_blk, nc, 1, n_stream), rows(cand_pin));
|
||||
cb(bmax, "cand_block_score", il);
|
||||
parts_sel.push_back(ggml_top_k(ctx0, bmax, std::min<int64_t>(hparams.dsv41_cand_topk_blocks, n_blk)));
|
||||
} else if (use_blk) {
|
||||
// own scores, restricted to the source's candidate blocks: 0 inside them, -inf elsewhere
|
||||
ggml_tensor * bsel = ggml_cont(ctx0, rows(cand_sel));
|
||||
const int64_t k_b = bsel->ne[0];
|
||||
ggml_tensor * bm = ggml_fill(ctx0, ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_blk, nc*n_stream, 1), -INFINITY);
|
||||
ggml_tensor * zeros = ggml_fill(ctx0, ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, k_b, nc*n_stream, 1), 0.0f);
|
||||
bm = ggml_set_rows(ctx0, bm, zeros, ggml_reshape_3d(ctx0, bsel, k_b, nc*n_stream, 1));
|
||||
bm = ggml_reshape_2d(ctx0, ggml_repeat_4d(ctx0, bm, bs, n_blk, nc*n_stream, 1), bs*n_blk, nc*n_stream);
|
||||
bm = ggml_reshape_4d(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, bm, n_comp, nc*n_stream, bm->nb[1], 0)), n_comp, nc, 1, n_stream);
|
||||
cb(bm, "cand_mask", il);
|
||||
score = ggml_add(ctx0, score, bm);
|
||||
}
|
||||
|
||||
parts_top_k.push_back(ggml_top_k(ctx0, score, std::min<int64_t>(n_comp, hparams.indexer_top_k)));
|
||||
}
|
||||
|
||||
const auto join = [&](std::vector<ggml_tensor *> v) {
|
||||
while (v.size() > 1) {
|
||||
std::vector<ggml_tensor *> w;
|
||||
for (size_t i = 0; i < v.size(); i += 2) {
|
||||
w.push_back(i + 1 < v.size() ? ggml_concat(ctx0, v[i], v[i + 1], 1) : v[i]);
|
||||
}
|
||||
v.swap(w);
|
||||
}
|
||||
return v[0];
|
||||
};
|
||||
if (pick_blk) {
|
||||
cand_sel = join(parts_sel);
|
||||
cand_n_comp = n_comp;
|
||||
}
|
||||
top_k = ggml_cont(ctx0, join(parts_top_k));
|
||||
cb(top_k, "idx_top_k", il);
|
||||
topk_of[il] = top_k;
|
||||
}
|
||||
GGML_ASSERT(top_k && "DeepSeek-V4.1: top-k source layer produced no top-k");
|
||||
|
||||
// a quantized K cache stores the window Hadamard-rotated (as build_raw_attention)
|
||||
ggml_tensor * k_rot = inp_attn->self_k_rot;
|
||||
if (k_rot) {
|
||||
q = llama_mul_mat_hadamard(ctx0, q, k_rot);
|
||||
kv = llama_mul_mat_hadamard(ctx0, kv, k_rot);
|
||||
}
|
||||
ggml_build_forward_expand(gf, q);
|
||||
ggml_build_forward_expand(gf, kv);
|
||||
ggml_build_forward_expand(gf, inp_attn->mctx->cpy_k(ctx0, kv, inp_attn->get_k_idxs(), il));
|
||||
|
||||
ggml_tensor * raw_k = inp_attn->mctx->get_k(ctx0, il);
|
||||
ggml_tensor * raw_mask = inp_attn->get_kq_mask();
|
||||
|
||||
const int32_t topk_src = hparams.dsv41_topk_src[il];
|
||||
const int64_t k_sel = top_k->ne[0];
|
||||
const int64_t n_raw = raw_k->ne[2];
|
||||
|
||||
// few queries per stream: gather each query's top-k rows and attend over [window | selected], one FA stream
|
||||
// per query, instead of a full n_comp mask per layer. same keys and mask as dense, different key order
|
||||
const bool gather = sparse && cparams.flash_attn && raw_mask->type == GGML_TYPE_F16 && ci.kq_mask->type == GGML_TYPE_F16 &&
|
||||
((n_stream == 1 && nt <= 32) || nt == n_stream);
|
||||
|
||||
if (gather) {
|
||||
// ggml_top_k returns the rows in no particular order (CUB DeviceTopK can vary between runs) and FA sums the
|
||||
// gathered keys in that order: sort the rows once per top-k source so the result is deterministic
|
||||
ggml_tensor *& sel = gather_sel_of[topk_src];
|
||||
if (!sel) {
|
||||
GGML_ASSERT(n_comp <= (INT64_C(1) << 24)); // row ids are exact in F32
|
||||
ggml_tensor * ids = ggml_reshape_2d(ctx0, top_k, k_sel, nt);
|
||||
ggml_tensor * ord = ggml_argsort(ctx0, ggml_cast(ctx0, ids, GGML_TYPE_F32), GGML_SORT_ORDER_ASC);
|
||||
sel = ggml_reshape_2d(ctx0, ggml_get_rows(ctx0, ggml_reshape_3d(ctx0, ids, 1, k_sel, nt), ord), k_sel, nt);
|
||||
cb(sel, "comp_top_k_sorted", il);
|
||||
}
|
||||
|
||||
// built once per (storage, top-k) pair
|
||||
ggml_tensor *& comp_k = gather_k_of[{rd, topk_src}];
|
||||
if (!comp_k) {
|
||||
if (n_stream == 1) {
|
||||
ggml_tensor * src_rows = ggml_view_3d(ctx0, comp_all, n_embd_head, n_comp, 1, comp_all->nb[2], comp_all->nb[3], 0);
|
||||
comp_k = ggml_get_rows(ctx0, src_rows, ggml_reshape_1d(ctx0, sel, k_sel*nt));
|
||||
} else {
|
||||
ggml_tensor * src_rows = ggml_view_3d(ctx0, comp_all, n_embd_head, n_comp, n_stream, comp_all->nb[2], comp_all->nb[3], 0);
|
||||
comp_k = ggml_get_rows(ctx0, src_rows, sel);
|
||||
}
|
||||
comp_k = ggml_reshape_4d(ctx0, comp_k, n_embd_head, 1, k_sel, nt); // F32 (get_rows dequantizes)
|
||||
if (k_rot) {
|
||||
// quantized K cache: the window is stored rotated, the compressed rows are not
|
||||
comp_k = llama_mul_mat_hadamard(ctx0, comp_k, k_rot);
|
||||
}
|
||||
if (comp_k->type != raw_k->type) {
|
||||
comp_k = ggml_cast(ctx0, comp_k, raw_k->type); // exact: rows are raw_k->type values
|
||||
}
|
||||
cb(comp_k, "comp_k_sel", il);
|
||||
}
|
||||
|
||||
if (n_stream == 1 && nt > 1) {
|
||||
if (ggml_is_quantized(raw_k->type)) {
|
||||
// no REPEAT kernel for quantized types: byte-exact copies by doubling concats
|
||||
ggml_tensor * rep = nullptr;
|
||||
ggml_tensor * pw = raw_k;
|
||||
for (int64_t n = nt; n > 0; n >>= 1) {
|
||||
if (n & 1) {
|
||||
rep = rep ? ggml_concat(ctx0, rep, pw, 3) : pw;
|
||||
}
|
||||
if (n > 1) {
|
||||
pw = ggml_concat(ctx0, pw, pw, 3);
|
||||
}
|
||||
}
|
||||
raw_k = rep;
|
||||
} else {
|
||||
raw_k = ggml_repeat_4d(ctx0, raw_k, n_embd_head, 1, n_raw, nt);
|
||||
}
|
||||
}
|
||||
ggml_tensor * k_all = ggml_concat(ctx0, raw_k, comp_k, 2);
|
||||
|
||||
// mask: window row, then comp mask at the selected rows (-inf where top-k picked an invisible row);
|
||||
// one per top-k source, shared by its layers
|
||||
ggml_tensor *& kq_mask = gather_mask_of[topk_src];
|
||||
if (!kq_mask) {
|
||||
ggml_tensor * sel_mask = ggml_get_rows(ctx0, ggml_reshape_3d(ctx0, ci.kq_mask, 1, n_comp, nt), sel);
|
||||
sel_mask = ggml_cast(ctx0, ggml_reshape_4d(ctx0, sel_mask, k_sel, 1, 1, nt), GGML_TYPE_F16);
|
||||
kq_mask = ggml_concat(ctx0, ggml_reshape_4d(ctx0, raw_mask, raw_mask->ne[0], 1, 1, nt), sel_mask, 0);
|
||||
cb(kq_mask, "comp_kq_mask_sel", il);
|
||||
}
|
||||
|
||||
out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, layer.attn_sinks, nullptr, 0, kq_scale, il);
|
||||
} else {
|
||||
// dense attention over [window | all compressed rows], selection applied as a mask
|
||||
ggml_tensor * comp_k = ggml_view_4d(ctx0, comp_all, n_embd_head, 1, n_comp, n_stream,
|
||||
comp_all->nb[1], comp_all->nb[2], comp_all->nb[3], 0);
|
||||
if (k_rot) {
|
||||
// compressed rows are stored unrotated: rotate in F32, then cast to the cache type (dequantized from
|
||||
// the strided view: CUDA has no same-type quantized copy)
|
||||
comp_k = comp_k->type != GGML_TYPE_F32 ? ggml_cast(ctx0, comp_k, GGML_TYPE_F32) : ggml_cont(ctx0, comp_k);
|
||||
comp_k = llama_mul_mat_hadamard(ctx0, comp_k, k_rot);
|
||||
if (comp_k->type != raw_k->type) {
|
||||
comp_k = ggml_cast(ctx0, comp_k, raw_k->type);
|
||||
}
|
||||
} else if (ggml_is_quantized(comp_k->type)) {
|
||||
// quantized cache without rotation (LLAMA_ATTN_ROT_DISABLE=1): via F32 for the same reason
|
||||
comp_k = ggml_cast(ctx0, ggml_cast(ctx0, comp_k, GGML_TYPE_F32), raw_k->type);
|
||||
} else {
|
||||
comp_k = comp_k->type != raw_k->type ? ggml_cast(ctx0, comp_k, raw_k->type) : ggml_cont(ctx0, comp_k);
|
||||
}
|
||||
cb(comp_k, "comp_k", il);
|
||||
ggml_tensor * k_all = ggml_concat(ctx0, raw_k, comp_k, 2);
|
||||
|
||||
ggml_tensor * sel_mask = build_top_k_mask(ci.kq_mask, top_k, "comp_top_k_mask", il);
|
||||
ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, sel_mask, 0);
|
||||
cb(kq_mask, "comp_kq_mask", il);
|
||||
|
||||
// at most window + top-k visible columns per query: lets CUDA FA gather only those
|
||||
const int64_t n_kv_max = sparse ? std::min<int64_t>(raw_mask->ne[0], hparams.n_swa) + k_sel : 0;
|
||||
out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, layer.attn_sinks, nullptr, n_kv_max, kq_scale, il);
|
||||
}
|
||||
if (k_rot) {
|
||||
out = llama_mul_mat_hadamard(ctx0, out, k_rot);
|
||||
}
|
||||
cb(out, "attn_comp", il);
|
||||
}
|
||||
|
||||
out = ggml_reshape_3d(ctx0, out, n_embd_head, n_head, nt);
|
||||
out = ggml_rope_ext_back(ctx0, out, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l,
|
||||
freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l);
|
||||
out = ggml_rope_set_offset(out, n_embd_head_nope);
|
||||
cb(out, "attn_derope", il);
|
||||
|
||||
out = ggml_permute(ctx0, ggml_reshape_3d(ctx0, out, n_head/n_groups*n_embd_head, n_groups, nt), 0, 2, 1, 3);
|
||||
ggml_tensor * oa = ggml_mul_mat(ctx0, layer.wo_a, out);
|
||||
cb(oa, "attn_wo_a", il);
|
||||
oa = ggml_cont_2d(ctx0, ggml_permute(ctx0, oa, 0, 2, 1, 3), o_lora_rank*n_groups, nt);
|
||||
|
||||
out = build_lora_mm(layer.wo_b, oa);
|
||||
cb(out, "attn_out", il);
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
llama_model_deepseek41::graph::graph(const llama_model & model_base, const llm_graph_params & params) :
|
||||
llama_model_deepseek4::graph(params) {
|
||||
const auto & model = static_cast<const llama_model_deepseek41 &>(model_base);
|
||||
|
||||
static const bool sparse_env = dsv41_env_on("LLAMA_DSV41_SPARSE");
|
||||
sparse = sparse_env;
|
||||
|
||||
topk_of.assign(n_layer, nullptr);
|
||||
comp_row_of.assign(n_layer, nullptr);
|
||||
|
||||
ggml_tensor * cur;
|
||||
|
||||
ggml_tensor * inp = build_inp_embd(model.tok_embd);
|
||||
ggml_tensor * inp_pos = build_inp_pos();
|
||||
ggml_tensor * inp_out_ids = build_inp_out_ids();
|
||||
llm_graph_input_dsv4 * inp_dsv4 = build_inp_dsv4();
|
||||
ggml_build_forward_expand(gf, inp_dsv4->get_raw()->self_kq_mask);
|
||||
|
||||
auto inp41 = std::make_unique<llm_graph_input_dsv41>(model, inp_dsv4->mctx);
|
||||
// embedding input has no token ids: it skips Engram, like the reference's masked (image) positions
|
||||
// the gathers go right after the token embeddings, so they share the first graph split (the table stays on the host)
|
||||
std::vector<ggml_tensor *> engram_rows;
|
||||
for (int il = 0; ubatch.token && hparams.ple_n_heads > 0 && il < n_layer; ++il) {
|
||||
if (!hparams.is_ple(il)) {
|
||||
continue;
|
||||
}
|
||||
ggml_tensor * ids = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, (int64_t) hparams.ple_n_heads*n_tokens);
|
||||
ggml_set_name(ids, ("engram_ids-" + std::to_string(il)).c_str());
|
||||
ggml_set_input(ids);
|
||||
inp41->engram_ids.push_back(ids);
|
||||
|
||||
ggml_tensor * rows = ggml_get_rows(ctx0, model.per_layer_tok_embd, ids);
|
||||
rows = ggml_reshape_2d(ctx0, rows, hparams.ple_head_dim*hparams.ple_n_heads, n_tokens);
|
||||
cb(rows, "engram_rows", il);
|
||||
ggml_build_forward_expand(gf, rows);
|
||||
engram_rows.push_back(rows);
|
||||
}
|
||||
ggml_tensor * engram_keep = nullptr;
|
||||
if (hparams.ple_n_heads > 0 && ubatch.token && ubatch.embd) {
|
||||
engram_keep = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, 1, n_tokens);
|
||||
ggml_set_input(engram_keep);
|
||||
ggml_set_name(engram_keep, "engram_keep");
|
||||
inp41->engram_keep = engram_keep;
|
||||
}
|
||||
ggml_tensor * cand_pin = nullptr;
|
||||
if (hparams.dsv41_cand_src >= 0) {
|
||||
inp41->cand_csa = hparams.dsv4_compress_ratios[hparams.dsv41_cand_src] == 2;
|
||||
const auto & ci = inp41->cand_csa ? inp_dsv4->get_csa() : inp_dsv4->get_hca();
|
||||
if (ci.kq_mask) {
|
||||
const int64_t bs = hparams.dsv41_cand_block_size;
|
||||
cand_pin = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, (ci.kq_mask->ne[0] + bs - 1)/bs, ci.kq_mask->ne[1], 1, ci.kq_mask->ne[3]);
|
||||
ggml_set_input(cand_pin);
|
||||
ggml_set_name(cand_pin, "cand_pin");
|
||||
inp41->cand_pin = cand_pin;
|
||||
}
|
||||
}
|
||||
res->add_input(std::move(inp41));
|
||||
|
||||
const int64_t hc = hparams.dsv4_hc_mult;
|
||||
ggml_tensor * inpL = ggml_repeat_4d(ctx0, ggml_reshape_3d(ctx0, inp, n_embd, 1, n_tokens), n_embd, hc, n_tokens, 1);
|
||||
cb(inpL, "hc_init", -1);
|
||||
|
||||
// single-pass mHC: the pre-mix used to collapse the stream lags one sublayer
|
||||
ggml_tensor * pre_prev = nullptr;
|
||||
|
||||
for (int il = 0; il < n_layer; ++il) {
|
||||
const auto & layer = model.layers[il];
|
||||
|
||||
const int e = model.engram_index(il);
|
||||
if (e >= 0 && !engram_rows.empty()) {
|
||||
inpL = build_engram(model, inpL, engram_rows[e], engram_keep, il);
|
||||
cb(inpL, "engram_out", il);
|
||||
}
|
||||
|
||||
// DSpark target features: mean over the hc copies of the layer input
|
||||
if ((size_t) il < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[il]) {
|
||||
res->t_layer_inp[il] = build_hc_mean(inpL);
|
||||
cb(res->t_layer_inp[il], "layer_inp", il);
|
||||
ggml_build_forward_expand(gf, res->t_layer_inp[il]);
|
||||
}
|
||||
|
||||
ggml_tensor * residual = inpL;
|
||||
ggml_tensor * post = nullptr;
|
||||
ggml_tensor * comb = nullptr;
|
||||
|
||||
ggml_tensor * attn_pre = build_hc_mixes(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, &post, &comb, true, il);
|
||||
|
||||
// the first layer's attention collapses with the identity pre-mix: copy 0
|
||||
cur = pre_prev ? build_hc_pre(inpL, pre_prev, il) : ggml_cont(ctx0, ggml_view_2d(ctx0, inpL, n_embd, n_tokens, inpL->nb[2], 0));
|
||||
cb(cur, "hc_attn_pre", il);
|
||||
|
||||
cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(cur, "attn_norm", il);
|
||||
|
||||
cur = build_attention41(model, inp_dsv4, cur, inp_pos, cand_pin, il);
|
||||
|
||||
inpL = build_hc_post(cur, residual, post, comb, il);
|
||||
cb(inpL, "hc_attn_post", il);
|
||||
|
||||
residual = inpL;
|
||||
pre_prev = build_hc_mixes(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, &post, &comb, true, il);
|
||||
|
||||
cur = build_hc_pre(inpL, attn_pre, il);
|
||||
cb(cur, "hc_ffn_pre", il);
|
||||
|
||||
ggml_build_forward_expand(gf, residual);
|
||||
ggml_build_forward_expand(gf, post);
|
||||
ggml_build_forward_expand(gf, comb);
|
||||
|
||||
cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(cur, "ffn_norm", il);
|
||||
|
||||
// F32 router inputs: expert selection is near-tie sensitive
|
||||
ggml_tensor * logits = build_lora_mm(layer.ffn_gate_inp, cur);
|
||||
ggml_prec_set_acc(logits, GGML_PREC_F32);
|
||||
if (!ggml_is_quantized(layer.ffn_gate_inp->type)) {
|
||||
ggml_prec_set_src(logits, GGML_PREC_F32, 1);
|
||||
}
|
||||
cb(logits, "ffn_moe_logits", il);
|
||||
|
||||
// image (embedding) tokens use the vision router bias, as DeepSeek-V4
|
||||
ggml_tensor * exp_probs_b = layer.ffn_exp_probs_b;
|
||||
if (ubatch.embd != nullptr && layer.ffn_exp_probs_b_vl) {
|
||||
exp_probs_b = layer.ffn_exp_probs_b_vl;
|
||||
}
|
||||
|
||||
// the reference normalizes only when it routes to more than one expert
|
||||
ggml_tensor * moe_out = build_moe_ffn(cur,
|
||||
layer.ffn_gate_inp,
|
||||
layer.ffn_up_exps,
|
||||
layer.ffn_gate_exps,
|
||||
layer.ffn_down_exps,
|
||||
exp_probs_b,
|
||||
n_expert, hparams.n_expert_used(),
|
||||
LLM_FFN_SILU, hparams.expert_weights_norm && hparams.n_expert_used() > 1,
|
||||
hparams.expert_weights_scale,
|
||||
(llama_expert_gating_func_type) hparams.expert_gating_func,
|
||||
il, logits);
|
||||
cb(moe_out, "ffn_moe_out", il);
|
||||
|
||||
ggml_tensor * ffn_shexp = build_ffn(cur,
|
||||
layer.ffn_up_shexp, nullptr, nullptr,
|
||||
layer.ffn_gate_shexp, nullptr, nullptr,
|
||||
layer.ffn_down_shexp, nullptr, nullptr,
|
||||
nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
|
||||
cb(ffn_shexp, "ffn_shexp", il);
|
||||
|
||||
cur = ggml_add(ctx0, moe_out, ffn_shexp);
|
||||
cb(cur, "ffn_out", il);
|
||||
|
||||
inpL = build_hc_post(cur, residual, post, comb, il);
|
||||
inpL = build_cvec(inpL, il);
|
||||
cb(inpL, "l_out", il);
|
||||
}
|
||||
|
||||
// no hc_head: collapse with the last FFN's pre-mix
|
||||
if (inp_out_ids) {
|
||||
inpL = ggml_reshape_3d(ctx0, ggml_get_rows(ctx0, ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens), inp_out_ids), n_embd, hc, inp_out_ids->ne[0]);
|
||||
pre_prev = ggml_get_rows(ctx0, pre_prev, inp_out_ids);
|
||||
}
|
||||
cur = build_hc_pre(inpL, pre_prev, -1);
|
||||
cb(cur, "hc_out", -1);
|
||||
|
||||
cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
|
||||
cb(cur, "result_norm", -1);
|
||||
res->t_embd = cur;
|
||||
|
||||
cur = ggml_mul_mat(ctx0, model.output, cur);
|
||||
cb(cur, "result_output", -1);
|
||||
res->t_logits = cur;
|
||||
}
|
||||
+260
-3
@@ -81,6 +81,7 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
|
||||
ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters);
|
||||
ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps);
|
||||
ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, false);
|
||||
ml.get_key(LLM_KV_HYPER_CONNECTION_SINGLE_PASS, hparams.dsv4_hc_single_pass, false);
|
||||
|
||||
GGML_ASSERT(hparams.dsv4_o_group_count > 0); // avoid div by zero
|
||||
|
||||
@@ -94,6 +95,11 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
|
||||
}
|
||||
|
||||
GGML_ASSERT(hparams.n_swa > 0);
|
||||
if (hparams.dsv4_hc_single_pass) {
|
||||
// DeepSeek-V4.1 DSpark: block queries see the main-KV ring plus the block via the mask row of the
|
||||
// first query, so that row's SWA window must reach window_size positions before the anchor
|
||||
hparams.n_swa += 1;
|
||||
}
|
||||
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
|
||||
hparams.set_swa_pattern(0);
|
||||
for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
|
||||
@@ -257,9 +263,12 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
|
||||
const int64_t hc_dim = hc_mult * n_embd;
|
||||
const int64_t hc_mix_dim = (2 + hc_mult) * hc_mult;
|
||||
|
||||
hc_head_fn = create_tensor(tn(LLM_TENSOR_HC_HEAD_FN, "weight"), {hc_dim, hc_mult}, 0);
|
||||
hc_head_base = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE, "weight"), {hc_mult}, 0);
|
||||
hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 0);
|
||||
// single-pass mHC: collapse with the last FFN pre-mix, no hc_head
|
||||
if (!hparams.dsv4_hc_single_pass) {
|
||||
hc_head_fn = create_tensor(tn(LLM_TENSOR_HC_HEAD_FN, "weight"), {hc_dim, hc_mult}, 0);
|
||||
hc_head_base = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE, "weight"), {hc_mult}, 0);
|
||||
hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 0);
|
||||
}
|
||||
|
||||
for (int i = 0; i < n_layer; ++i) {
|
||||
auto & layer = layers[i];
|
||||
@@ -1280,12 +1289,260 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
|
||||
}
|
||||
}
|
||||
|
||||
// DeepSeek-V4.1 DSpark decoder (mtp.0..2), dual-mode by batch type like graph_dsv4:
|
||||
// * embd batch -> writes every stage's kv_norm(wkv(main_proj(target features))) into its window ring
|
||||
// * token batch -> [anchor, noise...] block through the stages, single-pass mHC, no per-head q norm
|
||||
llama_model_dflash::graph_dsv41::graph_dsv41(const llama_model & model, const llm_graph_params & params) :
|
||||
llama_model_deepseek41::graph(params) {
|
||||
const int64_t n_embd_inp = hparams.n_embd_inp_enc();
|
||||
const int64_t n_embd_head = hparams.n_embd_head_k();
|
||||
const int64_t n_embd_head_rope = hparams.n_rot();
|
||||
const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;
|
||||
|
||||
GGML_ASSERT(hparams.dsv4_hc_single_pass);
|
||||
|
||||
ggml_tensor * inp_pos = build_inp_pos();
|
||||
|
||||
llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
|
||||
|
||||
// ratio-0 layers: plain rope at rope_theta, no YaRN (original_seq_len 0)
|
||||
const auto rope = [&](ggml_tensor * t) {
|
||||
t = ggml_rope_ext(ctx0, t, inp_pos, nullptr, n_embd_head_rope, rope_type, 0,
|
||||
freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
return ggml_rope_set_offset(t, n_embd_head_nope);
|
||||
};
|
||||
|
||||
ASSERT_EMBD_OR_TOKEN(ubatch);
|
||||
if (ubatch.embd) {
|
||||
auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
|
||||
|
||||
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);
|
||||
ggml_set_input(inp->embd);
|
||||
|
||||
ggml_tensor * inp_target = inp->embd;
|
||||
cb(inp_target, "inp_target_features", -1);
|
||||
|
||||
res->add_input(std::move(inp));
|
||||
|
||||
ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s);
|
||||
inp_g = build_norm(inp_g, model.output_norm_enc, nullptr, LLM_NORM_RMS, -1);
|
||||
cb(inp_g, "main_x", -1);
|
||||
|
||||
for (int il = 0; il < n_layer; ++il) {
|
||||
const auto & layer = model.layers[il];
|
||||
|
||||
ggml_tensor * kv = build_lora_mm(layer.wkv, inp_g);
|
||||
kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);
|
||||
kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens);
|
||||
kv = rope(kv);
|
||||
cb(kv, "kv_injected", il);
|
||||
|
||||
if (inp_attn->self_k_rot_swa) {
|
||||
kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa);
|
||||
}
|
||||
ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il));
|
||||
}
|
||||
|
||||
res->t_embd = inp_g;
|
||||
|
||||
ggml_build_forward_expand(gf, inp_g);
|
||||
return;
|
||||
}
|
||||
|
||||
auto * tok_embd = model.tok_embd;
|
||||
if (tok_embd == nullptr) {
|
||||
GGML_ASSERT(cparams.ctx_other != nullptr);
|
||||
const auto * model_other = llama_get_model(cparams.ctx_other);
|
||||
GGML_ASSERT(model_other->tok_embd != nullptr && "DSpark decoder requires the target model's token embeddings");
|
||||
tok_embd = model_other->tok_embd;
|
||||
}
|
||||
|
||||
auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
|
||||
|
||||
inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
|
||||
ggml_set_input(inp->tokens);
|
||||
|
||||
ggml_tensor * inp_tokens = inp->tokens;
|
||||
|
||||
ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens);
|
||||
cb(inpL, "inp_noise_embd", -1);
|
||||
|
||||
res->add_input(std::move(inp));
|
||||
|
||||
const int64_t hc = hparams.dsv4_hc_mult;
|
||||
const int64_t n_groups = hparams.dsv4_o_group_count;
|
||||
const int64_t n_heads_group = n_head / n_groups;
|
||||
const int64_t o_lora_rank = hparams.dsv4_o_lora_rank;
|
||||
const int64_t o_group_dim = n_heads_group*n_embd_head;
|
||||
const float kq_scale = 1.0f/sqrtf(float(n_embd_head));
|
||||
|
||||
GGML_ASSERT(n_embd_head == n_embd_head_v);
|
||||
GGML_ASSERT(n_head % n_groups == 0);
|
||||
|
||||
inpL = ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens);
|
||||
inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1);
|
||||
cb(inpL, "hc_init", -1);
|
||||
|
||||
// window of a block = window of its first query, see n_swa
|
||||
{
|
||||
ggml_tensor * m = inp_attn->self_kq_mask_swa_cnv;
|
||||
const int64_t n_blocks = ubatch.n_seqs_unq;
|
||||
const int64_t bd = n_blocks > 0 ? n_tokens / n_blocks : 0;
|
||||
if (m && bd > 1 && n_tokens % n_blocks == 0 && m->ne[2] == 1 && m->ne[1] % bd == 0 && ggml_is_contiguous(m)) {
|
||||
ggml_tensor * m4 = ggml_reshape_4d(ctx0, m, m->ne[0], bd, m->ne[1]/bd, m->ne[3]);
|
||||
ggml_tensor * r0 = ggml_view_4d(ctx0, m4, m4->ne[0], 1, m4->ne[2], m4->ne[3], m4->nb[1], m4->nb[2], m4->nb[3], 0);
|
||||
m4 = ggml_repeat(ctx0, r0, m4);
|
||||
inp_attn->self_kq_mask_swa_cnv = ggml_reshape_4d(ctx0, m4, m->ne[0], m->ne[1], 1, m->ne[3]);
|
||||
cb(inp_attn->self_kq_mask_swa_cnv, "dspark_block_mask", -1);
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor * pre_prev = nullptr;
|
||||
|
||||
for (int il = 0; il < n_layer; ++il) {
|
||||
const auto & layer = model.layers[il];
|
||||
|
||||
ggml_tensor * residual = inpL;
|
||||
ggml_tensor * post = nullptr;
|
||||
ggml_tensor * comb = nullptr;
|
||||
|
||||
ggml_tensor * attn_pre = build_hc_mixes(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, &post, &comb, false, il);
|
||||
|
||||
ggml_tensor * cur;
|
||||
if (pre_prev == nullptr) {
|
||||
// identity pre-mix of the first stage: copy 0
|
||||
cur = ggml_cont(ctx0, ggml_view_2d(ctx0, inpL, n_embd, n_tokens, inpL->nb[2], 0));
|
||||
} else {
|
||||
cur = build_hc_pre(inpL, pre_prev, il);
|
||||
}
|
||||
cb(cur, "hc_attn_pre", il);
|
||||
|
||||
cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(cur, "attn_norm", il);
|
||||
|
||||
// non-causal draft context: block queries see the main-KV ring plus every block key
|
||||
{
|
||||
ggml_tensor * qr = build_lora_mm(layer.wq_a, cur);
|
||||
qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(qr, "qr_norm", il);
|
||||
|
||||
ggml_tensor * q = build_lora_mm(layer.wq_b, qr);
|
||||
q = ggml_reshape_3d(ctx0, q, n_embd_head, n_head, n_tokens);
|
||||
q = rope(q);
|
||||
cb(q, "q", il);
|
||||
|
||||
ggml_tensor * kv = build_lora_mm(layer.wkv, cur);
|
||||
kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);
|
||||
kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens);
|
||||
kv = rope(kv);
|
||||
cb(kv, "kv", il);
|
||||
|
||||
ggml_tensor * out = build_attn(inp_attn,
|
||||
nullptr, nullptr, nullptr,
|
||||
q, kv, kv,
|
||||
nullptr, layer.attn_sinks, nullptr,
|
||||
kq_scale, il);
|
||||
cb(out, "attn_raw", il);
|
||||
|
||||
out = ggml_reshape_3d(ctx0, out, n_embd_head, n_head, n_tokens);
|
||||
out = ggml_rope_ext_back(ctx0, out, inp_pos, nullptr, n_embd_head_rope, rope_type, 0,
|
||||
freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
out = ggml_rope_set_offset(out, n_embd_head_nope);
|
||||
cb(out, "attn_derope", il);
|
||||
|
||||
out = ggml_reshape_3d(ctx0, out, o_group_dim, n_groups, n_tokens);
|
||||
out = ggml_permute(ctx0, out, 0, 2, 1, 3);
|
||||
ggml_tensor * oa = ggml_mul_mat(ctx0, layer.wo_a, out);
|
||||
cb(oa, "attn_wo_a", il);
|
||||
oa = ggml_permute(ctx0, oa, 0, 2, 1, 3);
|
||||
oa = ggml_cont_2d(ctx0, oa, o_lora_rank*n_groups, n_tokens);
|
||||
|
||||
cur = build_lora_mm(layer.wo_b, oa);
|
||||
cb(cur, "attn_out", il);
|
||||
}
|
||||
|
||||
inpL = build_hc_post(cur, residual, post, comb, il);
|
||||
cb(inpL, "hc_attn_post", il);
|
||||
|
||||
residual = inpL;
|
||||
ggml_tensor * ffn_pre = build_hc_mixes(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, &post, &comb, false, il);
|
||||
|
||||
cur = build_hc_pre(inpL, attn_pre, il);
|
||||
cb(cur, "hc_ffn_pre", il);
|
||||
|
||||
cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(cur, "ffn_norm", il);
|
||||
|
||||
ggml_tensor * moe_out = build_moe_ffn(cur,
|
||||
layer.ffn_gate_inp,
|
||||
layer.ffn_up_exps,
|
||||
layer.ffn_gate_exps,
|
||||
layer.ffn_down_exps,
|
||||
layer.ffn_exp_probs_b,
|
||||
n_expert, hparams.n_expert_used(),
|
||||
LLM_FFN_SILU, hparams.expert_weights_norm,
|
||||
hparams.expert_weights_scale,
|
||||
(llama_expert_gating_func_type) hparams.expert_gating_func,
|
||||
il);
|
||||
cb(moe_out, "ffn_moe_out", il);
|
||||
|
||||
ggml_tensor * ffn_shexp = build_ffn(cur,
|
||||
layer.ffn_up_shexp, nullptr, nullptr,
|
||||
layer.ffn_gate_shexp, nullptr, nullptr,
|
||||
layer.ffn_down_shexp, nullptr, nullptr,
|
||||
nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
|
||||
cb(ffn_shexp, "ffn_shexp", il);
|
||||
|
||||
cur = ggml_add(ctx0, moe_out, ffn_shexp);
|
||||
cb(cur, "ffn_out", il);
|
||||
|
||||
inpL = build_hc_post(cur, residual, post, comb, il);
|
||||
|
||||
ggml_tensor * l_out = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
|
||||
cb(l_out, "l_out", il);
|
||||
ggml_build_forward_expand(gf, l_out);
|
||||
|
||||
pre_prev = ffn_pre;
|
||||
}
|
||||
|
||||
// no hc_head: collapse with the last FFN's pre-mix
|
||||
ggml_tensor * cur = build_hc_pre(inpL, pre_prev, -1);
|
||||
cb(cur, "hc_out", -1);
|
||||
|
||||
// confidence head takes the un-normed collapsed state
|
||||
res->t_embd = cur;
|
||||
|
||||
cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
|
||||
cb(cur, "result_norm", -1);
|
||||
|
||||
auto * output = model.output;
|
||||
auto * output_s = model.output_s;
|
||||
if (output == nullptr) {
|
||||
GGML_ASSERT(cparams.ctx_other != nullptr);
|
||||
const auto * model_other = llama_get_model(cparams.ctx_other);
|
||||
GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection");
|
||||
output = model_other->output;
|
||||
output_s = model_other->output_s;
|
||||
}
|
||||
|
||||
cur = build_lora_mm(output, cur, output_s);
|
||||
cb(cur, "result_output", -1);
|
||||
res->t_logits = cur;
|
||||
|
||||
if (model.dspark_markov_w1) {
|
||||
build_dspark_markov_head(*this, model, inp_tokens);
|
||||
}
|
||||
}
|
||||
|
||||
std::unique_ptr<llm_graph_context> llama_model_dflash::build_arch_graph(const llm_graph_params & params) const {
|
||||
switch (params.gtype) {
|
||||
case LLM_GRAPH_TYPE_ENCODER:
|
||||
return std::make_unique<graph<true>>(*this, params);
|
||||
case LLM_GRAPH_TYPE_DEFAULT:
|
||||
case LLM_GRAPH_TYPE_DECODER:
|
||||
if (hparams.dsv4_hc_mult > 0 && hparams.dsv4_hc_single_pass) {
|
||||
return std::make_unique<graph_dsv41>(*this, params);
|
||||
}
|
||||
if (hparams.dsv4_hc_mult > 0) {
|
||||
return std::make_unique<graph_dsv4>(*this, params);
|
||||
}
|
||||
|
||||
@@ -550,9 +550,9 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
|
||||
if ((uint32_t) il < hparams.n_layer_dense_lead) {
|
||||
// leading dense FFN (swigluoai)
|
||||
cur = build_ffn(cur,
|
||||
model.layers[il].ffn_up, NULL, NULL,
|
||||
model.layers[il].ffn_gate, NULL, NULL,
|
||||
model.layers[il].ffn_down, NULL, NULL,
|
||||
model.layers[il].ffn_up, NULL, model.layers[il].ffn_up_s,
|
||||
model.layers[il].ffn_gate, NULL, model.layers[il].ffn_gate_s,
|
||||
model.layers[il].ffn_down, NULL, model.layers[il].ffn_down_s,
|
||||
NULL,
|
||||
LLM_FFN_SWIGLU_OAI_MOE, LLM_FFN_PAR, il);
|
||||
cb(cur, "ffn_out", il);
|
||||
@@ -568,14 +568,17 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
|
||||
LLM_FFN_SWIGLU_OAI_MOE, hparams.expert_weights_norm,
|
||||
hparams.expert_weights_scale,
|
||||
(llama_expert_gating_func_type) hparams.expert_gating_func,
|
||||
il);
|
||||
il, nullptr, nullptr,
|
||||
model.layers[il].ffn_up_exps_s,
|
||||
model.layers[il].ffn_gate_exps_s,
|
||||
model.layers[il].ffn_down_exps_s);
|
||||
cb(moe_out, "ffn_moe_out", il);
|
||||
|
||||
// shared expert (swigluoai)
|
||||
ggml_tensor * ffn_shexp = build_ffn(cur,
|
||||
model.layers[il].ffn_up_shexp, NULL, NULL,
|
||||
model.layers[il].ffn_gate_shexp, NULL, NULL,
|
||||
model.layers[il].ffn_down_shexp, NULL, NULL,
|
||||
model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s,
|
||||
model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s,
|
||||
model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s,
|
||||
NULL,
|
||||
LLM_FFN_SWIGLU_OAI_MOE, LLM_FFN_PAR, il);
|
||||
cb(ffn_shexp, "ffn_shexp", il);
|
||||
|
||||
@@ -1331,6 +1331,15 @@ struct llama_model_dots3note : public llama_model_base {
|
||||
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
|
||||
};
|
||||
|
||||
// YaRN attention factor that cancels ggml's rope mscale (the DeepSeek-V4 reference applies none)
|
||||
static inline float dsv4_rope_attn_factor(float freq_scale, float ext_factor) {
|
||||
if (ext_factor == 0.0f) {
|
||||
return 1.0f;
|
||||
}
|
||||
|
||||
return 1.0f / (1.0f + 0.1f*logf(1.0f/freq_scale));
|
||||
}
|
||||
|
||||
struct llama_model_deepseek4 : public llama_model_base {
|
||||
llama_model_deepseek4(const struct llama_model_params & params) : llama_model_base(params) {}
|
||||
void load_arch_hparams(llama_model_loader & ml) override;
|
||||
@@ -1359,6 +1368,17 @@ struct llama_model_deepseek4 : public llama_model_base {
|
||||
ggml_tensor ** comb,
|
||||
int il) const;
|
||||
|
||||
// returns the pre weights and fills post / comb
|
||||
ggml_tensor * build_hc_mixes(
|
||||
ggml_tensor * x,
|
||||
ggml_tensor * hc_fn,
|
||||
ggml_tensor * hc_scale,
|
||||
ggml_tensor * hc_base,
|
||||
ggml_tensor ** post,
|
||||
ggml_tensor ** comb,
|
||||
bool f32_src,
|
||||
int il) const;
|
||||
|
||||
ggml_tensor * build_hc_post(
|
||||
ggml_tensor * x,
|
||||
ggml_tensor * residual,
|
||||
@@ -1473,6 +1493,57 @@ struct llama_model_deepseek4 : public llama_model_base {
|
||||
};
|
||||
|
||||
|
||||
struct llama_model_deepseek41 : public llama_model_deepseek4 {
|
||||
llama_model_deepseek41(const struct llama_model_params & params) : llama_model_deepseek4(params) {}
|
||||
void load_arch_hparams(llama_model_loader & ml) override;
|
||||
void load_arch_tensors(llama_model_loader & ml) override;
|
||||
|
||||
// Engram: the hash constants are hparams.ple_* (one module per Engram layer), hashed on the host (ggml has no
|
||||
// int64 mul/xor/mod); the token map and its pad id are V4.1 only
|
||||
uint32_t engram_pad_id = 0;
|
||||
std::vector<int32_t> engram_token_map;
|
||||
|
||||
int engram_index(int il) const;
|
||||
|
||||
void engram_hash(int e, const llama_ubatch & ubatch, const std::vector<llama_token> & prev, std::vector<int32_t> & out) const;
|
||||
|
||||
struct graph : public llama_model_deepseek4::graph {
|
||||
graph(const llama_model & model, const llm_graph_params & params);
|
||||
// helpers only (no graph built): base of graph_dsv41
|
||||
explicit graph(const llm_graph_params & params) : llama_model_deepseek4::graph(params) {}
|
||||
|
||||
ggml_tensor * build_engram(
|
||||
const llama_model_deepseek41 & model,
|
||||
ggml_tensor * x,
|
||||
ggml_tensor * rows,
|
||||
ggml_tensor * keep,
|
||||
int il) const;
|
||||
|
||||
ggml_tensor * build_attention41(
|
||||
const llama_model_deepseek41 & model,
|
||||
llm_graph_input_dsv4 * inp_dsv4,
|
||||
ggml_tensor * cur,
|
||||
ggml_tensor * inp_pos,
|
||||
ggml_tensor * cand_pin,
|
||||
int il);
|
||||
|
||||
std::vector<ggml_tensor *> topk_of; // top-k of each index source layer
|
||||
std::vector<ggml_tensor *> comp_row_of; // [n_layer] new compressed rows of each kv source (for mirrors)
|
||||
ggml_tensor * cand_sel = nullptr; // candidate blocks of the candidate source layer
|
||||
int64_t cand_n_comp = 0; // compressed rows the candidate source scored
|
||||
|
||||
// sparse path (LLAMA_DSV41_SPARSE=0 disables): fused indexer, FA bounded to window + top-k, decode gather
|
||||
bool sparse = true;
|
||||
// gather path, shared by layers with the same top-k: selected K rows per (storage layer, top-k source),
|
||||
// the sorted top-k and the mask per top-k source (one input copy per split)
|
||||
std::map<std::pair<int32_t, int32_t>, ggml_tensor *> gather_k_of;
|
||||
std::map<int32_t, ggml_tensor *> gather_sel_of;
|
||||
std::map<int32_t, ggml_tensor *> gather_mask_of;
|
||||
};
|
||||
|
||||
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
|
||||
};
|
||||
|
||||
struct llama_model_deepseek2ocr : public llama_model_base {
|
||||
llama_model_deepseek2ocr(const struct llama_model_params & params) : llama_model_base(params) {}
|
||||
void load_arch_hparams(llama_model_loader & ml) override;
|
||||
@@ -1532,6 +1603,11 @@ struct llama_model_dflash : public llama_model_base {
|
||||
graph_dsv4(const llama_model & model, const llm_graph_params & params);
|
||||
};
|
||||
|
||||
// DeepSeek-V4.1 DSpark: single-pass mHC (lagged pre-mix, no hc_head), no per-head q norm
|
||||
struct graph_dsv41 : public llama_model_deepseek41::graph {
|
||||
graph_dsv41(const llama_model & model, const llm_graph_params & params);
|
||||
};
|
||||
|
||||
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
|
||||
};
|
||||
|
||||
|
||||
@@ -115,6 +115,9 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
|
||||
|
||||
hparams.ple_n_heads = (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram;
|
||||
hparams.ple_head_dim = hparams.n_embd_per_layer;
|
||||
hparams.ple_embd_token = hparams.ple_image_token_id != 0
|
||||
? (llama_token) hparams.ple_image_token_id
|
||||
: (llama_token) hparams.ple_eos_token_id;
|
||||
if (hparams.ple_ngram_size < 2 || hparams.ple_ngram_size > LLAMA_MAX_PLE_NGRAM) {
|
||||
throw std::runtime_error(format("PLE n-gram size %u is out of range", hparams.ple_ngram_size));
|
||||
}
|
||||
|
||||
@@ -49,6 +49,8 @@ deepseek32 ,0 ,any ,NORM+MUL+ADD , 2
|
||||
deepseek32 ,0 ,any ,RMS_NORM+MUL , 9
|
||||
deepseek4 ,0 ,any ,MUL+ADD , 8
|
||||
deepseek4 ,0 ,any ,RMS_NORM+MUL , 20
|
||||
deepseek41 ,0 ,any ,MUL+ADD , 12
|
||||
deepseek41 ,0 ,any ,RMS_NORM+MUL , 29
|
||||
dots1 ,0 ,any ,ADD+ADD , 1
|
||||
dots1 ,0 ,any ,MUL+ADD , 2
|
||||
dots1 ,0 ,any ,RMS_NORM+MUL , 9
|
||||
|
||||
|
+109
-40
@@ -1194,6 +1194,29 @@ static bool op_filter_entry_matches(std::string_view entry, std::string_view op_
|
||||
return std::regex_search(op_name.data(), op_name.data() + op_name.size(), re);
|
||||
}
|
||||
|
||||
// TODO: [TAG_MUL_MAT_ID_SKIP] see https://github.com/ggml-org/llama.cpp/pull/26631
|
||||
static bool backend_lacks_skipped_ids(ggml_backend_t backend, ggml_cgraph * gf) {
|
||||
const char * name = ggml_backend_reg_name(ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)));
|
||||
if (strcmp(name, "CUDA") != 0 && strcmp(name, "ROCm") != 0 && strcmp(name, "MUSA") != 0) {
|
||||
return false;
|
||||
}
|
||||
for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
|
||||
const ggml_tensor * node = ggml_graph_node(gf, i);
|
||||
if (node->op != GGML_OP_MUL_MAT_ID && node->op != GGML_OP_ADD_ID) {
|
||||
continue;
|
||||
}
|
||||
const ggml_tensor * ids = node->src[2];
|
||||
std::vector<int32_t> row(ids->ne[0]);
|
||||
for (int64_t i1 = 0; i1 < ids->ne[1]; i1++) {
|
||||
ggml_backend_tensor_get(ids, row.data(), i1*ids->nb[1], ids->ne[0]*sizeof(int32_t));
|
||||
if (std::find(row.begin(), row.end(), -1) != row.end()) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
struct test_case {
|
||||
virtual ~test_case() {}
|
||||
|
||||
@@ -1475,6 +1498,13 @@ struct test_case {
|
||||
initialize_tensors(ctx_weights.get());
|
||||
}
|
||||
|
||||
// TODO: [TAG_MUL_MAT_ID_SKIP] @ngxson : remove this once CUDA implements mul_mat_id with -1 = skipped
|
||||
if (backend_lacks_skipped_ids(backend1, gf)) {
|
||||
test_result result(ggml_backend_name(backend1), current_op_name, vars(), "test", false, false, "not supported");
|
||||
print_test_result_locked(output_printer, result);
|
||||
return test_status_t::NOT_SUPPORTED;
|
||||
}
|
||||
|
||||
// compare
|
||||
struct callback_userdata {
|
||||
bool ok;
|
||||
@@ -3564,6 +3594,16 @@ struct test_bin_bcast : public test_case {
|
||||
}
|
||||
};
|
||||
|
||||
// set some ids to -1, so that MUL_MAT_ID and ADD_ID skip those slots
|
||||
// every 4th row skips all of its slots, the other rows skip only some of them
|
||||
static void mask_skipped_ids(std::vector<int32_t> & ids, int64_t row) {
|
||||
for (size_t i = 0; i < ids.size(); i++) {
|
||||
if (row % 4 == 1 || (row + (int64_t) i) % 3 == 0) {
|
||||
ids[i] = -1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// GGML_OP_ADD_ID
|
||||
struct test_add_id : public test_case {
|
||||
const ggml_type type_a;
|
||||
@@ -3572,9 +3612,10 @@ struct test_add_id : public test_case {
|
||||
const int64_t n_experts;
|
||||
const int64_t n_experts_used;
|
||||
const int64_t n_token;
|
||||
const bool skip; // set some ids to a negative value
|
||||
|
||||
std::string vars() override {
|
||||
return VARS_TO_STR6(type_a, type_b, n_embd, n_experts, n_experts_used, n_token);
|
||||
return VARS_TO_STR7(type_a, type_b, n_embd, n_experts, n_experts_used, n_token, skip);
|
||||
}
|
||||
|
||||
size_t op_size(ggml_tensor * t) override {
|
||||
@@ -3586,9 +3627,10 @@ struct test_add_id : public test_case {
|
||||
int64_t n_embd = 128,
|
||||
int64_t n_experts = 16,
|
||||
int64_t n_experts_used = 8,
|
||||
int64_t n_token = 10)
|
||||
int64_t n_token = 10,
|
||||
bool skip = false)
|
||||
: type_a(type_a), type_b(type_b), n_embd(n_embd),
|
||||
n_experts(n_experts), n_experts_used(n_experts_used), n_token(n_token) {}
|
||||
n_experts(n_experts), n_experts_used(n_experts_used), n_token(n_token), skip(skip) {}
|
||||
|
||||
ggml_tensor * build_graph(ggml_context * ctx) override {
|
||||
ggml_tensor * a = ggml_new_tensor_3d(ctx, type_a, n_embd, n_experts_used, n_token);
|
||||
@@ -3617,6 +3659,9 @@ struct test_add_id : public test_case {
|
||||
data[i] = i % n_experts;
|
||||
}
|
||||
std::shuffle(data.begin(), data.end(), rng);
|
||||
if (skip) {
|
||||
mask_skipped_ids(data, r);
|
||||
}
|
||||
ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t));
|
||||
}
|
||||
} else {
|
||||
@@ -5550,7 +5595,8 @@ struct test_mul_mat_w4a4 : public test_mul_mat {
|
||||
}
|
||||
};
|
||||
|
||||
static void init_mul_mat_id_ids(ggml_context * ctx, int n_mats) {
|
||||
static void init_mul_mat_id_ids(ggml_context * ctx, int n_mats, bool skip = false) {
|
||||
|
||||
std::random_device rd;
|
||||
std::default_random_engine rng(rd());
|
||||
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
|
||||
@@ -5563,23 +5609,21 @@ static void init_mul_mat_id_ids(ggml_context * ctx, int n_mats) {
|
||||
data[i] = i % n_mats;
|
||||
}
|
||||
std::shuffle(data.begin(), data.end(), rng);
|
||||
if (skip) {
|
||||
mask_skipped_ids(data, r);
|
||||
}
|
||||
ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats, float amax = 1.0f) {
|
||||
static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats, bool skip = false) {
|
||||
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
|
||||
if (t->type == GGML_TYPE_I32) {
|
||||
continue;
|
||||
}
|
||||
if (amax != 1.0f && t->type == GGML_TYPE_F32) {
|
||||
init_tensor_uniform(t, -amax, amax);
|
||||
} else {
|
||||
if (t->type != GGML_TYPE_I32) {
|
||||
init_tensor_uniform(t);
|
||||
}
|
||||
}
|
||||
init_mul_mat_id_ids(ctx, n_mats);
|
||||
init_mul_mat_id_ids(ctx, n_mats, skip);
|
||||
}
|
||||
|
||||
// GGML_OP_MUL_MAT_ID
|
||||
@@ -5592,12 +5636,12 @@ struct test_mul_mat_id : public test_case {
|
||||
const int64_t m;
|
||||
const int64_t n;
|
||||
const int64_t k;
|
||||
const float amax; // magnitude of src1
|
||||
const int64_t m_v; // rows of as in memory, the experts of as are strided for m_v > m, no view for m_v == 0
|
||||
const bool ids_offset;
|
||||
const bool skip; // set some ids to a negative value
|
||||
|
||||
std::string vars() override {
|
||||
return VARS_TO_STR11(type_a, type_b, n_mats, n_used, b, m, n, k, amax, m_v, ids_offset);
|
||||
return VARS_TO_STR11(type_a, type_b, n_mats, n_used, b, m, n, k, m_v, ids_offset, skip);
|
||||
}
|
||||
|
||||
double max_nmse_err() override {
|
||||
@@ -5622,9 +5666,9 @@ struct test_mul_mat_id : public test_case {
|
||||
test_mul_mat_id(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32,
|
||||
int n_mats = 8, int n_used = 2, bool b = false,
|
||||
int64_t m = 32, int64_t n = 32, int64_t k = 32,
|
||||
float amax = 1.0f, int64_t m_v = 0, bool ids_offset = false)
|
||||
int64_t m_v = 0, bool ids_offset = false, bool skip = false)
|
||||
: type_a(type_a), type_b(type_b), n_mats(n_mats), n_used(n_used), b(b),
|
||||
m(m), n(n), k(k), amax(amax), m_v(m_v), ids_offset(ids_offset) {
|
||||
m(m), n(n), k(k), m_v(m_v), ids_offset(ids_offset), skip(skip) {
|
||||
GGML_ASSERT(n_used <= n_mats);
|
||||
GGML_ASSERT(m_v == 0 || m_v > m);
|
||||
}
|
||||
@@ -5650,20 +5694,15 @@ struct test_mul_mat_id : public test_case {
|
||||
ggml_tensor * out = ggml_mul_mat_id(ctx, as, b, ids);
|
||||
ggml_set_name(out, "out");
|
||||
|
||||
if (amax > 65504.0f) {
|
||||
// src1 exceeds F16 range
|
||||
ggml_prec_set_src(out, GGML_PREC_F32, 1);
|
||||
}
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
void initialize_tensors(ggml_context * ctx) override {
|
||||
init_mul_mat_id_tensors(ctx, n_mats, amax);
|
||||
init_mul_mat_id_tensors(ctx, n_mats, skip);
|
||||
}
|
||||
|
||||
void reinit_perf_iter(ggml_context * ctx) override {
|
||||
init_mul_mat_id_ids(ctx, n_mats);
|
||||
init_mul_mat_id_ids(ctx, n_mats, skip);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -5738,9 +5777,10 @@ struct test_mul_mat_id_fusion : public test_case {
|
||||
const int64_t k;
|
||||
const uint32_t o; // number of outputs
|
||||
const bool mul;
|
||||
const bool skip; // set some ids to a negative value
|
||||
|
||||
std::string vars() override {
|
||||
return VARS_TO_STR10(type_a, type_b, n_mats, n_used, b, m, n, k, o, mul);
|
||||
return VARS_TO_STR11(type_a, type_b, n_mats, n_used, b, m, n, k, o, mul, skip);
|
||||
}
|
||||
|
||||
double max_nmse_err() override {
|
||||
@@ -5754,9 +5794,9 @@ struct test_mul_mat_id_fusion : public test_case {
|
||||
|
||||
test_mul_mat_id_fusion(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32,
|
||||
int n_mats = 8, int n_used = 2, bool b = false,
|
||||
int64_t m = 32, int64_t n = 32, int64_t k = 32, uint32_t o = 1, bool mul = false)
|
||||
int64_t m = 32, int64_t n = 32, int64_t k = 32, uint32_t o = 1, bool mul = false, bool skip = false)
|
||||
: type_a(type_a), type_b(type_b), n_mats(n_mats), n_used(n_used), b(b),
|
||||
m(m), n(n), k(k), o(o), mul(mul) {
|
||||
m(m), n(n), k(k), o(o), mul(mul), skip(skip) {
|
||||
GGML_ASSERT(n_used <= n_mats);
|
||||
}
|
||||
|
||||
@@ -5796,7 +5836,7 @@ struct test_mul_mat_id_fusion : public test_case {
|
||||
}
|
||||
|
||||
void initialize_tensors(ggml_context * ctx) override {
|
||||
init_mul_mat_id_tensors(ctx, n_mats);
|
||||
init_mul_mat_id_tensors(ctx, n_mats, skip);
|
||||
}
|
||||
|
||||
bool run_whole_graph() override { return true; }
|
||||
@@ -7754,19 +7794,20 @@ struct test_mul_mat_vec_fusion : public test_case {
|
||||
const bool with_gate;
|
||||
const bool with_lane_scale;
|
||||
std::array<int64_t, 2> batch_dims;
|
||||
const bool skip; // set some ids to a negative value (only for use_id)
|
||||
|
||||
test_mul_mat_vec_fusion(ggml_type type, ggml_glu_op op, int64_t m, int64_t n, int64_t k,
|
||||
bool use_id = false, int n_mats = 1, int n_used = 1, bool b = false, bool with_bias = false, bool with_gate = true,
|
||||
bool with_lane_scale = false, std::array<int64_t, 2> batch_dims = {4, 2})
|
||||
bool with_lane_scale = false, std::array<int64_t, 2> batch_dims = {4, 2}, bool skip = false)
|
||||
: type(type), glu_op(op), m(m), n(n), k(k), use_id(use_id), n_mats(n_mats), n_used(n_used), b(b), with_bias(with_bias),
|
||||
with_gate(with_gate), with_lane_scale(with_lane_scale), batch_dims(batch_dims) {
|
||||
with_gate(with_gate), with_lane_scale(with_lane_scale), batch_dims(batch_dims), skip(skip) {
|
||||
if (use_id) {
|
||||
GGML_ASSERT(n_used <= n_mats);
|
||||
}
|
||||
}
|
||||
|
||||
std::string vars() override {
|
||||
return VARS_TO_STR13(type, glu_op, m, n, k, use_id, n_mats, n_used, b, with_bias, with_gate, with_lane_scale, batch_dims);
|
||||
return VARS_TO_STR14(type, glu_op, m, n, k, use_id, n_mats, n_used, b, with_bias, with_gate, with_lane_scale, batch_dims, skip);
|
||||
}
|
||||
|
||||
std::string op_desc(ggml_tensor * t) override {
|
||||
@@ -7917,7 +7958,7 @@ struct test_mul_mat_vec_fusion : public test_case {
|
||||
init_tensor_uniform(t);
|
||||
}
|
||||
} else {
|
||||
init_mul_mat_id_tensors(ctx, n_mats);
|
||||
init_mul_mat_id_tensors(ctx, n_mats, skip);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10242,6 +10283,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
|
||||
// im2col 1D
|
||||
test_cases.emplace_back(new test_im2col(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, {3000, 128, 1, 1}, {3, 128, 1280, 1}, 1, 0, 1, 0, 1, 0, false));
|
||||
// deepseek4v / deepseek41v aligner: 3x3 stride-3 unfold (F.unfold == im2col)
|
||||
test_cases.emplace_back(new test_im2col(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, {48, 33, 1024, 1}, {3, 3, 1024, 1}, 3, 3, 0, 0, 1, 1, true));
|
||||
test_cases.emplace_back(new test_im2col(GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_F32, {3000, 128, 1, 1}, {3, 128, 1280, 1}, 1, 0, 1, 0, 1, 0, false));
|
||||
test_cases.emplace_back(new test_im2col(GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_F16, {3000, 128, 1, 1}, {3, 128, 1280, 1}, 1, 0, 1, 0, 1, 0, false));
|
||||
test_cases.emplace_back(new test_im2col(GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_F16, {3000, 384, 1, 1}, {3, 384, 384, 1}, 1, 0, 1, 0, 1, 0, false));
|
||||
@@ -11285,8 +11328,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 8, 1, 64, {8, 1}, {1, 1}, {0, 1, 2, 3}, 0, 1, false, 16, 16));
|
||||
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 8, 16, 64, {8, 1}, {1, 1}, {0, 1, 2, 3}, 0, 1, false, 16, 16));
|
||||
// as is a view whose experts are strided by more rows than it uses
|
||||
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 4, 2, false, 8, 1, 64, 1.0f, 64));
|
||||
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 4, 2, false, 8, 16, 64, 1.0f, 64));
|
||||
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 4, 2, false, 8, 1, 64, 64));
|
||||
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 4, 2, false, 8, 16, 64, 64));
|
||||
|
||||
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 576, 512, 576, {1,1}, {1,1}));
|
||||
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 1, 2048, 8192, {1, 1}, {1, 1}));
|
||||
@@ -11400,7 +11443,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
}
|
||||
|
||||
for (ggml_type ta : {GGML_TYPE_F16, GGML_TYPE_Q4_0}) {
|
||||
test_cases.emplace_back(new test_mul_mat_id_dup(ta, GGML_TYPE_F32, 1025, 10, false, 64, 33, 256, 1.0f, 0, true));
|
||||
test_cases.emplace_back(new test_mul_mat_id_dup(ta, GGML_TYPE_F32, 1025, 10, false, 64, 33, 256, 0, true));
|
||||
test_cases.emplace_back(new test_mul_mat_id_dup(ta, GGML_TYPE_F32, 2050, 10, true, 64, 33, 256));
|
||||
}
|
||||
#endif
|
||||
@@ -11415,12 +11458,6 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 4, 4, false, 16, 10, 256));
|
||||
}
|
||||
|
||||
// test src1 f16 overflow
|
||||
for (int n : {16, 32, 64}) {
|
||||
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q4_K, GGML_TYPE_F32, 128, 4, false, 4096, n, 2048, 1e5f));
|
||||
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q8_0, GGML_TYPE_F32, 8, 2, false, 512, n, 256, 1e5f));
|
||||
}
|
||||
|
||||
for (ggml_type type_a : base_types) {
|
||||
for (ggml_type type_b : {GGML_TYPE_F32 /*, GGML_TYPE_F16 */}) {
|
||||
for (int n_mats : {4, 8}) {
|
||||
@@ -11453,6 +11490,21 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
}
|
||||
}
|
||||
|
||||
// skipped slots (negative ids)
|
||||
for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q4_0, GGML_TYPE_Q4_K, GGML_TYPE_Q8_0, GGML_TYPE_MXFP4}) {
|
||||
for (int n_used : {1, 2, 4}) {
|
||||
for (bool b : {false, true}) {
|
||||
for (int n : {1, 4, 33, 129}) {
|
||||
test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 8, n_used, b, 512, n, 256, 0, false, true));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// n_used == n_mats, so every expert can be skipped
|
||||
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 4, 4, false, 512, 64, 256, 0, false, true));
|
||||
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 1, 1, false, 8, 16, 1, 0, false, true));
|
||||
test_cases.emplace_back(new test_mul_mat_id_fusion(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, false, 32, 32, 32, 3, false, true));
|
||||
|
||||
for (int bs : {1, 4, 512}) {
|
||||
for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q4_0, GGML_TYPE_Q4_K}) {
|
||||
for (ggml_type type_b : {GGML_TYPE_F32}) {
|
||||
@@ -11500,6 +11552,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
for (int n_embd : {32, 129}) {
|
||||
for (int n_token : {1, 32, 129}) {
|
||||
test_cases.emplace_back(new test_add_id(type_a, type_b, n_embd, n_mats, n_used, n_token));
|
||||
test_cases.emplace_back(new test_add_id(type_a, type_b, n_embd, n_mats, n_used, n_token, true));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -12000,6 +12053,17 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
}
|
||||
}
|
||||
|
||||
// KV-head counts that do not divide the device count
|
||||
for (int hs : { 128, 256, 512, }) {
|
||||
for (int nh : { 1, 2, 6, }) {
|
||||
for (int nr2 : { 4, 8, }) {
|
||||
for (int nb : { 75, 1024, }) {
|
||||
test_cases.emplace_back(new test_flash_attn_ext(hs, hs, nh, {nr2, 1}, 1024, nb, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (int hsk : { 40, 64, 72, 80, 96, 128, 192, 256, 320, 512, 576 }) {
|
||||
for (int hsv : { 40, 64, 72, 80, 96, 128, 192, 256, 512 }) {
|
||||
if (hsk != 96 && hsk != 192 && hsk != 320 && hsk != 576 && hsk != hsv) continue;
|
||||
@@ -12238,6 +12302,11 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, m_batch, 32, 256,
|
||||
use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1}));
|
||||
}
|
||||
// the lane scale does a GET_ROWS on the ids, which does not accept -1
|
||||
if (use_id && !with_lane_scale) {
|
||||
test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, 8, 32, 256,
|
||||
use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1}, true));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -149,6 +149,21 @@ static common_chat_templates_ptr read_templates(const std::string & path) {
|
||||
return common_chat_templates_ptr(common_chat_templates_init(/* model= */ nullptr, read_file(path)));
|
||||
}
|
||||
|
||||
// DeepSeek V4.1 stand-in: the real template ships in the GGUF; this keeps only the turn markers,
|
||||
// the generation prompt, the V4.1 DSML block tag used for detection and the extra keys it reads
|
||||
static const char * DEEPSEEK_V4_1_MIN_TEMPLATE = R"({%- for m in messages -%}
|
||||
{%- if m.role == 'system' -%}<|System|>{{ m.content }}
|
||||
{%- elif m.role == 'user' -%}<|User|>
|
||||
{%- if m.content is string -%}{{ m.content }}{%- else -%}{%- for b in m.content -%}{%- if not loop.first -%}{{ '\n\n' }}{%- endif -%}{{ b.text }}{%- endfor -%}{%- endif -%}
|
||||
{%- if m.task -%}<|{{ m.task }}|>{%- endif -%}
|
||||
{%- elif m.role == 'assistant' -%}<|Assistant|></think>{{ m.content }}
|
||||
{%- for tc in m.tool_calls or [] -%}<|DSML| calls><|DSML| invoke name="{% if tc.namespace %}{{ tc.namespace }}::{% endif %}{{ tc.function.name }}"></|DSML| invoke></|DSML| calls>{%- endfor -%}
|
||||
<|end▁of▁sentence|>
|
||||
{%- endif -%}
|
||||
{%- endfor -%}
|
||||
{%- if reasoning_effort is defined and reasoning_effort is not none -%}<|effort={{ reasoning_effort }}|>{%- endif -%}
|
||||
{%- if add_generation_prompt -%}<|Assistant|>{%- if enable_thinking -%}<think>{%- else -%}</think>{%- endif -%}{%- endif -%})";
|
||||
|
||||
static std::unique_ptr<llama_grammar> build_grammar(const std::string & grammar_str) {
|
||||
return std::unique_ptr<llama_grammar>(
|
||||
llama_grammar_init_impl(nullptr, grammar_str.c_str(), "root", false, nullptr, 0, nullptr, 0));
|
||||
@@ -961,6 +976,7 @@ const common_chat_msg message_user{
|
||||
/* .reasoning_content = */ "",
|
||||
/* .tool_name = */ "",
|
||||
/* .tool_call_id = */ "",
|
||||
/* .extra_json = */ "",
|
||||
};
|
||||
|
||||
const common_chat_msg message_user_parts{
|
||||
@@ -977,6 +993,7 @@ const common_chat_msg message_user_parts{
|
||||
"",
|
||||
/* .tool_name = */ "",
|
||||
/* .tool_call_id = */ "",
|
||||
/* .extra_json = */ "",
|
||||
};
|
||||
|
||||
static common_chat_msg simple_assist_msg(const std::string & content,
|
||||
@@ -1486,6 +1503,11 @@ class peg_tester {
|
||||
template_path_(template_path),
|
||||
detailed_debug_(detailed_debug) {}
|
||||
|
||||
peg_tester(const std::string & name, common_chat_templates_ptr tmpls, const bool detailed_debug = false) :
|
||||
tmpls_(std::move(tmpls)),
|
||||
template_path_(name),
|
||||
detailed_debug_(detailed_debug) {}
|
||||
|
||||
const std::string & template_path() const { return template_path_; }
|
||||
|
||||
peg_test_builder test(const std::string & input);
|
||||
@@ -4381,6 +4403,103 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
|
||||
.run();
|
||||
}
|
||||
|
||||
// DeepSeek V4.1: V4 behaviour, DSML tag names with a leading space
|
||||
{
|
||||
auto tst = peg_tester("deepseek-v4.1-min",
|
||||
common_chat_templates_ptr(common_chat_templates_init(nullptr, DEEPSEEK_V4_1_MIN_TEMPLATE)),
|
||||
detailed_debug);
|
||||
|
||||
tst.test("Hello, world!\nWhat's up?")
|
||||
.enable_thinking(false)
|
||||
.reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
|
||||
.expect(message_assist)
|
||||
.run();
|
||||
|
||||
tst.test("I'm\nthinking</think>Hello, world!\nWhat's up?")
|
||||
.enable_thinking(true)
|
||||
.reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
|
||||
.expect(message_assist_thoughts)
|
||||
.run();
|
||||
|
||||
tst.test(
|
||||
"Let me check the time</think>\n\n"
|
||||
"<|DSML| calls>\n"
|
||||
"<|DSML| invoke name=\"get_time\">\n"
|
||||
"<|DSML| parameter name=\"city\" string=\"true\">Tokyo</|DSML| parameter>\n"
|
||||
"</|DSML| invoke>\n"
|
||||
"</|DSML| calls>")
|
||||
.enable_thinking(true)
|
||||
.reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
|
||||
.tools({ get_time_tool })
|
||||
.expect(message_with_tool_calls_and_reasoning("get_time", R"({"city": "Tokyo"})", "Let me check the time"))
|
||||
.run();
|
||||
|
||||
tst.test(
|
||||
"<|DSML| calls>\n"
|
||||
"<|DSML| invoke name=\"special_function\">\n"
|
||||
"<|DSML| parameter name=\"arg1\" string=\"false\">1</|DSML| parameter>\n"
|
||||
"</|DSML| invoke>\n"
|
||||
"</|DSML| calls>")
|
||||
.enable_thinking(false)
|
||||
.reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
|
||||
.tools({ special_function_tool })
|
||||
.expect(message_assist_call)
|
||||
.run();
|
||||
|
||||
tst.test(
|
||||
"Calling both</think>\n\n"
|
||||
"<|DSML| calls>\n"
|
||||
"<|DSML| invoke name=\"get_time\">\n"
|
||||
"<|DSML| parameter name=\"city\" string=\"true\">Paris</|DSML| parameter>\n"
|
||||
"</|DSML| invoke>\n"
|
||||
"<|DSML| invoke name=\"get_weather\">\n"
|
||||
"<|DSML| parameter name=\"city\" string=\"true\">Paris</|DSML| parameter>\n"
|
||||
"</|DSML| invoke>\n"
|
||||
"</|DSML| calls>")
|
||||
.enable_thinking(true)
|
||||
.reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
|
||||
.parallel_tool_calls(true)
|
||||
.tools({ get_time_tool, get_weather_tool })
|
||||
.expect(message_with_reasoning_content_and_multiple_tool_calls(
|
||||
"Calling both", "",
|
||||
{ { "get_time", R"({"city": "Paris"})" }, { "get_weather", R"({"city": "Paris"})" } }))
|
||||
.run();
|
||||
|
||||
tst.test("world!\nWhat's up?")
|
||||
.reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
|
||||
.enable_thinking(true)
|
||||
.messages({ message_user, message_assist_prefill_content })
|
||||
.add_generation_prompt(false)
|
||||
.continue_final_message(COMMON_CHAT_CONTINUATION_CONTENT)
|
||||
.expect_reasoning("I'm thinking")
|
||||
.expect_content("Hello, world!\nWhat's up?")
|
||||
.run();
|
||||
|
||||
tst.test(" thinking</think>Hello, world!\nWhat's up?")
|
||||
.reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
|
||||
.enable_thinking(true)
|
||||
.messages({ message_user, message_assist_prefill_reasoning })
|
||||
.add_generation_prompt(false)
|
||||
.continue_final_message(COMMON_CHAT_CONTINUATION_REASONING)
|
||||
.expect_reasoning("I'm thinking")
|
||||
.expect_content("Hello, world!\nWhat's up?")
|
||||
.run();
|
||||
|
||||
tst.test(
|
||||
"Let me check the time\n\n"
|
||||
"<|DSML| calls>\n"
|
||||
"<|DSML| invoke name=\"get_time\">\n"
|
||||
"<|DSML| parameter name=\"city\" string=\"true\">Tokyo</|DSML| parameter>\n"
|
||||
"</|DSML| invoke>\n"
|
||||
"</|DSML| calls>") // no </think> after the TC close because the grammar will immediately constrain it to end
|
||||
.enable_thinking(true)
|
||||
.reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
|
||||
.tools({ get_time_tool })
|
||||
.expect_reasoning("Let me check the time")
|
||||
.expect_tool_calls({ { "get_time", R"({"city": "Tokyo"})", {} } })
|
||||
.run();
|
||||
}
|
||||
|
||||
{
|
||||
// The DSML separator belongs to the tool call block, not assistant content.
|
||||
auto tst = peg_tester("models/templates/deepseek-ai-DeepSeek-V4-Flash-0731.jinja", detailed_debug);
|
||||
@@ -7825,6 +7944,62 @@ static void test_deepseek_v4_tool_result_ordering() {
|
||||
}
|
||||
}
|
||||
|
||||
// DeepSeek V4.1: extra message keys, "\n\n" around media markers, reasoning_effort mapping, system-message tools
|
||||
static void test_deepseek_v4_1_message_fields() {
|
||||
LOG_DBG("%s\n", __func__);
|
||||
|
||||
common_chat_templates_ptr tmpls(common_chat_templates_init(nullptr, DEEPSEEK_V4_1_MIN_TEMPLATE));
|
||||
auto render = [&](const char * messages_json, bool add_generation_prompt) {
|
||||
common_chat_templates_inputs inputs;
|
||||
inputs.messages = common_chat_msgs_parse_oaicompat(json::parse(messages_json));
|
||||
inputs.add_generation_prompt = add_generation_prompt;
|
||||
inputs.enable_thinking = false;
|
||||
return common_chat_templates_apply(tmpls.get(), inputs).prompt.str();
|
||||
};
|
||||
|
||||
assert_equals(std::string("<|User|>q<|action|>"),
|
||||
render(R"([{"role": "user", "content": "q", "task": "action"}])", false));
|
||||
|
||||
assert_contains(render(R"([
|
||||
{"role": "user", "content": "q"},
|
||||
{"role": "assistant", "content": "", "tool_calls": [{"type": "function", "namespace": "ns", "function": {"name": "f", "arguments": {"x": 1}}}]}
|
||||
])", false),
|
||||
"<|DSML| invoke name=\"ns::f\">");
|
||||
|
||||
assert_equals(std::string("<|User|>a\n\n<__media__>\n\nb<|Assistant|></think>"),
|
||||
render(R"([{"role": "user", "content": [{"type": "text", "text": "a"}, {"type": "media_marker", "text": "<__media__>"}, {"type": "text", "text": "b"}]}])", true));
|
||||
|
||||
for (const auto & [effort, expected] : std::vector<std::pair<std::string, std::string>>{
|
||||
{ "minimal", "<|effort=low|>" }, { "low", "<|effort=low|>" }, { "medium", "" },
|
||||
{ "high", "<|effort=high|>" }, { "xhigh", "<|effort=max|>" }, { "max", "<|effort=max|>" } }) {
|
||||
common_chat_templates_inputs inputs;
|
||||
inputs.messages = { message_user };
|
||||
inputs.add_generation_prompt = true;
|
||||
inputs.chat_template_kwargs["reasoning_effort"] = json(effort).dump();
|
||||
assert_equals("<|User|>Hey there!" + expected + "<|Assistant|><think>",
|
||||
common_chat_templates_apply(tmpls.get(), inputs).prompt.str());
|
||||
}
|
||||
|
||||
{
|
||||
common_chat_templates_inputs inputs;
|
||||
inputs.messages = common_chat_msgs_parse_oaicompat(json::parse(R"([
|
||||
{"role": "system", "content": "s", "tools": [{"type": "function", "function": {"name": "get_time",
|
||||
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}]},
|
||||
{"role": "user", "content": "q"}
|
||||
])"));
|
||||
inputs.add_generation_prompt = true;
|
||||
inputs.enable_thinking = false;
|
||||
inputs.reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK;
|
||||
common_chat_session session(tmpls.get(), nullptr, inputs, {});
|
||||
const auto & msg = session.finish(common_chat_input(std::string(
|
||||
"\n\n<|DSML| calls>\n<|DSML| invoke name=\"get_time\">\n"
|
||||
"<|DSML| parameter name=\"city\" string=\"true\">Tokyo</|DSML| parameter>\n</|DSML| invoke>\n</|DSML| calls>")));
|
||||
assert_equals(size_t(1), msg.tool_calls.size());
|
||||
assert_equals(std::string("get_time"), msg.tool_calls[0].name);
|
||||
assert_equals(std::string("{\"city\":\"Tokyo\"}"), msg.tool_calls[0].arguments);
|
||||
}
|
||||
}
|
||||
|
||||
static void test_prompt_parts() {
|
||||
LOG_DBG("%s\n", __func__);
|
||||
const std::string tmpl_src =
|
||||
@@ -8175,6 +8350,7 @@ int main(int argc, char ** argv) {
|
||||
test_translate_gemma();
|
||||
test_deepseek_v4_thinking_retention();
|
||||
test_deepseek_v4_tool_result_ordering();
|
||||
test_deepseek_v4_1_message_fields();
|
||||
test_template_generation_prompt();
|
||||
test_reasoning_effort_caps();
|
||||
test_prompt_parts();
|
||||
|
||||
@@ -1804,6 +1804,54 @@ static void test_array_methods(testing & t) {
|
||||
"z y "
|
||||
);
|
||||
|
||||
test_template(t, "array|select with test",
|
||||
"{{ items|select('odd')|join(',') }}",
|
||||
{{"items", json::array({1, 2, 3, 4, 5})}},
|
||||
"1,3,5"
|
||||
);
|
||||
|
||||
test_template(t, "array|reject with test",
|
||||
"{{ items|reject('odd')|join(',') }}",
|
||||
{{"items", json::array({1, 2, 3, 4, 5})}},
|
||||
"2,4"
|
||||
);
|
||||
|
||||
test_template(t, "array|select without test",
|
||||
"{{ items|select|join(',') }}",
|
||||
{{"items", json::array({0, 1, "", "a", nullptr, false})}},
|
||||
"1,a"
|
||||
);
|
||||
|
||||
test_template(t, "array|selectattr with test",
|
||||
"{% for p in items|selectattr('type', 'defined') %}{{ p.type }} {% endfor %}",
|
||||
{{"items", json::array({
|
||||
{{"type", "string"}},
|
||||
{{"enum", json::array({1})}},
|
||||
{{"type", "integer"}}
|
||||
})}},
|
||||
"string integer "
|
||||
);
|
||||
|
||||
test_template(t, "array|rejectattr with test",
|
||||
"{{ items|rejectattr('type', 'defined')|list|length }}",
|
||||
{{"items", json::array({
|
||||
{{"type", "string"}},
|
||||
{{"enum", json::array({1})}},
|
||||
{{"type", "integer"}}
|
||||
})}},
|
||||
"1"
|
||||
);
|
||||
|
||||
test_template(t, "array|selectattr with attr named like a test",
|
||||
"{% for p in items|selectattr('odd', 'even') %}{{ p.odd }} {% endfor %}",
|
||||
{{"items", json::array({
|
||||
{{"odd", 1}},
|
||||
{{"odd", 2}},
|
||||
{{"odd", 4}}
|
||||
})}},
|
||||
"2 4 "
|
||||
);
|
||||
|
||||
test_template(t, "array|tojson",
|
||||
"{{ arr|tojson }}",
|
||||
{{"arr", json::array({1, 2, 3})}},
|
||||
|
||||
@@ -195,6 +195,12 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
|
||||
n_head = 8;
|
||||
n_ff = 1024;
|
||||
n_layer = 4;
|
||||
} else if (arch == LLM_ARCH_DEEPSEEK41) {
|
||||
// window-only, Engram, ratio-2 and ratio-1 sources, top-k reuse and a candidate-masked index layer
|
||||
n_embd = 512;
|
||||
n_head = 8;
|
||||
n_ff = 1024;
|
||||
n_layer = 6;
|
||||
} else if (arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_LAGUNA) {
|
||||
n_embd = 160; // exercise per-head tensor split granularity with head size 80
|
||||
} else if (arch == LLM_ARCH_QWEN3 || arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE) {
|
||||
@@ -285,11 +291,11 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
|
||||
ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_per_layer);
|
||||
} else {
|
||||
ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head);
|
||||
ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(1) : n_head_kv);
|
||||
ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41) ? uint32_t(1) : n_head_kv);
|
||||
}
|
||||
|
||||
ms.add_kv(LLM_KV_ATTENTION_MAX_ALIBI_BIAS, 8.0f);
|
||||
if (arch == LLM_ARCH_DEEPSEEK4) {
|
||||
if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41) {
|
||||
ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, n_embd_head);
|
||||
ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, n_embd_head);
|
||||
ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, n_embd_head/2);
|
||||
@@ -337,7 +343,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
|
||||
ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, 1e-5f);
|
||||
ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_EPS, 1e-5f);
|
||||
ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_GROUPS, uint32_t(8));
|
||||
ms.add_kv(LLM_KV_ATTENTION_Q_LORA_RANK, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(64) : uint32_t(512));
|
||||
ms.add_kv(LLM_KV_ATTENTION_Q_LORA_RANK, (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41) ? uint32_t(64) : uint32_t(512));
|
||||
ms.add_kv(LLM_KV_ATTENTION_KV_LORA_RANK, uint32_t(512));
|
||||
ms.add_kv(LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT, uint32_t(8));
|
||||
ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW, n_ctx/8);
|
||||
@@ -476,6 +482,39 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
|
||||
ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true);
|
||||
}
|
||||
|
||||
if (arch == LLM_ARCH_DEEPSEEK41) {
|
||||
ms.add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, uint32_t(8));
|
||||
ms.add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, uint32_t(32));
|
||||
ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector<uint32_t>({0, 0, 2, 2, 1, 1}));
|
||||
ms.add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, 160000.0f);
|
||||
ms.add_kv(LLM_KV_ATTENTION_KV_SOURCE_LAYERS, std::vector<uint32_t>({2, 4}));
|
||||
ms.add_kv(LLM_KV_ATTENTION_INDEXER_SOURCE_LAYERS, std::vector<uint32_t>({2, 4, 5}));
|
||||
ms.add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_SOURCE_LAYER, int32_t(4));
|
||||
ms.add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS, uint32_t(4096));
|
||||
ms.add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE, uint32_t(4));
|
||||
ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4));
|
||||
ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(2));
|
||||
ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, 1.0e-6f);
|
||||
ms.add_kv(LLM_KV_HASH_LAYER_COUNT, uint32_t(0));
|
||||
ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, 10.0f);
|
||||
ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f);
|
||||
ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true);
|
||||
std::vector<int32_t> token_map(n_vocab);
|
||||
for (uint32_t i = 0; i < n_vocab; ++i) {
|
||||
token_map[i] = (int32_t) (i % 61);
|
||||
}
|
||||
ms.add_kv(LLM_KV_PLE_LAYERS, std::vector<uint32_t>({1, 3}));
|
||||
ms.add_kv(LLM_KV_PLE_NGRAM_SIZE, uint32_t(3));
|
||||
ms.add_kv(LLM_KV_PLE_HEADS_PER_NGRAM, uint32_t(2));
|
||||
ms.add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, uint32_t(64));
|
||||
ms.add_kv(LLM_KV_PLE_LAYER_MULTIPLIERS, std::vector<uint64_t>({ 76632096046245ull, 4839876093313ull, 35959672319349ull,
|
||||
67716810739261ull, 51510806800915ull, 30921347202721ull }));
|
||||
ms.add_kv(LLM_KV_PLE_HEAD_OFFSETS, std::vector<uint64_t>({0, 101, 204, 311, 420, 533, 660, 791}));
|
||||
ms.add_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES, std::vector<uint64_t>({101, 103, 107, 109, 113, 127, 131, 137}));
|
||||
ms.add_kv(LLM_KV_ENGRAM_PAD_TOKEN_ID, uint32_t(2));
|
||||
ms.add_kv(LLM_KV_ENGRAM_TOKEN_MAP, token_map);
|
||||
}
|
||||
|
||||
if (arch == LLM_ARCH_HRM_TEXT) {
|
||||
// 6 cache slots alias 2 physical blocks: 1 low-stack layer + 1 high-stack layer
|
||||
ms.add_kv(LLM_KV_HRM_LAYERS_PER_STACK, uint32_t(1));
|
||||
@@ -519,6 +558,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
|
||||
gating_func = 1;
|
||||
break;
|
||||
case LLM_ARCH_DEEPSEEK4:
|
||||
case LLM_ARCH_DEEPSEEK41:
|
||||
gating_func = 4;
|
||||
break;
|
||||
default:
|
||||
@@ -780,6 +820,7 @@ static bool moe_mandatory(const llm_arch arch) {
|
||||
case LLM_ARCH_DEEPSEEK32:
|
||||
case LLM_ARCH_DOTS3NOTE:
|
||||
case LLM_ARCH_DEEPSEEK4:
|
||||
case LLM_ARCH_DEEPSEEK41:
|
||||
case LLM_ARCH_GLM4_MOE:
|
||||
case LLM_ARCH_GLM_DSA:
|
||||
case LLM_ARCH_EXAONE_MOE:
|
||||
|
||||
@@ -60,10 +60,12 @@ static llama_tokens prepare_tokens(llama_model * model, const common_params & pa
|
||||
llama_tokens tokens;
|
||||
|
||||
if (params.prompt.empty()) {
|
||||
const int n_prompt = params.n_batch;
|
||||
// the prompt must not exceed n_batch and n_ctx of cases
|
||||
constexpr int n_prompt_max = 100;
|
||||
const int n_prompt = std::min(params.n_batch, n_prompt_max);
|
||||
|
||||
// this path is useful for model files that do not have a tokenizer
|
||||
TLOG_INF("%s: no prompt provided, generating %d (n_batch) random tokens\n", __func__, n_prompt);
|
||||
TLOG_INF("%s: no prompt provided, generating %d random tokens\n", __func__, n_prompt);
|
||||
|
||||
const auto * vocab = llama_model_get_vocab(model);
|
||||
const auto n_vocab = llama_vocab_n_tokens(vocab);
|
||||
@@ -205,7 +207,6 @@ int main(int argc, char ** argv) {
|
||||
|
||||
common_params params;
|
||||
params.prompt = "";
|
||||
params.n_batch = 100;
|
||||
params.out_file = "dump_state.bin";
|
||||
params.sampling.seed = 1234;
|
||||
|
||||
|
||||
@@ -529,6 +529,7 @@ enum projector_type {
|
||||
PROJECTOR_TYPE_DEEPSEEKOCR,
|
||||
PROJECTOR_TYPE_DEEPSEEKOCR2,
|
||||
PROJECTOR_TYPE_DEEPSEEK4V,
|
||||
PROJECTOR_TYPE_DEEPSEEK41V,
|
||||
PROJECTOR_TYPE_LFM2A,
|
||||
PROJECTOR_TYPE_D1OMNI_A,
|
||||
PROJECTOR_TYPE_GLM4V,
|
||||
@@ -606,6 +607,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
|
||||
{ PROJECTOR_TYPE_DEEPSEEKOCR, "deepseekocr"},
|
||||
{ PROJECTOR_TYPE_DEEPSEEKOCR2, "deepseekocr2"},
|
||||
{ PROJECTOR_TYPE_DEEPSEEK4V, "deepseek4v"},
|
||||
{ PROJECTOR_TYPE_DEEPSEEK41V, "deepseek41v"},
|
||||
{ PROJECTOR_TYPE_LFM2A, "lfm2a"},
|
||||
{ PROJECTOR_TYPE_D1OMNI_A, "d1omni_a"},
|
||||
{ PROJECTOR_TYPE_GLM4V, "glm4v"},
|
||||
@@ -879,6 +881,11 @@ static inline void clip_anyres_unpad(int cur_w, int cur_h, int orig_w, int orig_
|
||||
}
|
||||
}
|
||||
|
||||
// deepseek41v: [START] ([IMAGE]*n_llm_w [NEWLINE])*n_llm_h [END], ref: num_image_tokens() in image_processor.py
|
||||
static inline int dsv41_n_block_tokens(int n_llm_w, int n_llm_h) {
|
||||
return n_llm_h * (n_llm_w + 1) + 2;
|
||||
}
|
||||
|
||||
// deepseek4v: layout of the LLM token block built from the aligner grid
|
||||
struct dsv4_block_layout {
|
||||
int rows; // grid rows, padded to an even count
|
||||
|
||||
+73
-1
@@ -1111,6 +1111,7 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
|
||||
builder = std::make_unique<clip_graph_kimik25>(ctx, img);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_DEEPSEEK4V:
|
||||
case PROJECTOR_TYPE_DEEPSEEK41V:
|
||||
{
|
||||
builder = std::make_unique<clip_graph_deepseek4v>(ctx, img);
|
||||
} break;
|
||||
@@ -1773,6 +1774,30 @@ struct clip_model_loader {
|
||||
const int warmup_side = (int) std::sqrt((double) std::min(256, hparams.dsv4_max_n_token));
|
||||
hparams.set_warmup_n_tokens(warmup_side * warmup_side);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_DEEPSEEK41V:
|
||||
{
|
||||
// ref: load_image in inference/image_processor.py (V4.1), PIL ImageOps.pad on gray 127
|
||||
hparams.image_resize_algo = RESIZE_ALGO_BICUBIC;
|
||||
hparams.image_pad_color = {127, 127, 127};
|
||||
hparams.rope_theta = 10000.0f;
|
||||
get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge);
|
||||
get_u32(KEY_IMAGE_MIN_PIXELS, hparams.image_min_pixels);
|
||||
// vision_max_n_token = 1024, vision_max_wh_ratio = null (converter asserts both)
|
||||
hparams.dsv4_max_n_token = 1024;
|
||||
hparams.dsv4_max_wh_ratio = 0;
|
||||
const int patch_area = hparams.patch_size * hparams.patch_size * hparams.n_merge * hparams.n_merge;
|
||||
if (hparams.custom_image_min_tokens > 0) {
|
||||
hparams.image_min_pixels = hparams.custom_image_min_tokens * patch_area;
|
||||
}
|
||||
if (hparams.custom_image_max_tokens > 0) {
|
||||
// the cap is on the whole token block (newlines + start/end included)
|
||||
hparams.dsv4_max_n_token = std::max(hparams.custom_image_max_tokens, 16);
|
||||
}
|
||||
hparams.image_max_pixels = hparams.dsv4_max_n_token * patch_area;
|
||||
hparams.image_min_pixels = std::min(hparams.image_min_pixels, hparams.image_max_pixels);
|
||||
const int warmup_side = (int) std::sqrt((double) std::min(256, hparams.dsv4_max_n_token));
|
||||
hparams.set_warmup_n_tokens(warmup_side * warmup_side);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_GEMMA3:
|
||||
{
|
||||
// default value (used by all model sizes in gemma 3 family)
|
||||
@@ -3065,6 +3090,17 @@ struct clip_model_loader {
|
||||
model.token_embd_img_end = get_tensor(TN_TOK_IMG_END);
|
||||
model.token_embd_img_pad = get_tensor(TN_TOK_IMG_PAD);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_DEEPSEEK41V:
|
||||
{
|
||||
model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 1, "weight"));
|
||||
model.mm_1_b = get_tensor(string_format(TN_LLAVA_PROJ, 1, "bias"));
|
||||
model.mm_2_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight"));
|
||||
model.mm_2_b = get_tensor(string_format(TN_LLAVA_PROJ, 2, "bias"));
|
||||
// image_start / image_end / image_newline, written into the output block
|
||||
model.image_newline = get_tensor(TN_IMAGE_NEWLINE);
|
||||
model.token_embd_img_start = get_tensor(TN_TOK_IMG_START);
|
||||
model.token_embd_img_end = get_tensor(TN_TOK_IMG_END);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_PIXTRAL:
|
||||
{
|
||||
model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 1, "weight"));
|
||||
@@ -4377,7 +4413,7 @@ struct clip_cap clip_get_cap(const char * fname) {
|
||||
|
||||
int clip_get_image_max_tokens(const clip_ctx * ctx) {
|
||||
const auto & hparams = ctx->model.hparams;
|
||||
if (ctx->proj_type() == PROJECTOR_TYPE_DEEPSEEK4V) {
|
||||
if (ctx->proj_type() == PROJECTOR_TYPE_DEEPSEEK4V || ctx->proj_type() == PROJECTOR_TYPE_DEEPSEEK41V) {
|
||||
return hparams.dsv4_max_n_token;
|
||||
}
|
||||
if (hparams.image_max_pixels <= 0) {
|
||||
@@ -4585,6 +4621,13 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
|
||||
const int n_llm_h = CLIP_ALIGN(img->ny(), out_patch_size) / out_patch_size;
|
||||
n_patches = dsv4_get_block_layout(n_llm_w, n_llm_h, img->lead_pad).n_out;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_DEEPSEEK41V:
|
||||
{
|
||||
const int out_patch_size = params.patch_size * params.n_merge;
|
||||
const int n_llm_w = CLIP_ALIGN(img->nx(), out_patch_size) / out_patch_size;
|
||||
const int n_llm_h = CLIP_ALIGN(img->ny(), out_patch_size) / out_patch_size;
|
||||
n_patches = dsv41_n_block_tokens(n_llm_w, n_llm_h);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_PADDLEOCR:
|
||||
case PROJECTOR_TYPE_DOTS_OCR:
|
||||
case PROJECTOR_TYPE_DOTS3NOTE_V:
|
||||
@@ -5692,6 +5735,34 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
idx.push_back(idx_end);
|
||||
set_input_i32("layout_idx", idx);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_DEEPSEEK41V:
|
||||
{
|
||||
// mrope layout, only the first 2 channels are used
|
||||
const int n_patches_per_row = image_size_width / patch_size;
|
||||
std::vector<int32_t> positions(n_pos * 4, 0);
|
||||
for (int i = 0; i < n_pos; i++) {
|
||||
positions[i] = i / n_patches_per_row;
|
||||
positions[n_pos + i] = i % n_patches_per_row;
|
||||
}
|
||||
set_input_i32("positions", positions);
|
||||
|
||||
// reading-order block: rows [0, n_grid) are the aligner output, then start / end / newline
|
||||
const int n_merge = hparams.n_merge;
|
||||
const int n_llm_w = CLIP_ALIGN(pos_w, n_merge) / n_merge;
|
||||
const int n_llm_h = CLIP_ALIGN(pos_h, n_merge) / n_merge;
|
||||
const int n_grid = n_llm_w * n_llm_h;
|
||||
std::vector<int32_t> idx;
|
||||
idx.reserve(dsv41_n_block_tokens(n_llm_w, n_llm_h));
|
||||
idx.push_back(n_grid); // IMAGE_START
|
||||
for (int r = 0; r < n_llm_h; r++) {
|
||||
for (int c = 0; c < n_llm_w; c++) {
|
||||
idx.push_back(r * n_llm_w + c);
|
||||
}
|
||||
idx.push_back(n_grid + 2); // IMAGE_NEW_LINE
|
||||
}
|
||||
idx.push_back(n_grid + 1); // IMAGE_END
|
||||
set_input_i32("layout_idx", idx);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_GLM_EDGE:
|
||||
{
|
||||
// llava and other models
|
||||
@@ -6583,6 +6654,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
|
||||
case PROJECTOR_TYPE_KIMIK25:
|
||||
case PROJECTOR_TYPE_YASA2:
|
||||
case PROJECTOR_TYPE_DEEPSEEK4V:
|
||||
case PROJECTOR_TYPE_DEEPSEEK41V:
|
||||
case PROJECTOR_TYPE_COHERE2V:
|
||||
return ctx->model.mm_2_w->ne[1];
|
||||
case PROJECTOR_TYPE_HUNYUANVL:
|
||||
|
||||
@@ -14,6 +14,9 @@
|
||||
// the mapping is precomputed on CPU as the "layout_idx" input (see set_input in clip.cpp)
|
||||
//
|
||||
// ref: inference/vision.py and inference/image_processor.py in the HF repo
|
||||
//
|
||||
// DeepSeek-V4.1-Flash (deepseek41v) shares the encoder and aligner; its block is plain reading order
|
||||
// [START] ([IMAGE]*n_llm_w [NEWLINE])*n_llm_h [END], no lead pads, no pad row, no IMAGE_PAD tensor
|
||||
|
||||
ggml_cgraph * clip_graph_deepseek4v::build() {
|
||||
const int n_merge = hparams.n_merge;
|
||||
@@ -52,7 +55,10 @@ ggml_cgraph * clip_graph_deepseek4v::build() {
|
||||
cur = ggml_pad(ctx0, cur, pad_x, pad_y, 0, 0);
|
||||
}
|
||||
|
||||
ggml_tensor * kernel = ggml_view_3d(ctx0, cur, n_merge, n_merge, cur->ne[2], 0, 0, 0);
|
||||
// contiguous kernel view: CUDA rejects a zero-stride one and im2col falls back to CPU
|
||||
const size_t esz = ggml_element_size(cur);
|
||||
ggml_tensor * kernel = ggml_view_3d(ctx0, cur, n_merge, n_merge, cur->ne[2],
|
||||
n_merge*esz, n_merge*n_merge*esz, 0);
|
||||
cur = ggml_im2col(ctx0, kernel, cur, n_merge, n_merge, 0, 0, 1, 1, true, inp->type);
|
||||
cur = ggml_reshape_2d(ctx0, cur, cur->ne[0], cur->ne[1] * cur->ne[2]);
|
||||
|
||||
@@ -73,19 +79,23 @@ ggml_cgraph * clip_graph_deepseek4v::build() {
|
||||
const int64_t n_grid = cur->ne[1]; // n_llm_w * n_llm_h
|
||||
|
||||
// rows n_grid + 0..3, keep in sync with the index computation in set_input
|
||||
const bool is_v41 = proj_type == PROJECTOR_TYPE_DEEPSEEK41V;
|
||||
ggml_tensor * sentinels[] = {
|
||||
model.token_embd_img_start,
|
||||
model.token_embd_img_end,
|
||||
model.image_newline,
|
||||
model.token_embd_img_pad,
|
||||
is_v41 ? nullptr : model.token_embd_img_pad,
|
||||
};
|
||||
for (ggml_tensor * tok : sentinels) {
|
||||
cur = ggml_concat(ctx0, cur, ggml_reshape_2d(ctx0, tok, n_embd_out, 1), 1);
|
||||
if (tok) {
|
||||
cur = ggml_concat(ctx0, cur, ggml_reshape_2d(ctx0, tok, n_embd_out, 1), 1);
|
||||
}
|
||||
}
|
||||
|
||||
const int n_llm_w = CLIP_ALIGN(n_patches_x, n_merge) / n_merge;
|
||||
const int n_llm_h = CLIP_ALIGN(n_patches_y, n_merge) / n_merge;
|
||||
const int n_out = dsv4_get_block_layout(n_llm_w, n_llm_h, img.lead_pad).n_out;
|
||||
const int n_out = is_v41 ? dsv41_n_block_tokens(n_llm_w, n_llm_h)
|
||||
: dsv4_get_block_layout(n_llm_w, n_llm_h, img.lead_pad).n_out;
|
||||
GGML_ASSERT(n_grid == n_llm_w * n_llm_h);
|
||||
|
||||
ggml_tensor * layout_idx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_out);
|
||||
|
||||
@@ -1454,6 +1454,100 @@ mtmd_image_preproc_out mtmd_image_preprocessor_deepseek4v::preprocess(const clip
|
||||
return out;
|
||||
}
|
||||
|
||||
//
|
||||
// DeepSeek-V4.1-Flash (deepseek41v)
|
||||
//
|
||||
// port of plan_image_grid / load_image from inference/image_processor.py (V4.1)
|
||||
// differs from deepseek4v: block = h * (w + 1) + 2 tokens, no even-row rounding, a single solve, no stretch path
|
||||
//
|
||||
|
||||
void mtmd_image_preprocessor_deepseek41v::plan_grid(int width, int height, int p, int r, int max_n_token, int min_pixels,
|
||||
int & best_width, int & best_height) {
|
||||
auto n_tokens = [&](int bh, int bw) {
|
||||
const int n_llm_h = ((bh / p) + r - 1) / r;
|
||||
const int n_llm_w = ((bw / p) + r - 1) / r;
|
||||
return n_llm_h * (n_llm_w + 1) + 2;
|
||||
};
|
||||
if (min_pixels > 0 && (int64_t) width * height > 0 && (int64_t) width * height < min_pixels) {
|
||||
// python: ratio = (min_pixels / (w * h)) ** 0.5; int(w * ratio)
|
||||
const double ratio = std::sqrt((double) min_pixels / ((double) width * (double) height));
|
||||
width = (int) (width * ratio);
|
||||
height = (int) (height * ratio);
|
||||
}
|
||||
best_width = CLIP_ALIGN(width, p);
|
||||
best_height = CLIP_ALIGN(height, p);
|
||||
if (n_tokens(best_height, best_width) > max_n_token) {
|
||||
// solve_resize_ratio: largest aspect-preserving size whose token block fits
|
||||
const double ratio = (double) height / width;
|
||||
const double max_w_f = std::sqrt((max_n_token - 2) / ratio + 0.25) - 0.5;
|
||||
const double max_h_f = max_w_f * ratio;
|
||||
const int cell = p * r;
|
||||
if (max_w_f < 1.0) {
|
||||
best_height = (max_n_token - 2) / 2 * cell;
|
||||
best_width = cell;
|
||||
} else if (max_h_f < 1.0) {
|
||||
best_height = cell;
|
||||
best_width = (max_n_token - 3) * cell;
|
||||
} else {
|
||||
const double beta = std::min(std::floor(max_w_f) * cell / width, std::floor(max_h_f) * cell / height);
|
||||
best_height = (int) std::floor(height * beta / p) * p;
|
||||
best_width = (int) std::floor(width * beta / p) * p;
|
||||
}
|
||||
GGML_ASSERT(n_tokens(best_height, best_width) <= max_n_token);
|
||||
}
|
||||
}
|
||||
|
||||
mtmd_image_preproc_out mtmd_image_preprocessor_deepseek41v::preprocess(const clip_image_u8 & img) const {
|
||||
mtmd_image_preproc_out out;
|
||||
|
||||
const int p = hparams.patch_size;
|
||||
const clip_image_size orig = img.get_size();
|
||||
int best_width = 0, best_height = 0;
|
||||
plan_grid(orig.width, orig.height, p, hparams.n_merge, hparams.dsv4_max_n_token, hparams.image_min_pixels,
|
||||
best_width, best_height);
|
||||
|
||||
// PIL ImageOps.pad(color=127): ImageOps.contain then a centered paste; python round is half to even
|
||||
auto round_half_even = [](double x) { return (int) std::nearbyint(x); };
|
||||
int new_w = best_width;
|
||||
int new_h = best_height;
|
||||
const double im_ratio = (double) orig.width / orig.height;
|
||||
const double dest_ratio = (double) best_width / best_height;
|
||||
if (im_ratio != dest_ratio) {
|
||||
if (im_ratio > dest_ratio) {
|
||||
new_h = round_half_even((double) orig.height / orig.width * best_width);
|
||||
} else {
|
||||
new_w = round_half_even((double) orig.width / orig.height * best_height);
|
||||
}
|
||||
}
|
||||
clip_image_u8 resized;
|
||||
img_tool::resize(img, resized, {new_w, new_h}, hparams.image_resize_algo, PAD_NONE);
|
||||
|
||||
clip_image_u8 padded;
|
||||
if (new_w == best_width && new_h == best_height) {
|
||||
padded = std::move(resized);
|
||||
} else {
|
||||
padded.set_size({best_width, best_height}, img.is_placeholder());
|
||||
if (!img.is_placeholder()) {
|
||||
img_tool::fill(padded, hparams.image_pad_color);
|
||||
const int off_x = round_half_even((best_width - new_w) * 0.5);
|
||||
const int off_y = round_half_even((best_height - new_h) * 0.5);
|
||||
img_tool::composite(padded, resized, off_x, off_y);
|
||||
}
|
||||
}
|
||||
|
||||
out.append(hparams, padded);
|
||||
|
||||
// load_image rounds the normalized input to bf16 (round to nearest even)
|
||||
for (auto & e : out.entries) {
|
||||
std::vector<float> buf = e.get_ro_buf();
|
||||
for (float & v : buf) {
|
||||
v = ggml_bf16_to_fp32(ggml_fp32_to_bf16(v));
|
||||
}
|
||||
e.cpy_buf(buf);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
mtmd_image_preproc_out mtmd_image_preprocessor_deepseekocr::preprocess(const clip_image_u8 & img) const {
|
||||
mtmd_image_preproc_out output;
|
||||
int grid_w = 0;
|
||||
|
||||
@@ -166,6 +166,15 @@ private:
|
||||
static void safe_resize(int height, int width, int & best_height, int & best_width, int p, int r, int max_n_token);
|
||||
};
|
||||
|
||||
// ref: inference/image_processor.py in the HF repo (DeepSeek-V4.1-Flash): plan_image_grid / load_image
|
||||
struct mtmd_image_preprocessor_deepseek41v : mtmd_image_preprocessor {
|
||||
mtmd_image_preprocessor_deepseek41v(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {}
|
||||
mtmd_image_preproc_out preprocess(const clip_image_u8 & img) const override;
|
||||
// target pixel size (multiples of patch_size) for an image of the given size, exposed for tests
|
||||
static void plan_grid(int width, int height, int p, int r, int max_n_token, int min_pixels,
|
||||
int & best_width, int & best_height);
|
||||
};
|
||||
|
||||
// custom llava-uhd slicing logic for MiniCPM-V
|
||||
struct mtmd_image_preprocessor_minicpmv : mtmd_image_preprocessor_llava_uhd {
|
||||
using mtmd_image_preprocessor_llava_uhd::mtmd_image_preprocessor_llava_uhd;
|
||||
|
||||
@@ -903,6 +903,11 @@ struct mtmd_context {
|
||||
// no vocab tokens are added; the start/end/newline markers are learned embeddings emitted by the encoder
|
||||
image_preproc = std::make_unique<mtmd_image_preprocessor_deepseek4v>(ctx_v);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_DEEPSEEK41V:
|
||||
{
|
||||
// as deepseek4v: markers come from the encoder, the block is decoded causally
|
||||
image_preproc = std::make_unique<mtmd_image_preprocessor_deepseek41v>(ctx_v);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_DOTS_OCR:
|
||||
case PROJECTOR_TYPE_DOTS3NOTE_V:
|
||||
{
|
||||
|
||||
@@ -286,7 +286,7 @@ Instead of building everything from the ground up (like what most AI agents will
|
||||
|
||||
The flow for downloading a new model:
|
||||
- POST request comes in --> `post_router_models` --> validation
|
||||
- A new `llama-server` subprocess will be spawned with special `SERVER_CHILD_MODE_DOWNLOAD`
|
||||
- A new `llama-server` subprocess will be spawned with the `SERVER_ROLE_DOWNLOAD` role (see `server_role_type` in `server-common.h`)
|
||||
- Child process runs the download and report status back to router via stdin/out
|
||||
- If a stop request comes in, the router asks the child process to stop (same mechanism as running a model in child process)
|
||||
- Otherwise, upon completion, we call `load_models()` to refresh the list of models
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user